Spaces:
Runtime error
Runtime error
File size: 6,352 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 | """Interactive SQL execution facade over the split sql/* modules.
The heavy lifting now lives in validation/rewrite/registry/execution/writeback.
This module keeps the original `execute_sql_statements` entry point (and the
symbols the tests import) so existing callers and tests keep working while the
async core engine is built on top of the same pieces.
Author: mohamedgamal04
"""
from __future__ import annotations
from pathlib import Path
import duckdb
import pandas as pd
from rich.console import Console
from rich.prompt import Prompt
from .preview import print_dataframe_as_table
from .registry import (
SqlRewriteContext,
WorkbookRecord,
_build_execution_context,
_column_name_from_value,
_normalize_columns,
_table_name_from_file,
_table_name_from_sheet,
)
from .rewrite import (
_canonical_identifier,
_normalize_single_quoted_table_identifiers,
_normalize_sql_statement,
_prepare_statement,
_quote_known_identifiers,
_rewrite_to_normalized_identifiers,
_strip_identifier_quotes,
)
from .execution import (
_build_update_change_predicate,
_delete_statement_to_scope_query,
_execute_statement_safely,
_extract_delete_table_name,
_extract_update_set_clause,
_extract_update_table_name,
_extract_update_where_clause,
_handle_sql_execution_error,
_print_delete_preview,
_print_sql_preview,
_print_update_preview,
_split_top_level_csv,
_update_statement_to_scope_query,
)
from .validation import ALLOWED_SQL_COMMANDS, _is_dml_statement, _strip_leading_sql_noise, _validate_sql_allowlist
from .writeback import _extract_target_table_name, _save_dataframe_to_workbook
__all__ = [
"ALLOWED_SQL_COMMANDS",
"SqlRewriteContext",
"WorkbookRecord",
"execute_sql_statements",
"_build_execution_context",
"_build_update_change_predicate",
"_extract_update_table_name",
"_extract_update_where_clause",
"_normalize_single_quoted_table_identifiers",
"_prepare_statement",
"_quote_known_identifiers",
"_rewrite_to_normalized_identifiers",
"_strip_identifier_quotes",
"_update_statement_to_scope_query",
"_validate_sql_allowlist",
]
def execute_sql_statements(
sql_statements: list[str],
console: Console | None = None,
excel_dir: str | Path | None = None,
) -> None:
"""Execute SQL statements and optionally persist DML changes to Excel files."""
active_console = console or Console()
if not sql_statements:
active_console.print("No SQL statements to execute.")
return
for statement in sql_statements:
refusal_reason = _validate_sql_allowlist(statement)
if refusal_reason is not None:
active_console.print(
"[red]Refused:[/red] Only SELECT, INSERT, UPDATE, and DELETE statements are allowed. "
f"{refusal_reason}."
)
active_console.print(f"Skipped statement: {statement}")
return
_print_sql_preview(sql_statements, active_console)
connection = duckdb.connect()
rewrite_context, workbook_records, workbook_sheet_data = _build_execution_context(connection, excel_dir=excel_dir)
try:
wrote_data = False
for statement in sql_statements:
prepared_statement = _prepare_statement(statement, rewrite_context)
if statement.lstrip().lower().startswith("delete"):
_print_delete_preview(connection, prepared_statement, active_console)
wrote_data = True
continue
if statement.lstrip().lower().startswith("update"):
affected_rows = _print_update_preview(connection, prepared_statement, active_console)
if affected_rows > 0:
wrote_data = True
continue
cursor = _execute_statement_safely(connection, prepared_statement, active_console)
if cursor is None:
continue
if _is_dml_statement(statement):
wrote_data = True
if cursor.description is None:
if cursor.rowcount != -1:
active_console.print(cursor.rowcount)
continue
results = cursor.df()
print_dataframe_as_table(results, active_console)
if not wrote_data:
return
should_save = Prompt.ask(
"Save changes back to the Excel files?",
choices=["y", "n"],
default="n",
console=active_console,
).strip()
if should_save.lower() != "y":
return
saved_results: dict[str, pd.DataFrame] = {}
workbook_updates: dict[Path, dict[str, pd.DataFrame]] = {}
for statement in sql_statements:
if not _is_dml_statement(statement):
continue
prepared_statement = _prepare_statement(statement, rewrite_context)
target_table_name = _extract_target_table_name(prepared_statement)
if target_table_name is None:
continue
workbook_record = workbook_records.get(str(target_table_name))
if workbook_record is None:
continue
file_path = workbook_record["file_path"]
if file_path not in workbook_updates:
workbook_updates[file_path] = dict(workbook_sheet_data.get(file_path, {}))
current_df = connection.table(workbook_record["table_name"]).df()
workbook_updates[file_path][workbook_record["sheet_name"]] = current_df
saved_results[workbook_record["table_name"]] = current_df.copy()
for file_path, sheet_data in workbook_updates.items():
if not sheet_data:
continue
_save_dataframe_to_workbook(file_path, sheet_data)
if not workbook_updates:
return
active_console.print("Changes saved.")
for table_name, saved_df in saved_results.items():
preview_df = saved_df.head(20)
print_dataframe_as_table(preview_df, active_console, title=f"Query result ({table_name})")
if len(saved_df) > len(preview_df):
active_console.print(f"Showing first {len(preview_df)} of {len(saved_df)} rows.")
finally:
connection.close()
|