tests/test_sql_tool.py

"""Tests for src.tools.sql_tool — regex safety guards and static helpers.

These tests exercise the ``_clean``, ``_validate`` and ``_inject_limit``
methods without needing a real database connection.
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from src.tools.sql_tool import (
    _FORBIDDEN_KEYWORDS,
    _LIMIT_PATTERN,
    _MULTI_STATEMENT,
    _SELECT_PATTERN,
    SqlTool,
    SqlToolError,
)


# ---------------------------------------------------------------------------
# _clean
# ---------------------------------------------------------------------------


class TestClean:
    """``_clean`` normalises LLM output: fences, language tags, semicolons."""

    def test_strips_markdown_fence(self) -> None:
        raw = "```\nSELECT 1\n```"
        assert SqlTool._clean(raw) == "SELECT 1"

    def test_strips_sql_language_tag(self) -> None:
        raw = "```sql\nSELECT 1\n```"
        assert SqlTool._clean(raw) == "SELECT 1"

    def test_removes_trailing_semicolon(self) -> None:
        assert SqlTool._clean("SELECT 1;") == "SELECT 1"

    def test_noop_on_clean_input(self) -> None:
        assert SqlTool._clean("SELECT 1") == "SELECT 1"

    def test_strips_leading_whitespace(self) -> None:
        assert SqlTool._clean("   SELECT 1   ") == "SELECT 1"

    def test_empty_string(self) -> None:
        assert SqlTool._clean("") == ""

    def test_none_coerces_to_empty(self) -> None:
        assert SqlTool._clean(None) == ""  # type: ignore[arg-type]


# ---------------------------------------------------------------------------
# _validate — allowlist / forbidden-keyword guards
# ---------------------------------------------------------------------------


class TestValidate:
    """``_validate`` rejects anything that isn't a safe read-only statement."""

    def test_accepts_select(self) -> None:
        SqlTool._validate("SELECT * FROM customers")

    def test_accepts_with_cte(self) -> None:
        SqlTool._validate("WITH t AS (SELECT 1) SELECT * FROM t")

    def test_rejects_empty(self) -> None:
        with pytest.raises(SqlToolError) as exc:
            SqlTool._validate("")
        assert "Empty SQL" in str(exc.value)

    def test_rejects_insert(self) -> None:
        with pytest.raises(SqlToolError) as exc:
            SqlTool._validate("INSERT INTO t VALUES (1)")
        assert "Only SELECT" in str(exc.value)

    def test_rejects_update(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("UPDATE t SET x = 1")

    def test_rejects_delete(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("DELETE FROM t")

    def test_rejects_drop(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("DROP TABLE t")

    def test_rejects_alter(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("ALTER TABLE t ADD COLUMN x INT")

    def test_rejects_create(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("CREATE TABLE t (id INT)")

    def test_rejects_truncate(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("TRUNCATE TABLE t")

    def test_rejects_grant(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("GRANT SELECT ON t TO user")

    def test_rejects_copy(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("COPY t TO '/tmp/out.csv'")

    def test_rejects_multi_statement(self) -> None:
        with pytest.raises(SqlToolError) as exc:
            SqlTool._validate("SELECT 1; DROP TABLE t")
        assert "Multiple statements" in str(exc.value)

    def test_rejects_comment_on(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("COMMENT ON TABLE t IS 'desc'")

    def test_rejects_security_definer(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("CREATE FUNCTION f() SECURITY DEFINER")

    def test_case_insensitive_forbidden(self) -> None:
        with pytest.raises(SqlToolError):
            SqlTool._validate("insert into t values (1)")

    def test_select_in_string_literal_still_rejected(self) -> None:
        # The regex is intentionally simple (not a full SQL parser) and
        # matches forbidden keywords anywhere in the string. A literal
        # containing "DELETE" still triggers the guard.
        with pytest.raises(SqlToolError):
            SqlTool._validate("SELECT 'DELETE FROM t' AS warning")


# ---------------------------------------------------------------------------
# _inject_limit
# ---------------------------------------------------------------------------


class TestInjectLimit:
    """``_inject_limit`` appends LIMIT when absent."""

    @pytest.fixture
    def sql_tool(self) -> SqlTool:
        # Mock out engine creation and schema introspection so we don't need
        # a real Postgres driver.
        with (
            patch("src.tools.sql_tool.create_engine") as mock_engine,
            patch("src.tools.sql_tool.inspect") as mock_inspect,
        ):
            mock_engine.return_value = MagicMock()
            mock_inspector = MagicMock()
            mock_inspector.get_table_names.return_value = []
            mock_inspect.return_value = mock_inspector
            tool = SqlTool("postgresql://u:p@h/d")
            # Manually set the attributes _inject_limit depends on
            tool.row_limit = 200
            yield tool

    def test_adds_limit(self, sql_tool: SqlTool) -> None:
        result = sql_tool._inject_limit("SELECT * FROM t")
        assert result.endswith("LIMIT 200")

    def test_skips_when_present(self, sql_tool: SqlTool) -> None:
        original = "SELECT * FROM t LIMIT 10"
        result = sql_tool._inject_limit(original)
        assert result == original

    def test_case_insensitive_detection(self, sql_tool: SqlTool) -> None:
        original = "SELECT * FROM t limit 50"
        result = sql_tool._inject_limit(original)
        assert result == original


# ---------------------------------------------------------------------------
# Regex patterns (direct unit tests)
# ---------------------------------------------------------------------------


class TestRegexPatterns:
    """Direct tests for the compiled regex objects."""

    @pytest.mark.parametrize(
        "sql,matches",
        [
            ("SELECT 1", True),
            ("  SELECT 1", True),
            ("WITH t AS (SELECT 1) SELECT 1", True),
            ("INSERT INTO t VALUES (1)", False),
            ("UPDATE t SET x=1", False),
            ("", False),
        ],
    )
    def test_select_pattern(self, sql: str, matches: bool) -> None:
        assert bool(_SELECT_PATTERN.match(sql)) is matches

    @pytest.mark.parametrize(
        "sql,matches",
        [
            ("SELECT * FROM t LIMIT 10", True),
            ("SELECT * FROM t limit 5", True),
            ("SELECT * FROM t", False),
            ("SELECT LIMITLESS FROM t", False),
        ],
    )
    def test_limit_pattern(self, sql: str, matches: bool) -> None:
        assert bool(_LIMIT_PATTERN.search(sql)) is matches

    @pytest.mark.parametrize(
        "sql,matches",
        [
            ("SELECT 1; DROP TABLE t", True),
            ("SELECT 1; SELECT 2", True),
            ("SELECT 1", False),
            # String literals containing ';' are *not* multi-statement,
            # but the simple regex can't distinguish them. The regex
            # intentionally errs on the side of caution.
            ("SELECT ';' FROM t", True),
        ],
    )
    def test_multi_statement_pattern(self, sql: str, matches: bool) -> None:
        assert bool(_MULTI_STATEMENT.search(sql)) is matches

    @pytest.mark.parametrize(
        "keyword",
        [
            "INSERT",
            "UPDATE",
            "DELETE",
            "DROP",
            "ALTER",
            "CREATE",
            "TRUNCATE",
            "GRANT",
            "REVOKE",
            "COPY",
            "VACUUM",
            "ANALYZE",
            "REINDEX",
            "CLUSTER",
        ],
    )
    def test_forbidden_keywords(self, keyword: str) -> None:
        assert _FORBIDDEN_KEYWORDS.search(f"{keyword} something") is not None

    def test_forbidden_comment_on(self) -> None:
        assert _FORBIDDEN_KEYWORDS.search("COMMENT ON TABLE t IS 'x'") is not None

    def test_forbidden_security_definer(self) -> None:
        assert (
            _FORBIDDEN_KEYWORDS.search("CREATE FUNCTION f() SECURITY DEFINER")
            is not None
        )

    def test_forbidden_case_insensitive(self) -> None:
        assert _FORBIDDEN_KEYWORDS.search("insert into t") is not None