Spaces:
Runtime error
Runtime error
File size: 9,211 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 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 | 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()
|