Superset implementation with QueryDB

Author: Peng Ren

pymongosql/__init__.py

@@ -42,6 +42,26 @@ def connect(*args, **kwargs) -> "Connection":
     return Connection(*args, **kwargs)
 
 
+# Register superset execution strategy for mongodb+superset:// connections
+def _register_superset_executor() -> None:
+    """Register SupersetExecution strategy for superset mode.
+
+    This allows the executor and cursor to be unaware of superset -
+    the execution strategy is automatically selected based on the connection mode.
+    """
+    try:
+        from .executor import ExecutionPlanFactory
+        from .superset_mongodb.executor import SupersetExecution
+
+        ExecutionPlanFactory.register_strategy(SupersetExecution())
+    except ImportError:
+        # Superset module not available - skip registration
+        pass
+
+
+# Auto-register superset executor on module import
+_register_superset_executor()
+
 # SQLAlchemy integration (optional)
 # For SQLAlchemy functionality, import from pymongosql.sqlalchemy_mongodb:
 #   from pymongosql.sqlalchemy_mongodb import create_engine_url, create_engine_from_mongodb_uri

pymongosql/common.py

@@ -17,10 +17,12 @@ class BaseCursor(metaclass=ABCMeta):
     def __init__(
         self,
         connection: "Connection",
+        mode: str = "standard",
         **kwargs,
     ) -> None:
         super().__init__()
         self._connection = connection
+        self.mode = mode
 
     @property
     def connection(self) -> "Connection":

pymongosql/connection.py

@@ -12,6 +12,7 @@
 from .common import BaseCursor
 from .cursor import Cursor
 from .error import DatabaseError, NotSupportedError, OperationalError
+from .helper import ConnectionHelper
 
 _logger = logging.getLogger(__name__)
 
@@ -35,9 +36,17 @@ def __init__(
         to ensure full compatibility. All parameters are passed through directly
         to the underlying MongoClient.
 
+        Supports connection string patterns:
+        - mongodb://host:port/database - Core driver (no subquery support)
+        - mongodb+superset://host:port/database - Superset driver with subquery support
+
         See PyMongo MongoClient documentation for full parameter details.
         https://www.mongodb.com/docs/languages/python/pymongo-driver/current/connect/mongoclient/
         """
+        # Check if connection string specifies mode
+        connection_string = host if isinstance(host, str) else None
+        self._mode, host = ConnectionHelper.parse_connection_string(connection_string)
+
         # Extract commonly used parameters for backward compatibility
         self._host = host or "localhost"
         self._port = port or 27017
@@ -154,6 +163,11 @@ def database(self) -> Database:
             raise OperationalError("No database selected")
         return self._database
 
+    @property
+    def mode(self) -> str:
+        """Get the specified mode"""
+        return self._mode
+
     def use_database(self, database_name: str) -> None:
         """Switch to a different database"""
         if self._client is None:
@@ -267,6 +281,7 @@ def cursor(self, cursor: Optional[Type[BaseCursor]] = None, **kwargs) -> BaseCur
 
         new_cursor = cursor(
             connection=self,
+            mode=self._mode,
             **kwargs,
         )
         self.cursor_pool.append(new_cursor)

pymongosql/cursor.py

@@ -2,13 +2,11 @@
 import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, TypeVar
 
-from pymongo.errors import PyMongoError
-
 from .common import BaseCursor, CursorIterator
 from .error import DatabaseError, OperationalError, ProgrammingError, SqlSyntaxError
+from .executor import ExecutionContext, ExecutionPlanFactory
 from .result_set import DictResultSet, ResultSet
 from .sql.builder import ExecutionPlan
-from .sql.parser import SQLParser
 
 if TYPE_CHECKING:
     from .connection import Connection
@@ -22,9 +20,10 @@ class Cursor(BaseCursor, CursorIterator):
 
     NO_RESULT_SET = "No result set."
 
-    def __init__(self, connection: "Connection", **kwargs) -> None:
+    def __init__(self, connection: "Connection", mode: str = "standard", **kwargs) -> None:
         super().__init__(
             connection=connection,
+            mode=mode,
             **kwargs,
         )
         self._kwargs = kwargs
@@ -76,74 +75,6 @@ def _check_closed(self) -> None:
         if self._is_closed:
             raise ProgrammingError("Cursor is closed")
 
-    def _parse_sql(self, sql: str) -> ExecutionPlan:
-        """Parse SQL statement and return ExecutionPlan"""
-        try:
-            parser = SQLParser(sql)
-            execution_plan = parser.get_execution_plan()
-
-            if not execution_plan.validate():
-                raise SqlSyntaxError("Generated query plan is invalid")
-
-            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_execution_plan(self, execution_plan: ExecutionPlan) -> None:
-        """Execute an ExecutionPlan against MongoDB using db.command"""
-        try:
-            # Get database
-            if not execution_plan.collection:
-                raise ProgrammingError("No collection specified in query")
-
-            db = self.connection.database
-
-            # Build MongoDB find command
-            find_command = {"find": execution_plan.collection, "filter": execution_plan.filter_stage or {}}
-
-            # Apply projection if specified
-            if execution_plan.projection_stage:
-                find_command["projection"] = execution_plan.projection_stage
-
-            # Apply sort if specified
-            if execution_plan.sort_stage:
-                sort_spec = {}
-                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 execution_plan.skip_stage:
-                find_command["skip"] = execution_plan.skip_stage
-
-            # Apply limit if specified
-            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, execution_plan=execution_plan, database=db, **self._kwargs
-            )
-
-            _logger.info(f"Query executed successfully on collection '{execution_plan.collection}'")
-
-        except PyMongoError as e:
-            _logger.error(f"MongoDB command execution failed: {e}")
-            raise DatabaseError(f"Command execution failed: {e}")
-        except Exception as e:
-            _logger.error(f"Unexpected error during command execution: {e}")
-            raise OperationalError(f"Command execution error: {e}")
-
     def execute(self: _T, operation: str, parameters: Optional[Dict[str, Any]] = None) -> _T:
         """Execute a SQL statement
 
@@ -160,11 +91,25 @@ 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 ExecutionPlan
-            self._current_execution_plan = self._parse_sql(operation)
+            # Create execution context
+            context = ExecutionContext(operation, self.mode)
+
+            # Get appropriate execution strategy
+            strategy = ExecutionPlanFactory.get_strategy(context)
+
+            # Execute using selected strategy (Standard or Subquery)
+            result = strategy.execute(context, self.connection)
 
-            # Execute the execution plan
-            self._execute_execution_plan(self._current_execution_plan)
+            # Store execution plan for reference
+            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,
+            )
 
             return self