tests/test_sql_executor.py

import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch

import pandas as pd

from queryquest.sql.executor import (
    _build_update_change_predicate,
    _extract_update_table_name,
    _extract_update_where_clause,
    _prepare_statement,
    _normalize_single_quoted_table_identifiers,
    _quote_known_identifiers,
    _rewrite_to_normalized_identifiers,
    _update_statement_to_scope_query,
    _validate_sql_allowlist,
    _strip_identifier_quotes,
    execute_sql_statements,
)
from rich.console import Console


class SqlExecutorNormalizationTests(unittest.TestCase):
    def test_normalize_single_quoted_table_identifier_in_from_clause(self) -> None:
        statement = "SELECT * FROM 'Real Estate Listings' ORDER BY List Price DESC LIMIT 1"
        rewritten = _normalize_single_quoted_table_identifiers(statement, {"Real Estate Listings"})
        self.assertIn('FROM "Real Estate Listings"', rewritten)

    def test_does_not_requote_already_single_quoted_identifier(self) -> None:
        statement = "SELECT * FROM 'Real Estate Listings'"
        rewritten = _quote_known_identifiers(statement, {"Real Estate Listings"})
        self.assertEqual(rewritten, statement)

    def test_strip_identifier_quotes_handles_single_quotes(self) -> None:
        self.assertEqual(_strip_identifier_quotes("'Real Estate Listings'"), "Real Estate Listings")

    def test_rewrite_mixed_identifiers_to_normalized(self) -> None:
        statement = 'SELECT MAX("List Price") FROM Real_Estate_Listings'
        rewritten = _rewrite_to_normalized_identifiers(
            statement,
            table_name_map={"Real Estate Listings": "Real_Estate_Listings"},
            table_alias_map={"realestatelistings": "Real_Estate_Listings"},
            column_name_map={"List Price": "List_Price"},
            column_alias_map={"listprice": "List_Price"},
        )
        self.assertEqual(rewritten, "SELECT MAX(List_Price) FROM Real_Estate_Listings")

    def test_rewrite_original_table_and_columns_to_normalized(self) -> None:
        statement = 'SELECT MAX("List Price") FROM "Real Estate Listings"'
        rewritten = _rewrite_to_normalized_identifiers(
            statement,
            table_name_map={"Real Estate Listings": "Real_Estate_Listings"},
            table_alias_map={"realestatelistings": "Real_Estate_Listings"},
            column_name_map={"List Price": "List_Price"},
            column_alias_map={"listprice": "List_Price"},
        )
        self.assertEqual(rewritten, "SELECT MAX(List_Price) FROM Real_Estate_Listings")

    def test_rewrite_compact_table_alias_to_normalized(self) -> None:
        statement = 'SELECT MAX("List Price") FROM RealEstateListings'
        rewritten = _rewrite_to_normalized_identifiers(
            statement,
            table_name_map={"Real Estate Listings": "Real_Estate_Listings"},
            table_alias_map={"realestatelistings": "Real_Estate_Listings"},
            column_name_map={"List Price": "List_Price"},
            column_alias_map={"listprice": "List_Price"},
        )
        self.assertEqual(rewritten, "SELECT MAX(List_Price) FROM Real_Estate_Listings")

    def test_rewrite_qualified_compact_column_name_to_normalized(self) -> None:
        statement = "SELECT MC.CampaignName FROM Marketing_Campaigns MC"
        rewritten = _rewrite_to_normalized_identifiers(
            statement,
            table_name_map={},
            table_alias_map={"marketingcampaigns": "Marketing_Campaigns"},
            column_name_map={"Campaign Name": "Campaign_Name"},
            column_alias_map={"campaignname": "Campaign_Name"},
        )
        self.assertEqual(rewritten, "SELECT MC.Campaign_Name FROM Marketing_Campaigns MC")

    def test_rewrite_qualified_quoted_column_name_to_normalized(self) -> None:
        statement = 'SELECT MC."Campaign Name" FROM Marketing_Campaigns MC'
        rewritten = _rewrite_to_normalized_identifiers(
            statement,
            table_name_map={},
            table_alias_map={"marketingcampaigns": "Marketing_Campaigns"},
            column_name_map={"Campaign Name": "Campaign_Name"},
            column_alias_map={"campaignname": "Campaign_Name"},
        )
        self.assertEqual(rewritten, "SELECT MC.Campaign_Name FROM Marketing_Campaigns MC")

    def test_prepare_statement_rewrites_spaced_table_name_for_delete(self) -> None:
        statement = "DELETE FROM Real Estate Listings WHERE Listing_ID IS NULL"
        rewrite_context = {
            "table_identifiers": {"Real Estate Listings", "Real_Estate_Listings"},
            "column_identifiers": {"Listing ID", "Listing_ID"},
            "table_name_map": {"Real Estate Listings": "Real_Estate_Listings"},
            "table_alias_map": {"realestatelistings": "Real_Estate_Listings"},
            "column_name_map": {"Listing ID": "Listing_ID"},
            "column_alias_map": {"listingid": "Listing_ID"},
        }
        prepared = _prepare_statement(statement, rewrite_context)
        self.assertEqual(prepared, "DELETE FROM Real_Estate_Listings WHERE Listing_ID IS NULL")

    def test_validate_sql_allowlist_rejects_schema_changes(self) -> None:
        self.assertEqual(_validate_sql_allowlist("CREATE TABLE x (id INT)"), "'CREATE' is not allowed")

    def test_validate_sql_allowlist_rejects_multiple_statements(self) -> None:
        self.assertEqual(
            _validate_sql_allowlist("SELECT 1; DROP TABLE users"),
            "multiple SQL statements are not allowed",
        )

    def test_validate_sql_allowlist_allows_select_join(self) -> None:
        self.assertIsNone(_validate_sql_allowlist("SELECT * FROM a JOIN b ON a.id = b.id"))

    def test_validate_sql_allowlist_rejects_join_in_dml(self) -> None:
        self.assertEqual(
            _validate_sql_allowlist("UPDATE a SET x = 1 FROM b JOIN c ON b.id = c.id"),
            "'JOIN' is not allowed",
        )

    def test_extract_update_table_name(self) -> None:
        statement = "UPDATE listings SET price = 2000 WHERE price = 1000"
        self.assertEqual(_extract_update_table_name(statement), "listings")

    def test_extract_update_where_clause(self) -> None:
        statement = "UPDATE listings SET price = 2000 WHERE price = 1000"
        self.assertEqual(_extract_update_where_clause(statement), "price = 1000")

    def test_extract_update_where_clause_missing(self) -> None:
        statement = "UPDATE listings SET price = 2000"
        self.assertIsNone(_extract_update_where_clause(statement))

    def test_update_statement_to_scope_query_with_where(self) -> None:
        statement = "UPDATE listings SET price = 2000 WHERE price = 1000"
        self.assertEqual(
            _update_statement_to_scope_query(statement),
            "SELECT * FROM listings WHERE (price = 1000) AND ((price IS DISTINCT FROM (2000)))",
        )

    def test_update_statement_to_scope_query_without_where(self) -> None:
        statement = "UPDATE listings SET price = 2000"
        self.assertEqual(
            _update_statement_to_scope_query(statement),
            "SELECT * FROM listings WHERE ((price IS DISTINCT FROM (2000)))",
        )

    def test_build_update_change_predicate_multiple_assignments(self) -> None:
        statement = "UPDATE listings SET price = 2000, status = 'active' WHERE id = 1"
        self.assertEqual(
            _build_update_change_predicate(statement),
            "(price IS DISTINCT FROM (2000)) OR (status IS DISTINCT FROM ('active'))",
        )

    def test_execute_sql_statements_refuses_disallowed_sql(self) -> None:
        console = Console(record=True)

        execute_sql_statements(["DROP TABLE users"], console=console)

        output = console.export_text()
        self.assertIn("Refused", output)
        self.assertIn("DROP", output)
        self.assertNotIn("SQL statements to execute", output)

    @patch("queryquest.sql.executor.Prompt.ask", return_value="y")
    def test_execute_sql_statements_updates_non_first_sheet_and_saves(self, _mock_prompt) -> None:
        with TemporaryDirectory() as tmp_dir:
            excel_path = Path(tmp_dir) / "inventory.xlsx"

            current_df = pd.DataFrame({"Item": ["A", "B"], "Qty": [1, 2]})
            archive_df = pd.DataFrame({"Item": ["A", "B"], "Qty": [10, 20]})

            with pd.ExcelWriter(excel_path, engine="openpyxl") as writer:
                current_df.to_excel(writer, sheet_name="Current", index=False)
                archive_df.to_excel(writer, sheet_name="Archive Data", index=False)

            console = Console(record=True)
            execute_sql_statements(
                ["UPDATE inventory__Archive_Data SET Qty = 99 WHERE Item = 'A'"],
                console=console,
                excel_dir=tmp_dir,
            )

            reloaded = pd.read_excel(excel_path, sheet_name=None)
            self.assertEqual(int(reloaded["Archive Data"].loc[0, "Qty"]), 99)
            self.assertEqual(int(reloaded["Archive Data"].loc[1, "Qty"]), 20)
            self.assertEqual(int(reloaded["Current"].loc[0, "Qty"]), 1)
            self.assertEqual(int(reloaded["Current"].loc[1, "Qty"]), 2)

if __name__ == "__main__":
    unittest.main()