# 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.
from __future__ import annotations

import asyncio
import json
from unittest.mock import MagicMock, PropertyMock, patch

import pytest
from pydantic_ai.exceptions import ModelRetry

from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA
from airflow.providers.common.sql.hooks.sql import DbApiHook


def _make_mock_db_hook(
    table_names: list[str] | None = None,
    table_schema: list[dict[str, str]] | None = None,
    records: list[tuple] | None = None,
    last_description: list[tuple] | None = None,
):
    """Create a mock DbApiHook with sensible defaults."""
    mock = MagicMock(spec=DbApiHook)
    mock.inspector = MagicMock()
    mock.inspector.get_table_names.return_value = table_names or ["users", "orders"]
    mock.get_table_schema.return_value = table_schema or [
        {"name": "id", "type": "INTEGER"},
        {"name": "name", "type": "VARCHAR"},
    ]
    mock.get_records.return_value = records or [(1, "Alice"), (2, "Bob")]
    type(mock).last_description = PropertyMock(return_value=last_description or [("id",), ("name",)])
    return mock


class TestSQLToolsetInit:
    def test_id_includes_conn_id(self):
        ts = SQLToolset("my_pg")
        assert ts.id == "sql-my_pg"


class TestSQLToolsetGetTools:
    def test_returns_four_tools(self):
        ts = SQLToolset("pg_default")
        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
        assert set(tools.keys()) == {"list_tables", "get_schema", "query", "check_query"}

    def test_tool_definitions_have_descriptions(self):
        ts = SQLToolset("pg_default")
        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
        for tool in tools.values():
            assert tool.tool_def.description

    @pytest.mark.skipif(
        not _SUPPORTS_RETURN_SCHEMA, reason="pydantic-ai too old for ToolDefinition.return_schema"
    )
    def test_tools_declare_string_return_schema(self):
        # Every tool returns a JSON-encoded string, so code mode should see `-> str`.
        ts = SQLToolset("pg_default")
        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
        for tool in tools.values():
            assert tool.tool_def.return_schema == {"type": "string"}


class TestSQLToolsetListTables:
    def test_returns_all_tables(self):
        ts = SQLToolset("pg_default")
        mock_hook = _make_mock_db_hook(table_names=["users", "orders", "products"])
        ts._hook = mock_hook

        result = asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))
        tables = json.loads(result)
        assert tables == ["users", "orders", "products"]

    def test_filters_by_allowed_tables(self):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        mock_hook = _make_mock_db_hook(table_names=["users", "orders", "products"])
        ts._hook = mock_hook

        result = asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))
        tables = json.loads(result)
        assert tables == ["orders"]

    def test_introspection_error_raises_model_retry(self):
        """A failure while listing tables is returned to the agent as a retry."""
        ts = SQLToolset("pg_default")
        mock_hook = _make_mock_db_hook()
        mock_hook.inspector.get_table_names.side_effect = Exception("could not connect to server")
        ts._hook = mock_hook

        with pytest.raises(ModelRetry) as exc_info:
            asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))
        assert "could not connect to server" in exc_info.value.message


class TestSQLToolsetGetSchema:
    def test_returns_column_info(self):
        ts = SQLToolset("pg_default")
        mock_hook = _make_mock_db_hook()
        ts._hook = mock_hook

        result = asyncio.run(
            ts.call_tool("get_schema", {"table_name": "users"}, ctx=MagicMock(), tool=MagicMock())
        )
        columns = json.loads(result)
        assert columns == [{"name": "id", "type": "INTEGER"}, {"name": "name", "type": "VARCHAR"}]
        mock_hook.get_table_schema.assert_called_once_with("users", schema=None)

    def test_blocks_table_not_in_allowed_list(self):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        result = asyncio.run(
            ts.call_tool("get_schema", {"table_name": "secrets"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert "error" in data
        assert "secrets" in data["error"]

    def test_introspection_error_raises_model_retry(self):
        """A failure while reading a table's schema is returned to the agent as a retry."""
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()
        ts._hook.get_table_schema.side_effect = Exception('relation "users" does not exist')

        with pytest.raises(ModelRetry) as exc_info:
            asyncio.run(
                ts.call_tool("get_schema", {"table_name": "users"}, ctx=MagicMock(), tool=MagicMock())
            )
        assert "does not exist" in exc_info.value.message


class TestSQLToolsetQuery:
    def test_returns_rows_as_json(self):
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook(
            records=[(1, "Alice"), (2, "Bob")],
            last_description=[("id",), ("name",)],
        )

        result = asyncio.run(
            ts.call_tool("query", {"sql": "SELECT id, name FROM users"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert data["rows"] == [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]
        assert data["count"] == 2

    def test_truncates_at_max_rows(self):
        ts = SQLToolset("pg_default", max_rows=1)
        ts._hook = _make_mock_db_hook(
            records=[(1, "Alice"), (2, "Bob"), (3, "Charlie")],
            last_description=[("id",), ("name",)],
        )

        result = asyncio.run(
            ts.call_tool("query", {"sql": "SELECT id, name FROM users"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert len(data["rows"]) == 1
        assert data["truncated"] is True
        assert data["count"] == 3

    def test_unsafe_sql_raises_model_retry(self):
        """An unsafe statement is surfaced to the agent as a retry so it can switch to a SELECT."""
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()

        with pytest.raises(ModelRetry) as exc_info:
            asyncio.run(ts.call_tool("query", {"sql": "DROP TABLE users"}, ctx=MagicMock(), tool=MagicMock()))
        assert "not allowed" in exc_info.value.message

    def test_allows_writes_when_enabled(self):
        ts = SQLToolset("pg_default", allow_writes=True)
        ts._hook = _make_mock_db_hook(
            records=[(1,)],
            last_description=[("count",)],
        )

        # Should not raise even with INSERT
        result = asyncio.run(
            ts.call_tool(
                "query", {"sql": "INSERT INTO users VALUES (3, 'Eve')"}, ctx=MagicMock(), tool=MagicMock()
            )
        )
        # The mock doesn't actually execute, just returns mocked records
        data = json.loads(result)
        assert "rows" in data

    @pytest.mark.parametrize(
        "error",
        [
            Exception("001003 (42000): SQL compilation error: unexpected 'rows'"),
            RuntimeError("type mismatch"),
            ConnectionError("could not connect to server"),
        ],
    )
    def test_query_error_is_returned_to_agent_as_model_retry(self, error):
        """Any error from the query, whatever its type, is handed back to the agent as a retry with
        the database's own message. The toolset never inspects the error type or text; pydantic-ai's
        max_retries bounds the loop, so an unrecoverable error still fails the task."""
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()
        ts._hook.get_records.side_effect = error

        with pytest.raises(ModelRetry) as exc_info:
            asyncio.run(
                ts.call_tool("query", {"sql": "SELECT foo FROM bar"}, ctx=MagicMock(), tool=MagicMock())
            )
        message = exc_info.value.message
        assert str(error) in message
        assert "list_tables" in message
        assert "get_schema" in message


class TestSQLToolsetCheckQuery:
    def test_valid_select(self):
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()

        result = asyncio.run(
            ts.call_tool("check_query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert data["valid"] is True

    def test_invalid_sql(self):
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()

        result = asyncio.run(
            ts.call_tool("check_query", {"sql": "DROP TABLE users"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert data["valid"] is False
        assert "error" in data


class TestSQLToolsetHookResolution:
    @patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True)
    def test_lazy_resolves_db_hook(self, mock_base_hook):
        mock_hook = MagicMock(spec=DbApiHook)
        mock_conn = MagicMock(spec=["get_hook"])
        mock_conn.get_hook.return_value = mock_hook
        mock_base_hook.get_connection.return_value = mock_conn

        ts = SQLToolset("pg_default")
        hook = ts._get_db_hook()

        assert hook is mock_hook
        mock_base_hook.get_connection.assert_called_once_with("pg_default")

    @patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True)
    def test_raises_for_non_dbapi_hook(self, mock_base_hook):
        mock_conn = MagicMock(spec=["get_hook"])
        mock_conn.get_hook.return_value = MagicMock()  # Not a DbApiHook
        mock_base_hook.get_connection.return_value = mock_conn

        ts = SQLToolset("bad_conn")

        with pytest.raises(ValueError, match="does not provide a DbApiHook"):
            ts._get_db_hook()

    @patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True)
    def test_caches_hook_after_first_resolution(self, mock_base_hook):
        mock_hook = MagicMock(spec=DbApiHook)
        mock_conn = MagicMock(spec=["get_hook"])
        mock_conn.get_hook.return_value = mock_hook
        mock_base_hook.get_connection.return_value = mock_conn

        ts = SQLToolset("pg_default")
        ts._get_db_hook()
        ts._get_db_hook()

        # Only called once because result is cached.
        mock_base_hook.get_connection.assert_called_once()


class TestSQLToolsetMultiSchema:
    """Schema-qualified allowed_tables span multiple schemas in one database."""

    @staticmethod
    def _schema_aware_hook(tables_by_schema: dict[str | None, list[str]]):
        hook = MagicMock(spec=DbApiHook)
        hook.inspector = MagicMock()
        hook.inspector.get_table_names.side_effect = lambda schema=None: tables_by_schema.get(schema, [])
        hook.get_table_schema.return_value = [{"name": "id", "type": "INTEGER"}]
        return hook

    def test_list_tables_spans_multiple_schemas(self):
        ts = SQLToolset(
            "sf",
            allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"],
        )
        ts._hook = self._schema_aware_hook(
            {
                "MODEL_ASTRO": ["DEPLOYMENT_IMAGE_DETAILS", "OTHER_TABLE"],
                "MODEL_CRM": ["SF_ASTRO_ORGS"],
            }
        )

        result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())))
        assert result == ["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"]

    def test_list_tables_never_introspects_none_schema_when_all_qualified(self):
        """Regression for the 'SHOW TABLES IN SCHEMA "DB"."None"' failure."""
        ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.X", "MODEL_CRM.Y"])
        ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["X"], "MODEL_CRM": ["Y"]})

        asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))

        called_schemas = {c.kwargs.get("schema") for c in ts._hook.inspector.get_table_names.call_args_list}
        assert called_schemas == {"MODEL_ASTRO", "MODEL_CRM"}
        assert None not in called_schemas

    def test_list_tables_mixed_qualified_and_default(self):
        ts = SQLToolset("pg", allowed_tables=["users", "MODEL_ASTRO.X"], schema="public")
        ts._hook = self._schema_aware_hook({"public": ["users", "orders"], "MODEL_ASTRO": ["X", "Z"]})

        result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())))
        # Qualified schemas listed first (sorted), then the default schema.
        assert result == ["MODEL_ASTRO.X", "users"]

    def test_get_schema_routes_to_qualified_schema(self):
        ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"])
        ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["DEPLOYMENT_IMAGE_DETAILS"]})

        result = json.loads(
            asyncio.run(
                ts.call_tool(
                    "get_schema",
                    {"table_name": "MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"},
                    ctx=MagicMock(),
                    tool=MagicMock(),
                )
            )
        )
        assert result == [{"name": "id", "type": "INTEGER"}]
        ts._hook.get_table_schema.assert_called_once_with("DEPLOYMENT_IMAGE_DETAILS", schema="MODEL_ASTRO")

    def test_get_schema_blocks_table_outside_allowed_schema(self):
        ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.X"])
        ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["X"]})

        result = json.loads(
            asyncio.run(
                ts.call_tool(
                    "get_schema", {"table_name": "SECRETS.PASSWORDS"}, ctx=MagicMock(), tool=MagicMock()
                )
            )
        )
        assert "error" in result
        ts._hook.get_table_schema.assert_not_called()

    def test_get_schema_unqualified_uses_default_schema(self):
        ts = SQLToolset("pg", schema="public")
        ts._hook = self._schema_aware_hook({"public": ["users"]})

        asyncio.run(ts.call_tool("get_schema", {"table_name": "users"}, ctx=MagicMock(), tool=MagicMock()))
        ts._hook.get_table_schema.assert_called_once_with("users", schema="public")

    def test_list_tables_matches_case_insensitively(self):
        """Snowflake reflects unquoted names lowercased; uppercase allowed_tables still match."""
        ts = SQLToolset(
            "sf",
            allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"],
        )
        ts._hook = self._schema_aware_hook(
            {
                "MODEL_ASTRO": ["deployment_image_details", "other"],
                "MODEL_CRM": ["sf_astro_orgs"],
            }
        )

        result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())))
        assert result == ["MODEL_ASTRO.deployment_image_details", "MODEL_CRM.sf_astro_orgs"]

    def test_get_schema_matches_case_insensitively(self):
        ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"])
        ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["deployment_image_details"]})

        result = json.loads(
            asyncio.run(
                ts.call_tool(
                    "get_schema",
                    {"table_name": "MODEL_ASTRO.deployment_image_details"},
                    ctx=MagicMock(),
                    tool=MagicMock(),
                )
            )
        )
        assert "error" not in result
        ts._hook.get_table_schema.assert_called_once_with("deployment_image_details", schema="MODEL_ASTRO")

    def test_list_tables_deduplicates_same_table(self):
        """A table listed both qualified and unqualified appears once."""
        ts = SQLToolset("pg", allowed_tables=["public.users", "users"], schema="public")
        ts._hook = self._schema_aware_hook({"public": ["users"]})

        result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())))
        assert result == ["public.users"]


class TestSQLToolsetMetadataStatements:
    """Read-only metadata statements (DESCRIBE/SHOW) flow through the query tool."""

    def test_describe_allowed_through_query(self):
        """DESCRIBE is read-only metadata and should not be rejected as unsafe."""
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook(
            records=[("id", "INTEGER"), ("name", "VARCHAR")],
            last_description=[("column_name",), ("data_type",)],
        )

        result = asyncio.run(
            ts.call_tool("query", {"sql": "DESCRIBE TABLE users"}, ctx=MagicMock(), tool=MagicMock())
        )
        data = json.loads(result)
        assert "rows" in data
        ts._hook.get_records.assert_called_once_with("DESCRIBE TABLE users")

    def test_show_allowed_with_snowflake_dialect(self):
        """SHOW parses to a metadata statement once the hook's dialect is passed through."""
        ts = SQLToolset("sf_default")
        ts._hook = _make_mock_db_hook(records=[("USERS",)], last_description=[("name",)])
        ts._hook.dialect_name = "snowflake"

        result = asyncio.run(ts.call_tool("query", {"sql": "SHOW TABLES"}, ctx=MagicMock(), tool=MagicMock()))
        data = json.loads(result)
        assert "rows" in data
        ts._hook.get_records.assert_called_once_with("SHOW TABLES")

    @pytest.mark.parametrize(
        "sql",
        # SHOW falls back to Command on Postgres (no SHOW support); DELETE is a write.
        ["SHOW TABLES", "DELETE FROM users"],
        ids=["show_without_dialect_support", "write"],
    )
    def test_query_blocks_disallowed_statements(self, sql):
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = "postgresql"

        # The statement is rejected before execution and surfaced to the agent as a
        # retry; get_records is never reached, so the guardrail still holds.
        with pytest.raises(ModelRetry) as exc_info:
            asyncio.run(ts.call_tool("query", {"sql": sql}, ctx=MagicMock(), tool=MagicMock()))
        assert "not allowed" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_check_query_accepts_describe(self):
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook()

        result = asyncio.run(
            ts.call_tool("check_query", {"sql": "DESCRIBE TABLE users"}, ctx=MagicMock(), tool=MagicMock())
        )
        assert json.loads(result)["valid"] is True

    def test_check_query_handles_unresolvable_connection(self):
        """check_query stays usable (dialect-agnostic) when the connection can't be resolved."""
        ts = SQLToolset("missing_conn")
        with patch.object(ts, "_get_db_hook", side_effect=RuntimeError("no such connection")):
            result = asyncio.run(
                ts.call_tool("check_query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock())
            )
        assert json.loads(result)["valid"] is True


def _run_query(ts: SQLToolset, sql: str):
    return asyncio.run(ts.call_tool("query", {"sql": sql}, ctx=MagicMock(), tool=MagicMock()))


def _run_check(ts: SQLToolset, sql: str):
    return json.loads(
        asyncio.run(ts.call_tool("check_query", {"sql": sql}, ctx=MagicMock(), tool=MagicMock()))
    )


class TestSQLToolsetAllowedTablesQueryEnforcement:
    """``allowed_tables`` is enforced on the query/check_query tools, not just on discovery."""

    def test_query_allows_table_on_the_list(self):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])

        result = _run_query(ts, "SELECT id FROM orders")

        assert "rows" in json.loads(result)
        ts._hook.get_records.assert_called_once_with("SELECT id FROM orders")

    def test_query_blocks_table_off_the_list(self):
        """The headline escape: querying a table that is not on the allow-list is refused."""
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "SELECT * FROM secret_salaries")

        assert "not in the allowed tables list" in exc_info.value.message
        assert "secret_salaries" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    @pytest.mark.parametrize(
        "sql",
        [
            "SELECT * FROM (SELECT * FROM secret_salaries) x",
            "WITH s AS (SELECT * FROM secret_salaries) SELECT * FROM s",
            "SELECT * FROM orders JOIN secret_salaries ON orders.id = secret_salaries.id",
            "SELECT * FROM orders UNION SELECT * FROM secret_salaries",
            "SELECT * FROM secret_salaries WHERE id IN (SELECT id FROM orders)",
        ],
        ids=["subquery", "cte_body", "join", "union", "where_subquery"],
    )
    def test_query_blocks_disallowed_table_reached_indirectly(self, sql):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, sql)

        assert "secret_salaries" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_query_blocks_catalog_enumeration(self):
        """information_schema/pg_catalog are ordinary tables, so the allow-list blocks them too."""
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "SELECT table_name FROM information_schema.tables")

        assert "information_schema.tables" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_query_allows_cte_reference_not_mistaken_for_table(self):
        """A CTE whose name is not on the list is fine as long as its body stays allowed."""
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])

        result = _run_query(ts, "WITH ranked AS (SELECT * FROM orders) SELECT * FROM ranked")

        assert "rows" in json.loads(result)

    def test_query_blocks_table_valued_function(self):
        """dblink reaches data through a path the list can't describe, so it is refused."""
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = "postgresql"

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "SELECT * FROM dblink('host=evil', 'SELECT 1') AS t(x int)")

        assert "cannot be checked against allowed_tables" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_query_blocks_show_when_allowlist_active(self):
        ts = SQLToolset("sf_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = "snowflake"

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "SHOW TABLES")

        assert "cannot be checked against allowed_tables" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_query_blocks_describe_of_disallowed_table(self):
        ts = SQLToolset("sf_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = "snowflake"

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "DESCRIBE TABLE secret_salaries")

        assert "secret_salaries" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_query_allows_describe_of_allowed_table(self):
        ts = SQLToolset("sf_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook(records=[("id", "INT")], last_description=[("name",), ("type",)])
        ts._hook.dialect_name = "snowflake"

        result = _run_query(ts, "DESCRIBE TABLE orders")

        assert "rows" in json.loads(result)
        ts._hook.get_records.assert_called_once_with("DESCRIBE TABLE orders")

    def test_query_allows_schema_qualified_table_on_list(self):
        ts = SQLToolset("sf", allowed_tables=["MODEL_CRM.SF_ASTRO_ORGS"])
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])
        ts._hook.dialect_name = "snowflake"

        result = _run_query(ts, "SELECT * FROM MODEL_CRM.SF_ASTRO_ORGS")

        assert "rows" in json.loads(result)

    def test_query_unqualified_resolves_to_default_schema(self):
        """``public.orders`` and ``orders`` denote the same table when schema='public'."""
        ts = SQLToolset("pg", allowed_tables=["orders"], schema="public")
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])

        # Qualifying with the default schema must still match the bare allow-list entry.
        result = _run_query(ts, "SELECT * FROM public.orders")
        assert "rows" in json.loads(result)

    def test_no_allowlist_leaves_queries_unrestricted(self):
        """Without allowed_tables the query tool behaves exactly as before (allow-all)."""
        ts = SQLToolset("pg_default")
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])

        result = _run_query(ts, "SELECT * FROM anything_at_all")

        assert "rows" in json.loads(result)
        ts._hook.get_records.assert_called_once_with("SELECT * FROM anything_at_all")

    def test_check_query_reports_disallowed_table_as_invalid(self):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        data = _run_check(ts, "SELECT * FROM secret_salaries")

        assert data["valid"] is False
        assert "secret_salaries" in data["error"]

    def test_check_query_valid_for_allowed_table(self):
        ts = SQLToolset("pg_default", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook()

        assert _run_check(ts, "SELECT * FROM orders")["valid"] is True

    def test_writes_still_bounded_by_allowed_tables(self):
        """allow_writes widens the statement types, but the allow-list still scopes the target."""
        ts = SQLToolset("pg_default", allowed_tables=["orders"], allow_writes=True)
        ts._hook = _make_mock_db_hook(records=[], last_description=None)

        # An allowed target is written.
        _run_query(ts, "INSERT INTO orders (id) VALUES (1)")
        ts._hook.get_records.assert_called_once_with("INSERT INTO orders (id) VALUES (1)")

        # A disallowed target is refused before execution.
        ts._hook.get_records.reset_mock()
        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "INSERT INTO secret_salaries (id) VALUES (1)")
        assert "secret_salaries" in exc_info.value.message
        ts._hook.get_records.assert_not_called()

    def test_writes_reject_dynamic_sql_the_parser_cannot_inspect(self):
        """allow_writes skips the read-only validator, so the allow-list must still
        refuse dynamic SQL (EXEC/EXECUTE) whose table access is opaque."""
        ts = SQLToolset("mssql_default", allowed_tables=["orders"], allow_writes=True)
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = "mssql"

        with pytest.raises(ModelRetry) as exc_info:
            _run_query(ts, "EXEC sp_who")

        assert "cannot be checked against allowed_tables" in exc_info.value.message
        ts._hook.get_records.assert_not_called()


class TestSQLToolsetAllowedTablesBypassRegressions:
    """Regression tests for bypasses found by adversarial red-teaming of the allow-list."""

    @pytest.mark.parametrize(
        ("sql", "dialect", "allow_writes"),
        [
            # CTE scope: a same-named CTE in an inner/sibling scope must not hide the real table.
            (
                "SELECT * FROM secret_salaries WHERE id IN "
                "(WITH secret_salaries AS (SELECT 1 id) SELECT id FROM secret_salaries)",
                "postgresql",
                False,
            ),
            # Non-recursive CTE is not in scope within its own body.
            (
                "WITH secret_salaries AS (SELECT * FROM secret_salaries) SELECT * FROM secret_salaries",
                "postgresql",
                False,
            ),
            # A CTE may only reference earlier siblings; a later-defined name is the real table.
            (
                "WITH a AS (SELECT * FROM secret_salaries), secret_salaries AS (SELECT 1 id) SELECT * FROM a",
                "postgresql",
                False,
            ),
            # Cross-database / catalog qualifier the schema.table allow-list cannot describe.
            ("SELECT * FROM secretdb.public.orders", "snowflake", False),
            ("SELECT * FROM secret_salaries..orders", "mssql", False),
            # MySQL executable comments execute on the engine but sqlglot treats them as inert.
            ("SELECT * FROM orders/*!UNION SELECT * FROM secret_salaries*/", "mysql", False),
            ("SELECT id FROM orders /*!50000 UNION SELECT id FROM secret_salaries */", "mysql", False),
            # TABLE <name> shorthand (mis-parsed) and TABLE('name') row source (string-named).
            ("TABLE secret_salaries UNION SELECT * FROM orders", "postgresql", False),
            ("SELECT * FROM TABLE('secret_salaries')", "snowflake", False),
            # Write-mode CTE shadowing the DML target.
            ("WITH secret_salaries AS (SELECT 1) DELETE FROM secret_salaries", "postgresql", True),
            # Quoted identifier is case-distinct on the engine but case-folds into the list.
            ('SELECT * FROM "Orders"', "postgresql", False),
            # A DML source CTE whose body reads an off-list table is still caught.
            (
                "WITH src AS (SELECT * FROM secret_salaries) INSERT INTO orders SELECT * FROM src",
                "postgresql",
                True,
            ),
        ],
        ids=[
            "cte_inner_shadow",
            "cte_self_body",
            "cte_forward_ref",
            "catalog_cross_db",
            "mssql_empty_middle",
            "mysql_exec_comment",
            "mysql_versioned_comment",
            "table_shorthand",
            "table_row_source",
            "write_cte_target",
            "quoted_case_distinct",
            "dml_cte_body_reads_offlist",
        ],
    )
    def test_known_bypasses_are_rejected(self, sql, dialect, allow_writes):
        ts = SQLToolset("c", allowed_tables=["orders"], allow_writes=allow_writes)
        ts._hook = _make_mock_db_hook()
        ts._hook.dialect_name = dialect

        with pytest.raises(ModelRetry):
            _run_query(ts, sql)
        ts._hook.get_records.assert_not_called()

    def test_legit_cte_over_allowed_table_still_runs(self):
        """The scope-aware fix must not false-reject a genuine CTE over an allowed table."""
        ts = SQLToolset("c", allowed_tables=["orders"])
        ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("id",)])

        result = _run_query(ts, "WITH ranked AS (SELECT * FROM orders) SELECT * FROM ranked")

        assert "rows" in json.loads(result)
        ts._hook.get_records.assert_called_once()

    def test_dml_with_cte_source_over_allowed_table_runs(self):
        """A CTE used as a DML source must not be mistaken for a disallowed base table."""
        ts = SQLToolset("c", allowed_tables=["orders"], allow_writes=True)
        ts._hook = _make_mock_db_hook(records=[], last_description=None)

        sql = "WITH src AS (SELECT * FROM orders) INSERT INTO orders SELECT * FROM src"
        _run_query(ts, sql)

        ts._hook.get_records.assert_called_once_with(sql)
