"""Run prepared statements and build DELETE/UPDATE previews. I keep the affected-row previews here so the caller can show the user exactly what a destructive statement will touch before it commits. Author: mohamedgamal04 """ from __future__ import annotations import re import duckdb from rich.console import Console from .preview import print_dataframe_as_table, print_sql_statements_table from .rewrite import _strip_identifier_quotes def _extract_delete_table_name(statement: str) -> str | None: """Extract target table name from a DELETE statement when possible.""" match = re.match( r'^\s*delete\s+from\s+("[^"]+"|\[[^\]]+\]|`[^`]+`|\'[^\']+\'|[A-Za-z_][A-Za-z0-9_]*)', statement, flags=re.IGNORECASE, ) if not match: return None return _strip_identifier_quotes(match.group(1)) def _extract_update_table_name(statement: str) -> str | None: """Extract target table name from an UPDATE statement when possible.""" match = re.match( r'^\s*update\s+("[^"]+"|\[[^\]]+\]|`[^`]+`|\'[^\']+\'|[A-Za-z_][A-Za-z0-9_]*)\s+set\b', statement, flags=re.IGNORECASE, ) if not match: return None return _strip_identifier_quotes(match.group(1)) def _extract_update_where_clause(statement: str) -> str | None: """Extract WHERE clause text from an UPDATE statement, when present.""" trimmed = statement.strip().rstrip(";") match = re.search(r"\bwhere\b\s+(.+)$", trimmed, flags=re.IGNORECASE | re.DOTALL) if not match: return None where_clause = match.group(1).strip() return where_clause or None def _extract_update_set_clause(statement: str) -> str | None: """Extract SET clause text from an UPDATE statement.""" trimmed = statement.strip().rstrip(";") match = re.search(r"\bset\b\s+(.+?)(?:\bwhere\b\s+.+)?$", trimmed, flags=re.IGNORECASE | re.DOTALL) if not match: return None set_clause = match.group(1).strip() return set_clause or None def _split_top_level_csv(text: str) -> list[str]: """Split comma-separated SQL expressions while respecting quotes/parens.""" parts: list[str] = [] current: list[str] = [] depth = 0 in_single_quote = False in_double_quote = False escape = False for char in text: if in_single_quote: current.append(char) if escape: escape = False elif char == "\\": escape = True elif char == "'": in_single_quote = False continue if in_double_quote: current.append(char) if escape: escape = False elif char == "\\": escape = True elif char == '"': in_double_quote = False continue if char == "'": in_single_quote = True current.append(char) continue if char == '"': in_double_quote = True current.append(char) continue if char == "(": depth += 1 current.append(char) continue if char == ")" and depth > 0: depth -= 1 current.append(char) continue if char == "," and depth == 0: part = "".join(current).strip() if part: parts.append(part) current = [] continue current.append(char) tail = "".join(current).strip() if tail: parts.append(tail) return parts def _build_update_change_predicate(statement: str) -> str | None: """Build predicate matching only rows whose values will actually change.""" set_clause = _extract_update_set_clause(statement) if not set_clause: return None comparisons: list[str] = [] for assignment in _split_top_level_csv(set_clause): lhs, sep, rhs = assignment.partition("=") if not sep: continue column_expr = lhs.strip() value_expr = rhs.strip() if not column_expr or not value_expr: continue comparisons.append(f"({column_expr} IS DISTINCT FROM ({value_expr}))") if not comparisons: return None return " OR ".join(comparisons) def _update_statement_to_scope_query(statement: str) -> str | None: """Convert an UPDATE statement into a SELECT query for rows that will change.""" table_name = _extract_update_table_name(statement) if table_name is None: return None where_clause = _extract_update_where_clause(statement) change_predicate = _build_update_change_predicate(statement) predicates: list[str] = [] if where_clause: predicates.append(f"({where_clause})") if change_predicate: predicates.append(f"({change_predicate})") if predicates: return f"SELECT * FROM {table_name} WHERE {' AND '.join(predicates)}" return f"SELECT * FROM {table_name}" def _delete_statement_to_scope_query(statement: str) -> str | None: """Convert a DELETE statement into an equivalent SELECT scope query.""" scope_query = re.sub( r"^\s*delete\s+from\s+", "SELECT * FROM ", statement.strip().rstrip(";"), count=1, flags=re.IGNORECASE, ) if scope_query.lower().startswith("select * from "): return scope_query return None def _handle_sql_execution_error(error: Exception, statement: str, console: Console) -> None: """Render SQL execution errors without aborting the interactive CLI session.""" if isinstance(error, duckdb.Error): console.print(f"[red]SQL execution failed:[/red] {error}") else: console.print(f"[red]Unexpected SQL execution error:[/red] {error}") console.print(f"Skipped statement: {statement}") def _execute_statement_safely( connection: duckdb.DuckDBPyConnection, statement: str, console: Console, ) -> duckdb.DuckDBPyConnection | None: """Execute one SQL statement and return None when execution fails.""" try: return connection.execute(statement) except Exception as error: # Broad on purpose to keep the CLI alive. _handle_sql_execution_error(error, statement, console) return None def _print_sql_preview(sql_statements: list[str], console: Console) -> None: """Show statements that are about to execute.""" console.print() print_sql_statements_table(sql_statements, console) def _print_delete_preview( connection: duckdb.DuckDBPyConnection, prepared_statement: str, console: Console, ) -> None: """Show rows targeted by DELETE before execution.""" scope_query = _delete_statement_to_scope_query(prepared_statement) if scope_query is None: return table_name = _extract_delete_table_name(prepared_statement) or "target_table" limited_scope_query = f"SELECT * FROM ({scope_query}) AS delete_scope_preview LIMIT 20" try: before_count_row = connection.execute( f"SELECT COUNT(*) FROM ({scope_query}) AS delete_scope" ).fetchone() before_count = int(before_count_row[0]) if before_count_row is not None else 0 before_df = connection.execute( limited_scope_query ).df() console.print(f"[yellow]Affected rows:[/yellow] {before_count} row(s) match the filter.") if before_count == 0: console.print("[yellow]DELETE skipped:[/yellow] no rows matched the filter.") return print_dataframe_as_table(before_df, console, title=f"Rows to delete ({table_name})") except Exception as error: console.print(f"[yellow]DELETE preview skipped:[/yellow] {error}") _execute_statement_safely(connection, prepared_statement, console) return cursor = _execute_statement_safely(connection, prepared_statement, console) if cursor is None: return def _print_update_preview( connection: duckdb.DuckDBPyConnection, prepared_statement: str, console: Console, ) -> int: """Show before/after rows for UPDATE statements and return affected row count.""" scope_query = _update_statement_to_scope_query(prepared_statement) table_name = _extract_update_table_name(prepared_statement) or "target_table" if scope_query is None: cursor = _execute_statement_safely(connection, prepared_statement, console) if cursor is None or cursor.description is None: return max(cursor.rowcount, 0) if cursor is not None and cursor.rowcount is not None else 0 after_df = cursor.df() return len(after_df) limited_scope_query = f"SELECT * FROM ({scope_query}) AS update_scope_preview LIMIT 20" try: before_count_row = connection.execute( f"SELECT COUNT(*) FROM ({scope_query}) AS update_scope" ).fetchone() before_count = int(before_count_row[0]) if before_count_row is not None else 0 before_df = connection.execute(limited_scope_query).df() except Exception as error: console.print(f"[yellow]UPDATE preview skipped:[/yellow] {error}") cursor = _execute_statement_safely(connection, prepared_statement, console) if cursor is None: return 0 if cursor.description is None: return max(cursor.rowcount, 0) if cursor.rowcount is not None else 0 return len(cursor.df()) console.print(f"[yellow]Affected rows:[/yellow] {before_count} row(s) match the filter.") if before_count == 0: console.print("[yellow]UPDATE skipped:[/yellow] no rows matched the filter.") return 0 print_dataframe_as_table(before_df, console, title=f"Rows before update ({table_name})") if before_count > len(before_df): console.print(f"Showing first {len(before_df)} of {before_count} matched rows before update.") statement_without_semicolon = prepared_statement.strip().rstrip(";") update_with_returning = statement_without_semicolon if " returning " not in statement_without_semicolon.lower(): update_with_returning = f"{statement_without_semicolon} RETURNING *" cursor = _execute_statement_safely(connection, update_with_returning, console) if cursor is None: return 0 if cursor.description is None: # Fallback for engines that execute UPDATE without RETURNING result sets. rowcount = max(cursor.rowcount, 0) if cursor.rowcount is not None else 0 console.print(f"[yellow]Updated rows:[/yellow] {rowcount}") return rowcount after_df = cursor.df() after_count = len(after_df) print_dataframe_as_table(after_df.head(20), console, title=f"Rows after update ({table_name})") if after_count > 20: console.print(f"Showing first 20 of {after_count} rows after update.") return after_count