Added transaction support

Author: Peng Ren

pymongosql/connection.py

@@ -11,7 +11,7 @@
 
 from .common import BaseCursor
 from .cursor import Cursor
-from .error import DatabaseError, NotSupportedError, OperationalError
+from .error import DatabaseError, OperationalError
 from .helper import ConnectionHelper
 
 _logger = logging.getLogger(__name__)
@@ -212,8 +212,8 @@ def in_transaction(self) -> bool:
         return self._in_transaction
 
     @in_transaction.setter
-    def in_transaction(self, value: bool) -> bool:
-        self._in_transaction = False
+    def in_transaction(self, value: bool) -> None:
+        self._in_transaction = value
 
     @property
     def host(self) -> str:
@@ -407,24 +407,54 @@ def _with_transaction(self, callback, **kwargs):
         return self._session.with_transaction(callback, **kwargs)
 
     def begin(self) -> None:
-        """Begin transaction (DB-API 2.0 standard method)"""
+        """Begin transaction (DB-API 2.0 standard method)
+
+        Starts an explicit transaction. After calling begin(), operations
+        are executed within the transaction context until commit() or
+        rollback() is called. Requires MongoDB 4.0+ for multi-document
+        transactions on replica sets or sharded clusters.
+
+        Example:
+            conn.begin()
+            try:
+                cursor.execute("INSERT INTO users VALUES (...)")
+                cursor.execute("UPDATE accounts SET balance = balance - 100")
+                conn.commit()
+            except Exception:
+                conn.rollback()
+
+        Raises:
+            OperationalError: If unable to start transaction
+        """
         self._start_transaction()
 
     def commit(self) -> None:
-        """Commit transaction (DB-API 2.0 standard method)"""
+        """Commit transaction (DB-API 2.0 standard method)
+
+        Commits the current transaction to the database. All operations
+        executed since begin() will be atomically persisted. If no
+        transaction is active, this is a no-op (DB-API 2.0 compliant).
+
+        Raises:
+            OperationalError: If commit fails
+        """
         if self._session and self._session.in_transaction:
             self._commit_transaction()
-        else:
-            # Fallback for non-session based operations
-            self._in_transaction = False
-            self._autocommit = True
+        # If no transaction, this is a no-op (DB-API 2.0 compliant)
 
     def rollback(self) -> None:
-        """Rollback transaction (DB-API 2.0 standard method)"""
+        """Rollback transaction (DB-API 2.0 standard method)
+
+        Rolls back (aborts) the current transaction, undoing all operations
+        executed since begin(). If no transaction is active, this is a no-op
+        (DB-API 2.0 compliant).
+
+        Raises:
+            OperationalError: If rollback fails
+        """
         if self._session and self._session.in_transaction:
             self._abort_transaction()
-        else:
-            raise NotSupportedError("MongoDB doesn't support rollback without an active transaction")
+        # If no transaction, this is a no-op (DB-API 2.0 compliant)
 
     def test_connection(self) -> bool:
         """Test if the connection is alive"""

pymongosql/executor.py

@@ -101,11 +101,23 @@ def _replace_placeholders(self, obj: Any, parameters: Sequence[Any]) -> Any:
     def _execute_execution_plan(
         self,
         execution_plan: QueryExecutionPlan,
-        db: Any,
+        connection: Any = None,
         parameters: Optional[Sequence[Any]] = None,
     ) -> Optional[Dict[str, Any]]:
-        """Execute a QueryExecutionPlan against MongoDB using db.command"""
+        """Execute a QueryExecutionPlan against MongoDB using db.command
+
+        Args:
+            execution_plan: QueryExecutionPlan to execute
+            connection: Connection object (for session and database access)
+            parameters: Parameters for placeholder replacement
+        """
         try:
+            # Get database from connection
+            if not connection:
+                raise OperationalError("No connection provided")
+
+            db = connection.database
+
             # Get database
             if not execution_plan.collection:
                 raise ProgrammingError("No collection specified in query")
@@ -144,8 +156,11 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB command: {find_command}")
 
-            # Execute find command directly
-            result = db.command(find_command)
+            # Execute find command with session if in transaction
+            if connection and connection.session and connection.session.in_transaction:
+                result = db.command(find_command, session=connection.session)
+            else:
+                result = db.command(find_command)
 
             # Create command result
             return result
@@ -182,7 +197,7 @@ def execute(
         # Parse the query
         self._execution_plan = self._parse_sql(processed_query)
 
-        return self._execute_execution_plan(self._execution_plan, connection.database, processed_params)
+        return self._execute_execution_plan(self._execution_plan, connection, processed_params)
 
 
 class InsertExecution(ExecutionStrategy):
@@ -224,10 +239,16 @@ def _replace_placeholders(
     def _execute_execution_plan(
         self,
         execution_plan: InsertExecutionPlan,
-        db: Any,
+        connection: Any = None,
         parameters: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
     ) -> Optional[Dict[str, Any]]:
         try:
+            # Get database from connection
+            if not connection:
+                raise OperationalError("No connection provided")
+
+            db = connection.database
+
             if not execution_plan.collection:
                 raise ProgrammingError("No collection specified in insert")
 
@@ -238,7 +259,11 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB insert command: {command}")
 
-            return db.command(command)
+            # Execute with session if in transaction
+            if connection and connection.session and connection.session.in_transaction:
+                return db.command(command, session=connection.session)
+            else:
+                return db.command(command)
         except PyMongoError as e:
             _logger.error(f"MongoDB insert failed: {e}")
             raise DatabaseError(f"Insert execution failed: {e}")
@@ -259,7 +284,7 @@ def execute(
 
         self._execution_plan = self._parse_sql(context.query)
 
-        return self._execute_execution_plan(self._execution_plan, connection.database, parameters)
+        return self._execute_execution_plan(self._execution_plan, connection, parameters)
 
 
 class DeleteExecution(ExecutionStrategy):
@@ -293,10 +318,16 @@ def _parse_sql(self, sql: str) -> Any:
     def _execute_execution_plan(
         self,
         execution_plan: Any,
-        db: Any,
+        connection: Any = None,
         parameters: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
     ) -> Optional[Dict[str, Any]]:
         try:
+            # Get database from connection
+            if not connection:
+                raise OperationalError("No connection provided")
+
+            db = connection.database
+
             if not execution_plan.collection:
                 raise ProgrammingError("No collection specified in delete")
 
@@ -312,7 +343,11 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB delete command: {command}")
 
-            return db.command(command)
+            # Execute with session if in transaction
+            if connection and connection.session and connection.session.in_transaction:
+                return db.command(command, session=connection.session)
+            else:
+                return db.command(command)
         except PyMongoError as e:
             _logger.error(f"MongoDB delete failed: {e}")
             raise DatabaseError(f"Delete execution failed: {e}")
@@ -333,7 +368,7 @@ def execute(
 
         self._execution_plan = self._parse_sql(context.query)
 
-        return self._execute_execution_plan(self._execution_plan, connection.database, parameters)
+        return self._execute_execution_plan(self._execution_plan, connection, parameters)
 
 
 class UpdateExecution(ExecutionStrategy):
@@ -367,10 +402,16 @@ def _parse_sql(self, sql: str) -> Any:
     def _execute_execution_plan(
         self,
         execution_plan: Any,
-        db: Any,
+        connection: Any = None,
         parameters: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
     ) -> Optional[Dict[str, Any]]:
         try:
+            # Get database from connection
+            if not connection:
+                raise OperationalError("No connection provided")
+
+            db = connection.database
+
             if not execution_plan.collection:
                 raise ProgrammingError("No collection specified in update")
 
@@ -406,7 +447,11 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB update command: {command}")
 
-            return db.command(command)
+            # Execute with session if in transaction
+            if connection and connection.session and connection.session.in_transaction:
+                return db.command(command, session=connection.session)
+            else:
+                return db.command(command)
         except PyMongoError as e:
             _logger.error(f"MongoDB update failed: {e}")
             raise DatabaseError(f"Update execution failed: {e}")
@@ -427,7 +472,7 @@ def execute(
 
         self._execution_plan = self._parse_sql(context.query)
 
-        return self._execute_execution_plan(self._execution_plan, connection.database, parameters)
+        return self._execute_execution_plan(self._execution_plan, connection, parameters)
 
 
 class ExecutionPlanFactory:

pymongosql/superset_mongodb/executor.py

@@ -63,7 +63,7 @@ def execute(
         _logger.debug(f"Stage 1: Executing MongoDB subquery: {mongo_query}")
 
         mongo_execution_plan = self._parse_sql(mongo_query)
-        mongo_result = self._execute_execution_plan(mongo_execution_plan, connection.database)
+        mongo_result = self._execute_execution_plan(mongo_execution_plan, connection)
 
         # Extract result set from MongoDB
         mongo_result_set = ResultSet(

pyproject.toml

@@ -89,6 +89,9 @@ testpaths = ["tests"]
 python_files = ["test_*.py"]
 python_classes = ["Test*"]
 python_functions = ["test_*"]
+markers = [
+    "transactional: marks tests that require MongoDB transaction support (requires replica set or sharded cluster)",
+]
 
 [tool.coverage.run]
 source = ["pymongosql"]