superset2/tests/unit_tests/connectors/sqla/models_test.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.

import pandas as pd
import pytest
from pytest_mock import MockerFixture
from sqlalchemy import create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.session import Session

from superset.connectors.sqla.models import (
    SqlaTable,
    TableColumn,
    validate_stored_expression,
)
from superset.daos.dataset import DatasetDAO
from superset.daos.exceptions import DatasourceNotFound
from superset.exceptions import (
    OAuth2RedirectError,
    SupersetDisallowedSQLFunctionException,
    SupersetDisallowedSQLTableException,
    SupersetSecurityException,
)
from superset.models.core import Database
from superset.sql.parse import Table
from superset.superset_typing import QueryObjectDict


def test_query_bubbles_errors(mocker: MockerFixture) -> None:
    """
    Test that the `query` method bubbles exceptions correctly.

    When a user needs to authenticate via OAuth2 to access data, a custom exception is
    raised. The exception needs to bubble up all the way to the frontend as a SIP-40
    compliant payload with the error type `DATABASE_OAUTH2_REDIRECT_URI` so that the
    frontend can initiate the OAuth2 authentication.

    This tests verifies that the method does not capture these exceptions; otherwise the
    user will be never be prompted to authenticate via OAuth2.
    """
    database = mocker.MagicMock()
    database.get_df.side_effect = OAuth2RedirectError(
        url="http://example.com",
        tab_id="1234",
        redirect_uri="http://redirect.example.com",
    )

    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
    )
    mocker.patch.object(
        sqla_table,
        "get_query_str_extended",
        return_value=mocker.MagicMock(sql="SELECT * FROM my_sqla_table"),
    )
    query_obj: QueryObjectDict = {
        "granularity": None,
        "from_dttm": None,
        "to_dttm": None,
        "groupby": ["id", "username", "email"],
        "metrics": [],
        "is_timeseries": False,
        "filter": [],
    }
    with pytest.raises(OAuth2RedirectError):
        sqla_table.query(query_obj)


def _query_obj() -> QueryObjectDict:
    return {
        "granularity": None,
        "from_dttm": None,
        "to_dttm": None,
        "groupby": ["id"],
        "metrics": [],
        "is_timeseries": False,
        "filter": [],
    }


def _build_sqla_table_for_query(
    mocker: MockerFixture, sql: str, engine: str = "postgresql"
) -> SqlaTable:
    db_engine_spec = mocker.MagicMock()
    db_engine_spec.engine = engine
    database = mocker.MagicMock()
    database.db_engine_spec = db_engine_spec
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
    )
    mocker.patch.object(
        SqlaTable,
        "db_engine_spec",
        new=property(lambda self: db_engine_spec),
    )
    mocker.patch.object(
        sqla_table,
        "get_query_str_extended",
        return_value=mocker.MagicMock(sql=sql, labels_expected=[]),
    )
    return sqla_table


def test_query_blocks_disallowed_function_on_chart_data_path(
    mocker: MockerFixture,
) -> None:
    mocker.patch.dict(
        "flask.current_app.config",
        {
            "DISALLOWED_SQL_FUNCTIONS": {"postgresql": {"version"}},
            "DISALLOWED_SQL_TABLES": {},
        },
        clear=False,
    )
    sqla_table = _build_sqla_table_for_query(mocker, "SELECT version()")
    with pytest.raises(SupersetDisallowedSQLFunctionException):
        sqla_table.query(_query_obj())
    sqla_table.database.get_df.assert_not_called()  # type: ignore[attr-defined]


def test_query_blocks_disallowed_table_on_chart_data_path(
    mocker: MockerFixture,
) -> None:
    mocker.patch.dict(
        "flask.current_app.config",
        {
            "DISALLOWED_SQL_FUNCTIONS": {},
            "DISALLOWED_SQL_TABLES": {"postgresql": {"pg_authid"}},
        },
        clear=False,
    )
    sqla_table = _build_sqla_table_for_query(mocker, "SELECT rolname FROM pg_authid")
    with pytest.raises(SupersetDisallowedSQLTableException):
        sqla_table.query(_query_obj())
    sqla_table.database.get_df.assert_not_called()  # type: ignore[attr-defined]


def test_query_disallowed_table_error_reports_only_matched_tables(
    mocker: MockerFixture,
) -> None:
    mocker.patch.dict(
        "flask.current_app.config",
        {
            "DISALLOWED_SQL_FUNCTIONS": {},
            "DISALLOWED_SQL_TABLES": {
                "postgresql": {"pg_authid", "pg_shadow", "pg_stat_activity"}
            },
        },
        clear=False,
    )
    sqla_table = _build_sqla_table_for_query(mocker, "SELECT rolname FROM pg_authid")
    with pytest.raises(SupersetDisallowedSQLTableException) as excinfo:
        sqla_table.query(_query_obj())
    message = str(excinfo.value)
    assert "pg_authid" in message
    assert "pg_shadow" not in message
    assert "pg_stat_activity" not in message


def test_query_allows_benign_sql_on_chart_data_path(mocker: MockerFixture) -> None:
    mocker.patch.dict(
        "flask.current_app.config",
        {
            "DISALLOWED_SQL_FUNCTIONS": {"postgresql": {"version"}},
            "DISALLOWED_SQL_TABLES": {"postgresql": {"pg_authid"}},
        },
        clear=False,
    )
    sqla_table = _build_sqla_table_for_query(mocker, "SELECT id FROM my_sqla_table")
    sqla_table.database.get_df.return_value = pd.DataFrame()  # type: ignore[attr-defined]
    result = sqla_table.query(_query_obj())
    sqla_table.database.get_df.assert_called_once()  # type: ignore[attr-defined]
    assert result is not None


def test_permissions_without_catalog() -> None:
    """
    Test permissions when the table has no catalog.
    """
    database = Database(database_name="my_db")
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
        schema="schema1",
        catalog=None,
        id=1,
    )

    assert sqla_table.get_perm() == "[my_db].[my_sqla_table](id:1)"
    assert sqla_table.get_catalog_perm() is None
    assert sqla_table.get_schema_perm() == "[my_db].[schema1]"


def test_permissions_with_catalog() -> None:
    """
    Test permissions when the table with a catalog set.
    """
    database = Database(database_name="my_db")
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
        schema="schema1",
        catalog="db1",
        id=1,
    )

    assert sqla_table.get_perm() == "[my_db].[my_sqla_table](id:1)"
    assert sqla_table.get_catalog_perm() == "[my_db].[db1]"
    assert sqla_table.get_schema_perm() == "[my_db].[db1].[schema1]"


def test_query_datasources_by_name(mocker: MockerFixture) -> None:
    """
    Test the `query_datasources_by_name` method.
    """
    db = mocker.patch("superset.connectors.sqla.models.db")

    database = Database(database_name="my_db", id=1)
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
    )

    sqla_table.query_datasources_by_name(database, "my_table")
    db.session.query().filter_by.assert_called_with(
        database_id=1,
        table_name="my_table",
    )

    sqla_table.query_datasources_by_name(database, "my_table", "db1", "schema1")
    db.session.query().filter_by.assert_called_with(
        database_id=1,
        table_name="my_table",
        catalog="db1",
        schema="schema1",
    )


def test_query_datasources_by_permissions(mocker: MockerFixture) -> None:
    """
    Test the `query_datasources_by_permissions` method.
    """
    db = mocker.patch("superset.connectors.sqla.models.db")

    engine = create_engine("sqlite://")
    database = Database(database_name="my_db", id=1)
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
    )

    sqla_table.query_datasources_by_permissions(database, set(), set(), set())
    db.session.query().filter_by.assert_called_with(database_id=1)
    clause = db.session.query().filter_by().filter.mock_calls[0].args[0]
    assert str(clause.compile(engine, compile_kwargs={"literal_binds": True})) == ""


def test_query_datasources_by_permissions_with_catalog_schema(
    mocker: MockerFixture,
) -> None:
    """
    Test the `query_datasources_by_permissions` method passing a catalog and schema.
    """
    db = mocker.patch("superset.connectors.sqla.models.db")

    engine = create_engine("sqlite://")
    database = Database(database_name="my_db", id=1)
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=database,
    )
    sqla_table.query_datasources_by_permissions(
        database,
        {"[my_db].[table1](id:1)"},
        {"[my_db].[db1]"},
        # pass as list to have deterministic order for test
        ["[my_db].[db1].[schema1]", "[my_other_db].[schema]"],  # type: ignore
    )
    clause = db.session.query().filter_by().filter.mock_calls[0].args[0]
    assert str(clause.compile(engine, compile_kwargs={"literal_binds": True})) == (
        "tables.perm IN ('[my_db].[table1](id:1)') OR "
        "tables.schema_perm IN ('[my_db].[db1].[schema1]', '[my_other_db].[schema]') OR "  # noqa: E501
        "tables.catalog_perm IN ('[my_db].[db1]')"
    )


def test_dataset_uniqueness(session: Session) -> None:
    """
    Test dataset uniqueness constraints.
    """
    Database.metadata.create_all(session.bind)

    database = Database(database_name="my_db", sqlalchemy_uri="sqlite://")

    # add prod.schema.table
    dataset = SqlaTable(
        database=database,
        catalog="prod",
        schema="schema",
        table_name="table",
    )
    session.add(dataset)
    session.commit()

    # add dev.schema.table
    dataset = SqlaTable(
        database=database,
        catalog="dev",
        schema="schema",
        table_name="table",
    )
    session.add(dataset)
    session.commit()

    # try to add dev.schema.table again, fails
    dataset = SqlaTable(
        database=database,
        catalog="dev",
        schema="schema",
        table_name="table",
    )
    session.add(dataset)
    with pytest.raises(IntegrityError):
        session.commit()
    session.rollback()

    # add schema.table
    dataset = SqlaTable(
        database=database,
        catalog=None,
        schema="schema",
        table_name="table",
    )
    session.add(dataset)
    session.commit()

    # add schema.table again, works because in SQL `NULlL != NULL`
    dataset = SqlaTable(
        database=database,
        catalog=None,
        schema="schema",
        table_name="table",
    )
    session.add(dataset)
    session.commit()

    # but the DAO enforces application logic for uniqueness
    assert not DatasetDAO.validate_uniqueness(
        database,
        Table("table", "schema", None),
    )

    assert DatasetDAO.validate_uniqueness(
        database,
        Table("table", "schema", "some_catalog"),
    )


def test_normalize_prequery_result_type_custom_sql() -> None:
    """
    Test that the `_normalize_prequery_result_type` can handle custom SQL.
    """
    sqla_table = SqlaTable(
        table_name="my_sqla_table",
        columns=[],
        metrics=[],
        database=Database(database_name="my_db", sqlalchemy_uri="sqlite://"),
    )
    row: pd.Series = {
        "custom_sql": "Car",
    }
    dimension: str = "custom_sql"
    columns_by_name: dict[str, TableColumn] = {
        "product_line": TableColumn(column_name="product_line"),
    }
    assert (
        sqla_table._normalize_prequery_result_type(row, dimension, columns_by_name)
        == "Car"
    )


def test_fetch_metadata_with_comment_field_new_columns(mocker: MockerFixture) -> None:
    """Test that fetch_metadata correctly assigns comment field to description
    for new columns
    """
    # Mock database
    database = mocker.MagicMock()
    database.get_metrics.return_value = []

    # Mock db_engine_spec
    mock_db_engine_spec = mocker.MagicMock()
    mock_db_engine_spec.alter_new_orm_column = mocker.MagicMock()
    database.db_engine_spec = mock_db_engine_spec

    # Create table
    table = SqlaTable(
        table_name="test_table",
        database=database,
    )

    # Mock external_metadata to return columns with comment fields
    mock_columns = [
        {
            "column_name": "id",
            "type": "INTEGER",
            "comment": "Primary key identifier",
        },
        {
            "column_name": "name",
            "type": "VARCHAR",
            "comment": "Full name of the user",
        },
        {
            "column_name": "status",
            "type": "VARCHAR",
            # No comment field for this column
        },
    ]

    # Mock dependencies
    mocker.patch.object(table, "external_metadata", return_value=mock_columns)
    mocker.patch("superset.connectors.sqla.models.db.session")
    mocker.patch(
        "superset.connectors.sqla.models.config", {"SQLA_TABLE_MUTATOR": lambda x: None}
    )

    # Execute fetch_metadata
    result = table.fetch_metadata()

    # Verify results
    assert len(result.added) == 3
    assert set(result.added) == {"id", "name", "status"}

    # Check that descriptions were set correctly from comments
    columns_by_name = {col.column_name: col for col in table.columns}

    assert columns_by_name["id"].description == "Primary key identifier"
    assert columns_by_name["name"].description == "Full name of the user"
    # Column without comment should have None description
    assert columns_by_name["status"].description is None


def test_fetch_metadata_with_comment_field_existing_columns(
    mocker: MockerFixture,
) -> None:
    """Test that fetch_metadata correctly updates description for existing columns"""
    # Mock database
    database = mocker.MagicMock()
    database.get_metrics.return_value = []

    # Mock db_engine_spec
    mock_db_engine_spec = mocker.MagicMock()
    mock_db_engine_spec.alter_new_orm_column = mocker.MagicMock()
    database.db_engine_spec = mock_db_engine_spec

    # Create table with existing columns
    table = SqlaTable(
        table_name="test_table_existing",
        database=database,
    )
    table.id = 1  # Set ID so it's treated as existing table

    # Create existing columns
    existing_col1 = TableColumn(
        column_name="id",
        type="INTEGER",
        table=table,
        description="Old description",
    )
    existing_col2 = TableColumn(
        column_name="name",
        type="VARCHAR",
        table=table,
    )
    table.columns = [existing_col1, existing_col2]

    # Mock external_metadata to return updated columns with comments
    mock_columns = [
        {
            "column_name": "id",
            "type": "INTEGER",
            "comment": "Updated primary key description",
        },
        {
            "column_name": "name",
            "type": "VARCHAR",
            "comment": "Updated name description",
        },
    ]

    # Mock dependencies
    mock_session = mocker.patch("superset.connectors.sqla.models.db.session")
    mock_session.query.return_value.filter.return_value.all.return_value = [
        existing_col1,
        existing_col2,
    ]
    mocker.patch.object(table, "external_metadata", return_value=mock_columns)
    mocker.patch(
        "superset.connectors.sqla.models.config", {"SQLA_TABLE_MUTATOR": lambda x: None}
    )

    # Execute fetch_metadata
    result = table.fetch_metadata()

    # Verify no new columns were added
    assert len(result.added) == 0

    # Check that descriptions were updated from comments
    columns_by_name = {col.column_name: col for col in table.columns}

    assert columns_by_name["id"].description == "Updated primary key description"
    assert columns_by_name["name"].description == "Updated name description"


def test_fetch_metadata_mixed_comment_scenarios(mocker: MockerFixture) -> None:
    """Test fetch_metadata with mix of new/existing columns and with/without
    comments
    """
    # Mock database
    database = mocker.MagicMock()
    database.get_metrics.return_value = []

    # Mock db_engine_spec
    mock_db_engine_spec = mocker.MagicMock()
    mock_db_engine_spec.alter_new_orm_column = mocker.MagicMock()
    database.db_engine_spec = mock_db_engine_spec

    # Create table with one existing column
    table = SqlaTable(
        table_name="test_table_mixed",
        database=database,
    )
    table.id = 1

    existing_col = TableColumn(
        column_name="existing_col",
        type="INTEGER",
        table=table,
        description="Existing description",
    )
    table.columns = [existing_col]

    # Mock external_metadata with mixed scenarios
    mock_columns = [
        {
            "column_name": "existing_col",
            "type": "INTEGER",
            "comment": "Updated existing column comment",
        },
        {
            "column_name": "new_with_comment",
            "type": "VARCHAR",
            "comment": "New column with comment",
        },
        {
            "column_name": "new_without_comment",
            "type": "VARCHAR",
            # No comment field
        },
    ]

    # Mock dependencies
    mock_session = mocker.patch("superset.connectors.sqla.models.db.session")
    mock_session.query.return_value.filter.return_value.all.return_value = [
        existing_col
    ]
    mocker.patch.object(table, "external_metadata", return_value=mock_columns)
    mocker.patch(
        "superset.connectors.sqla.models.config", {"SQLA_TABLE_MUTATOR": lambda x: None}
    )

    # Execute fetch_metadata
    result = table.fetch_metadata()

    # Check added columns
    assert len(result.added) == 2
    assert set(result.added) == {"new_with_comment", "new_without_comment"}

    # Check all column descriptions
    columns_by_name = {col.column_name: col for col in table.columns}

    # Existing column should have updated description
    assert (
        columns_by_name["existing_col"].description == "Updated existing column comment"
    )

    # New column with comment should have description set
    assert columns_by_name["new_with_comment"].description == "New column with comment"

    # New column without comment should have None description
    assert columns_by_name["new_without_comment"].description is None


def test_fetch_metadata_no_comment_field_safe_handling(
    mocker: MockerFixture,
) -> None:
    """Test that fetch_metadata safely handles columns with no comment field"""
    # Mock database
    database = mocker.MagicMock()
    database.get_metrics.return_value = []

    # Mock db_engine_spec
    mock_db_engine_spec = mocker.MagicMock()
    mock_db_engine_spec.alter_new_orm_column = mocker.MagicMock()
    database.db_engine_spec = mock_db_engine_spec

    # Create table
    table = SqlaTable(
        table_name="test_table_no_comments",
        database=database,
    )

    # Mock external_metadata with columns that have no comment fields
    mock_columns = [
        {"column_name": "col1", "type": "INTEGER"},
        {"column_name": "col2", "type": "VARCHAR"},
    ]

    # Mock dependencies
    mocker.patch.object(table, "external_metadata", return_value=mock_columns)
    mocker.patch("superset.connectors.sqla.models.db.session")
    mocker.patch(
        "superset.connectors.sqla.models.config", {"SQLA_TABLE_MUTATOR": lambda x: None}
    )

    # Execute fetch_metadata - should not raise any exceptions
    result = table.fetch_metadata()

    # Check that columns were added successfully
    assert len(result.added) == 2
    assert set(result.added) == {"col1", "col2"}

    # Check that descriptions are None (not set)
    columns_by_name = {col.column_name: col for col in table.columns}
    assert columns_by_name["col1"].description is None
    assert columns_by_name["col2"].description is None


def test_fetch_metadata_empty_comment_field_handling(mocker: MockerFixture) -> None:
    """Test that fetch_metadata handles empty comment fields correctly"""
    # Mock database
    database = mocker.MagicMock()
    database.get_metrics.return_value = []

    # Mock db_engine_spec
    mock_db_engine_spec = mocker.MagicMock()
    mock_db_engine_spec.alter_new_orm_column = mocker.MagicMock()
    database.db_engine_spec = mock_db_engine_spec

    # Create table
    table = SqlaTable(
        table_name="test_table_empty_comments",
        database=database,
    )

    # Mock external_metadata with empty comment fields
    mock_columns = [
        {
            "column_name": "col_with_empty_comment",
            "type": "INTEGER",
            "comment": "",  # Empty string comment
        },
        {
            "column_name": "col_with_none_comment",
            "type": "VARCHAR",
            "comment": None,  # None comment
        },
        {
            "column_name": "col_with_valid_comment",
            "type": "VARCHAR",
            "comment": "Valid comment",
        },
    ]

    # Mock dependencies
    mocker.patch.object(table, "external_metadata", return_value=mock_columns)
    mocker.patch("superset.connectors.sqla.models.db.session")
    mocker.patch(
        "superset.connectors.sqla.models.config", {"SQLA_TABLE_MUTATOR": lambda x: None}
    )

    # Execute fetch_metadata
    result = table.fetch_metadata()

    # Check that all columns were added
    assert len(result.added) == 3

    columns_by_name = {col.column_name: col for col in table.columns}

    # Empty string comment should not be set (falsy)
    assert columns_by_name["col_with_empty_comment"].description is None

    # None comment should not be set
    assert columns_by_name["col_with_none_comment"].description is None

    # Valid comment should be set
    assert columns_by_name["col_with_valid_comment"].description == "Valid comment"


@pytest.mark.parametrize(
    "supports_cross_catalog,table_name,catalog,schema,expected_name,expected_schema",
    [
        # Database supports cross-catalog queries (like BigQuery)
        (
            True,
            "test_table",
            "test_project",
            "test_dataset",
            '"test_project"."test_dataset"."test_table"',
            None,
        ),
        # Database supports cross-catalog queries, catalog only (no schema)
        (
            True,
            "test_table",
            "test_project",
            None,
            '"test_project"."test_table"',
            None,
        ),
        # Database supports cross-catalog queries, schema only (no catalog)
        (
            True,
            "test_table",
            None,
            "test_schema",
            "test_table",
            "test_schema",
        ),
        # Database supports cross-catalog queries, no catalog or schema
        (
            True,
            "test_table",
            None,
            None,
            "test_table",
            None,
        ),
        # Database doesn't support cross-catalog queries, catalog ignored
        (
            False,
            "test_table",
            "test_catalog",
            "test_schema",
            "test_table",
            "test_schema",
        ),
        # Database doesn't support cross-catalog queries, no schema
        (
            False,
            "test_table",
            "test_catalog",
            None,
            "test_table",
            None,
        ),
    ],
)
def test_get_sqla_table_with_catalog(
    mocker: MockerFixture,
    supports_cross_catalog: bool,
    table_name: str,
    catalog: str | None,
    schema: str | None,
    expected_name: str,
    expected_schema: str | None,
) -> None:
    """
    Test that `get_sqla_table` handles catalog inclusion correctly.
    """
    # Mock database with specified cross-catalog support
    database = mocker.MagicMock()
    database.db_engine_spec.supports_cross_catalog_queries = supports_cross_catalog
    # Provide a simple quote_identifier
    database.quote_identifier = lambda x: f'"{x}"'

    # Create table with specified parameters
    table = SqlaTable(
        table_name=table_name,
        database=database,
        schema=schema,
        catalog=catalog,
    )

    # Get the SQLAlchemy table representation
    sqla_table = table.get_sqla_table()

    # Verify expected table name and schema
    assert sqla_table.name == expected_name
    assert sqla_table.schema == expected_schema


@pytest.mark.parametrize(
    "table_name, catalog, schema, expected_in_sql, not_expected_in_sql",
    [
        (
            "My-Table",
            "My-DB",
            "My-Schema",
            '"My-DB"."My-Schema"."My-Table"',
            '"My-DB.My-Schema.My-Table"',  # Should NOT be one quoted string
        ),
        (
            "ORDERS",
            "PROD_DB",
            "SALES",
            '"PROD_DB"."SALES"."ORDERS"',
            '"PROD_DB.SALES.ORDERS"',  # Should NOT be one quoted string
        ),
        (
            "My Table",
            "My DB",
            "My Schema",
            '"My DB"."My Schema"."My Table"',
            '"My DB.My Schema.My Table"',  # Should NOT be one quoted string
        ),
    ],
)
def test_get_sqla_table_quoting_for_cross_catalog(
    mocker: MockerFixture,
    table_name: str,
    catalog: str | None,
    schema: str | None,
    expected_in_sql: str,
    not_expected_in_sql: str,
) -> None:
    """
    Test that `get_sqla_table` properly quotes each component of the identifier.
    """
    from sqlalchemy import create_engine, select

    # Create a Postgres-like engine to test proper quoting
    engine = create_engine("postgresql://user:pass@host/db")

    # Mock database with cross-catalog support and proper quote_identifier
    database = mocker.MagicMock()
    database.db_engine_spec.supports_cross_catalog_queries = True
    database.quote_identifier = engine.dialect.identifier_preparer.quote

    # Create table
    table = SqlaTable(
        table_name=table_name,
        database=database,
        schema=schema,
        catalog=catalog,
    )

    # Get the SQLAlchemy table representation
    sqla_table = table.get_sqla_table()
    query = select(sqla_table)
    compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True}))

    # The compiled SQL should contain each part quoted separately
    assert expected_in_sql in compiled, f"Expected {expected_in_sql} in SQL: {compiled}"
    # Should NOT have the entire identifier quoted as one string
    assert not_expected_in_sql not in compiled, (
        f"Should not have {not_expected_in_sql} in SQL: {compiled}"
    )


def test_get_sqla_table_without_cross_catalog_ignores_catalog(
    mocker: MockerFixture,
) -> None:
    """
    Test that databases without cross-catalog support ignore the catalog field.
    """
    from sqlalchemy import create_engine, select

    # Create a PostgreSQL engine (doesn't support cross-catalog queries)
    engine = create_engine("postgresql://user:pass@localhost/db")

    # Mock database without cross-catalog support
    database = mocker.MagicMock()
    database.db_engine_spec.supports_cross_catalog_queries = False
    database.quote_identifier = engine.dialect.identifier_preparer.quote

    # Create table with catalog - should be ignored
    table = SqlaTable(
        table_name="my_table",
        database=database,
        schema="my_schema",
        catalog="my_catalog",
    )

    # Get the SQLAlchemy table representation
    sqla_table = table.get_sqla_table()

    # Compile to SQL
    query = select(sqla_table)
    compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True}))

    # Should only have schema.table, not catalog.schema.table
    assert "my_schema" in compiled
    assert "my_table" in compiled
    assert "my_catalog" not in compiled


def test_quoted_name_prevents_double_quoting(mocker: MockerFixture) -> None:
    """
    Test that `quoted_name(..., quote=False)` does not cause double quoting.
    """
    from sqlalchemy import create_engine, select

    engine = create_engine("postgresql://user:pass@host/db")

    # Mock database
    database = mocker.MagicMock()
    database.db_engine_spec.supports_cross_catalog_queries = True
    database.quote_identifier = engine.dialect.identifier_preparer.quote

    # Use uppercase table name to force quoting
    table = SqlaTable(
        table_name="MY_TABLE",
        database=database,
        schema="MY_SCHEMA",
        catalog="MY_DB",
    )

    # Get the SQLAlchemy table representation
    sqla_table = table.get_sqla_table()

    # Compile to SQL
    query = select(sqla_table)
    compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True}))

    # Should NOT have the entire identifier quoted as one:
    # BAD:  '"MY_DB.MY_SCHEMA.MY_TABLE"'
    # This would cause: SQL compilation error: Object '"MY_DB.MY_SCHEMA.MY_TABLE"'
    # does not exist
    assert '"MY_DB.MY_SCHEMA.MY_TABLE"' not in compiled

    # Should have each part quoted separately:
    # GOOD: "MY_DB"."MY_SCHEMA"."MY_TABLE"
    assert '"MY_DB"."MY_SCHEMA"."MY_TABLE"' in compiled


def test_sqla_table_currency_code_column_property() -> None:
    """
    Test currency_code_column property on SqlaTable.
    """
    database = Database(database_name="my_db")
    table = SqlaTable(
        table_name="sales",
        database=database,
        currency_code_column="currency",
    )
    assert table.currency_code_column == "currency"


def test_sqla_table_data_includes_currency_code_column(mocker: MockerFixture) -> None:
    """
    Test that data property includes currency_code_column.
    """
    database = mocker.MagicMock()
    database.get_sqla_engine.return_value.__enter__ = mocker.MagicMock()
    database.get_sqla_engine.return_value.__exit__ = mocker.MagicMock()

    table = SqlaTable(
        table_name="sales",
        database=database,
        currency_code_column="currency_code",
        main_dttm_col="ds",
    )
    table.columns = []
    table.metrics = []

    # Mock the columns property to return empty list
    mocker.patch.object(SqlaTable, "columns", [])
    mocker.patch.object(SqlaTable, "metrics", [])

    data = table.data
    assert data["currency_code_column"] == "currency_code"
    assert data["main_dttm_col"] == "ds"


def test_sqla_table_link_escapes_url(mocker: MockerFixture) -> None:
    """
    Test that link property properly escapes URL to prevent XSS.
    """
    database = Database(database_name="my_db")
    table = SqlaTable(
        table_name='test<script>alert("xss")</script>',
        database=database,
        id=1,
    )

    # Mock explore_url to return a URL with special characters
    mocker.patch.object(
        SqlaTable,
        "explore_url",
        new_callable=mocker.PropertyMock,
        return_value='/explore/?datasource_type=table&datasource_id=1&name=<script>alert("xss")</script>',
    )

    link = table.link
    # Verify that special characters are escaped in both name and URL
    assert "&lt;script&gt;" in str(link)
    assert "<script>" not in str(link)


def test_data_for_slices_handles_missing_datasource(mocker: MockerFixture) -> None:
    """
    Test that data_for_slices gracefully handles a chart whose query_context
    references a datasource that no longer exists.

    When a chart's query_context references a deleted datasource, get_query_context()
    raises DatasourceNotFound. The fix ensures this exception is caught and logged,
    allowing the dashboard to load normally instead of returning a 404.
    """
    database = mocker.MagicMock()
    database.id = 1

    table = SqlaTable(
        table_name="test_table",
        database=database,
        columns=[],
        metrics=[],
    )

    # Create a mock slice whose get_query_context raises DatasourceNotFound
    mock_slice = mocker.MagicMock()
    mock_slice.id = 1
    mock_slice.slice_name = "Test Chart"
    mock_slice.form_data = {}
    mock_slice.get_query_context.side_effect = DatasourceNotFound()

    # Mock the columns and metrics properties to return empty lists
    mocker.patch.object(SqlaTable, "columns", [])
    mocker.patch.object(SqlaTable, "metrics", [])

    # This should not raise an exception - the fix catches DatasourceNotFound
    result = table.data_for_slices([mock_slice])

    # Verify the method returns a valid data structure
    assert "columns" in result
    assert "metrics" in result
    assert "verbose_map" in result


def test_owners_data_includes_email(mocker: MockerFixture) -> None:
    """Test that the owners_data property includes the email field."""
    database = mocker.MagicMock()

    table = SqlaTable(
        table_name="test_table",
        database=database,
        columns=[],
        metrics=[],
    )

    mock_owner = mocker.MagicMock()
    mock_owner.first_name = "John"
    mock_owner.last_name = "Doe"
    mock_owner.username = "johndoe"
    mock_owner.id = 1
    mock_owner.email = "john@example.com"

    table.owners = [mock_owner]

    owners_data = table.owners_data
    assert len(owners_data) == 1
    assert owners_data[0] == {
        "first_name": "John",
        "last_name": "Doe",
        "username": "johndoe",
        "id": 1,
        "email": "john@example.com",
    }


def _database_for_expression(mocker: MockerFixture) -> Database:
    database = mocker.MagicMock(spec=Database)
    database.backend = "sqlite"
    database.allow_multi_catalog = False
    return database


def test_validate_stored_expression_rejects_multi_statement(
    mocker: MockerFixture,
) -> None:
    database = _database_for_expression(mocker)
    with pytest.raises(SupersetSecurityException):
        validate_stored_expression(database, None, None, "1; DROP TABLE users")


def test_validate_stored_expression_rejects_set_operation(
    mocker: MockerFixture,
) -> None:
    database = _database_for_expression(mocker)
    with pytest.raises(SupersetSecurityException):
        validate_stored_expression(
            database, None, None, "1 UNION SELECT password FROM ab_user"
        )


def test_validate_stored_expression_accepts_case_expression(
    mocker: MockerFixture,
) -> None:
    database = _database_for_expression(mocker)
    validate_stored_expression(
        database, None, None, "CASE WHEN amount > 0 THEN 'a' ELSE 'b' END"
    )


def test_validate_stored_expression_rejects_subquery(
    mocker: MockerFixture,
) -> None:
    """
    With ``ALLOW_ADHOC_SUBQUERY=False`` (the default), a stored
    expression that contains a sub-query is rejected by the same
    ``validate_adhoc_subquery`` gate that already covers adhoc SQL.
    Locks in the sub-query branch so a future refactor that
    removes the ``validate_adhoc_subquery`` call gets a red test.
    """
    database = _database_for_expression(mocker)
    mocker.patch("superset.models.helpers.is_feature_enabled", return_value=False)
    with pytest.raises(SupersetSecurityException):
        validate_stored_expression(
            database,
            None,
            None,
            "(SELECT password FROM ab_user LIMIT 1)",
        )


@pytest.mark.parametrize(
    "expression",
    [
        "case when '{{ current_username() }}' = 'abc' then 'yes' else 'no' end",
        "SUM(price) * {{ url_param('multiplier') }}",
        "{# comment #} amount",
        "{% if 1 %}amount{% endif %}",
    ],
)
def test_validate_stored_expression_accepts_jinja(
    mocker: MockerFixture, expression: str
) -> None:
    """
    Stored expressions can contain Jinja templating. Balanced Jinja blocks
    are replaced with a placeholder so the surrounding SQL is still parsed;
    skeletons whose control flow leaves them unparseable defer to runtime.
    """
    database = _database_for_expression(mocker)
    validate_stored_expression(database, None, None, expression)


def test_validate_stored_expression_rejects_set_op_around_jinja(
    mocker: MockerFixture,
) -> None:
    """
    A ``UNION`` smuggled around a Jinja block must still be rejected: the
    Jinja substitution leaves the set operator visible to the parser.
    """
    database = _database_for_expression(mocker)
    with pytest.raises(SupersetSecurityException):
        validate_stored_expression(
            database,
            None,
            None,
            "'{{ current_username() }}' UNION SELECT password FROM ab_user",
        )


def test_validate_stored_expression_rejects_subquery_around_jinja(
    mocker: MockerFixture,
) -> None:
    """
    Sub-queries combined with a Jinja comment block must still be rejected:
    stripping the ``{# ... #}`` block leaves the sub-query visible to the
    ``validate_adhoc_subquery`` gate.
    """
    database = _database_for_expression(mocker)
    mocker.patch("superset.models.helpers.is_feature_enabled", return_value=False)
    with pytest.raises(SupersetSecurityException):
        validate_stored_expression(
            database,
            None,
            None,
            "(SELECT password FROM ab_user LIMIT 1) {# x #}",
        )