Spaces:
Sleeping
Sleeping
File size: 8,789 Bytes
256d433 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | """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
|