Spaces:
Running
Running
File size: 2,947 Bytes
1e10174 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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()
|