superset2/tests/unit_tests/extensions/test_sqlalchemy.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-outer-name, import-outside-toplevel, unused-argument

import os
import re
from collections.abc import Iterator
from typing import TYPE_CHECKING

import pytest
from flask import current_app
from pytest_mock import MockerFixture
from sqlalchemy import text
from sqlalchemy.engine import create_engine
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm.session import Session

from superset import db
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetSecurityException
from tests.conftest import with_config
from tests.unit_tests.conftest import with_feature_flags

if TYPE_CHECKING:
    from superset.models.core import Database


def _normalize_sqla_doc_link(message: str) -> str:
    """
    Replace the SQLAlchemy error-doc URL's version segment (e.g. ``/e/20/``)
    with a placeholder so assertions don't need updating every time
    SQLAlchemy bumps its minor version.
    """
    return re.sub(r"/e/\d+/", "/e/XX/", message)


@pytest.fixture
def database1(session: Session) -> Iterator["Database"]:
    from superset.models.core import Database

    engine = db.session.connection().engine
    Database.metadata.create_all(engine)  # pylint: disable=no-member

    database = Database(
        database_name="database1",
        sqlalchemy_uri="sqlite:///database1.db",
        allow_dml=True,
    )
    db.session.add(database)
    db.session.commit()

    yield database

    db.session.delete(database)
    db.session.commit()
    if os.path.exists("database1.db"):
        os.unlink("database1.db")


@pytest.fixture
def table1(session: Session, database1: "Database") -> Iterator[None]:
    with database1.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text("CREATE TABLE table1 (a INTEGER NOT NULL PRIMARY KEY, b INTEGER)")
            )
            conn.execute(text("INSERT INTO table1 (a, b) VALUES (1, 10), (2, 20)"))
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table1"))
        db.session.commit()


@pytest.fixture
def database2(session: Session) -> Iterator["Database"]:
    from superset.models.core import Database

    database = Database(
        database_name="database2",
        sqlalchemy_uri="sqlite:///database2.db",
        allow_dml=False,
    )
    db.session.add(database)
    db.session.commit()

    yield database

    db.session.delete(database)
    db.session.commit()
    if os.path.exists("database2.db"):
        os.unlink("database2.db")


@pytest.fixture
def table2(session: Session, database2: "Database") -> Iterator[None]:
    with database2.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text("CREATE TABLE table2 (a INTEGER NOT NULL PRIMARY KEY, b TEXT)")
            )
            conn.execute(
                text("INSERT INTO table2 (a, b) VALUES (1, 'ten'), (2, 'twenty')")
            )
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table2"))
        db.session.commit()


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset(mocker: MockerFixture, app_context: None, table1: None) -> None:
    """
    Simple test querying a table.
    """
    # Mock the security_manager.raise_for_access to allow access
    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    # Mock Flask g.user for security checks
    # In Python 3.8+, we can't directly patch flask.g
    # Instead, we need to ensure g.user exists in the context
    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://")
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10), (2, 20)]


@with_config(
    {
        "DB_SQLA_URI_VALIDATOR": None,
        "SUPERSET_META_DB_LIMIT": 1,
        "DATABASE_OAUTH2_CLIENTS": {},
        "SQLALCHEMY_CUSTOM_PASSWORD_STORE": None,
    }
)
@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_limit(mocker: MockerFixture, app_context: None, table1: None) -> None:
    """
    Simple that limit is applied when querying a table.
    """
    # Note: We don't patch flask.current_app.config directly anymore
    # The @with_config decorator handles the config patching

    # Mock the security_manager.raise_for_access to allow access
    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    # Mock Flask g.user for security checks
    # In Python 3.8+, we can't directly patch flask.g
    # Instead, we need to ensure g.user exists in the context
    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://")
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10)]


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_joins(
    mocker: MockerFixture,
    app_context: None,
    table1: None,
    table2: None,
) -> None:
    """
    A test joining across databases.
    """
    # Mock the security_manager.raise_for_access to allow access
    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    # Mock Flask g.user for security checks
    # In Python 3.8+, we can't directly patch flask.g
    # Instead, we need to ensure g.user exists in the context
    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://")
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(
            text("""
            SELECT t1.b, t2.b
            FROM "database1.table1" AS t1
            JOIN "database2.table2" AS t2
            ON t1.a = t2.a
            """)
        )
        assert list(results) == [(10, "ten"), (20, "twenty")]


@pytest.mark.parametrize(
    ("statement", "expected"),
    [
        # A single table reference is not a multi-table statement...
        ('SELECT * FROM "database1.table1"', 1),
        # ...even when it has a dotted, double-quoted column alias, which a
        # naive `"[^"]*\.[^"]*"`-shaped regex would also match, misidentifying
        # a single-table statement as multi-table and silently skipping
        # SUPERSET_META_DB_LIMIT for it.
        (
            'SELECT COUNT(id) AS "metric.value" FROM "database1.table1"',
            1,
        ),
        (
            'SELECT t1.b, t2.b FROM "database1.table1" AS t1 '
            'JOIN "database2.table2" AS t2 ON t1.a = t2.a',
            2,
        ),
        (
            'SELECT * FROM "database1.table1", "database2.table2" WHERE t1.a = t2.a',
            2,
        ),
        # Statements the parser can't handle fall back to the safe default
        # (treat as single-table, so the app-wide limit still applies).
        ("this is not valid sql (((", 1),
    ],
)
def test_count_referenced_tables(statement: str, expected: int) -> None:
    """
    Regression for a review comment on #42598/#36304: the multi-table
    detection used to gate SUPERSET_META_DB_LIMIT must count actual table
    references via the real SQL parser, not pattern-match dotted quoted
    identifiers, which also matches dotted column aliases.
    """
    from superset.extensions.metadb import _count_referenced_tables

    assert _count_referenced_tables(statement) == expected


@pytest.fixture
def table1_large(session: Session, database1: "Database") -> Iterator[None]:
    with database1.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE table1_large (a INTEGER NOT NULL PRIMARY KEY, "
                    "b INTEGER)"
                )
            )
            conn.execute(
                text("INSERT INTO table1_large (a, b) VALUES (1, 10), (2, 20), (3, 30)")
            )
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table1_large"))
        db.session.commit()


@pytest.fixture
def table2_late_match(session: Session, database2: "Database") -> Iterator[None]:
    with database2.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE table2_late_match (a INTEGER NOT NULL PRIMARY KEY, "
                    "b TEXT)"
                )
            )
            conn.execute(
                text("INSERT INTO table2_late_match (a, b) VALUES (3, 'thirty')")
            )
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table2_late_match"))
        db.session.commit()


@pytest.fixture
def table2_multi_late_match(session: Session, database2: "Database") -> Iterator[None]:
    with database2.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE table2_multi_late_match "
                    "(a INTEGER NOT NULL PRIMARY KEY, b TEXT)"
                )
            )
            conn.execute(
                text(
                    "INSERT INTO table2_multi_late_match (a, b) "
                    "VALUES (2, 'twenty'), (3, 'thirty')"
                )
            )
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table2_multi_late_match"))
        db.session.commit()


@pytest.fixture
def table2_fanout_match(session: Session, database2: "Database") -> Iterator[None]:
    with database2.get_sqla_engine() as engine:
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE table2_fanout_match "
                    "(id INTEGER NOT NULL PRIMARY KEY, a INTEGER, b TEXT)"
                )
            )
            # `a` is deliberately not unique (unlike table2_late_match, where
            # it's the primary key): a single outer row matching on `a=3`
            # fans out into two inner rows here, so reading the match
            # requires pulling more than one row through the cursor per
            # outer probe, instead of a single unique-index lookup.
            conn.execute(
                text(
                    "INSERT INTO table2_fanout_match (a, b) "
                    "VALUES (3, 'thirty-x'), (3, 'thirty-y')"
                )
            )
        db.session.commit()

        yield

        with engine.begin() as conn:
            conn.execute(text("DROP TABLE table2_fanout_match"))
        db.session.commit()


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_joins_with_limit_drops_fanout_matches(
    mocker: MockerFixture,
    monkeypatch: pytest.MonkeyPatch,
    app_context: None,
    table1_large: None,
    table2_fanout_match: None,
) -> None:
    """
    Coverage note from review of #42598: the other join regression tests
    match on a primary key on both sides, so each inner lookup returns at
    most one row. Here `table1_large`'s single match (a=3) fans out into two
    rows in `table2_fanout_match`, so satisfying it means pulling more than
    one row through the cursor for a single outer probe, rather than a single
    unique-index lookup.
    """
    monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
    monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
    monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
    monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)

    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", future=True)
    except Exception as e:
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(
            text("""
            SELECT t1.b, t2.b
            FROM "database1.table1_large" AS t1
            JOIN "database2.table2_fanout_match" AS t2
            ON t1.a = t2.a
            ORDER BY t2.b
            """)
        )
        assert list(results) == [(30, "thirty-x"), (30, "thirty-y")]


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_joins_with_limit_multiple_late_matches(
    mocker: MockerFixture,
    monkeypatch: pytest.MonkeyPatch,
    app_context: None,
    table1_large: None,
    table2_multi_late_match: None,
) -> None:
    """
    Diagnostic probe raised in review of #42598: does the per-table
    ``SUPERSET_META_DB_LIMIT`` skip (keyed on the
    ``_executing_multi_table_query`` ContextVar) hold for every row of a
    multi-row join result, or only the first? ``table1_large`` has two
    genuine matches in ``table2_multi_late_match`` (a=2 and a=3), both of
    which fall past SUPERSET_META_DB_LIMIT=2 in table1_large's own row
    order for a naive per-table truncation.
    """
    monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
    monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
    monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
    monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)

    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", future=True)
    except Exception as e:
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(
            text("""
            SELECT t1.b, t2.b
            FROM "database1.table1_large" AS t1
            JOIN "database2.table2_multi_late_match" AS t2
            ON t1.a = t2.a
            ORDER BY t1.b
            """)
        )
        assert list(results) == [(20, "twenty"), (30, "thirty")]


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_joins_with_limit_drops_matches(
    mocker: MockerFixture,
    monkeypatch: pytest.MonkeyPatch,
    app_context: None,
    table1_large: None,
    table2_late_match: None,
) -> None:
    """
    Regression for #36304: SUPERSET_META_DB_LIMIT is applied to each
    underlying table independently, before the in-memory join runs. A row
    that has a genuine match on the other side of the join but falls past
    the per-table limit is silently dropped from the join result, with no
    error or truncation warning.
    """
    # Use monkeypatch (rather than the `@with_config` decorator) so the
    # config overrides are guaranteed to be undone even though this test is
    # expected to fail its assertion until the underlying bug is fixed.
    # `@with_config` only restores the original values after the wrapped
    # test function returns normally, so an assertion failure here would
    # otherwise leak SUPERSET_META_DB_LIMIT=2 into later tests.
    monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
    monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
    monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
    monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)

    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", future=True)
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(
            text("""
            SELECT t1.b, t2.b
            FROM "database1.table1_large" AS t1
            JOIN "database2.table2_late_match" AS t2
            ON t1.a = t2.a
            """)
        )
        # table2_late_match's only row (a=3) has a genuine match in
        # table1_large (a=3, b=30), but SUPERSET_META_DB_LIMIT=2 truncates
        # table1_large to its first two rows (a=1, a=2) before the join
        # runs, so the join comes back empty instead of finding the match.
        assert list(results) == [(30, "thirty")]


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_comma_join_with_limit_drops_matches(
    mocker: MockerFixture,
    monkeypatch: pytest.MonkeyPatch,
    app_context: None,
    table1_large: None,
    table2_late_match: None,
) -> None:
    """
    Regression for #36304: an implicit comma join (``FROM a, b WHERE ...``)
    references two tables just like an explicit ``JOIN``, but doesn't contain
    the literal `JOIN` keyword. Multi-table detection has to catch this shape
    too, or the per-table limit still gets applied and silently drops matches.
    """
    # See test_superset_joins_with_limit_drops_matches for why monkeypatch is
    # used here instead of `@with_config`.
    monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
    monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
    monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
    monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)

    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", future=True)
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(
            text("""
            SELECT t1.b, t2.b
            FROM "database1.table1_large" AS t1, "database2.table2_late_match" AS t2
            WHERE t1.a = t2.a
            """)
        )
        # Same scenario as test_superset_joins_with_limit_drops_matches, but
        # using a comma join instead of the `JOIN` keyword.
        assert list(results) == [(30, "thirty")]


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_superset_joins_via_raw_cursor_drops_matches(
    mocker: MockerFixture,
    monkeypatch: pytest.MonkeyPatch,
    app_context: None,
    table1_large: None,
    table2_late_match: None,
) -> None:
    """
    Regression for #36304: SQL Lab executes statements through a raw DBAPI
    cursor (``engine.raw_connection().cursor()``), not through SQLAlchemy's
    ``Connection.execute()``. That path never reaches
    ``SupersetAPSWDialect.do_execute*``, so a fix keyed only on those hooks
    leaves the per-table ``SUPERSET_META_DB_LIMIT`` skip blind to exactly the
    statements SQL Lab runs, and the same join match SQL Lab users see would
    still be silently dropped even though
    ``test_superset_joins_with_limit_drops_matches`` (which goes through
    ``Connection.execute()``) passes.
    """
    monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
    monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
    monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
    monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)

    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", future=True)
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    raw_connection = engine.raw_connection()
    try:
        cursor = raw_connection.cursor()
        cursor.execute(
            """
            SELECT t1.b, t2.b
            FROM "database1.table1_large" AS t1
            JOIN "database2.table2_late_match" AS t2
            ON t1.a = t2.a
            """
        )
        # Same scenario as test_superset_joins_with_limit_drops_matches, but
        # executed the way SQL Lab actually runs queries: a raw DBAPI cursor
        # obtained from `engine.raw_connection()`, bypassing `do_execute*`.
        assert list(cursor) == [(30, "thirty")]
    finally:
        raw_connection.close()


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_dml(
    mocker: MockerFixture,
    app_context: None,
    table1: None,
    table2: None,
) -> None:
    """
    DML tests.

    Test that we can update/delete data, only if DML is enabled.
    """
    # Mock the security_manager.raise_for_access to allow access
    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    # Mock Flask g.user for security checks
    # In Python 3.8+, we can't directly patch flask.g
    # Instead, we need to ensure g.user exists in the context
    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://")
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.begin() as conn:
        conn.execute(text('INSERT INTO "database1.table1" (a, b) VALUES (3, 30)'))
    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10), (2, 20), (3, 30)]
    with engine.begin() as conn:
        conn.execute(text('UPDATE "database1.table1" SET b=35 WHERE a=3'))
    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10), (2, 20), (3, 35)]
    with engine.begin() as conn:
        conn.execute(text('DELETE FROM "database1.table1" WHERE b>20'))
    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10), (2, 20)]

    with engine.begin() as conn:
        with pytest.raises(ProgrammingError) as excinfo:
            conn.execute(
                text("""INSERT INTO "database2.table2" (a, b) VALUES (3, 'thirty')""")
            )
    assert _normalize_sqla_doc_link(str(excinfo.value).strip()) == (
        "(shillelagh.exceptions.ProgrammingError) DML not enabled in database "
        '"database2"\n[SQL: INSERT INTO "database2.table2" (a, b) '
        "VALUES (3, 'thirty')]\n(Background on this error at: "
        "https://sqlalche.me/e/XX/f405)"
    )


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_security_manager(
    mocker: MockerFixture, app_context: None, table1: None
) -> None:
    """
    Test that we use the security manager to check for permissions.
    """
    # Skip this test if metadb dependencies are not available
    try:
        import superset.extensions.metadb  # noqa: F401
    except ImportError:
        pytest.skip("metadb dependencies not available")

    # Mock Flask g.user first to avoid AttributeError
    # We need to mock the actual g object that's imported by security.manager
    mock_user = mocker.MagicMock()
    mock_user.is_anonymous = False
    mocker.patch("superset.security.manager.g", mocker.MagicMock(user=mock_user))

    # Then patch the security_manager to raise an exception
    security_manager = mocker.MagicMock()
    # Patch it in the metadb module where it's actually used
    mocker.patch(
        "superset.extensions.metadb.security_manager",
        new=security_manager,
    )
    security_manager.raise_for_access.side_effect = SupersetSecurityException(
        SupersetError(
            error_type=SupersetErrorType.TABLE_SECURITY_ACCESS_ERROR,
            message=(
                "You need access to the following tables: `table1`,\n            "
                "`all_database_access` or `all_datasource_access` permission"
            ),
            level=ErrorLevel.ERROR,
        )
    )

    try:
        engine = create_engine("superset://")
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        with pytest.raises(SupersetSecurityException) as excinfo:
            conn.execute(text('SELECT * FROM "database1.table1"'))
    assert str(excinfo.value) == (
        "You need access to the following tables: `table1`,\n            "
        "`all_database_access` or `all_datasource_access` permission"
    )


@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
def test_allowed_dbs(mocker: MockerFixture, app_context: None, table1: None) -> None:
    """
    Test that DBs can be restricted.
    """
    # Mock the security_manager.raise_for_access to allow access
    mocker.patch(
        "superset.extensions.metadb.security_manager.raise_for_access",
        return_value=None,
    )

    # Mock Flask g.user for security checks
    # In Python 3.8+, we can't directly patch flask.g
    # Instead, we need to ensure g.user exists in the context
    from flask import g

    g.user = mocker.MagicMock()
    g.user.is_anonymous = False

    try:
        engine = create_engine("superset://", allowed_dbs=["database1"])
    except Exception as e:
        # Skip test if superset:// dialect can't be loaded (common in Docker)
        pytest.skip(f"Superset dialect not available: {e}")

    with engine.connect() as conn:
        results = conn.execute(text('SELECT * FROM "database1.table1"'))
        assert list(results) == [(1, 10), (2, 20)]

    with engine.connect() as conn:
        with pytest.raises(ProgrammingError) as excinfo:
            conn.execute(text('SELECT * FROM "database2.table2"'))
    assert _normalize_sqla_doc_link(str(excinfo.value)) == (
        """
(shillelagh.exceptions.ProgrammingError) Unsupported table: database2.table2
[SQL: SELECT * FROM "database2.table2"]
(Background on this error at: https://sqlalche.me/e/XX/f405)
        """.strip()
    )