AgentLab/backend/tests/test_sql_guard.py

"""SQL 守卫。

这是数据源接入里唯一不能出错的地方：SQL 由模型生成，不是人写的。
守卫放松一分，模型某次抽风就可能打到生产库上——所以这组用例只增不减。
"""
from __future__ import annotations

import pytest

from app.data.guard import SqlRejected, check, first_verb, is_write, split_statements


def rejected(sql: str, *, readonly: bool = True) -> str | None:
    try:
        check(sql, readonly=readonly, source_name="测试库")
        return None
    except SqlRejected as e:
        return str(e)


@pytest.mark.parametrize("sql", [
    "SELECT * FROM orders",
    "  select 1  ",
    "WITH t AS (SELECT 1) SELECT * FROM t",
    "(SELECT a FROM b)",
    "SHOW TABLES",
    "DESCRIBE orders",
    "EXPLAIN SELECT * FROM orders",
    "/* 注释 */ SELECT 1",
    "-- 行注释\nSELECT 2",
    "SELECT 1;",                       # 尾分号不算第二条语句
])
def test_readonly_allows_queries(sql):
    assert rejected(sql) is None, f"不该拒绝：{sql!r}"


@pytest.mark.parametrize("sql", [
    "UPDATE orders SET x=1",
    "DELETE FROM orders",
    "INSERT INTO orders VALUES(1)",
    "MERGE INTO a USING b ON (1=1)",   # 白名单式判定才拦得住的方言私货
    "CALL some_proc()",
    "REPLACE INTO t VALUES(1)",
])
def test_readonly_blocks_writes(sql):
    reason = rejected(sql)
    assert reason and "只读" in reason


@pytest.mark.parametrize("sql", [
    # 首关键字是查询，正文在写。只看首关键字的话全都能混过去——
    # 第一条在 SQLite 上实测过：守卫放行，数据真的删了
    "WITH k AS (SELECT 1) DELETE FROM t WHERE x < 3",
    "WITH x AS (DELETE FROM t RETURNING *) SELECT * FROM x",      # PG 的数据修改 CTE
    "WITH k AS (SELECT 1) UPDATE t SET x = 1",
    "WITH k AS (SELECT 1) INSERT INTO t SELECT * FROM k",
    "SELECT * INTO backup FROM orders",                           # PG：建一张新表
    "SELECT * FROM orders INTO OUTFILE '/tmp/o.csv'",             # MySQL：写服务器上的文件
    "EXPLAIN ANALYZE DELETE FROM orders",                         # PG：ANALYZE 会真的执行
    "SELECT * FROM orders FOR UPDATE",                            # 行锁
    "SELECT set_config('default_transaction_read_only', 'off', false)",  # 关掉连接层只读
    "SELECT \"set_config\"('default_transaction_read_only', 'off', false)",  # 函数名加引号
    "SELECT pg_catalog.set_config('default_transaction_read_only', 'off', false)",
    "SELECT * FROM dblink('dbname=x', 'DELETE FROM t RETURNING *') AS r(x int)",  # 另开连接
    "SELECT nextval('order_seq')",
])
def test_readonly_blocks_writes_hidden_behind_a_query_verb(sql):
    reason = rejected(sql)
    assert reason and "只读" in reason, f"漏了：{sql!r}"


@pytest.mark.parametrize("sql", [
    "PRAGMA user_version = 7",
    "PRAGMA user_version(7)",            # 函数写法同样是赋值
    "PRAGMA query_only = 0",             # 想先把只读关掉
    "PRAGMA main.writable_schema = 1",
    "PRAGMA optimize",                   # 会跑 ANALYZE，写统计表
    "PRAGMA wal_checkpoint(TRUNCATE)",
])
def test_readonly_blocks_pragmas_that_write(sql):
    reason = rejected(sql)
    assert reason and "只读" in reason, f"漏了：{sql!r}"


@pytest.mark.parametrize("sql", [
    "PRAGMA table_info(orders)",
    "PRAGMA main.table_info('orders')",
    "PRAGMA index_list(orders)",
    "PRAGMA foreign_key_list(orders)",
    "PRAGMA user_version",
    "SELECT updated_at, deleted_flag, created_by FROM orders",   # 标识符里带关键字不算
    "SELECT 'please delete me' AS note",                         # 字符串里的不算
    'SELECT "update" FROM t',                                    # 引号里的标识符不算
    "SELECT REPLACE(name, 'a', 'b'), INSERT('abc', 1, 1, 'x'), TRUNCATE(price, 2) FROM t",
    "SHOW CREATE TABLE orders",                                  # MySQL 看建表语句
    "WITH t AS (SELECT 1 AS into_count) SELECT * FROM t",
])
def test_readonly_does_not_misfire_on_harmless_queries(sql):
    """正文扫描收紧了，误杀也得守住：同名函数、标识符、字符串都不是在写。"""
    assert rejected(sql) is None, f"误杀了：{sql!r}"


def test_is_write_sees_writes_hidden_behind_a_query_verb():
    """可写源靠 is_write 决定要不要人工审批——WITH 开头的写不能被当成查询放过去。"""
    assert is_write("WITH k AS (SELECT 1) DELETE FROM t")
    assert is_write("SELECT * INTO backup FROM t")
    assert is_write("UPDATE t SET x = 1")
    assert not is_write("SELECT * FROM t")
    assert not is_write("WITH k AS (SELECT 1) SELECT * FROM k")


@pytest.mark.parametrize("sql", [
    "DROP TABLE orders",
    "TRUNCATE TABLE orders",
    "ALTER TABLE t ADD c INT",
    "GRANT ALL ON db.* TO u",
    "CREATE TABLE t(a INT)",
])
def test_destructive_blocked_even_on_writable_source(sql):
    """结构变更和权限操作在可写源上也不放行——不可逆的事不交给模型发起。"""
    reason = rejected(sql, readonly=False)
    assert reason and "任何数据源" in reason


@pytest.mark.parametrize("sql", [
    "SELECT 1; DROP TABLE users",
    "SELECT 1;DROP TABLE users;",
    "SELECT 1 /*x*/; DELETE FROM t",
])
def test_multi_statement_blocked(sql):
    """分号拼接必须在到达驱动之前断掉，不赌驱动默认不跑第二条。"""
    reason = rejected(sql)
    assert reason and "一条语句" in reason


@pytest.mark.parametrize("sql", [
    "SELECT ';' AS x",
    "SELECT 'a;b' FROM t WHERE c = 'd;e'",
    "SELECT '它说：''你好'';' AS s",      # 连续两个引号是转义，不是字符串结束
])
def test_semicolon_inside_string_is_not_a_second_statement(sql):
    """误杀比漏杀更让人无从下手：字符串字面量里的分号不算语句分隔。"""
    assert rejected(sql) is None, f"误杀了：{sql!r}"


def test_writable_source_allows_dml():
    for sql in ("UPDATE orders SET x=1", "INSERT INTO t VALUES(1)", "DELETE FROM t WHERE id=1"):
        assert rejected(sql, readonly=False) is None


@pytest.mark.parametrize("sql", ["", "   ", "-- 只有注释", "!!!"])
def test_empty_and_malformed_rejected(sql):
    assert rejected(sql) is not None


def test_check_returns_statement_without_trailing_semicolon():
    """调用方必须执行返回值而不是原始输入，否则守卫等于没做。"""
    assert check("SELECT 1;", readonly=True) == "SELECT 1"


def test_first_verb_sees_through_comments_and_parens():
    assert first_verb("/* x */ (SELECT 1)") == "select"
    assert first_verb("-- c\nDROP TABLE t") == "drop"


def test_split_keeps_quoted_semicolons_together():
    assert len(split_statements("SELECT ';'")) == 1
    assert len(split_statements("SELECT 1; SELECT 2")) == 2