WL13847: Add support for context managers

Context managers in Python, allows the allocation and releasing of
resources precisely when it's needed.
The most widely used example of context managers is the with statement.

This worklog implements this functionality in the classic protocol for
Connection objects and in the X DevAPI protocol for Session and Client
objects.

Author: nmariz

lib/mysql/connector/abstracts.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2014, 2019, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2014, 2020, Oracle and/or its affiliates. All rights reserved.
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License, version 2.0, as
@@ -111,6 +111,12 @@ def __init__(self, **kwargs):
 
         self._consume_results = False
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def _get_self(self):
         """Return self for weakref.proxy
 
@@ -1232,6 +1238,12 @@ def __init__(self):
         self._warnings = None
         self.arraysize = 1
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     @abstractmethod
     def callproc(self, procname, args=()):
         """Calls a stored procedure with the given arguments

lib/mysqlx/connection.py

@@ -1845,6 +1845,12 @@ def __init__(self, settings):
                     else "Default schema '{}' does not exists".format(schema)
                 raise InterfaceError(errmsg, err.errno)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def _init_attributes(self):
         """Setup default and user defined connection-attributes."""
         if os.name == "nt":
@@ -2125,6 +2131,12 @@ def __init__(self, connection_dict, options_dict=None):
         self.settings["max_size"] = self.max_size
         self.settings["client_id"] = self.client_id
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def _set_pool_size(self, pool_size):
         """Set the size of the pool.
 

tests/test_connection.py

@@ -46,6 +46,7 @@
 
 try:
     from mysql.connector.connection_cext import HAVE_CMYSQL, CMySQLConnection
+    from mysql.connector import cursor_cext
 except ImportError:
     # Test without C Extension
     CMySQLConnection = None
@@ -2001,7 +2002,6 @@ def test_get_connection_with_tls_version(self):
                 cur = cnx.cursor()
                 cur.execute("SHOW STATUS LIKE 'Ssl_version%'")
                 res = cur.fetchall()
-    
                 self.assertEqual(res[0][1], expected_ssl_version,
                                  err_msg.format(expected_ssl_version, res))
 
@@ -2195,6 +2195,29 @@ def test_dns_srv(self):
         self.assertRaises(InterfaceError, connect, **config)
         del config["unix_socket"]
 
+    @unittest.skipIf(not HAVE_CMYSQL, "C Extension not available")
+    def test_context_manager_cext(self):
+        """Test connection and cursor context manager using the C extension."""
+        config = tests.get_mysql_config().copy()
+        config["use_pure"] = False
+        with connect(**config) as conn:
+            self.assertTrue(conn.is_connected())
+            with conn.cursor() as cur:
+                self.assertIsInstance(cur, cursor_cext.CMySQLCursor)
+            self.assertIsNone(cur._cnx)
+        self.assertFalse(conn.is_connected())
+
+    def test_context_manager_pure(self):
+        """Test connection and cursor context manager using pure Python."""
+        config = tests.get_mysql_config().copy()
+        config["use_pure"] = True
+        with connect(**config) as conn:
+            self.assertTrue(conn.is_connected())
+            with conn.cursor() as cur:
+                self.assertIsInstance(cur, cursor.MySQLCursor)
+            self.assertIsNone(cur._connection)
+        self.assertFalse(conn.is_connected())
+
 
 class WL13335(tests.MySQLConnectorTests):
     """WL#13335: Avoid set config values whit flag CAN_HANDLE_EXPIRED_PASSWORDS

tests/test_mysqlx_connection.py

@@ -1583,6 +1583,13 @@ def test_dns_srv(self):
         uri = "mysqlx+srv://root:@localhost/myschema?dns-srv=true"
         self.assertRaises(InterfaceError, mysqlx.get_session, uri)
 
+    def test_context_manager(self):
+        """Test mysqlx.get_session() context manager."""
+        with mysqlx.get_session(self.connect_kwargs) as session:
+            self.assertIsInstance(session, mysqlx.Session)
+            self.assertTrue(session.is_open())
+        self.assertFalse(session.is_open())
+
 
 @unittest.skipIf(tests.MYSQL_VERSION < (8, 0, 20), "XPlugin not compatible")
 class MySQLxInnitialNoticeTests(tests.MySQLxTests):

tests/test_mysqlx_pooling.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2018, 2020, Oracle and/or its affiliates. All rights reserved.
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License, version 2.0, as
@@ -593,6 +593,21 @@ def test_get_client_with_tls_version(self):
         session.close()
         client.close()
 
+    def test_context_manager(self):
+        """Test mysqlx.get_client() context manager."""
+        settings = self.connect_kwargs.copy()
+        pooling_dict = {"enabled": True, "max_size": 5}
+        cnx_options = {"pooling": pooling_dict}
+        with mysqlx.get_client(settings, cnx_options) as client:
+            with client.get_session() as session:
+                self.assertIsInstance(session, mysqlx.Session)
+                self.assertTrue(session.is_open())
+            self.assertFalse(session.is_open())
+            # Create one more session
+            _ = client.get_session()
+        for session in client.sessions:
+            self.assertFalse(session.is_open())
+
 
 @unittest.skipIf(tests.MYSQL_VERSION < (5, 7, 12), "XPlugin not compatible")
 class MySQLxClientPoolingTests(tests.MySQLxTests):