Added basic insert support

Author: Peng Ren

pymongosql/cursor.py

@@ -29,7 +29,7 @@ def __init__(self, connection: "Connection", mode: str = "standard", **kwargs) -
         self._kwargs = kwargs
         self._result_set: Optional[ResultSet] = None
         self._result_set_class = ResultSet
-        self._current_execution_plan: Optional[QueryExecutionPlan] = None
+        self._current_execution_plan: Optional[Any] = None
         self._is_closed = False
 
     @property
@@ -103,12 +103,32 @@ def execute(self: _T, operation: str, parameters: Optional[Any] = None) -> _T:
             self._current_execution_plan = strategy.execution_plan
 
             # Create result set from command result
-            self._result_set = self._result_set_class(
-                command_result=result,
-                execution_plan=self._current_execution_plan,
-                database=self.connection.database,
-                **self._kwargs,
-            )
+            # For SELECT/QUERY operations, use the execution plan directly
+            if isinstance(self._current_execution_plan, QueryExecutionPlan):
+                execution_plan_for_rs = self._current_execution_plan
+                self._result_set = self._result_set_class(
+                    command_result=result,
+                    execution_plan=execution_plan_for_rs,
+                    database=self.connection.database,
+                    **self._kwargs,
+                )
+            else:
+                # For INSERT and other non-query operations, create a minimal synthetic result
+                # since INSERT commands don't return a cursor structure
+                stub_plan = QueryExecutionPlan(collection=self._current_execution_plan.collection)
+                self._result_set = self._result_set_class(
+                    command_result={
+                        "cursor": {
+                            "id": 0,
+                            "firstBatch": [],
+                        }
+                    },
+                    execution_plan=stub_plan,
+                    database=self.connection.database,
+                    **self._kwargs,
+                )
+                # Store the actual insert result for reference
+                self._result_set._insert_result = result
 
             return self
 

pymongosql/executor.py

@@ -7,6 +7,8 @@
 from pymongo.errors import PyMongoError
 
 from .error import DatabaseError, OperationalError, ProgrammingError, SqlSyntaxError
+from .helper import SQLHelper
+from .sql.insert_builder import InsertExecutionPlan
 from .sql.parser import SQLParser
 from .sql.query_builder import QueryExecutionPlan
 
@@ -30,7 +32,7 @@ class ExecutionStrategy(ABC):
 
     @property
     @abstractmethod
-    def execution_plan(self) -> QueryExecutionPlan:
+    def execution_plan(self) -> Union[QueryExecutionPlan, InsertExecutionPlan]:
         """Name of the execution plan"""
         pass
 
@@ -60,7 +62,7 @@ def supports(self, context: ExecutionContext) -> bool:
         pass
 
 
-class StandardExecution(ExecutionStrategy):
+class StandardQueryExecution(ExecutionStrategy):
     """Standard execution strategy for simple SELECT queries without subqueries"""
 
     @property
@@ -70,7 +72,8 @@ def execution_plan(self) -> QueryExecutionPlan:
 
     def supports(self, context: ExecutionContext) -> bool:
         """Support simple queries without subqueries"""
-        return "standard" in context.execution_mode.lower()
+        normalized = context.query.lstrip().upper()
+        return "standard" in context.execution_mode.lower() and normalized.startswith("SELECT")
 
     def _parse_sql(self, sql: str) -> QueryExecutionPlan:
         """Parse SQL statement and return QueryExecutionPlan"""
@@ -91,29 +94,7 @@ def _parse_sql(self, sql: str) -> QueryExecutionPlan:
 
     def _replace_placeholders(self, obj: Any, parameters: Sequence[Any]) -> Any:
         """Recursively replace ? placeholders with parameter values in filter/projection dicts"""
-        param_index = [0]  # Use list to allow modification in nested function
-
-        def replace_recursive(value: Any) -> Any:
-            if isinstance(value, str):
-                # Replace ? with the next parameter value
-                if value == "?":
-                    if param_index[0] < len(parameters):
-                        result = parameters[param_index[0]]
-                        param_index[0] += 1
-                        return result
-                    else:
-                        raise ProgrammingError(
-                            f"Not enough parameters provided: expected at least {param_index[0] + 1}"
-                        )
-                return value
-            elif isinstance(value, dict):
-                return {k: replace_recursive(v) for k, v in value.items()}
-            elif isinstance(value, list):
-                return [replace_recursive(item) for item in value]
-            else:
-                return value
-
-        return replace_recursive(obj)
+        return SQLHelper.replace_placeholders_generic(obj, parameters, "qmark")
 
     def _execute_execution_plan(
         self,
@@ -202,10 +183,87 @@ def execute(
         return self._execute_execution_plan(self._execution_plan, connection.database, processed_params)
 
 
+class InsertExecution(ExecutionStrategy):
+    """Execution strategy for INSERT statements."""
+
+    @property
+    def execution_plan(self) -> InsertExecutionPlan:
+        return self._execution_plan
+
+    def supports(self, context: ExecutionContext) -> bool:
+        return context.query.lstrip().upper().startswith("INSERT")
+
+    def _parse_sql(self, sql: str) -> InsertExecutionPlan:
+        try:
+            parser = SQLParser(sql)
+            plan = parser.get_execution_plan()
+
+            if not isinstance(plan, InsertExecutionPlan):
+                raise SqlSyntaxError("Expected INSERT execution plan")
+
+            if not plan.validate():
+                raise SqlSyntaxError("Generated insert 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 _replace_placeholders(
+        self,
+        documents: Sequence[Dict[str, Any]],
+        parameters: Optional[Union[Sequence[Any], Dict[str, Any]]],
+        style: Optional[str],
+    ) -> Sequence[Dict[str, Any]]:
+        return SQLHelper.replace_placeholders_generic(documents, parameters, style)
+
+    def _execute_execution_plan(
+        self,
+        execution_plan: InsertExecutionPlan,
+        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 insert")
+
+            docs = execution_plan.insert_documents or []
+            docs = self._replace_placeholders(docs, parameters, execution_plan.parameter_style)
+
+            command = {"insert": execution_plan.collection, "documents": docs}
+
+            _logger.debug(f"Executing MongoDB insert command: {command}")
+
+            return db.command(command)
+        except PyMongoError as e:
+            _logger.error(f"MongoDB insert failed: {e}")
+            raise DatabaseError(f"Insert 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 insert execution: {e}")
+            raise OperationalError(f"Insert 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 insert 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 = [StandardExecution()]
+    _strategies = [InsertExecution(), StandardQueryExecution()]
 
     @classmethod
     def get_strategy(cls, context: ExecutionContext) -> ExecutionStrategy:
@@ -216,7 +274,7 @@ def get_strategy(cls, context: ExecutionContext) -> ExecutionStrategy:
                 return strategy
 
         # Fallback to standard execution
-        return StandardExecution()
+        return StandardQueryExecution()
 
     @classmethod
     def register_strategy(cls, strategy: ExecutionStrategy) -> None:

pymongosql/helper.py

@@ -6,9 +6,11 @@
 """
 
 import logging
-from typing import Optional, Tuple
+from typing import Any, Optional, Sequence, Tuple
 from urllib.parse import parse_qs, urlparse
 
+from .error import ProgrammingError
+
 _logger = logging.getLogger(__name__)
 
 
@@ -95,3 +97,54 @@ def parse_connection_string(connection_string: Optional[str]) -> Tuple[Optional[
         except Exception as e:
             _logger.error(f"Failed to parse connection string: {e}")
             raise ValueError(f"Invalid connection string format: {e}")
+
+
+class SQLHelper:
+    """SQL-related helper utilities."""
+
+    @staticmethod
+    def replace_placeholders_generic(value: Any, parameters: Any, style: Optional[str]) -> Any:
+        """Recursively replace placeholders in nested structures for qmark or named styles."""
+        if style is None or parameters is None:
+            return value
+
+        if style == "qmark":
+            if not isinstance(parameters, Sequence) or isinstance(parameters, (str, bytes, dict)):
+                raise ProgrammingError("Positional parameters must be provided as a sequence")
+
+            idx = [0]
+
+            def replace(val: Any) -> Any:
+                if isinstance(val, str) and val == "?":
+                    if idx[0] >= len(parameters):
+                        raise ProgrammingError("Not enough parameters provided")
+                    out = parameters[idx[0]]
+                    idx[0] += 1
+                    return out
+                if isinstance(val, dict):
+                    return {k: replace(v) for k, v in val.items()}
+                if isinstance(val, list):
+                    return [replace(v) for v in val]
+                return val
+
+            return replace(value)
+
+        if style == "named":
+            if not isinstance(parameters, dict):
+                raise ProgrammingError("Named parameters must be provided as a mapping")
+
+            def replace(val: Any) -> Any:
+                if isinstance(val, str) and val.startswith(":"):
+                    key = val[1:]
+                    if key not in parameters:
+                        raise ProgrammingError(f"Missing named parameter: {key}")
+                    return parameters[key]
+                if isinstance(val, dict):
+                    return {k: replace(v) for k, v in val.items()}
+                if isinstance(val, list):
+                    return [replace(v) for v in val]
+                return val
+
+            return replace(value)
+
+        return value

pymongosql/sql/ast.py

@@ -46,6 +46,7 @@ def _initialize_handlers(self) -> Dict[str, BaseHandler]:
             "select": HandlerFactory.get_visitor_handler("select"),
             "from": HandlerFactory.get_visitor_handler("from"),
             "where": HandlerFactory.get_visitor_handler("where"),
+            "insert": HandlerFactory.get_visitor_handler("insert"),
         }
 
     @property
@@ -74,6 +75,9 @@ def _build_query_plan(self) -> QueryExecutionPlan:
 
     def _build_insert_plan(self) -> InsertExecutionPlan:
         """Build an INSERT execution plan from INSERT parsing."""
+        if self._insert_parse_result.has_errors:
+            raise SqlSyntaxError(self._insert_parse_result.error_message or "INSERT parsing failed")
+
         builder = BuilderFactory.create_insert_builder().collection(self._insert_parse_result.collection)
 
         documents = self._insert_parse_result.insert_documents or []
@@ -145,6 +149,24 @@ def visitWhereClauseSelect(self, ctx: PartiQLParser.WhereClauseSelectContext) ->
             _logger.warning(f"Error processing WHERE clause: {e}")
             return self.visitChildren(ctx)
 
+    def visitInsertStatement(self, ctx: PartiQLParser.InsertStatementContext) -> Any:
+        """Handle INSERT statements via the insert handler."""
+        _logger.debug("Processing INSERT statement")
+        self._current_operation = "insert"
+        handler = self._handlers.get("insert")
+        if handler:
+            return handler.handle_visitor(ctx, self._insert_parse_result)
+        return self.visitChildren(ctx)
+
+    def visitInsertStatementLegacy(self, ctx: PartiQLParser.InsertStatementLegacyContext) -> Any:
+        """Handle legacy INSERT statements."""
+        _logger.debug("Processing INSERT legacy statement")
+        self._current_operation = "insert"
+        handler = self._handlers.get("insert")
+        if handler:
+            return handler.handle_visitor(ctx, self._insert_parse_result)
+        return self.visitChildren(ctx)
+
     def visitOrderByClause(self, ctx: PartiQLParser.OrderByClauseContext) -> Any:
         """Handle ORDER BY clause for sorting"""
         _logger.debug("Processing ORDER BY clause")

pymongosql/sql/insert_builder.py

@@ -80,7 +80,7 @@ def insert_documents(self, documents: List[Dict[str, Any]]) -> "MongoInsertBuild
 
     def parameter_style(self, style: Optional[str]) -> "MongoInsertBuilder":
         """Set parameter binding style for tracking."""
-        if style and style not in ["qmark"]:
+        if style and style not in ["qmark", "named"]:
             self._add_error(f"Invalid parameter style: {style}")
             return self