from sql_security import validate_sql def test_accepts_select_cte(): result = validate_sql( """ WITH ranked AS ( SELECT id, ROW_NUMBER() OVER (ORDER BY created_at DESC) AS rn FROM customers ) SELECT * FROM ranked WHERE rn <= 10 """ ) assert result.valid is True def test_rejects_drop_table(): result = validate_sql("DROP TABLE transactions") assert result.valid is False assert "SELECT" in result.reason or "DROP" in result.reason def test_rejects_update_statement(): result = validate_sql("UPDATE customers SET lifetime_value = 0") assert result.valid is False def test_rejects_multiple_statements_after_select(): result = validate_sql("SELECT * FROM customers; DELETE FROM customers;") assert result.valid is False def test_rejects_unsafe_functions_inside_select(): result = validate_sql("SELECT pg_read_file('/etc/passwd')") assert result.valid is False assert "pg_read_file" in result.reason