Added classic insert statement with column list and works in SQLAlchemy

Author: Peng Ren

pymongosql/sql/ast.py

@@ -137,7 +137,10 @@ def visitInsertStatement(self, ctx: PartiQLParser.InsertStatementContext) -> Any
         self._insert_parse_result = InsertParseResult.for_visitor()
         handler = self._handlers.get("insert")
         if handler:
-            return handler.handle_visitor(ctx, self._insert_parse_result)
+            handler.handle_visitor(ctx, self._insert_parse_result)
+            # Continue visiting children to process columnList and values
+            self.visitChildren(ctx)
+            return self._insert_parse_result
         return self.visitChildren(ctx)
 
     def visitInsertStatementLegacy(self, ctx: PartiQLParser.InsertStatementLegacyContext) -> Any:
@@ -151,6 +154,22 @@ def visitInsertStatementLegacy(self, ctx: PartiQLParser.InsertStatementLegacyCon
             return handler.handle_visitor(ctx, self._insert_parse_result)
         return self.visitChildren(ctx)
 
+    def visitColumnList(self, ctx: PartiQLParser.ColumnListContext) -> Any:
+        """Handle column list in INSERT statements."""
+        if self._current_operation == "insert":
+            handler = self._handlers.get("insert")
+            if handler:
+                return handler.handle_column_list(ctx, self._insert_parse_result)
+        return self.visitChildren(ctx)
+
+    def visitValues(self, ctx: PartiQLParser.ValuesContext) -> Any:
+        """Handle VALUES clause in INSERT statements."""
+        if self._current_operation == "insert":
+            handler = self._handlers.get("insert")
+            if handler:
+                return handler.handle_values(ctx, self._insert_parse_result)
+        return self.visitChildren(ctx)
+
     def visitFromClauseSimpleExplicit(self, ctx: PartiQLParser.FromClauseSimpleExplicitContext) -> Any:
         """Handle FROM clause (explicit form) in DELETE statements."""
         if self._current_operation == "delete":

pymongosql/sql/insert_handler.py

@@ -39,8 +39,17 @@ def can_handle(self, ctx: Any) -> bool:
     def handle_visitor(self, ctx: Any, parse_result: InsertParseResult) -> InsertParseResult:
         try:
             collection = self._extract_collection(ctx)
-            value_text = self._extract_value_text(ctx)
 
+            # Check if this is a VALUES clause INSERT (new syntax)
+            if hasattr(ctx, "values") and ctx.values():
+                _logger.debug("Processing INSERT with VALUES clause")
+                parse_result.collection = collection
+                parse_result.insert_type = "values"
+                # Return parse_result - visitor will call handle_column_list and handle_values
+                return parse_result
+
+            # Otherwise, handle legacy value expression INSERT
+            value_text = self._extract_value_text(ctx)
             documents = self._parse_value_expr(value_text)
             param_style, param_count = self._detect_parameter_style(documents)
 
@@ -58,6 +67,137 @@ def handle_visitor(self, ctx: Any, parse_result: InsertParseResult) -> InsertPar
             parse_result.error_message = str(exc)
             return parse_result
 
+    def handle_column_list(self, ctx: Any, parse_result: InsertParseResult) -> Optional[List[str]]:
+        """Extract column names from columnList context."""
+        _logger.debug("InsertHandler processing column list")
+        try:
+            columns = []
+            column_names = ctx.columnName()
+            if column_names:
+                if not isinstance(column_names, list):
+                    column_names = [column_names]
+                for col_name_ctx in column_names:
+                    # columnName contains a symbolPrimitive
+                    symbol = col_name_ctx.symbolPrimitive()
+                    if symbol:
+                        columns.append(symbol.getText())
+            parse_result.insert_columns = columns
+            _logger.debug(f"Extracted columns for INSERT: {columns}")
+            return columns
+        except Exception as e:
+            _logger.warning(f"Error processing column list: {e}")
+            parse_result.has_errors = True
+            parse_result.error_message = str(e)
+        return None
+
+    def handle_values(self, ctx: Any, parse_result: InsertParseResult) -> Optional[List[List[Any]]]:
+        """Extract value rows from VALUES clause."""
+        _logger.debug("InsertHandler processing VALUES clause")
+        try:
+            rows = []
+            value_rows = ctx.valueRow()
+            if value_rows:
+                if not isinstance(value_rows, list):
+                    value_rows = [value_rows]
+                for value_row_ctx in value_rows:
+                    row_values = self._extract_value_row(value_row_ctx)
+                    rows.append(row_values)
+
+            parse_result.insert_values = rows
+
+            # Convert rows to documents
+            columns = parse_result.insert_columns
+            documents = self._convert_rows_to_documents(columns, rows)
+            parse_result.insert_documents = documents
+
+            # Detect parameter style
+            param_style, param_count = self._detect_parameter_style(documents)
+            parse_result.parameter_style = param_style
+            parse_result.parameter_count = param_count
+            parse_result.has_errors = False
+            parse_result.error_message = None
+
+            _logger.debug(f"Extracted {len(rows)} value rows for INSERT")
+            return rows
+        except Exception as e:
+            _logger.warning(f"Error processing VALUES clause: {e}")
+            parse_result.has_errors = True
+            parse_result.error_message = str(e)
+        return None
+
+    def _extract_value_row(self, value_row_ctx: Any) -> List[Any]:
+        """Extract values from a single valueRow."""
+        row_values = []
+        exprs = value_row_ctx.expr()
+        if exprs:
+            if not isinstance(exprs, list):
+                exprs = [exprs]
+            for expr_ctx in exprs:
+                value = self._parse_expression_value(expr_ctx)
+                row_values.append(value)
+        return row_values
+
+    def _parse_expression_value(self, expr_ctx: Any) -> Any:
+        """Parse a single expression value from the parse tree."""
+        if not expr_ctx:
+            return None
+
+        text = expr_ctx.getText()
+
+        # Handle NULL
+        if text.upper() == "NULL":
+            return None
+
+        # Handle boolean literals
+        if text.upper() == "TRUE":
+            return True
+        if text.upper() == "FALSE":
+            return False
+
+        # Handle string literals (quoted)
+        if (text.startswith("'") and text.endswith("'")) or (text.startswith('"') and text.endswith('"')):
+            return text[1:-1]
+
+        # Handle numeric literals
+        try:
+            if "." in text:
+                return float(text)
+            return int(text)
+        except ValueError:
+            pass
+
+        # Handle parameters (? or :name)
+        if text == "?":
+            return "?"
+        if text.startswith(":"):
+            return text
+
+        # For complex expressions, return as-is
+        return text
+
+    def _convert_rows_to_documents(self, columns: Optional[List[str]], rows: List[List[Any]]) -> List[Dict[str, Any]]:
+        """Convert rows to MongoDB documents."""
+        documents = []
+
+        for row in rows:
+            doc = {}
+
+            if columns:
+                # Use explicit column names
+                if len(row) != len(columns):
+                    raise ValueError(f"Column count ({len(columns)}) does not match value count ({len(row)})")
+
+                for col, val in zip(columns, row):
+                    doc[col] = val
+            else:
+                # Generate automatic column names (col0, col1, ...)
+                for idx, val in enumerate(row):
+                    doc[f"col{idx}"] = val
+
+            documents.append(doc)
+
+        return documents
+
     def _extract_collection(self, ctx: Any) -> str:
         if hasattr(ctx, "symbolPrimitive") and ctx.symbolPrimitive():
             return ctx.symbolPrimitive().getText()

pymongosql/sql/partiql/PartiQLLexer.py

@@ -1,4 +1,4 @@
-# Generated from PartiQLLexer.g4 by ANTLR 4.13.2
+# Generated from PartiQLLexer.g4 by ANTLR 4.13.0
 from antlr4 import *
 from io import StringIO
 import sys
@@ -1678,7 +1678,7 @@ class PartiQLLexer(Lexer):
 
     def __init__(self, input=None, output:TextIO = sys.stdout):
         super().__init__(input, output)
-        self.checkVersion("4.13.2")
+        self.checkVersion("4.13.0")
         self._interp = LexerATNSimulator(self, self.atn, self.decisionsToDFA, PredictionContextCache())
         self._actions = None
         self._predicates = None

pymongosql/sql/partiql/PartiQLParser.g4

@@ -166,8 +166,8 @@ insertCommandReturning
 
 // See the Grammar at https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md#2-proposed-grammar-and-semantics
 insertStatement
-    : INSERT INTO symbolPrimitive asIdent? value=expr onConflict?
-    | INSERT INTO symbolPrimitive columnList? values onConflict?
+    : INSERT INTO symbolPrimitive columnList? values onConflict?
+    | INSERT INTO symbolPrimitive asIdent? value=expr onConflict?
     ;
 
 columnList

pymongosql/sql/partiql/PartiQLParser.py

@@ -1,4 +1,4 @@
-# Generated from PartiQLParser.g4 by ANTLR 4.13.2
+# Generated from PartiQLParser.g4 by ANTLR 4.13.0
 # encoding: utf-8
 from antlr4 import *
 from io import StringIO
@@ -303,21 +303,21 @@ def serializedATN():
         513,511,1,0,0,0,513,514,1,0,0,0,514,516,1,0,0,0,515,517,3,66,33,
         0,516,515,1,0,0,0,516,517,1,0,0,0,517,519,1,0,0,0,518,520,3,86,43,
         0,519,518,1,0,0,0,519,520,1,0,0,0,520,57,1,0,0,0,521,522,5,112,0,
-        0,522,523,5,117,0,0,523,525,3,12,6,0,524,526,3,6,3,0,525,524,1,0,
-        0,0,525,526,1,0,0,0,526,527,1,0,0,0,527,529,3,188,94,0,528,530,3,
-        62,31,0,529,528,1,0,0,0,529,530,1,0,0,0,530,542,1,0,0,0,531,532,
-        5,112,0,0,532,533,5,117,0,0,533,535,3,12,6,0,534,536,3,60,30,0,535,
-        534,1,0,0,0,535,536,1,0,0,0,536,537,1,0,0,0,537,539,3,220,110,0,
-        538,540,3,62,31,0,539,538,1,0,0,0,539,540,1,0,0,0,540,542,1,0,0,
-        0,541,521,1,0,0,0,541,531,1,0,0,0,542,59,1,0,0,0,543,544,5,294,0,
-        0,544,549,3,24,12,0,545,546,5,270,0,0,546,548,3,24,12,0,547,545,
-        1,0,0,0,548,551,1,0,0,0,549,547,1,0,0,0,549,550,1,0,0,0,550,552,
-        1,0,0,0,551,549,1,0,0,0,552,553,5,295,0,0,553,61,1,0,0,0,554,555,
-        5,147,0,0,555,557,5,244,0,0,556,558,3,68,34,0,557,556,1,0,0,0,557,
-        558,1,0,0,0,558,559,1,0,0,0,559,560,3,72,36,0,560,63,1,0,0,0,561,
-        562,5,112,0,0,562,563,5,117,0,0,563,564,3,46,23,0,564,565,5,218,
-        0,0,565,568,3,188,94,0,566,567,5,13,0,0,567,569,3,188,94,0,568,566,
-        1,0,0,0,568,569,1,0,0,0,569,571,1,0,0,0,570,572,3,66,33,0,571,570,
+        0,522,523,5,117,0,0,523,525,3,12,6,0,524,526,3,60,30,0,525,524,1,
+        0,0,0,525,526,1,0,0,0,526,527,1,0,0,0,527,529,3,220,110,0,528,530,
+        3,62,31,0,529,528,1,0,0,0,529,530,1,0,0,0,530,542,1,0,0,0,531,532,
+        5,112,0,0,532,533,5,117,0,0,533,535,3,12,6,0,534,536,3,6,3,0,535,
+        534,1,0,0,0,535,536,1,0,0,0,536,537,1,0,0,0,537,539,3,188,94,0,538,
+        540,3,62,31,0,539,538,1,0,0,0,539,540,1,0,0,0,540,542,1,0,0,0,541,
+        521,1,0,0,0,541,531,1,0,0,0,542,59,1,0,0,0,543,544,5,294,0,0,544,
+        549,3,24,12,0,545,546,5,270,0,0,546,548,3,24,12,0,547,545,1,0,0,
+        0,548,551,1,0,0,0,549,547,1,0,0,0,549,550,1,0,0,0,550,552,1,0,0,
+        0,551,549,1,0,0,0,552,553,5,295,0,0,553,61,1,0,0,0,554,555,5,147,
+        0,0,555,557,5,244,0,0,556,558,3,68,34,0,557,556,1,0,0,0,557,558,
+        1,0,0,0,558,559,1,0,0,0,559,560,3,72,36,0,560,63,1,0,0,0,561,562,
+        5,112,0,0,562,563,5,117,0,0,563,564,3,46,23,0,564,565,5,218,0,0,
+        565,568,3,188,94,0,566,567,5,13,0,0,567,569,3,188,94,0,568,566,1,
+        0,0,0,568,569,1,0,0,0,569,571,1,0,0,0,570,572,3,66,33,0,571,570,
         1,0,0,0,571,572,1,0,0,0,572,65,1,0,0,0,573,574,5,147,0,0,574,575,
         5,244,0,0,575,576,5,225,0,0,576,577,3,188,94,0,577,578,5,245,0,0,
         578,579,5,250,0,0,579,67,1,0,0,0,580,581,5,294,0,0,581,586,3,12,
@@ -1384,7 +1384,7 @@ class PartiQLParser ( Parser ):
 
     def __init__(self, input:TokenStream, output:TextIO = sys.stdout):
         super().__init__(input, output)
-        self.checkVersion("4.13.2")
+        self.checkVersion("4.13.0")
         self._interp = ParserATNSimulator(self, self.atn, self.decisionsToDFA, self.sharedContextCache)
         self._predicates = None
 
@@ -3969,24 +3969,24 @@ def symbolPrimitive(self):
             return self.getTypedRuleContext(PartiQLParser.SymbolPrimitiveContext,0)
 
 
-        def expr(self):
-            return self.getTypedRuleContext(PartiQLParser.ExprContext,0)
+        def values(self):
+            return self.getTypedRuleContext(PartiQLParser.ValuesContext,0)
 
 
-        def asIdent(self):
-            return self.getTypedRuleContext(PartiQLParser.AsIdentContext,0)
+        def columnList(self):
+            return self.getTypedRuleContext(PartiQLParser.ColumnListContext,0)
 
 
         def onConflict(self):
             return self.getTypedRuleContext(PartiQLParser.OnConflictContext,0)
 
 
-        def values(self):
-            return self.getTypedRuleContext(PartiQLParser.ValuesContext,0)
+        def expr(self):
+            return self.getTypedRuleContext(PartiQLParser.ExprContext,0)
 
 
-        def columnList(self):
-            return self.getTypedRuleContext(PartiQLParser.ColumnListContext,0)
+        def asIdent(self):
+            return self.getTypedRuleContext(PartiQLParser.AsIdentContext,0)
 
 
         def getRuleIndex(self):
@@ -4029,13 +4029,13 @@ def insertStatement(self):
                 self.state = 525
                 self._errHandler.sync(self)
                 _la = self._input.LA(1)
-                if _la==10:
+                if _la==294:
                     self.state = 524
-                    self.asIdent()
+                    self.columnList()
 
 
                 self.state = 527
-                localctx.value = self.expr()
+                self.values()
                 self.state = 529
                 self._errHandler.sync(self)
                 _la = self._input.LA(1)
@@ -4057,13 +4057,13 @@ def insertStatement(self):
                 self.state = 535
                 self._errHandler.sync(self)
                 _la = self._input.LA(1)
-                if _la==294:
+                if _la==10:
                     self.state = 534
-                    self.columnList()
+                    self.asIdent()
 
 
                 self.state = 537
-                self.values()
+                localctx.value = self.expr()
                 self.state = 539
                 self._errHandler.sync(self)
                 _la = self._input.LA(1)

pymongosql/sql/partiql/PartiQLParserListener.py

@@ -1,4 +1,4 @@
-# Generated from PartiQLParser.g4 by ANTLR 4.13.2
+# Generated from PartiQLParser.g4 by ANTLR 4.13.0
 from antlr4 import *
 if "." in __name__:
     from .PartiQLParser import PartiQLParser

pymongosql/sql/partiql/PartiQLParserVisitor.py

@@ -1,4 +1,4 @@
-# Generated from PartiQLParser.g4 by ANTLR 4.13.2
+# Generated from PartiQLParser.g4 by ANTLR 4.13.0
 from antlr4 import *
 if "." in __name__:
     from .PartiQLParser import PartiQLParser