backend/test_sql_security.py

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