Add connection finalizer

Adds a finalizer to the PySQL Connection class. This helps mitigate common application errors causing the session to not be closed correctly.

Added tests to ensure the finalizer is called when there are no more references to the connection.

Author: CJStuart

cmdexec/clients/python/src/databricks/sql/client.py

@@ -110,6 +110,12 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
+    def __del__(self):
+        if self.open:
+            logger.debug("Closing unclosed connection for session "
+                         "{}".format(self.get_session_id()))
+            self._close(close_cursors=False)
+
     def get_session_id(self):
         return self.thrift_backend.handle_to_id(self._session_handle)
 
@@ -134,11 +140,15 @@ def cursor(self,
 
     def close(self) -> None:
         """Close the underlying session and mark all associated cursors as closed."""
+        self._close()
+
+    def _close(self, close_cursors=True) -> None:
         self.thrift_backend.close_session(self._session_handle)
         self.open = False
 
-        for cursor in self._cursors:
-            cursor.close()
+        if close_cursors:
+            for cursor in self._cursors:
+                cursor.close()
 
     def commit(self):
         """No-op because Databricks does not support transactions"""

cmdexec/clients/python/tests/tests.py

@@ -1,3 +1,4 @@
+import gc
 import re
 import sys
 import unittest
@@ -507,6 +508,34 @@ def test_column_name_api(self):
                 "third_col": expected[2]
             })
 
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_finalizer_closes_abandoned_connection(self, mock_client_class):
+        instance = mock_client_class.return_value
+        instance.open_session.return_value = b'\x22'
+
+        databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
+
+        # not strictly necessary as the refcount is 0, but just to be sure
+        gc.collect()
+
+        # Check the close session request has an id of x22
+        close_session_id = instance.close_session.call_args[0][0]
+        self.assertEqual(close_session_id, b'\x22')
+
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_cursor_keeps_connection_alive(self, mock_client_class):
+        instance = mock_client_class.return_value
+        instance.open_session.return_value = b'\x22'
+
+        connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
+        cursor = connection.cursor()
+        del connection
+
+        gc.collect()
+
+        self.assertEqual(instance.close_session.call_count, 0)
+        cursor.close()
+
 
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])