Fix for ODBC resource leakage

Author: subrata-ms

mssql_python/__init__.py

@@ -4,8 +4,10 @@
 This module initializes the mssql_python package.
 """
 
+import atexit
 import sys
 import types
+import weakref
 from typing import Dict
 
 # Import settings from helpers to avoid circular imports
@@ -67,6 +69,40 @@
 # Pooling
 from .pooling import PoolingManager
 
+# Global registry for tracking active connections (using weak references)
+_active_connections = weakref.WeakSet()
+
+
+def _register_connection(conn):
+    """Register a connection for cleanup before shutdown."""
+    _active_connections.add(conn)
+
+
+def _cleanup_connections():
+    """
+    Cleanup function called by atexit to close all active connections.
+
+    This prevents resource leaks during interpreter shutdown by ensuring
+    all ODBC handles are freed in the correct order before Python finalizes.
+    """
+    # Make a copy of the connections to avoid modification during iteration
+    connections_to_close = list(_active_connections)
+
+    for conn in connections_to_close:
+        try:
+            # Check if connection is still valid and not closed
+            if hasattr(conn, "_closed") and not conn._closed:
+                # Close will handle both cursors and the connection
+                conn.close()
+        except Exception:
+            # Silently ignore errors during shutdown cleanup
+            # We're prioritizing crash prevention over error reporting
+            pass
+
+
+# Register cleanup function to run before Python exits
+atexit.register(_cleanup_connections)
+
 # GLOBALS
 # Read-Only
 apilevel: str = "2.0"

mssql_python/connection.py

@@ -239,6 +239,17 @@ def __init__(
         )
         self.setautocommit(autocommit)
 
+        # Register this connection for cleanup before Python shutdown
+        # This ensures ODBC handles are freed in correct order, preventing leaks
+        try:
+            import mssql_python
+
+            if hasattr(mssql_python, "_register_connection"):
+                mssql_python._register_connection(self)
+        except (ImportError, AttributeError):
+            # If registration fails, continue - cleanup will still happen via __del__
+            pass
+
     def _construct_connection_string(self, connection_str: str = "", **kwargs: Any) -> str:
         """
         Construct the connection string by parsing, validating, and merging parameters.

mssql_python/pybind/ddbc_bindings.cpp

@@ -1130,6 +1130,15 @@ void SqlHandle::free() {
         // Type 3 = SQL_HANDLE_STMT (parent: DBC)
         // Type 2 = SQL_HANDLE_DBC (parent: ENV, which is static and may destruct first)
         // Type 1 = SQL_HANDLE_ENV (no parent)
+        //
+        // RESOURCE LEAK MITIGATION:
+        // When handles are skipped during shutdown, they are not freed, which could
+        // cause resource leaks. However, this is mitigated by:
+        // 1. Python-side atexit cleanup (in __init__.py) that explicitly closes all
+        //    connections before shutdown, ensuring handles are freed in correct order
+        // 2. OS-level cleanup at process termination recovers any remaining resources
+        // 3. This tradeoff prioritizes crash prevention over resource cleanup, which
+        //    is appropriate since we're already in shutdown sequence
         if (pythonShuttingDown && (_type == 3 || _type == 2)) {
             _handle = nullptr;  // Mark as freed to prevent double-free attempts
             return;

tests/test_013_SqlHandle_free_shutdown.py

@@ -259,7 +259,7 @@ def test_stmt_handle_cleanup_at_shutdown(self, conn_str):
         assert result.returncode == 0, f"Process crashed. stderr: {result.stderr}"
         assert "STMT handle cleanup test: Exiting without explicit cleanup" in result.stdout
         assert "Query result: [(1,)]" in result.stdout
-        print(f"✓ STMT handle (Type 3) cleanup during shutdown: PASSED")
+        print(f"PASS: STMT handle (Type 3) cleanup during shutdown")
 
     def test_dbc_handle_cleanup_at_shutdown(self, conn_str):
         """
@@ -308,7 +308,7 @@ def test_dbc_handle_cleanup_at_shutdown(self, conn_str):
         assert "Connection 0: created and cursor closed" in result.stdout
         assert "Connection 1: created and cursor closed" in result.stdout
         assert "Connection 2: created and cursor closed" in result.stdout
-        print(f"✓ DBC handle (Type 2) cleanup during shutdown: PASSED")
+        print(f"PASS: DBC handle (Type 2) cleanup during shutdown")
 
     def test_env_handle_cleanup_at_shutdown(self, conn_str):
         """
@@ -356,7 +356,7 @@ def test_env_handle_cleanup_at_shutdown(self, conn_str):
         assert "Connection 0: properly closed" in result.stdout
         assert "Connection 1: properly closed" in result.stdout
         assert "Connection 2: properly closed" in result.stdout
-        print(f"✓ ENV handle (Type 1) cleanup during shutdown: PASSED")
+        print(f"PASS: ENV handle (Type 1) cleanup during shutdown")
 
     def test_mixed_handle_cleanup_at_shutdown(self, conn_str):
         """
@@ -430,7 +430,7 @@ def test_mixed_handle_cleanup_at_shutdown(self, conn_str):
         assert "Connection 1: cursors left open" in result.stdout
         assert "Connection 2: cursors closed, connection left open" in result.stdout
         assert "Connection 3: everything properly closed" in result.stdout
-        print(f"✓ Mixed handle cleanup during shutdown: PASSED")
+        print(f"PASS: Mixed handle cleanup during shutdown")
 
     def test_rapid_connection_churn_with_shutdown(self, conn_str):
         """
@@ -483,7 +483,7 @@ def test_rapid_connection_churn_with_shutdown(self, conn_str):
         assert result.returncode == 0, f"Process crashed. stderr: {result.stderr}"
         assert "Created 10 connections, closed 5 explicitly" in result.stdout
         assert "Rapid churn test: Exiting with mixed cleanup" in result.stdout
-        print(f"✓ Rapid connection churn with shutdown: PASSED")
+        print(f"PASS: Rapid connection churn with shutdown")
 
     def test_exception_during_query_with_shutdown(self, conn_str):
         """
@@ -524,7 +524,7 @@ def test_exception_during_query_with_shutdown(self, conn_str):
         assert result.returncode == 0, f"Process crashed. stderr: {result.stderr}"
         assert "Expected error occurred: ProgrammingError" in result.stdout
         assert "Exception test: Exiting after exception without cleanup" in result.stdout
-        print(f"✓ Exception during query with shutdown: PASSED")
+        print(f"PASS: Exception during query with shutdown")
 
     def test_weakref_cleanup_at_shutdown(self, conn_str):
         """
@@ -578,7 +578,7 @@ def callback(ref):
 
         assert result.returncode == 0, f"Process crashed. stderr: {result.stderr}"
         assert "Weakref test: Exiting with weakrefs active" in result.stdout
-        print(f"✓ Weakref cleanup at shutdown: PASSED")
+        print(f"PASS: Weakref cleanup at shutdown")
 
     def test_gc_during_shutdown_with_circular_refs(self, conn_str):
         """
@@ -638,7 +638,7 @@ def execute_query(self):
 
         assert result.returncode == 0, f"Process crashed. stderr: {result.stderr}"
         assert "Circular ref test: Exiting after GC with cycles" in result.stdout
-        print(f"✓ GC during shutdown with circular refs: PASSED")
+        print(f"PASS: GC during shutdown with circular refs")
 
     def test_all_handle_types_comprehensive(self, conn_str):
         """
@@ -717,7 +717,7 @@ def test_all_handle_types_comprehensive(self, conn_str):
         assert "Scenario 3: Both cursor and connection left open" in result.stdout
         assert "Scenario 4: Multiple cursors per connection left open" in result.stdout
         assert "=== Exiting ===" in result.stdout
-        print(f"✓ Comprehensive all handle types test: PASSED")
+        print(f"PASS: Comprehensive all handle types test")
 
 
 if __name__ == "__main__":
@@ -755,7 +755,7 @@ def test_all_handle_types_comprehensive(self, conn_str):
         test.test_all_handle_types_comprehensive(conn_str)
 
         print("\n" + "=" * 70)
-        print("✓ ALL TESTS PASSED - No segfaults detected")
+        print("PASS: ALL TESTS PASSED - No segfaults detected")
         print("=" * 70 + "\n")
     except AssertionError as e:
         print(f"\n✗ TEST FAILED: {e}")