tests/test_sql_guard.py

"""Security tests for the LLM SQL guard.

This is the highest-risk surface in the project: text from a language model is
executed against a database. These tests assert the guard blocks the obvious
attacks and, equally importantly, does not block legitimate analytics.
"""

from __future__ import annotations

import pytest

from gridpulse.agent.text2sql import SQLGuardError, guard_sql


class TestAllowsLegitimateQueries:
    def test_plain_select(self):
        assert "LIMIT" in guard_sql("SELECT ba_code FROM fact_demand_hourly")

    def test_cte(self):
        sql = (
            "WITH averages AS (SELECT ba_code, avg(demand_mwh) d FROM fact_demand_hourly GROUP BY 1) "
            "SELECT * FROM averages ORDER BY d DESC LIMIT 5"
        )
        assert guard_sql(sql).startswith("WITH")

    def test_join_between_allowed_tables(self):
        sql = "SELECT * FROM fact_forecast_accuracy JOIN dim_ba USING (ba_code) LIMIT 10"
        assert guard_sql(sql)

    def test_markdown_fences_are_stripped(self):
        assert not guard_sql("```sql\nSELECT 1 FROM dim_ba LIMIT 1\n```").startswith("`")

    def test_limit_is_injected_when_missing(self):
        assert "LIMIT" in guard_sql("SELECT * FROM dim_ba")

    def test_existing_limit_is_preserved(self):
        assert guard_sql("SELECT * FROM dim_ba LIMIT 7").count("LIMIT") == 1


class TestBlocksAttacks:
    @pytest.mark.parametrize(
        "sql",
        [
            "DROP TABLE dim_ba",
            "DELETE FROM fact_demand_hourly",
            "INSERT INTO dim_ba VALUES (1)",
            "UPDATE dim_ba SET ba_code = 'X'",
            "CREATE TABLE evil AS SELECT 1",
            "ALTER TABLE dim_ba ADD COLUMN x INT",
        ],
    )
    def test_ddl_and_dml_are_refused(self, sql):
        with pytest.raises(SQLGuardError):
            guard_sql(sql)

    def test_stacked_statement_is_refused(self):
        with pytest.raises(SQLGuardError, match="Multiple SQL statements"):
            guard_sql("SELECT * FROM dim_ba; DROP TABLE dim_ba;")

    def test_attach_is_refused(self):
        with pytest.raises(SQLGuardError):
            guard_sql("ATTACH '/etc/passwd' AS leak")

    @pytest.mark.parametrize(
        "sql",
        [
            "SELECT * FROM duckdb_settings()",
            "SELECT * FROM sqlite_master",
            "SELECT * FROM information_schema.tables",
            "SELECT * FROM secret_table",
        ],
    )
    def test_unknown_tables_are_refused(self, sql):
        with pytest.raises(SQLGuardError, match="unknown table"):
            guard_sql(sql)

    def test_empty_input_is_refused(self):
        with pytest.raises(SQLGuardError):
            guard_sql("   ")

    def test_comment_hidden_payload_is_neutralised(self):
        """A DELETE hidden behind a comment must not survive into the executed SQL."""
        result = guard_sql("SELECT * FROM dim_ba -- ; DELETE FROM dim_ba")
        assert "delete" not in result.lower()