FEAT: Context manager support for connection and cursor class

mssql_python/connection.py

@@ -29,13 +29,27 @@ 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 commit the transaction when
+    exiting the context (if autocommit is False), but will NOT close the connection.
+    This behavior matches pyodbc for compatibility.
+
+    Example usage:
+        with connect(connection_string) as conn:
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO table VALUES (?)", [value])
+            # Transaction is automatically committed when exiting the with block
+        # Connection remains open after exiting the with block
+
     Methods:
         __init__(database: str) -> None:
         connect_to_db() -> None:
         cursor() -> Cursor:
         commit() -> None:
         rollback() -> None:
         close() -> None:
+        __enter__() -> Connection:
+        __exit__(exc_type, exc_val, exc_tb) -> None:
     """
 
     def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_before: dict = None, **kwargs) -> None:
@@ -147,7 +161,7 @@ def autocommit(self, value: bool) -> None:
         self.setautocommit(value)
         log('info', "Autocommit mode set to %s.", value)
 
-    def setautocommit(self, value: bool = True) -> None:
+    def setautocommit(self, value: bool = False) -> None:
         """
         Set the autocommit mode of the connection.
         Args:
@@ -259,6 +273,14 @@ def close(self) -> None:
         # Close the connection even if cursor cleanup had issues
         try:
             if self._conn:
+                if not self.autocommit:
+                    # 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
                 self._conn.close()
                 self._conn = None
         except Exception as e:
@@ -271,6 +293,66 @@ 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, exc_type, exc_val, exc_tb) -> None:
+        """
+        Exit the context manager.
+
+        This method is called when exiting the 'with' statement. It follows pyodbc
+        behavior where:
+        - If autocommit is False, it commits the current transaction
+        - If an exception occurred, it rolls back instead of committing
+        - The connection is NOT closed (matches pyodbc behavior)
+
+        Args:
+            exc_type: The exception type if an exception occurred, None otherwise
+            exc_val: The exception value if an exception occurred, None otherwise  
+            exc_tb: The exception traceback if an exception occurred, None otherwise
+
+        Note:
+            This method does not return True, so exceptions are not suppressed
+            and will propagate normally.
+        """
+        try:
+            if exc_type is not None:
+                # An exception occurred in the with block
+                if not self.autocommit:
+                    log('info', "Exception occurred in context manager, rolling back transaction.")
+                    self._conn.rollback()
+                else:
+                    log('info', "Exception occurred in context manager, but autocommit is enabled.")
+            else:
+                # No exception occurred
+                if not self.autocommit:
+                    log('info', "Exiting connection context manager, committing transaction.")
+                    self._conn.commit()
+                else:
+                    log('info', "Exiting connection context manager, autocommit is enabled (no commit needed).")
+        except Exception as e:
+            log('error', f"Error during context manager exit: {e}")
+            # Let the exception propagate - don't suppress it
+            raise
+        finally:
+            log('info', "Exited connection context manager (connection remains open).")
+
     def __del__(self):
         """
         Destructor to ensure the connection is closed when the connection object is no longer needed.

mssql_python/cursor.py

@@ -789,6 +789,46 @@ 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, exc_type, exc_val, exc_tb):
+        """
+        Exit the runtime context for the cursor.
+
+        Following pyodbc behavior:
+        - If autocommit is False, commit the transaction
+        - The cursor is NOT closed (unlike connection context manager)
+        - Returns None to propagate any exceptions
+
+        Args:
+            exc_type: Exception type (if any)
+            exc_val: Exception value (if any) 
+            exc_tb: Exception traceback (if any)
+        """
+        try:
+            # Only commit if autocommit is False, following pyodbc behavior
+            if not self.connection.autocommit:
+                self.connection.commit()
+                log('debug', "Transaction committed in cursor context manager exit")
+        except Exception as e:
+            log('error', "Error committing transaction in cursor context manager: %s", e)
+            # Re-raise the exception to maintain proper error handling
+            raise
+
+        # Note: Unlike connection context manager, cursor is NOT closed here
+        # This matches pyodbc behavior exactly
+
+        # Return None to propagate any exception that occurred in the with block
+        return None
+
     def __del__(self):
         """
         Destructor to ensure the cursor is closed when it is no longer needed.

mssql_python/db_connection.py

@@ -5,7 +5,7 @@
 """
 from mssql_python.connection import Connection
 
-def connect(connection_str: str = "", autocommit: bool = True, attrs_before: dict = None, **kwargs) -> Connection:
+def connect(connection_str: str = "", autocommit: bool = False, attrs_before: dict = None, **kwargs) -> Connection:
     """
     Constructor for creating a connection to the database.
 

mssql_python/pybind/connection/connection.cpp

@@ -132,9 +132,14 @@ void Connection::setAutocommit(bool enable) {
         ThrowStdException("Connection handle not allocated");
     }
     SQLINTEGER value = enable ? SQL_AUTOCOMMIT_ON : SQL_AUTOCOMMIT_OFF;
-    LOG("Set SQL Connection Attribute");
+    LOG("Setting SQL Connection Attribute");
     SQLRETURN ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_AUTOCOMMIT, reinterpret_cast<SQLPOINTER>(static_cast<SQLULEN>(value)), 0);
     checkError(ret);
+    if(value == SQL_AUTOCOMMIT_ON) {
+        LOG("SQL Autocommit set to True");
+    } else {
+        LOG("SQL Autocommit set to False");
+    }
     _autocommit = enable;
 }