This is an automated email from the ASF dual-hosted git repository.
msyavuz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 838376dbc28 fix(mongodb): treat the Mongo database as the SQL Lab
schema so Data Preview returns rows (#44141)
838376dbc28 is described below
commit 838376dbc281f518ca1108a8a224c0793bd62240
Author: Mehmet Salih Yavuz <[email protected]>
AuthorDate: Wed Sep 23 16:23:21 2026 +0300
fix(mongodb): treat the Mongo database as the SQL Lab schema so Data
Preview returns rows (#44141)
---
superset/db_engine_specs/mongodb.py | 63 ++++++-
tests/unit_tests/db_engine_specs/test_mongodb.py | 205 ++++++++++++++++++++++-
2 files changed, 266 insertions(+), 2 deletions(-)
diff --git a/superset/db_engine_specs/mongodb.py
b/superset/db_engine_specs/mongodb.py
index 03ea7dc4b92..28c41759525 100644
--- a/superset/db_engine_specs/mongodb.py
+++ b/superset/db_engine_specs/mongodb.py
@@ -23,12 +23,19 @@ to enable SQL queries on MongoDB collections.
from __future__ import annotations
from datetime import datetime
-from typing import Any, Optional
+from typing import Any, Optional, TYPE_CHECKING
from sqlalchemy import types
+from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.engine.url import URL
from superset.constants import TimeGrain
+from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import BaseEngineSpec, DatabaseCategory
+from superset.sql.parse import Table
+
+if TYPE_CHECKING:
+ from superset.models.core import Database
class MongoDBEngineSpec(BaseEngineSpec):
@@ -38,6 +45,13 @@ class MongoDBEngineSpec(BaseEngineSpec):
engine_name = "MongoDB"
force_column_alias_quotes = False
+ # A MongoDB database is exposed as a SQLAlchemy schema. PyMongoSQL treats
the
+ # whole FROM reference as the collection name, so the schema is selected
through
+ # the driver's ``database`` connect argument rather than by qualifying the
+ # collection. The URI database is left untouched because MongoDB also uses
it as
+ # the default ``authSource``.
+ supports_dynamic_schema = True
+
metadata = {
"description": ("MongoDB is a document-oriented, operational NoSQL
database."),
"logo": "mongodb.png",
@@ -102,6 +116,53 @@ class MongoDBEngineSpec(BaseEngineSpec):
),
}
+ @classmethod
+ def adjust_engine_params(
+ cls,
+ uri: URL,
+ connect_args: dict[str, Any],
+ catalog: Optional[str] = None,
+ schema: Optional[str] = None,
+ ) -> tuple[URL, dict[str, Any]]:
+ uri, new_connect_args = super().adjust_engine_params(
+ uri,
+ connect_args,
+ catalog,
+ schema,
+ )
+
+ if schema:
+ new_connect_args = {**new_connect_args, "database": schema}
+
+ return uri, new_connect_args
+
+ @classmethod
+ def get_schema_from_engine_params(
+ cls,
+ sqlalchemy_uri: URL,
+ connect_args: dict[str, Any],
+ ) -> Optional[str]:
+ # PyMongoSQL gives the ``database`` connect argument precedence over
the URI.
+ return connect_args.get("database") or sqlalchemy_uri.database or None
+
+ @classmethod
+ def get_default_schema(
+ cls,
+ database: Database,
+ catalog: Optional[str],
+ ) -> Optional[str]:
+ return cls.get_schema_from_engine_params(
+ make_url_safe(database.sqlalchemy_uri), database.connect_args
+ )
+
+ @classmethod
+ def quote_table(cls, table: Table, dialect: Dialect) -> str:
+ """
+ Quote only the collection name. PyMongoSQL treats the whole FROM
reference
+ as the collection, so the schema is applied via the ``database``
connect arg.
+ """
+ return dialect.identifier_preparer.quote(table.table)
+
@classmethod
def epoch_to_dttm(cls) -> str:
return "datetime({col}, 'unixepoch')"
diff --git a/tests/unit_tests/db_engine_specs/test_mongodb.py
b/tests/unit_tests/db_engine_specs/test_mongodb.py
index 2272fd51610..8a481daa082 100644
--- a/tests/unit_tests/db_engine_specs/test_mongodb.py
+++ b/tests/unit_tests/db_engine_specs/test_mongodb.py
@@ -15,11 +15,13 @@
# specific language governing permissions and limitations
# under the License.
from datetime import datetime
-from typing import Optional
+from typing import Any, Optional
import pytest
+from pytest_mock import MockerFixture
from superset.constants import TimeGrain
+from superset.utils import json
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
from tests.unit_tests.fixtures.common import dttm # noqa: F401
@@ -123,3 +125,204 @@ def test_engine_metadata() -> None:
assert spec.engine == "mongodb"
assert spec.engine_name == "MongoDB"
assert spec.force_column_alias_quotes is False
+ assert spec.supports_dynamic_schema is True
+
+
[email protected](
+ "connect_args,schema,expected",
+ [
+ ({"foo": "bar"}, "dbtwo", {"foo": "bar", "database": "dbtwo"}),
+ ({"foo": "bar"}, None, {"foo": "bar"}),
+ (
+ {"database": "dbone", "authSource": "dbone"},
+ "dbtwo",
+ {"database": "dbtwo", "authSource": "dbone"},
+ ),
+ ],
+)
+def test_adjust_engine_params(
+ connect_args: dict[str, Any],
+ schema: Optional[str],
+ expected: dict[str, Any],
+) -> None:
+ """
+ The selected schema is applied through the ``database`` connect argument
and
+ the URI is left untouched, since its database is the default
``authSource``.
+ """
+ from sqlalchemy.engine.url import make_url
+
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+
+ uri = "mongodb://user:pass@host:27017/dbone?mode=superset"
+ original = dict(connect_args)
+
+ adjusted, new_connect_args = MongoDBEngineSpec.adjust_engine_params(
+ make_url(uri), connect_args, schema=schema
+ )
+
+ assert adjusted.render_as_string(hide_password=False) == uri
+ assert new_connect_args == expected
+ assert connect_args == original
+
+
[email protected](
+ "uri,connect_args,expected_auth_source",
+ [
+ ("mongodb://user:pass@host:27017/dbone?mode=superset", {}, "dbone"),
+ (
+
"mongodb://user:pass@host:27017/dbone?mode=superset&authSource=admin",
+ {},
+ "admin",
+ ),
+ (
+ "mongodb://user:pass@host:27017/dbone?mode=superset",
+ {"database": "dbone", "authSource": "dbone"},
+ "dbone",
+ ),
+ ],
+)
+def test_adjust_engine_params_driver(
+ uri: str,
+ connect_args: dict[str, Any],
+ expected_auth_source: str,
+) -> None:
+ """
+ With the real driver, selecting ``dbtwo`` queries ``dbtwo`` while keeping
the
+ original authentication database.
+ """
+ pytest.importorskip("pymongosql")
+
+ from sqlalchemy import create_engine
+ from sqlalchemy.engine.url import make_url
+
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+
+ adjusted, new_connect_args = MongoDBEngineSpec.adjust_engine_params(
+ make_url(uri), connect_args, schema="dbtwo"
+ )
+ engine = create_engine(
+ adjusted, connect_args={**new_connect_args, "connect": False}
+ )
+ connection = engine.raw_connection().driver_connection
+ try:
+ credentials = connection.client.options.pool_options._credentials
+ assert connection.database_name == "dbtwo"
+ assert credentials.source == expected_auth_source
+ finally:
+ connection.close()
+
+
+def test_get_schema_from_engine_params() -> None:
+ from sqlalchemy.engine.url import make_url
+
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+
+ assert (
+ MongoDBEngineSpec.get_schema_from_engine_params(
+ make_url("mongodb://user:pass@host:27017/dbone?mode=superset"), {}
+ )
+ == "dbone"
+ )
+ assert (
+ MongoDBEngineSpec.get_schema_from_engine_params(
+ make_url("mongodb://user:pass@host:27017"), {}
+ )
+ is None
+ )
+ assert (
+ MongoDBEngineSpec.get_schema_from_engine_params(
+ make_url("mongodb://user:pass@host:27017/dbone?mode=superset"),
+ {"database": "dbtwo"},
+ )
+ == "dbtwo"
+ )
+
+
+def test_get_default_schema() -> None:
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+ from superset.models.core import Database
+
+ database = Database(
+ database_name="mongo",
+ sqlalchemy_uri="mongodb://user:pass@host:27017/dbone?mode=superset",
+ )
+
+ assert MongoDBEngineSpec.get_default_schema(database, None) == "dbone"
+
+ database.extra = json.dumps(
+ {"engine_params": {"connect_args": {"database": "dbtwo"}}}
+ )
+ assert MongoDBEngineSpec.get_default_schema(database, None) == "dbtwo"
+
+
+def test_get_default_schema_for_query(mocker: MockerFixture) -> None:
+ """
+ Access checks must resolve unqualified collections against the query
schema,
+ which is the database the connection is bound to.
+ """
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+ from superset.models.core import Database
+
+ database = Database(
+ database_name="mongo",
+ sqlalchemy_uri="mongodb://user:pass@host:27017/dbone?mode=superset",
+ )
+ query = mocker.MagicMock(schema="dbtwo", catalog=None)
+
+ assert MongoDBEngineSpec.get_default_schema_for_query(database, query) ==
"dbtwo"
+
+
+def test_select_star_does_not_qualify_collection(mocker: MockerFixture) ->
None:
+ """
+ PyMongoSQL treats ``schema.collection`` as a literal collection name, so
the
+ preview query must reference the bare collection and rely on the schema
+ being applied to the connection instead.
+ """
+ from sqlalchemy.engine.default import DefaultDialect
+
+ from superset.db_engine_specs.mongodb import MongoDBEngineSpec
+ from superset.sql.parse import Table
+
+ database = mocker.MagicMock()
+ database.compile_sqla_query.side_effect = lambda qry, catalog, schema: str(
+ qry.compile(dialect=DefaultDialect(), compile_kwargs={"literal_binds":
True})
+ )
+
+ sql = MongoDBEngineSpec.select_star(
+ database,
+ Table("orders", "testdb"),
+ DefaultDialect(),
+ limit=10,
+ show_cols=False,
+ latest_partition=False,
+ )
+
+ assert sql == "SELECT\n *\nFROM orders\nLIMIT 10"
+ database.compile_sqla_query.assert_called_once()
+ assert database.compile_sqla_query.call_args.args[1:] == (None, "testdb")
+
+
+def test_get_sqla_engine_applies_selected_schema() -> None:
+ """
+ The ``database`` connect argument returned by ``adjust_engine_params`` must
+ reach the driver when the engine is built through ``Database``.
+ """
+ pytest.importorskip("pymongosql")
+
+ from superset.models.core import Database
+
+ database = Database(
+ database_name="mongo",
+ sqlalchemy_uri="mongodb://user:pass@host:27017/dbone?mode=superset",
+ extra=json.dumps({"engine_params": {"connect_args": {"connect":
False}}}),
+ )
+
+ with database.get_sqla_engine(schema="dbtwo") as engine:
+ raw_connection = engine.raw_connection()
+ try:
+ connection = raw_connection.driver_connection
+ credentials = connection.client.options.pool_options._credentials
+ assert connection.database_name == "dbtwo"
+ assert credentials.source == "dbone"
+ finally:
+ raw_connection.close()
