FEAT: Context manager support for connection and cursor class (#160)

jahnvi480 · bewithgaurav · Copilot · web-flow · commit f4c899e059b0 · 2025-08-27T18:38:03.000+05:30
### Work Item / Issue Reference   > [AB#38074](https://sqlclientdrivers.visualstudio.com/c6d89619-62de-46a0-8b46-70b92a84d85e/_workitems/edit/38074) > GitHub Issue: #23 ------------------------------------------------------------------- ### Summary This pull request introduces context manager support for the `Connection` and `Cursor` classes in the `mssql_python` module, aligning their behavior with `pyodbc`. It also adds extensive tests to ensure correct functionality, covering various scenarios such as normal exits, exceptions, nested transactions, and manual commit/rollback. The key changes are grouped below: ### Context Manager Support * Added `__enter__` and `__exit__` methods to the `Connection` class to enable usage with the `with` statement. The connection commits transactions on normal exit and rolls back on exceptions, without closing the connection. * Added `__enter__` and `__exit__` methods to the `Cursor` class to allow usage with the `with` statement. The cursor commits transactions on normal exit but does not close itself, aligning with `pyodbc` behavior. ### Logging Enhancements * Enhanced the `Connection.close` method to log when uncommitted changes are rolled back before closing. ### Test Coverage * Added new tests in `tests/test_003_connection.py` to verify context manager behavior for connections, including: - Committing transactions on normal exit. - Rolling back transactions on exceptions. - Handling autocommit mode. - Ensuring connections remain open after context exit. - Supporting nested context managers. - Testing manual commit/rollback within a context manager. * Added a test for using `contextlib.closing` with connections to ensure proper closure after context exit. * Updated `tests/test_004_cursor.py` to include `contextlib.closing` for cursor tests. --------- Co-authored-by: Jahnvi Thakkar <jathakkar@microsoft.com> Co-authored-by: Gaurav Sharma <sharmag@microsoft.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
diff --git a/mssql_python/connection.py b/mssql_python/connection.py
@@ -42,13 +42,32 @@ class Connection:
     to be used in a context where database operations are required, such as executing queries
     and fetching results.
 
+    The Connection class supports the Python context manager protocol (with statement).
+    When used as a context manager, it will automatically close the connection when
+    exiting the context, ensuring proper resource cleanup.
+
+    Example usage:
+        with connect(connection_string) as conn:
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO table VALUES (?)", [value])
+        # Connection is automatically closed when exiting the with block
+        
+    For long-lived connections, use without context manager:
+        conn = connect(connection_string)
+        try:
+            # Multiple operations...
+        finally:
+            conn.close()
+
     Methods:
         __init__(database: str) -> None:
         connect_to_db() -> None:
         cursor() -> Cursor:
         commit() -> None:
         rollback() -> None:
         close() -> None:
+        __enter__() -> Connection:
+        __exit__() -> None:
     """
 
     # DB-API 2.0 Exception attributes
@@ -289,6 +308,7 @@ def close(self) -> None:
                     # If autocommit is disabled, rollback any uncommitted changes
                     # This is important to ensure no partial transactions remain
                     # For autocommit True, this is not necessary as each statement is committed immediately
+                    log('info', "Rolling back uncommitted changes before closing connection.")
                     self._conn.rollback()
                 # TODO: Check potential race conditions in case of multithreaded scenarios
                 # Close the connection
@@ -304,6 +324,35 @@ def close(self) -> None:
         
         log('info', "Connection closed successfully.")
 
+    def __enter__(self) -> 'Connection':
+        """
+        Enter the context manager.
+        
+        This method enables the Connection to be used with the 'with' statement.
+        When entering the context, it simply returns the connection object itself.
+        
+        Returns:
+            Connection: The connection object itself.
+            
+        Example:
+            with connect(connection_string) as conn:
+                cursor = conn.cursor()
+                cursor.execute("INSERT INTO table VALUES (?)", [value])
+                # Transaction will be committed automatically when exiting
+        """
+        log('info', "Entering connection context manager.")
+        return self
+
+    def __exit__(self, *args) -> None:
+        """
+        Exit the context manager.
+        
+        Closes the connection when exiting the context, ensuring proper resource cleanup.
+        This follows the modern standard used by most database libraries.
+        """
+        if not self._closed:
+            self.close()
+
     def __del__(self):
         """
         Destructor to ensure the connection is closed when the connection object is no longer needed.
diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py
@@ -16,7 +16,7 @@
 from mssql_python.constants import ConstantsDDBC as ddbc_sql_const
 from mssql_python.helpers import check_error, log
 from mssql_python import ddbc_bindings
-from mssql_python.exceptions import InterfaceError
+from mssql_python.exceptions import InterfaceError, ProgrammingError
 from .row import Row
 
 
@@ -796,6 +796,22 @@ def nextset(self) -> Union[bool, None]:
             return False
         return True
 
+    def __enter__(self):
+        """
+        Enter the runtime context for the cursor.
+        
+        Returns:
+            The cursor instance itself.
+        """
+        self._check_closed()
+        return self
+    
+    def __exit__(self, *args):
+        """Closes the cursor when exiting the context, ensuring proper resource cleanup."""
+        if not self.closed:
+            self.close()
+        return None
+
     def __del__(self):
         """
         Destructor to ensure the cursor is closed when it is no longer needed.
diff --git a/tests/test_003_connection.py b/tests/test_003_connection.py
@@ -16,12 +16,16 @@
 - test_connection_string_with_attrs_before: Check if the connection string is constructed correctly with attrs_before.
 - test_connection_string_with_odbc_param: Check if the connection string is constructed correctly with ODBC parameters.
 - test_rollback_on_close: Test that rollback occurs on connection close if autocommit is False.
+- test_context_manager_commit: Test that context manager commits transaction on normal exit.
+- test_context_manager_autocommit_mode: Test context manager behavior with autocommit enabled.
+- test_context_manager_connection_closes: Test that context manager closes the connection.
 """
 
 from mssql_python.exceptions import InterfaceError
 import pytest
 import time
 from mssql_python import Connection, connect, pooling
+from contextlib import closing
 import threading
 
 # Import all exception classes for testing
@@ -500,6 +504,113 @@ def test_connection_pooling_basic(conn_str):
     conn1.close()
     conn2.close()
 
+def test_context_manager_commit(conn_str):
+    """Test that context manager closes connection on normal exit"""
+    # Create a permanent table for testing across connections
+    setup_conn = connect(conn_str)
+    setup_cursor = setup_conn.cursor()
+    drop_table_if_exists(setup_cursor, "pytest_context_manager_test")
+    
+    try:
+        setup_cursor.execute("CREATE TABLE pytest_context_manager_test (id INT PRIMARY KEY, value VARCHAR(50));")
+        setup_conn.commit()
+        setup_conn.close()
+        
+        # Test context manager closes connection
+        with connect(conn_str) as conn:
+            assert conn.autocommit is False, "Autocommit should be False by default"
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO pytest_context_manager_test (id, value) VALUES (1, 'context_test');")
+            conn.commit()  # Manual commit now required
+        # Connection should be closed here
+        
+        # Verify data was committed manually
+        verify_conn = connect(conn_str)
+        verify_cursor = verify_conn.cursor()
+        verify_cursor.execute("SELECT * FROM pytest_context_manager_test WHERE id = 1;")
+        result = verify_cursor.fetchone()
+        assert result is not None, "Manual commit failed: No data found"
+        assert result[1] == 'context_test', "Manual commit failed: Incorrect data"
+        verify_conn.close()
+        
+    except Exception as e:
+        pytest.fail(f"Context manager test failed: {e}")
+    finally:
+        # Cleanup
+        cleanup_conn = connect(conn_str)
+        cleanup_cursor = cleanup_conn.cursor()
+        drop_table_if_exists(cleanup_cursor, "pytest_context_manager_test")
+        cleanup_conn.commit()
+        cleanup_conn.close()
+
+def test_context_manager_connection_closes(conn_str):
+    """Test that context manager closes the connection"""
+    conn = None
+    try:
+        with connect(conn_str) as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT 1")
+            result = cursor.fetchone()
+            assert result[0] == 1, "Connection should work inside context manager"
+        
+        # Connection should be closed after exiting context manager
+        assert conn._closed, "Connection should be closed after exiting context manager"
+        
+        # Should not be able to use the connection after closing
+        with pytest.raises(InterfaceError):
+            conn.cursor()
+            
+    except Exception as e:
+        pytest.fail(f"Context manager connection close test failed: {e}")
+
+def test_close_with_autocommit_true(conn_str):
+    """Test that connection.close() with autocommit=True doesn't trigger rollback."""
+    cursor = None
+    conn = None
+    
+    try:
+        # Create a temporary table for testing
+        setup_conn = connect(conn_str)
+        setup_cursor = setup_conn.cursor()
+        drop_table_if_exists(setup_cursor, "pytest_autocommit_close_test")
+        setup_cursor.execute("CREATE TABLE pytest_autocommit_close_test (id INT PRIMARY KEY, value VARCHAR(50));")
+        setup_conn.commit()
+        setup_conn.close()
+        
+        # Create a connection with autocommit=True
+        conn = connect(conn_str)
+        conn.autocommit = True
+        assert conn.autocommit is True, "Autocommit should be True"
+        
+        # Insert data
+        cursor = conn.cursor()
+        cursor.execute("INSERT INTO pytest_autocommit_close_test (id, value) VALUES (1, 'test_autocommit');")
+        
+        # Close the connection without explicitly committing
+        conn.close()
+        
+        # Verify the data was committed automatically despite connection.close()
+        verify_conn = connect(conn_str)
+        verify_cursor = verify_conn.cursor()
+        verify_cursor.execute("SELECT * FROM pytest_autocommit_close_test WHERE id = 1;")
+        result = verify_cursor.fetchone()
+        
+        # Data should be present if autocommit worked and wasn't affected by close()
+        assert result is not None, "Autocommit failed: Data not found after connection close"
+        assert result[1] == 'test_autocommit', "Autocommit failed: Incorrect data after connection close"
+        
+        verify_conn.close()
+        
+    except Exception as e:
+        pytest.fail(f"Test failed: {e}")
+    finally:
+        # Clean up
+        cleanup_conn = connect(conn_str)
+        cleanup_cursor = cleanup_conn.cursor()
+        drop_table_if_exists(cleanup_cursor, "pytest_autocommit_close_test")
+        cleanup_conn.commit()
+        cleanup_conn.close()
+        
 # DB-API 2.0 Exception Attribute Tests
 def test_connection_exception_attributes_exist(db_connection):
     """Test that all DB-API 2.0 exception classes are available as Connection attributes"""
@@ -684,3 +795,4 @@ def test_connection_exception_attributes_comprehensive_list():
         exc_class = getattr(Connection, exc_name)
         assert isinstance(exc_class, type), f"Connection.{exc_name} should be a class"
         assert issubclass(exc_class, Exception), f"Connection.{exc_name} should be an Exception subclass"
+
diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py
