Added basic update support

Author: Peng Ren

pymongosql/executor.py

@@ -8,9 +8,11 @@
 
 from .error import DatabaseError, OperationalError, ProgrammingError, SqlSyntaxError
 from .helper import SQLHelper
+from .sql.delete_builder import DeleteExecutionPlan
 from .sql.insert_builder import InsertExecutionPlan
 from .sql.parser import SQLParser
 from .sql.query_builder import QueryExecutionPlan
+from .sql.update_builder import UpdateExecutionPlan
 
 _logger = logging.getLogger(__name__)
 
@@ -263,10 +265,6 @@ def execute(
 class DeleteExecution(ExecutionStrategy):
     """Strategy for executing DELETE statements."""
 
-    def __init__(self) -> None:
-        """Initialize DELETE execution strategy."""
-        self._execution_plan: Optional[Any] = None
-
     @property
     def execution_plan(self) -> Any:
         return self._execution_plan
@@ -275,8 +273,6 @@ def supports(self, context: ExecutionContext) -> bool:
         return context.query.lstrip().upper().startswith("DELETE")
 
     def _parse_sql(self, sql: str) -> Any:
-        from .sql.delete_builder import DeleteExecutionPlan
-
         try:
             parser = SQLParser(sql)
             plan = parser.get_execution_plan()
@@ -340,10 +336,104 @@ def execute(
         return self._execute_execution_plan(self._execution_plan, connection.database, parameters)
 
 
+class UpdateExecution(ExecutionStrategy):
+    """Strategy for executing UPDATE statements."""
+
+    @property
+    def execution_plan(self) -> Any:
+        return self._execution_plan
+
+    def supports(self, context: ExecutionContext) -> bool:
+        return context.query.lstrip().upper().startswith("UPDATE")
+
+    def _parse_sql(self, sql: str) -> Any:
+        try:
+            parser = SQLParser(sql)
+            plan = parser.get_execution_plan()
+
+            if not isinstance(plan, UpdateExecutionPlan):
+                raise SqlSyntaxError("Expected UPDATE execution plan")
+
+            if not plan.validate():
+                raise SqlSyntaxError("Generated update plan is invalid")
+
+            return plan
+        except SqlSyntaxError:
+            raise
+        except Exception as e:
+            _logger.error(f"SQL parsing failed: {e}")
+            raise SqlSyntaxError(f"Failed to parse SQL: {e}")
+
+    def _execute_execution_plan(
+        self,
+        execution_plan: Any,
+        db: Any,
+        parameters: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
+    ) -> Optional[Dict[str, Any]]:
+        try:
+            if not execution_plan.collection:
+                raise ProgrammingError("No collection specified in update")
+
+            if not execution_plan.update_fields:
+                raise ProgrammingError("No fields to update specified")
+
+            filter_conditions = execution_plan.filter_conditions or {}
+            update_fields = execution_plan.update_fields or {}
+
+            # Replace placeholders if parameters provided
+            # Note: We need to replace both update_fields and filter_conditions in one pass
+            # to maintain correct parameter ordering (SET clause first, then WHERE clause)
+            if parameters:
+                # Combine structures for replacement in correct order
+                combined = {"update_fields": update_fields, "filter_conditions": filter_conditions}
+                replaced = SQLHelper.replace_placeholders_generic(combined, parameters, execution_plan.parameter_style)
+                update_fields = replaced["update_fields"]
+                filter_conditions = replaced["filter_conditions"]
+
+            # MongoDB update command format
+            # https://www.mongodb.com/docs/manual/reference/command/update/
+            command = {
+                "update": execution_plan.collection,
+                "updates": [
+                    {
+                        "q": filter_conditions,  # query filter
+                        "u": {"$set": update_fields},  # update document using $set operator
+                        "multi": True,  # update all matching documents (like SQL UPDATE)
+                        "upsert": False,  # don't insert if no match
+                    }
+                ],
+            }
+
+            _logger.debug(f"Executing MongoDB update command: {command}")
+
+            return db.command(command)
+        except PyMongoError as e:
+            _logger.error(f"MongoDB update failed: {e}")
+            raise DatabaseError(f"Update execution failed: {e}")
+        except (ProgrammingError, DatabaseError, OperationalError):
+            # Re-raise our own errors without wrapping
+            raise
+        except Exception as e:
+            _logger.error(f"Unexpected error during update execution: {e}")
+            raise OperationalError(f"Update execution error: {e}")
+
+    def execute(
+        self,
+        context: ExecutionContext,
+        connection: Any,
+        parameters: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
+    ) -> Optional[Dict[str, Any]]:
+        _logger.debug(f"Using update execution for query: {context.query[:100]}")
+
+        self._execution_plan = self._parse_sql(context.query)
+
+        return self._execute_execution_plan(self._execution_plan, connection.database, parameters)
+
+
 class ExecutionPlanFactory:
     """Factory for creating appropriate execution strategy based on query context"""
 
-    _strategies = [DeleteExecution(), InsertExecution(), StandardQueryExecution()]
+    _strategies = [StandardQueryExecution(), InsertExecution(), UpdateExecution(), DeleteExecution()]
 
     @classmethod
     def get_strategy(cls, context: ExecutionContext) -> ExecutionStrategy:

pymongosql/result_set.py

@@ -198,7 +198,21 @@ def errors(self) -> List[Dict[str, str]]:
 
     @property
     def rowcount(self) -> int:
-        """Return number of rows fetched so far (not total available)"""
+        """Return number of rows fetched/affected"""
+        # Check for write operation results (UPDATE, DELETE, INSERT)
+        if hasattr(self, "_insert_result") and self._insert_result:
+            # INSERT operation - return number of inserted documents
+            return self._insert_result.get("n", 0)
+
+        # Check command result for write operations
+        if self._command_result:
+            # For UPDATE/DELETE operations, check 'n' (modified count) or 'nModified'
+            if "n" in self._command_result:
+                return self._command_result.get("n", 0)
+            if "nModified" in self._command_result:
+                return self._command_result.get("nModified", 0)
+
+        # For SELECT/QUERY operations, return number of fetched rows
         return self._total_fetched
 
     @property

pymongosql/sql/ast.py

@@ -14,6 +14,8 @@
 from .partiql.PartiQLParserVisitor import PartiQLParserVisitor
 from .query_builder import QueryExecutionPlan
 from .query_handler import QueryParseResult
+from .update_builder import UpdateExecutionPlan
+from .update_handler import UpdateParseResult
 
 _logger = logging.getLogger(__name__)
 
@@ -38,6 +40,7 @@ def __init__(self) -> None:
         self._parse_result = QueryParseResult.for_visitor()
         self._insert_parse_result = InsertParseResult.for_visitor()
         self._delete_parse_result = DeleteParseResult.for_visitor()
+        self._update_parse_result = UpdateParseResult.for_visitor()
         # Track current statement kind generically so UPDATE/DELETE can reuse this
         self._current_operation: str = "select"  # expected values: select | insert | update | delete
         self._handlers = self._initialize_handlers()
@@ -50,6 +53,7 @@ def _initialize_handlers(self) -> Dict[str, BaseHandler]:
             "from": HandlerFactory.get_visitor_handler("from"),
             "where": HandlerFactory.get_visitor_handler("where"),
             "insert": HandlerFactory.get_visitor_handler("insert"),
+            "update": HandlerFactory.get_visitor_handler("update"),
             "delete": HandlerFactory.get_visitor_handler("delete"),
         }
 
@@ -58,12 +62,16 @@ def parse_result(self) -> QueryParseResult:
         """Get the current parse result"""
         return self._parse_result
 
-    def parse_to_execution_plan(self) -> Union[QueryExecutionPlan, InsertExecutionPlan, DeleteExecutionPlan]:
+    def parse_to_execution_plan(
+        self,
+    ) -> Union[QueryExecutionPlan, InsertExecutionPlan, DeleteExecutionPlan, UpdateExecutionPlan]:
         """Convert the parse result to an execution plan using BuilderFactory."""
         if self._current_operation == "insert":
             return self._build_insert_plan()
         elif self._current_operation == "delete":
             return self._build_delete_plan()
+        elif self._current_operation == "update":
+            return self._build_update_plan()
 
         return self._build_query_plan()
 
@@ -110,6 +118,23 @@ def _build_delete_plan(self) -> DeleteExecutionPlan:
 
         return builder.build()
 
+    def _build_update_plan(self) -> UpdateExecutionPlan:
+        """Build an UPDATE execution plan from UPDATE parsing."""
+        _logger.debug(
+            f"Building UPDATE plan with collection: {self._update_parse_result.collection}, "
+            f"update_fields: {self._update_parse_result.update_fields}, "
+            f"filters: {self._update_parse_result.filter_conditions}"
+        )
+        builder = BuilderFactory.create_update_builder().collection(self._update_parse_result.collection)
+
+        if self._update_parse_result.update_fields:
+            builder.update_fields(self._update_parse_result.update_fields)
+
+        if self._update_parse_result.filter_conditions:
+            builder.filter_conditions(self._update_parse_result.filter_conditions)
+
+        return builder.build()
+
     def visitRoot(self, ctx: PartiQLParser.RootContext) -> Any:
         """Visit root node and process child nodes"""
         _logger.debug("Starting to parse SQL query")
@@ -212,6 +237,12 @@ def visitWhereClause(self, ctx: PartiQLParser.WhereClauseContext) -> Any:
                 if handler:
                     return handler.handle_where_clause(ctx, self._delete_parse_result)
                 return {}
+            # For UPDATE, use the update handler
+            elif self._current_operation == "update":
+                handler = self._handlers.get("update")
+                if handler:
+                    return handler.handle_where_clause(ctx, self._update_parse_result)
+                return {}
             else:
                 # For other operations, use the where handler
                 handler = self._handlers["where"]
@@ -293,3 +324,29 @@ def visitOffsetByClause(self, ctx: PartiQLParser.OffsetByClauseContext) -> Any:
         except Exception as e:
             _logger.warning(f"Error processing OFFSET clause: {e}")
             return self.visitChildren(ctx)
+
+    def visitUpdateClause(self, ctx: PartiQLParser.UpdateClauseContext) -> Any:
+        """Handle UPDATE clause to extract collection/table name."""
+        _logger.debug("Processing UPDATE clause")
+        self._current_operation = "update"
+        # Reset update parse result for this statement
+        self._update_parse_result = UpdateParseResult.for_visitor()
+
+        handler = self._handlers.get("update")
+        if handler:
+            handler.handle_visitor(ctx, self._update_parse_result)
+
+        # Visit children to process SET and WHERE clauses
+        return self.visitChildren(ctx)
+
+    def visitSetCommand(self, ctx: PartiQLParser.SetCommandContext) -> Any:
+        """Handle SET command for UPDATE statements."""
+        _logger.debug("Processing SET command")
+
+        if self._current_operation == "update":
+            handler = self._handlers.get("update")
+            if handler:
+                handler.handle_set_command(ctx, self._update_parse_result)
+                return None
+
+        return self.visitChildren(ctx)

pymongosql/sql/builder.py

@@ -57,6 +57,14 @@ def create_delete_builder():
 
         return MongoDeleteBuilder()
 
+    @staticmethod
+    def create_update_builder():
+        """Create a builder for UPDATE queries"""
+        # Local import to avoid circular dependency during module import
+        from .update_builder import MongoUpdateBuilder
+
+        return MongoUpdateBuilder()
+
 
 __all__ = [
     "ExecutionPlan",

pymongosql/sql/handler.py

@@ -809,13 +809,15 @@ def _initialize_visitor_handlers(cls):
             from .delete_handler import DeleteHandler
             from .insert_handler import InsertHandler
             from .query_handler import FromHandler, SelectHandler, WhereHandler
+            from .update_handler import UpdateHandler
 
             cls._visitor_handlers = {
                 "select": SelectHandler(),
                 "from": FromHandler(),
                 "where": WhereHandler(),
                 "insert": InsertHandler(),
                 "delete": DeleteHandler(),
+                "update": UpdateHandler(),
             }
         return cls._visitor_handlers
 

pymongosql/sql/parser.py

@@ -11,6 +11,7 @@
 from .delete_builder import DeleteExecutionPlan
 from .insert_builder import InsertExecutionPlan
 from .query_builder import QueryExecutionPlan
+from .update_builder import UpdateExecutionPlan
 
 _logger = logging.getLogger(__name__)
 
@@ -128,8 +129,10 @@ def _validate_ast(self) -> None:
 
         _logger.debug("AST validation successful")
 
-    def get_execution_plan(self) -> Union[QueryExecutionPlan, InsertExecutionPlan, DeleteExecutionPlan]:
-        """Parse SQL and return an execution plan (SELECT, INSERT, or DELETE)."""
+    def get_execution_plan(
+        self,
+    ) -> Union[QueryExecutionPlan, InsertExecutionPlan, DeleteExecutionPlan, UpdateExecutionPlan]:
+        """Parse SQL and return an execution plan (SELECT, INSERT, DELETE, or UPDATE)."""
         if self._ast is None:
             raise SqlSyntaxError("No AST available - parsing may have failed")