Update connection string pattern

Author: Peng Ren

pymongosql/__init__.py

@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     from .connection import Connection
 
-__version__: str = "0.2.3"
+__version__: str = "0.2.4"
 
 # Globals https://www.python.org/dev/peps/pep-0249/#globals
 apilevel: str = "2.0"

pymongosql/connection.py

@@ -38,14 +38,18 @@ def __init__(
 
         Supports connection string patterns:
         - mongodb://host:port/database - Core driver (no subquery support)
-        - mongodb+superset://host:port/database - Superset driver with subquery support
+        - mongodb+srv://host:port/database - Cloud/SRV connection string
+        - mongodb://host:port/database?mode=superset - Superset driver with subquery support
+        - mongodb+srv://host:port/database?mode=superset - Cloud SRV with superset mode
+
+        Mode is specified via the ?mode= query parameter. If not specified, defaults to "standard".
 
         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
-        mode, host = ConnectionHelper.parse_connection_string(connection_string)
+        mode, db_from_uri, host = ConnectionHelper.parse_connection_string(connection_string)
 
         self._mode = kwargs.pop("mode", None)
         if not self._mode and mode:
@@ -56,7 +60,10 @@ def __init__(
         self._port = port or 27017
 
         # Handle database parameter separately (not a MongoClient parameter)
+        # Explicit 'database' parameter takes precedence over database in URI
         self._database_name = kwargs.pop("database", None)
+        if not self._database_name and db_from_uri:
+            self._database_name = db_from_uri
 
         # Store all PyMongo parameters to pass through directly
         self._pymongo_params = kwargs.copy()

pymongosql/helper.py

@@ -7,7 +7,7 @@
 
 import logging
 from typing import Optional, Tuple
-from urllib.parse import urlparse
+from urllib.parse import parse_qs, urlparse
 
 _logger = logging.getLogger(__name__)
 
@@ -17,52 +17,80 @@ class ConnectionHelper:
 
     Supports connection string patterns:
     - mongodb://host:port/database - Core driver (no subquery support)
-    - mongodb+superset://host:port/database - Superset driver with subquery support
+    - mongodb+srv://host:port/database - Cloud/SRV connection string
+    - mongodb://host:port/database?mode=superset - Superset driver with subquery support
+    - mongodb+srv://host:port/database?mode=superset - Cloud SRV with superset mode
+
+    Mode is specified via query parameter (?mode=superset) and defaults to "standard" if not specified.
     """
 
     @staticmethod
-    def parse_connection_string(connection_string: str) -> Tuple[str, str, Optional[str], int, Optional[str]]:
+    def parse_connection_string(connection_string: Optional[str]) -> Tuple[Optional[str], Optional[str], Optional[str]]:
         """
-        Parse PyMongoSQL connection string and determine driver mode.
+        Parse MongoDB connection string and extract driver mode from query parameters.
+
+        Mode is extracted from the 'mode' query parameter and removed from the normalized
+        connection string. Database name is extracted from the path. If mode is not specified,
+        it defaults to "standard".
+
+        Supports all standard MongoDB connection string patterns:
+        mongodb://[username:password@]host1[:port1][,host2[:port2]...][/[defaultauthdb]?options]
+
+        Args:
+            connection_string: MongoDB connection string
+
+        Returns:
+            Tuple of (mode, database_name, normalized_connection_string)
+            - mode: "standard" (default) or other mode values specified via ?mode= parameter
+            - database_name: extracted database name from path, or None if not specified
+            - normalized_connection_string: connection string without the mode parameter
         """
         try:
             if not connection_string:
-                return "standard", None
+                return "standard", None, None
 
             parsed = urlparse(connection_string)
-            scheme = parsed.scheme
 
             if not parsed.scheme:
-                return "standard", connection_string
-
-            base_scheme = "mongodb"
-            mode = "standard"
-
-            # Determine mode from scheme
-            if "+" in scheme:
-                base_scheme = scheme.split("+")[0].lower()
-                mode = scheme.split("+")[-1].lower()
-
-            host = parsed.hostname or "localhost"
-            port = parsed.port or 27017
-            database = parsed.path.lstrip("/") if parsed.path else None
-
-            # Build normalized connection string with mongodb scheme (removing any +mode)
-            # Reconstruct netloc with credentials if present
-            netloc = host
-            if parsed.username:
-                creds = parsed.username
-                if parsed.password:
-                    creds += f":{parsed.password}"
-                netloc = f"{creds}@{host}"
-            netloc += f":{port}"
-
-            query_part = f"?{parsed.query}" if parsed.query else ""
-            normalized_connection_string = f"{base_scheme}://{netloc}/{database or ''}{query_part}"
-
-            _logger.debug(f"Parsed connection string - Mode: {mode}, Host: {host}, Port: {port}, Database: {database}")
-
-            return mode, normalized_connection_string
+                return "standard", None, connection_string
+
+            # Extract mode from query parameters (defaults to "standard" if not specified)
+            query_params = parse_qs(parsed.query, keep_blank_values=True) if parsed.query else {}
+            mode = query_params.get("mode", ["standard"])[0]
+
+            # Extract database name from path
+            database_name = None
+            if parsed.path:
+                # Remove leading slash and trailing slashes
+                path_parts = parsed.path.strip("/").split("/")
+                if path_parts and path_parts[0]:  # Get the first path segment as database name
+                    database_name = path_parts[0]
+
+            # Remove mode from query parameters
+            query_params.pop("mode", None)
+
+            # Rebuild query string without mode parameter
+            query_string = (
+                "&".join(f"{k}={v}" if v else k for k, v_list in query_params.items() for v in v_list)
+                if query_params
+                else ""
+            )
+
+            # Reconstruct the connection string without mode parameter
+            if query_string:
+                if parsed.path:
+                    normalized_connection_string = f"{parsed.scheme}://{parsed.netloc}{parsed.path}?{query_string}"
+                else:
+                    normalized_connection_string = f"{parsed.scheme}://{parsed.netloc}?{query_string}"
+            else:
+                if parsed.path:
+                    normalized_connection_string = f"{parsed.scheme}://{parsed.netloc}{parsed.path}"
+                else:
+                    normalized_connection_string = f"{parsed.scheme}://{parsed.netloc}"
+
+            _logger.debug(f"Parsed connection string - Mode: {mode}, Database: {database_name}")
+
+            return mode, database_name, normalized_connection_string
 
         except Exception as e:
             _logger.error(f"Failed to parse connection string: {e}")

pymongosql/sqlalchemy_mongodb/__init__.py

@@ -52,100 +52,23 @@ def create_engine_url(
         >>> url = create_engine_url("localhost", 27017, "mydb", mode="superset")
         >>> engine = sqlalchemy.create_engine(url)
     """
-    scheme = "mongodb+superset" if mode == "superset" else "mongodb"
+    scheme = "mongodb"
 
     params = []
     for key, value in kwargs.items():
         params.append(f"{key}={value}")
 
+    # Add mode parameter if not standard
+    if mode != "standard":
+        params.append(f"mode={mode}")
+
     param_str = "&".join(params)
     if param_str:
         param_str = "?" + param_str
 
     return f"{scheme}://{host}:{port}/{database}{param_str}"
 
 
-def create_mongodb_url(mongodb_uri: str) -> str:
-    """Convert a standard MongoDB URI to work with PyMongoSQL SQLAlchemy dialect.
-
-    Args:
-        mongodb_uri: Standard MongoDB connection string
-                    (e.g., 'mongodb://localhost:27017/mydb' or 'mongodb+srv://...')
-
-    Returns:
-        SQLAlchemy-compatible URL for PyMongoSQL
-
-    Example:
-        >>> url = create_mongodb_url("mongodb://user:pass@localhost:27017/mydb")
-        >>> engine = sqlalchemy.create_engine(url)
-    """
-    # Return the MongoDB URI as-is since the dialect now handles MongoDB URLs directly
-    return mongodb_uri
-
-
-def create_engine_from_mongodb_uri(mongodb_uri: str, **engine_kwargs):
-    """Create a SQLAlchemy engine from any MongoDB connection string.
-
-    This function handles mongodb://, mongodb+srv://, and mongodb+superset:// URIs properly.
-    Use this instead of create_engine() directly for special URI schemes.
-
-    Args:
-        mongodb_uri: MongoDB connection string (supports standard, SRV, and superset modes)
-        **engine_kwargs: Additional arguments passed to create_engine
-
-    Returns:
-        SQLAlchemy Engine object
-
-    Example:
-        >>> # For SRV records (Atlas/Cloud)
-        >>> engine = create_engine_from_mongodb_uri("mongodb+srv://user:pass@cluster.net/db")
-        >>> # For standard MongoDB
-        >>> engine = create_engine_from_mongodb_uri("mongodb://localhost:27017/mydb")
-        >>> # For superset mode (with subquery support)
-        >>> engine = create_engine_from_mongodb_uri("mongodb+superset://localhost:27017/mydb")
-    """
-    try:
-        from sqlalchemy import create_engine
-
-        if mongodb_uri.startswith("mongodb+srv://"):
-            # For MongoDB+SRV, convert to standard mongodb:// for SQLAlchemy compatibility
-            # SQLAlchemy doesn't handle the + character in scheme names well
-            converted_uri = mongodb_uri.replace("mongodb+srv://", "mongodb://")
-
-            # Create engine with converted URI
-            engine = create_engine(converted_uri, **engine_kwargs)
-
-            def custom_create_connect_args(url):
-                # Use original SRV URI for actual MongoDB connection
-                opts = {"host": mongodb_uri}
-                return [], opts
-
-            engine.dialect.create_connect_args = custom_create_connect_args
-            return engine
-        elif mongodb_uri.startswith("mongodb+superset://"):
-            # For MongoDB+Superset, convert to standard mongodb:// for SQLAlchemy compatibility
-            # but preserve the superset mode by passing it through connection options
-            converted_uri = mongodb_uri.replace("mongodb+superset://", "mongodb://")
-
-            # Create engine with converted URI
-            engine = create_engine(converted_uri, **engine_kwargs)
-
-            def custom_create_connect_args(url):
-                # Use original superset URI for actual MongoDB connection
-                # This preserves the superset mode for subquery support
-                opts = {"host": mongodb_uri}
-                return [], opts
-
-            engine.dialect.create_connect_args = custom_create_connect_args
-            return engine
-        else:
-            # Standard mongodb:// URLs work fine with SQLAlchemy
-            return create_engine(mongodb_uri, **engine_kwargs)
-
-    except ImportError:
-        raise ImportError("SQLAlchemy is required for engine creation")
-
-
 def register_dialect():
     """Register the PyMongoSQL dialect with SQLAlchemy.
 
@@ -166,20 +89,7 @@ def register_dialect():
             registry.register("mongodb+srv", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
             registry.register("mongodb.srv", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
         except Exception:
-            # If registration fails we fall back to handling SRV URIs in
-            # create_engine_from_mongodb_uri by converting 'mongodb+srv' to 'mongodb'.
-            pass
-
-        try:
-            registry.register(
-                "mongodb+superset", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect"
-            )
-            registry.register(
-                "mongodb.superset", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect"
-            )
-        except Exception:
-            # If registration fails we fall back to handling Superset URIs in
-            # create_engine_from_mongodb_uri by converting 'mongodb+superset' to 'mongodb'.
+            # If registration fails, users can convert URIs to standard mongodb:// format
             pass
 
         return True
@@ -197,8 +107,6 @@ def register_dialect():
 # Export all SQLAlchemy-related functionality
 __all__ = [
     "create_engine_url",
-    "create_mongodb_url",
-    "create_engine_from_mongodb_uri",
     "register_dialect",
     "__sqlalchemy_version__",
     "__supports_sqlalchemy__",

tests/conftest.py

@@ -28,15 +28,18 @@ def make_conn(**kwargs):
 def make_superset_conn(**kwargs):
     """Create a superset-mode Connection using TEST_URI if provided, otherwise use a local default."""
     if TEST_URI:
-        # Convert test URI to superset mode by replacing mongodb:// with mongodb+superset://
-        superset_uri = TEST_URI.replace("mongodb://", "mongodb+superset://", 1)
+        # Convert test URI to superset mode by adding ?mode=superset query parameter
+        if "?" in TEST_URI:
+            superset_uri = TEST_URI + "&mode=superset"
+        else:
+            superset_uri = TEST_URI + "?mode=superset"
         if "database" not in kwargs:
             kwargs["database"] = TEST_DB
         return Connection(host=superset_uri, **kwargs)
 
     # Default local connection parameters with superset mode
     defaults = {
-        "host": "mongodb+superset://testuser:testpass@localhost:27017/test_db?authSource=test_db",
+        "host": "mongodb://testuser:testpass@localhost:27017/test_db?authSource=test_db&mode=superset",
         "database": "test_db",
     }
     for k, v in defaults.items():

tests/test_connection.py

@@ -172,3 +172,16 @@ def test_no_database_param_uses_client_default_database(self):
         assert conn.database is not None
         assert conn.database.name == "test_db"
         conn.close()
+
+    def test_connection_string_with_mode_query_param(self):
+        """Test that connection string with ?mode parameter is parsed correctly"""
+        if TEST_URI:
+            # Test with mode parameter in query string
+            test_url = f"{TEST_URI.rstrip('/')}/test_db?mode=superset"
+        else:
+            test_url = "mongodb://localhost:27017/test_db?mode=superset"
+
+        conn = Connection(host=test_url)
+        assert conn.mode == "superset"
+        assert conn.database_name == "test_db"
+        conn.close()