# 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.
"""
Tests for virtual dataset SQL handling with RLS and sqlglot formatting.

Two issues are covered:

1. **Unnecessary sqlglot round-trip** – ``get_from_clause()`` used to call
   ``format()`` even when no RLS rules applied, which could silently rewrite
   dialect-specific SQL (e.g. ``NVL`` → ``COALESCE`` on Redshift).

2. **RLS subquery alias mismatch** – ``RLSAsSubqueryTransformer`` used the
   fully-qualified table name (``"schema.table"``) as the subquery alias,
   which broke column references that used just the table name
   (``table.column``), producing Redshift errors like
   ``column "X" does not exist in virtual_table``.
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from flask import Flask
from sqlalchemy.sql.elements import TextClause

from superset.exceptions import QueryObjectValidationError
from superset.models.helpers import ExploreMixin
from superset.sql.parse import RLSMethod, SQLStatement, Table

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def virtual_datasource() -> MagicMock:
    """
    Create a mock datasource that behaves like a virtual dataset.
    """
    datasource = MagicMock(spec=ExploreMixin)

    # Wire up real methods from ExploreMixin
    datasource.get_from_clause = ExploreMixin.get_from_clause.__get__(datasource)
    datasource.text = lambda sql: TextClause(sql)

    # Mock the database and db_engine_spec
    datasource.db_engine_spec.engine = "redshift"
    datasource.db_engine_spec.get_cte_query.return_value = None
    datasource.db_engine_spec.cte_alias = "__cte"
    datasource.database.get_default_schema.return_value = "public"
    datasource.catalog = None
    datasource.schema = "public"

    return datasource


def _set_virtual_sql(datasource: MagicMock, sql: str) -> None:
    """
    Configure the mock datasource to return the given SQL.
    """
    datasource.get_rendered_sql.return_value = sql


def _get_subquery_sql(datasource: MagicMock) -> str:
    """
    Run get_from_clause and extract the inner SQL from the virtual_table alias.
    """
    from_clause, _ = datasource.get_from_clause(template_processor=None)
    return str(from_clause.element).strip()


# ---------------------------------------------------------------------------
# 1. format() should only run when RLS predicates were actually applied
# ---------------------------------------------------------------------------


class TestVirtualDatasetNoRLS:
    """
    When no RLS predicates apply, the virtual dataset SQL must not be
    round-tripped through sqlglot's formatter.
    """

    @patch("superset.models.helpers.apply_rls", return_value=False)
    def test_sql_preserved_when_no_rls(
        self,
        mock_apply_rls: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        """
        The original SQL should be used verbatim when apply_rls returns False.
        """
        original_sql = "SELECT pen_id, is_green FROM public.pens"
        _set_virtual_sql(virtual_datasource, original_sql)

        inner_sql = _get_subquery_sql(virtual_datasource)
        assert inner_sql == original_sql

    @patch("superset.models.helpers.apply_rls", return_value=False)
    def test_redshift_nvl_preserved_when_no_rls(
        self,
        mock_apply_rls: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        """
        sqlglot rewrites ``NVL(a, b)`` → ``COALESCE(a, b)`` for Redshift.
        Without the fix the inner SQL would be silently rewritten.
        """
        original_sql = "SELECT pen_id, NVL(burn_flag, false) FROM pens"
        _set_virtual_sql(virtual_datasource, original_sql)

        inner_sql = _get_subquery_sql(virtual_datasource)
        assert "NVL" in inner_sql
        assert "COALESCE" not in inner_sql

    @patch("superset.models.helpers.apply_rls", return_value=False)
    def test_redshift_current_timestamp_preserved_when_no_rls(
        self,
        mock_apply_rls: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        """
        sqlglot rewrites ``current_timestamp`` → ``GETDATE()`` for Redshift.
        """
        original_sql = "SELECT pen_id, current_timestamp FROM pens"
        _set_virtual_sql(virtual_datasource, original_sql)

        inner_sql = _get_subquery_sql(virtual_datasource)
        assert "current_timestamp" in inner_sql
        assert "GETDATE" not in inner_sql

    @patch("superset.models.helpers.apply_rls", return_value=False)
    def test_redshift_cast_syntax_preserved_when_no_rls(
        self,
        mock_apply_rls: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        """
        sqlglot rewrites Redshift ``::`` cast syntax to ``CAST(... AS ...)``.
        """
        original_sql = (
            "SELECT pen_name::varchar(256) AS name, is_green FROM public.pens"
        )
        _set_virtual_sql(virtual_datasource, original_sql)

        inner_sql = _get_subquery_sql(virtual_datasource)
        assert "::varchar(256)" in inner_sql


class TestVirtualDatasetWithRLS:
    """
    When RLS predicates are applied, the SQL must be regenerated via
    sqlglot to serialize the AST modifications.
    """

    @patch("superset.models.helpers.apply_rls", return_value=True)
    def test_sql_reformatted_when_rls_applied(
        self,
        mock_apply_rls: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        """
        When apply_rls returns True, the SQL should be regenerated through
        sqlglot's format() to serialize the AST modifications.
        """
        original_sql = "SELECT pen_id, is_green FROM public.pens"
        _set_virtual_sql(virtual_datasource, original_sql)

        inner_sql = _get_subquery_sql(virtual_datasource)

        # After format(), sqlglot pretty-prints (adds newlines, etc.)
        assert inner_sql != original_sql
        # But the column must still be present
        assert "is_green" in inner_sql


# ---------------------------------------------------------------------------
# 2. apply_rls() return value
# ---------------------------------------------------------------------------


class TestApplyRlsReturnValue:
    """
    Test that apply_rls correctly reports whether predicates were applied.
    """

    def test_returns_false_when_no_tables(self, app: Flask) -> None:
        """
        apply_rls should return False when the statement has no tables.
        """
        from superset.utils.rls import apply_rls

        database = MagicMock()
        database.db_engine_spec.get_rls_method.return_value = MagicMock()
        database.get_default_catalog.return_value = None

        statement = MagicMock()
        statement.tables = []

        result = apply_rls(
            database=database,
            catalog=None,
            schema="public",
            parsed_statement=statement,
        )
        assert result is False

    @patch("superset.utils.rls.get_predicates_for_table")
    def test_returns_false_when_predicates_empty(
        self,
        mock_get_predicates: MagicMock,
        app: Flask,
    ) -> None:
        """
        apply_rls should return False when tables exist but have no RLS rules.
        """
        from superset.utils.rls import apply_rls

        mock_get_predicates.return_value = []

        database = MagicMock()
        database.db_engine_spec.get_rls_method.return_value = MagicMock()
        database.get_default_catalog.return_value = None

        mock_table = MagicMock()
        mock_table.qualify.return_value = Table("pens", "public", None)

        statement = MagicMock()
        statement.tables = [mock_table]

        result = apply_rls(
            database=database,
            catalog=None,
            schema="public",
            parsed_statement=statement,
        )
        assert result is False

    @patch("superset.utils.rls.get_predicates_for_table")
    def test_returns_true_when_predicates_exist(
        self,
        mock_get_predicates: MagicMock,
        app: Flask,
    ) -> None:
        """
        apply_rls should return True when RLS predicates are found.
        """
        from superset.utils.rls import apply_rls

        mock_get_predicates.return_value = ["user_id = 42"]

        database = MagicMock()
        database.db_engine_spec.get_rls_method.return_value = MagicMock()
        database.get_default_catalog.return_value = None

        mock_table = MagicMock()
        mock_table.qualify.return_value = Table("pens", "public", None)

        statement = MagicMock()
        statement.tables = [mock_table]
        statement.parse_predicate.return_value = MagicMock()

        result = apply_rls(
            database=database,
            catalog=None,
            schema="public",
            parsed_statement=statement,
        )
        assert result is True
        statement.apply_rls.assert_called_once()


# ---------------------------------------------------------------------------
# 3. RLS subquery alias must use just the table name, not schema-qualified
# ---------------------------------------------------------------------------


class TestRLSSubqueryAlias:
    """
    When RLSAsSubqueryTransformer replaces a table with a filtered subquery,
    the subquery alias must match what column references in the query expect.

    If the SQL says ``SELECT pens.col FROM public.pens``, the subquery
    must be aliased as ``"pens"`` (not ``"public.pens"``), otherwise
    ``pens.col`` won't resolve and Redshift returns::

        column "col" does not exist in virtual_table
    """

    def test_table_qualified_columns_no_alias(self, app: Flask) -> None:
        """
        Column references like ``pens.col`` must resolve after RLS wraps
        ``public.pens`` in a subquery.
        """
        sql = "SELECT pens.pen_id, pens.is_green FROM public.pens"
        statement = SQLStatement(sql, engine="redshift")
        predicate = statement.parse_predicate("user_id = 1")
        statement.apply_rls(
            None,
            "public",
            {Table("pens", "public", None): [predicate]},
            RLSMethod.AS_SUBQUERY,
        )
        result = statement.format()

        # The subquery alias must be just "pens", not "public.pens"
        assert 'AS "pens"' in result
        assert 'AS "public.pens"' not in result
        assert "pens.is_green" in result

    def test_catalog_schema_qualified_table_no_alias(self, app: Flask) -> None:
        """
        Even with a catalog-qualified table, the subquery alias should be
        just the table name so that ``table.col`` references still work.
        """
        sql = "SELECT pens.pen_id, pens.is_green FROM mycat.public.pens"
        statement = SQLStatement(sql, engine="redshift")
        predicate = statement.parse_predicate("user_id = 1")
        statement.apply_rls(
            None,
            "public",
            {Table("pens", "public", "mycat"): [predicate]},
            RLSMethod.AS_SUBQUERY,
        )
        result = statement.format()

        assert 'AS "pens"' in result
        assert 'AS "mycat.public.pens"' not in result

    def test_explicit_alias_preserved(self, app: Flask) -> None:
        """
        When the table already has an explicit alias, it should be reused.
        """
        sql = "SELECT p.pen_id, p.is_green FROM public.pens p"
        statement = SQLStatement(sql, engine="redshift")
        predicate = statement.parse_predicate("user_id = 1")
        statement.apply_rls(
            None,
            "public",
            {Table("pens", "public", None): [predicate]},
            RLSMethod.AS_SUBQUERY,
        )
        result = statement.format()

        assert "AS p" in result
        assert "p.is_green" in result

    def test_unqualified_columns_work(self, app: Flask) -> None:
        """
        Unqualified column references should work regardless of the alias.
        """
        sql = "SELECT pen_id, is_green FROM public.pens"
        statement = SQLStatement(sql, engine="redshift")
        predicate = statement.parse_predicate("user_id = 1")
        statement.apply_rls(
            None,
            "public",
            {Table("pens", "public", None): [predicate]},
            RLSMethod.AS_SUBQUERY,
        )
        result = statement.format()

        assert "is_green" in result
        assert "WHERE" in result  # RLS predicate applied


# ---------------------------------------------------------------------------
# 4. RLS injection failures must fail closed when predicates apply
# ---------------------------------------------------------------------------


class TestVirtualDatasetRLSFailClosed:
    """
    When RLS predicates exist for the underlying tables but cannot be
    injected into the virtual dataset SQL, the query must be aborted
    instead of running against the unfiltered inner SQL.
    """

    @patch(
        "superset.models.helpers.get_predicates_for_table",
        return_value=["user_id = 42"],
    )
    @patch(
        "superset.models.helpers.apply_rls",
        side_effect=NotImplementedError("engine cannot apply RLS"),
    )
    def test_raises_when_rls_predicates_cannot_be_applied(
        self,
        mock_apply_rls: MagicMock,
        mock_get_predicates: MagicMock,
        virtual_datasource: MagicMock,
        app: Flask,
    ) -> None:
        _set_virtual_sql(virtual_datasource, "SELECT pen_id FROM public.pens")

        with pytest.raises(QueryObjectValidationError):
            virtual_datasource.get_from_clause(template_processor=None)