Spaces:
Runtime error
Runtime error
File size: 4,099 Bytes
0cac9cf | 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 | """Tests for the static SQL safety sandbox.
Author: mohamedgamal04
"""
import unittest
from queryquest.sql.validation import (
validate_statement_static,
validate_statement_tables,
)
class ValidateStatementStaticTests(unittest.TestCase):
def test_allows_plain_select(self) -> None:
self.assertIsNone(validate_statement_static("SELECT * FROM listings"))
def test_allows_update_set_clause(self) -> None:
# SET must not be treated as a blocked admin token.
self.assertIsNone(validate_statement_static("UPDATE listings SET price = 10 WHERE id = 1"))
def test_blocks_read_csv(self) -> None:
self.assertEqual(
validate_statement_static("SELECT * FROM read_csv('/etc/passwd')"),
"'READ_CSV' is not allowed",
)
def test_blocks_read_parquet(self) -> None:
self.assertEqual(
validate_statement_static("SELECT * FROM read_parquet('/tmp/secret.parquet')"),
"'READ_PARQUET' is not allowed",
)
def test_blocks_glob(self) -> None:
self.assertEqual(
validate_statement_static("SELECT * FROM glob('/home/*')"),
"'GLOB' is not allowed",
)
def test_blocks_attach(self) -> None:
self.assertEqual(
validate_statement_static("ATTACH 'evil.db' AS e"),
"'ATTACH' is not allowed",
)
def test_blocks_copy(self) -> None:
self.assertEqual(
validate_statement_static("COPY listings TO '/tmp/out.csv'"),
"'COPY' is not allowed",
)
def test_blocks_pragma(self) -> None:
self.assertEqual(
validate_statement_static("PRAGMA database_list"),
"'PRAGMA' is not allowed",
)
def test_blocks_comments_anywhere(self) -> None:
self.assertEqual(
validate_statement_static("SELECT * FROM listings -- sneaky"),
"SQL comments are not allowed",
)
def test_blocks_block_comments(self) -> None:
self.assertEqual(
validate_statement_static("SELECT /* hi */ * FROM listings"),
"SQL comments are not allowed",
)
def test_semicolon_inside_string_is_not_multistatement(self) -> None:
# The historic bug: a literal ';' tripped the multi-statement guard.
self.assertIsNone(
validate_statement_static("SELECT * FROM listings WHERE note = 'a;b'")
)
def test_rejects_multiple_statements(self) -> None:
self.assertEqual(
validate_statement_static("SELECT 1; DROP TABLE users"),
"multiple SQL statements are not allowed",
)
def test_trailing_semicolon_allowed(self) -> None:
self.assertIsNone(validate_statement_static("SELECT * FROM listings;"))
def test_rejects_schema_change(self) -> None:
self.assertEqual(
validate_statement_static("CREATE TABLE x (id INT)"),
"'CREATE' is not allowed",
)
class ValidateStatementTablesTests(unittest.TestCase):
def test_accepts_registered_table(self) -> None:
self.assertIsNone(
validate_statement_tables("SELECT * FROM listings", {"listings"})
)
def test_rejects_unregistered_table(self) -> None:
self.assertEqual(
validate_statement_tables("SELECT * FROM secrets", {"listings"}),
"table 'secrets' is not a registered workbook table",
)
def test_accepts_registered_join(self) -> None:
self.assertIsNone(
validate_statement_tables(
"SELECT * FROM orders JOIN customers ON orders.cid = customers.id",
{"orders", "customers"},
)
)
def test_rejects_unregistered_join_table(self) -> None:
self.assertEqual(
validate_statement_tables(
"SELECT * FROM orders JOIN secrets ON orders.id = secrets.id",
{"orders", "customers"},
),
"table 'secrets' is not a registered workbook table",
)
if __name__ == "__main__":
unittest.main()
|