Merge branch 'main' of https://github.com/microsoft/mssql-python into release/0.7.0

Author: bewithgaurav

mssql_python/connection.py

@@ -4,14 +4,20 @@
 This module defines the Connection class, which is used to manage a connection to a database.
 The class provides methods to establish a connection, create cursors, commit transactions, 
 roll back transactions, and close the connection.
+Resource Management:
+- All cursors created from this connection are tracked internally.
+- When close() is called on the connection, all open cursors are automatically closed.
+- Do not use any cursor after the connection is closed; doing so will raise an exception.
+- Cursors are also cleaned up automatically when no longer referenced, to prevent memory leaks.
 """
-import ctypes
+import weakref
 from mssql_python.cursor import Cursor
 from mssql_python.logging_config import get_logger, ENABLE_LOGGING
 from mssql_python.constants import ConstantsDDBC as ddbc_sql_const
 from mssql_python.helpers import add_driver_to_connection_str, check_error
 from mssql_python import ddbc_bindings
 from mssql_python.pooling import PoolingManager
+from mssql_python.exceptions import DatabaseError, InterfaceError
 
 logger = get_logger()
 
@@ -58,6 +64,15 @@ def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_bef
             connection_str, **kwargs
         )
         self._attrs_before = attrs_before or {}
+        self._closed = False
+        
+        # Using WeakSet which automatically removes cursors when they are no longer in use
+        # It is a set that holds weak references to its elements.
+        # When an object is only weakly referenced, it can be garbage collected even if it's still in the set.
+        # It prevents memory leaks by ensuring that cursors are cleaned up when no longer in use without requiring explicit deletion.
+        # TODO: Think and implement scenarios for multi-threaded access to cursors
+        self._cursors = weakref.WeakSet()
+
         # Auto-enable pooling if user never called
         if not PoolingManager.is_initialized():
             PoolingManager.enable()
@@ -152,7 +167,17 @@ def cursor(self) -> Cursor:
             DatabaseError: If there is an error while creating the cursor.
             InterfaceError: If there is an error related to the database interface.
         """
-        return Cursor(self)
+        """Return a new Cursor object using the connection."""
+        if self._closed:
+            # raise InterfaceError
+            raise InterfaceError(
+                driver_error="Cannot create cursor on closed connection",
+                ddbc_error="Cannot create cursor on closed connection",
+            )
+
+        cursor = Cursor(self)
+        self._cursors.add(cursor)  # Track the cursor
+        return cursor
 
     def commit(self) -> None:
         """
@@ -201,6 +226,45 @@ def close(self) -> None:
             DatabaseError: If there is an error while closing the connection.
         """
         # Close the connection
-        self._conn.close()
+        if self._closed:
+            return
+        
+        # Close all cursors first, but don't let one failure stop the others
+        if hasattr(self, '_cursors'):
+            # Convert to list to avoid modification during iteration
+            cursors_to_close = list(self._cursors)
+            close_errors = []
+            
+            for cursor in cursors_to_close:
+                try:
+                    if not cursor.closed:
+                        cursor.close()
+                except Exception as e:
+                    # Collect errors but continue closing other cursors
+                    close_errors.append(f"Error closing cursor: {e}")
+                    if ENABLE_LOGGING:
+                        logger.warning(f"Error closing cursor: {e}")
+            
+            # If there were errors closing cursors, log them but continue
+            if close_errors and ENABLE_LOGGING:
+                logger.warning(f"Encountered {len(close_errors)} errors while closing cursors")
+
+            # Clear the cursor set explicitly to release any internal references
+            self._cursors.clear()
+
+        # Close the connection even if cursor cleanup had issues
+        try:
+            if self._conn:
+                self._conn.close()
+                self._conn = None
+        except Exception as e:
+            if ENABLE_LOGGING:
+                logger.error(f"Error closing database connection: {e}")
+            # Re-raise the connection close error as it's more critical
+            raise
+        finally:
+            # Always mark as closed, even if there were errors
+            self._closed = True
+        
         if ENABLE_LOGGING:
             logger.info("Connection closed successfully.")

mssql_python/cursor.py

@@ -2,6 +2,11 @@
 Copyright (c) Microsoft Corporation.
 Licensed under the MIT license.
 This module contains the Cursor class, which represents a database cursor.
+Resource Management:
+- Cursors are tracked by their parent connection.
+- Closing the connection will automatically close all open cursors.
+- Do not use a cursor after it is closed, or after its parent connection is closed.
+- Use close() to release resources held by the cursor as soon as it is no longer needed.
 """
 import ctypes
 import decimal

tests/test_003_connection.py

@@ -10,14 +10,15 @@
 Note: The cursor function is not yet implemented, so related tests are commented out.
 """
 
+from mssql_python.exceptions import InterfaceError
 import pytest
 import time
 from mssql_python import Connection, connect, pooling
 
 def drop_table_if_exists(cursor, table_name):
     """Drop the table if it exists"""
     try:
-        cursor.execute(f"IF OBJECT_ID('{table_name}', 'U') IS NOT NULL DROP TABLE {table_name}")
+        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
     except Exception as e:
         pytest.fail(f"Failed to drop table {table_name}: {e}")
 
@@ -223,4 +224,172 @@ def test_connection_pooling_basic(conn_str):
 
     conn1.close()
     conn2.close()
-    
+
+# Add these tests at the end of the file
+
+def test_cursor_cleanup_on_connection_close(conn_str):
+    """Test that cursors are properly cleaned up when connection is closed"""
+    # Create a new connection for this test
+    conn = connect(conn_str)
+    
+    # Create multiple cursors
+    cursor1 = conn.cursor()
+    cursor2 = conn.cursor()
+    cursor3 = conn.cursor()
+    
+    # Execute something on each cursor to ensure they have statement handles
+    # Option 1: Fetch results immediately to free the connection
+    cursor1.execute("SELECT 1")
+    cursor1.fetchall() 
+    
+    cursor2.execute("SELECT 2")
+    cursor2.fetchall()
+    
+    cursor3.execute("SELECT 3")
+    cursor3.fetchall()
+
+    # Close one cursor explicitly
+    cursor1.close()
+    assert cursor1.closed is True, "Cursor1 should be closed"
+    
+    # Close the connection (should clean up remaining cursors)
+    conn.close()
+    
+    # Verify all cursors are closed
+    assert cursor1.closed is True, "Cursor1 should remain closed"
+    assert cursor2.closed is True, "Cursor2 should be closed by connection.close()"
+    assert cursor3.closed is True, "Cursor3 should be closed by connection.close()"
+
+def test_cursor_weakref_cleanup(conn_str):
+    """Test that WeakSet properly removes garbage collected cursors"""
+    conn = connect(conn_str)
+    
+    # Create cursors
+    cursor1 = conn.cursor()
+    cursor2 = conn.cursor()
+    
+    # Check initial cursor count
+    assert len(conn._cursors) == 2, "Should have 2 cursors"
+    
+    # Delete reference to cursor1 (should be garbage collected)
+    cursor1_id = id(cursor1)
+    del cursor1
+    
+    # Force garbage collection
+    import gc
+    gc.collect()
+    
+    # Check cursor count after garbage collection
+    assert len(conn._cursors) == 1, "Should have 1 cursor after garbage collection"
+    
+    # Verify cursor2 is still there
+    assert cursor2 in conn._cursors, "Cursor2 should still be in the set"
+    
+    conn.close()
+
+def test_cursor_cleanup_order_no_segfault(conn_str):
+    """Test that proper cleanup order prevents segfaults"""
+    # This test ensures cursors are cleaned before connection
+    conn = connect(conn_str)
+    
+    # Create multiple cursors with active statements
+    cursors = []
+    for i in range(5):
+        cursor = conn.cursor()
+        cursor.execute(f"SELECT {i}")
+        cursor.fetchall()
+        cursors.append(cursor)
+    
+    # Don't close any cursors explicitly
+    # Just close the connection - it should handle cleanup properly
+    conn.close()
+    
+    # Verify all cursors were closed
+    for cursor in cursors:
+        assert cursor.closed is True, "All cursors should be closed"
+
+def test_cursor_close_removes_from_connection(conn_str):
+    """Test that closing a cursor properly cleans up references"""
+    conn = connect(conn_str)
+    
+    # Create cursors
+    cursor1 = conn.cursor()
+    cursor2 = conn.cursor()
+    cursor3 = conn.cursor()
+    
+    assert len(conn._cursors) == 3, "Should have 3 cursors"
+    
+    # Close cursor2
+    cursor2.close()
+    
+    # cursor2 should still be in the WeakSet (until garbage collected)
+    # but it should be marked as closed
+    assert cursor2.closed is True, "Cursor2 should be closed"
+    
+    # Delete the reference and force garbage collection
+    del cursor2
+    import gc
+    gc.collect()
+    
+    # Now should have 2 cursors
+    assert len(conn._cursors) == 2, "Should have 2 cursors after closing and GC"
+    
+    conn.close()
+
+def test_connection_close_idempotent(conn_str):
+    """Test that calling close() multiple times is safe"""
+    conn = connect(conn_str)
+    cursor = conn.cursor()
+    cursor.execute("SELECT 1")
+    
+    # First close
+    conn.close()
+    assert conn._closed is True, "Connection should be closed"
+    
+    # Second close (should not raise exception)
+    conn.close()
+    assert conn._closed is True, "Connection should remain closed"
+    
+    # Cursor should also be closed
+    assert cursor.closed is True, "Cursor should be closed"
+
+def test_cursor_after_connection_close(conn_str):
+    """Test that creating cursor after connection close raises error"""
+    conn = connect(conn_str)
+    conn.close()
+    
+    # Should raise exception when trying to create cursor on closed connection
+    with pytest.raises(InterfaceError) as excinfo:
+        cursor = conn.cursor()
+    
+    assert "closed connection" in str(excinfo.value).lower(), "Should mention closed connection"
+
+def test_multiple_cursor_operations_cleanup(conn_str):
+    """Test cleanup with multiple cursor operations"""
+    conn = connect(conn_str)
+    
+    # Create table for testing
+    cursor_setup = conn.cursor()
+    drop_table_if_exists(cursor_setup, "#test_cleanup")
+    cursor_setup.execute("CREATE TABLE #test_cleanup (id INT, value VARCHAR(50))")
+    cursor_setup.close()
+    
+    # Create multiple cursors doing different operations
+    cursor_insert = conn.cursor()
+    cursor_insert.execute("INSERT INTO #test_cleanup VALUES (1, 'test1'), (2, 'test2')")
+    
+    cursor_select1 = conn.cursor()
+    cursor_select1.execute("SELECT * FROM #test_cleanup WHERE id = 1")
+    cursor_select1.fetchall()
+    
+    cursor_select2 = conn.cursor()
+    cursor_select2.execute("SELECT * FROM #test_cleanup WHERE id = 2")
+    cursor_select2.fetchall()
+
+    # Close connection without closing cursors
+    conn.close()
+    
+    # All cursors should be closed
+    assert cursor_insert.closed is True
+    assert cursor_select1.closed is True
+    assert cursor_select2.closed is True