Add more cases for query

Author: Peng Ren

pymongosql/cursor.py

@@ -8,7 +8,7 @@
 from .common import BaseCursor, CursorIterator
 from .error import DatabaseError, OperationalError, ProgrammingError, SqlSyntaxError
 from .result_set import ResultSet
-from .sql.builder import QueryPlan
+from .sql.builder import ExecutionPlan
 from .sql.parser import SQLParser
 
 if TYPE_CHECKING:
@@ -31,7 +31,7 @@ def __init__(self, connection: "Connection", **kwargs) -> None:
         self._kwargs = kwargs
         self._result_set: Optional[ResultSet] = None
         self._result_set_class = ResultSet
-        self._current_query_plan: Optional[QueryPlan] = None
+        self._current_execution_plan: Optional[ExecutionPlan] = None
         self._mongo_cursor: Optional[MongoCursor] = None
         self._is_closed = False
 
@@ -78,65 +78,66 @@ def _check_closed(self) -> None:
         if self._is_closed:
             raise ProgrammingError("Cursor is closed")
 
-    def _parse_sql(self, sql: str) -> QueryPlan:
-        """Parse SQL statement and return QueryPlan"""
+    def _parse_sql(self, sql: str) -> ExecutionPlan:
+        """Parse SQL statement and return ExecutionPlan"""
         try:
             parser = SQLParser(sql)
-            query_plan = parser.get_query_plan()
+            execution_plan = parser.get_execution_plan()
 
-            if not query_plan.validate():
+            if not execution_plan.validate():
                 raise SqlSyntaxError("Generated query plan is invalid")
 
-            return query_plan
+            return execution_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_query_plan(self, query_plan: QueryPlan) -> None:
-        """Execute a QueryPlan against MongoDB using db.command"""
+    def _execute_execution_plan(self, execution_plan: ExecutionPlan) -> None:
+        """Execute an ExecutionPlan against MongoDB using db.command"""
         try:
             # Get database
-            if not query_plan.collection:
+            if not execution_plan.collection:
                 raise ProgrammingError("No collection specified in query")
 
             db = self.connection.database
 
             # Build MongoDB find command
-            find_command = {"find": query_plan.collection, "filter": query_plan.filter_stage or {}}
+            find_command = {"find": execution_plan.collection, "filter": execution_plan.filter_stage or {}}
 
-            # Convert projection stage from alias mapping to MongoDB format
-            if query_plan.projection_stage:
-                # Convert {"field": "alias"} to {"field": 1} for MongoDB
-                find_command["projection"] = {field: 1 for field in query_plan.projection_stage.keys()}
+            # Apply projection if specified (already in MongoDB format)
+            if execution_plan.projection_stage:
+                find_command["projection"] = execution_plan.projection_stage
 
             # Apply sort if specified
-            if query_plan.sort_stage:
+            if execution_plan.sort_stage:
                 sort_spec = {}
-                for sort_dict in query_plan.sort_stage:
+                for sort_dict in execution_plan.sort_stage:
                     for field, direction in sort_dict.items():
                         sort_spec[field] = direction
                 find_command["sort"] = sort_spec
 
             # Apply skip if specified
-            if query_plan.skip_stage:
-                find_command["skip"] = query_plan.skip_stage
+            if execution_plan.skip_stage:
+                find_command["skip"] = execution_plan.skip_stage
 
             # Apply limit if specified
-            if query_plan.limit_stage:
-                find_command["limit"] = query_plan.limit_stage
+            if execution_plan.limit_stage:
+                find_command["limit"] = execution_plan.limit_stage
 
             _logger.debug(f"Executing MongoDB command: {find_command}")
 
             # Execute find command directly
             result = db.command(find_command)
 
             # Create result set from command result
-            self._result_set = self._result_set_class(command_result=result, query_plan=query_plan, **self._kwargs)
+            self._result_set = self._result_set_class(
+                command_result=result, execution_plan=execution_plan, **self._kwargs
+            )
 
-            _logger.info(f"Query executed successfully on collection '{query_plan.collection}'")
+            _logger.info(f"Query executed successfully on collection '{execution_plan.collection}'")
 
         except PyMongoError as e:
             _logger.error(f"MongoDB command execution failed: {e}")
@@ -161,11 +162,11 @@ def execute(self: _T, operation: str, parameters: Optional[Dict[str, Any]] = Non
             _logger.warning("Parameter substitution not yet implemented, ignoring parameters")
 
         try:
-            # Parse SQL to QueryPlan
-            self._current_query_plan = self._parse_sql(operation)
+            # Parse SQL to ExecutionPlan
+            self._current_execution_plan = self._parse_sql(operation)
 
-            # Execute the query plan
-            self._execute_query_plan(self._current_query_plan)
+            # Execute the execution plan
+            self._execute_execution_plan(self._current_execution_plan)
 
             return self
 

pymongosql/result_set.py

@@ -7,7 +7,7 @@
 
 from .common import CursorIterator
 from .error import DatabaseError, ProgrammingError
-from .sql.builder import QueryPlan
+from .sql.builder import ExecutionPlan
 
 _logger = logging.getLogger(__name__)
 
@@ -19,7 +19,7 @@ def __init__(
         self,
         command_result: Optional[Dict[str, Any]] = None,
         mongo_cursor: Optional[MongoCursor] = None,
-        query_plan: QueryPlan = None,
+        execution_plan: ExecutionPlan = None,
         arraysize: int = None,
         **kwargs,
     ) -> None:
@@ -32,7 +32,7 @@ def __init__(
             # Extract cursor info from command result
             self._result_cursor = command_result.get("cursor", {})
             self._raw_results = self._result_cursor.get("firstBatch", [])
-            self._cached_results: List[Dict[str, Any]] = []  # Will be populated after query_plan is set
+            self._cached_results: List[Dict[str, Any]] = []
         elif mongo_cursor is not None:
             self._mongo_cursor = mongo_cursor
             self._command_result = None
@@ -41,33 +41,33 @@ def __init__(
         else:
             raise ProgrammingError("Either command_result or mongo_cursor must be provided")
 
-        self._query_plan = query_plan
+        self._execution_plan = execution_plan
         self._is_closed = False
         self._cache_exhausted = False
         self._total_fetched = 0
         self._description: Optional[List[Tuple[str, str, None, None, None, None, None]]] = None
         self._errors: List[Dict[str, str]] = []
 
-        # Apply projection mapping for command results now that query_plan is set
+        # Apply projection mapping for command results now that execution_plan is set
         if command_result is not None and self._raw_results:
             self._cached_results = [self._process_document(doc) for doc in self._raw_results]
 
         # Build description from projection
         self._build_description()
 
     def _build_description(self) -> None:
-        """Build column description from query plan projection"""
-        if not self._query_plan.projection_stage:
+        """Build column description from execution plan projection"""
+        if not self._execution_plan.projection_stage:
             # No projection specified, description will be built dynamically
             self._description = None
             return
 
-        # Build description from projection
+        # Build description from projection (now in MongoDB format {field: 1})
         description = []
-        for field_name, alias in self._query_plan.projection_stage.items():
+        for field_name, include_flag in self._execution_plan.projection_stage.items():
             # SQL cursor description format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
-            column_name = alias if alias != field_name else field_name
-            description.append((column_name, "VARCHAR", None, None, None, None, None))
+            if include_flag == 1:  # Field is included in projection
+                description.append((field_name, "VARCHAR", None, None, None, None, None))
 
         self._description = description
 
@@ -111,20 +111,19 @@ def _ensure_results_available(self, count: int = 1) -> None:
 
     def _process_document(self, doc: Dict[str, Any]) -> Dict[str, Any]:
         """Process a MongoDB document according to projection mapping"""
-        if not self._query_plan.projection_stage:
+        if not self._execution_plan.projection_stage:
             # No projection, return document as-is (including _id)
             return dict(doc)
 
-        # Apply projection mapping
+        # Apply projection mapping (now using MongoDB format {field: 1})
         processed = {}
-        for field_name, alias in self._query_plan.projection_stage.items():
-            if field_name in doc:
-                output_name = alias if alias != field_name else field_name
-                processed[output_name] = doc[field_name]
-            elif field_name != "_id":  # _id might be excluded by MongoDB
-                # Field not found, set to None
-                output_name = alias if alias != field_name else field_name
-                processed[output_name] = None
+        for field_name, include_flag in self._execution_plan.projection_stage.items():
+            if include_flag == 1:  # Field is included in projection
+                if field_name in doc:
+                    processed[field_name] = doc[field_name]
+                elif field_name != "_id":  # _id might be excluded by MongoDB
+                    # Field not found, set to None
+                    processed[field_name] = None
 
         return processed
 

pymongosql/sql/ast.py

@@ -3,7 +3,7 @@
 from typing import Any, Dict
 
 from ..error import SqlSyntaxError
-from .builder import QueryPlan
+from .builder import ExecutionPlan
 from .handler import BaseHandler, HandlerFactory, ParseResult
 from .partiql.PartiQLLexer import PartiQLLexer
 from .partiql.PartiQLParser import PartiQLParser
@@ -46,9 +46,9 @@ def parse_result(self) -> ParseResult:
         """Get the current parse result"""
         return self._parse_result
 
-    def parse_to_query_plan(self) -> QueryPlan:
-        """Convert the parse result to a QueryPlan"""
-        return QueryPlan(
+    def parse_to_execution_plan(self) -> ExecutionPlan:
+        """Convert the parse result to an ExecutionPlan"""
+        return ExecutionPlan(
             collection=self._parse_result.collection,
             filter_stage=self._parse_result.filter_conditions,
             projection_stage=self._parse_result.projection,
@@ -114,3 +114,43 @@ def visitWhereClauseSelect(self, ctx: PartiQLParser.WhereClauseSelectContext) ->
         except Exception as e:
             _logger.warning(f"Error processing WHERE clause: {e}")
             return self.visitChildren(ctx)
+
+    def visitOrderByClause(self, ctx: PartiQLParser.OrderByClauseContext) -> Any:
+        """Handle ORDER BY clause for sorting"""
+        _logger.debug("Processing ORDER BY clause")
+
+        try:
+            sort_specs = []
+            if hasattr(ctx, "orderSortSpec") and ctx.orderSortSpec():
+                for sort_spec in ctx.orderSortSpec():
+                    field_name = sort_spec.expr().getText() if sort_spec.expr() else "_id"
+                    # Check for ASC/DESC (default is ASC = 1)
+                    direction = 1  # ASC
+                    if hasattr(sort_spec, "DESC") and sort_spec.DESC():
+                        direction = -1  # DESC
+                    # Convert to the expected format: List[Dict[str, int]]
+                    sort_specs.append({field_name: direction})
+
+            self._parse_result.sort_fields = sort_specs
+            _logger.debug(f"Extracted sort specifications: {sort_specs}")
+            return self.visitChildren(ctx)
+        except Exception as e:
+            _logger.warning(f"Error processing ORDER BY clause: {e}")
+            return self.visitChildren(ctx)
+
+    def visitLimitClause(self, ctx: PartiQLParser.LimitClauseContext) -> Any:
+        """Handle LIMIT clause for result limiting"""
+        _logger.debug("Processing LIMIT clause")
+        try:
+            if hasattr(ctx, "exprSelect") and ctx.exprSelect():
+                limit_text = ctx.exprSelect().getText()
+                try:
+                    limit_value = int(limit_text)
+                    self._parse_result.limit_value = limit_value
+                    _logger.debug(f"Extracted limit value: {limit_value}")
+                except ValueError as e:
+                    _logger.warning(f"Invalid LIMIT value '{limit_text}': {e}")
+            return self.visitChildren(ctx)
+        except Exception as e:
+            _logger.warning(f"Error processing LIMIT clause: {e}")
+            return self.visitChildren(ctx)