432 lines
18 KiB
Python
432 lines
18 KiB
Python
# #region Plugin.SqlGenerator.SQLGenerator [C:3] [TYPE Module] [SEMANTICS clickhouse, translate, sql, insert, generate]
|
|
# @defgroup Translate Module group.
|
|
# @BRIEF Dialect-aware safe SQL generation for INSERT/UPSERT operations.
|
|
# @LAYER Domain
|
|
# @RELATION DEPENDS_ON -> [Models.Translate.TranslationJob]
|
|
# @RELATION DEPENDS_ON -> [Models.Translate.TranslationRun]
|
|
# @PRE Job has target_schema and target_table configured. Dialect is one of supported SUPPORTED_DIALECTS.
|
|
# @POST Returns safe SQL strings for the target dialect.
|
|
# @SIDE_EFFECT None — pure code generation.
|
|
# @RATIONALE Dialect-aware SQL uses ON CONFLICT for PostgreSQL; plain INSERT for ClickHouse with documented limitations.
|
|
# @REJECTED UPDATE statements — source is append-only; UPSERT covers overwrite case.
|
|
# @REJECTED ORM-based insert bypasses Superset's SQL Lab audit trail.
|
|
|
|
from datetime import UTC, datetime
|
|
from decimal import Decimal
|
|
import math
|
|
from typing import Any
|
|
|
|
from ...core.logger import belief_scope, logger
|
|
|
|
# PostgreSQL-family dialects that support ON CONFLICT (UPSERT)
|
|
POSTGRESQL_DIALECTS = {"postgresql", "redshift"}
|
|
# Dialects that support UPSERT via ON CONFLICT (subset of all supported dialects)
|
|
UPSERT_SUPPORTED_DIALECTS = {"postgresql", "redshift"}
|
|
# Dialects that use backtick quoting
|
|
CLICKHOUSE_DIALECTS = {"clickhouse", "clickhousedb"}
|
|
# Dialects that use backtick quoting (MySQL native, ClickHouse)
|
|
BACKTICK_DIALECTS = {"clickhouse", "clickhousedb", "mysql"}
|
|
|
|
|
|
# #region Plugin.SqlGenerator.NormalizeTimestampValue [C:2] [TYPE Function] [SEMANTICS translate,sql,timestamp]
|
|
# @BRIEF Convert Unix timestamp values to a ClickHouse temporal literal only when the target type is explicit.
|
|
# @RATIONALE Value magnitude alone cannot distinguish a timestamp from a numeric primary key, so bare
|
|
# numeric values are preserved unless the target column type is explicitly temporal. String
|
|
# timestamp values (as delivered by Superset chart data) are still normalized without type
|
|
# metadata. Epoch 0 with an explicit Date type maps to 1970-01-01.
|
|
# @REJECTED Dialect-wide timestamp heuristics — they corrupt legitimate IDs and discard DateTime precision.
|
|
def _normalize_timestamp_value(value: Any, target_type: str | None = None) -> str | None:
|
|
"""Convert a Unix timestamp to a ClickHouse temporal literal.
|
|
|
|
Rules:
|
|
- An explicit temporal target_type (date/datetime*) always triggers normalization,
|
|
including epoch 0 (1970-01-01) and millisecond precision.
|
|
- A non-temporal or absent target_type preserves numeric values unchanged — a bare
|
|
integer like 1705320000 is indistinguishable from a primary key.
|
|
- String values that parse as Unix timestamps are still normalized without a target
|
|
type: Superset chart data delivers ClickHouse temporal values as timestamp strings
|
|
with no column-type metadata.
|
|
"""
|
|
if target_type:
|
|
normalized_type = target_type.strip().lower()
|
|
if not normalized_type.startswith(("date", "datetime")):
|
|
return None
|
|
elif not isinstance(value, str):
|
|
return None
|
|
else:
|
|
normalized_type = "date"
|
|
try:
|
|
ts = float(value)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
if 0 <= ts < 1e12:
|
|
pass
|
|
elif 1e12 <= ts < 1e15:
|
|
ts /= 1000.0
|
|
else:
|
|
return None
|
|
try:
|
|
dt = datetime.fromtimestamp(ts, tz=UTC)
|
|
except (OSError, OverflowError, ValueError):
|
|
return None
|
|
if normalized_type == "date":
|
|
return dt.strftime("%Y-%m-%d")
|
|
if normalized_type.startswith("datetime64"):
|
|
return dt.strftime("%Y-%m-%d %H:%M:%S.%f").rstrip("0").rstrip(".")
|
|
return dt.strftime("%Y-%m-%d %H:%M:%S")
|
|
# #endregion Plugin.SqlGenerator.NormalizeTimestampValue
|
|
|
|
|
|
# #region Plugin.SqlGenerator.QuoteIdentifier [C:4] [TYPE Function]
|
|
# @BRIEF Quote an identifier per dialect rules. PostgreSQL uses double quotes; ClickHouse uses backticks.
|
|
# @PRE identifier is a non-empty string.
|
|
# @POST Returns safely quoted identifier.
|
|
# @RATIONALE Native quoting per dialect prevents syntax errors.
|
|
# PostgreSQL/Redshift require double quotes for case-sensitive/reserved identifiers.
|
|
# ClickHouse/MySQL use backticks natively.
|
|
# ANSI double-quote fallback works with any modern SQL engine.
|
|
# @REJECTED Single-quote quoting — would break all dialects for identifiers.
|
|
# No-op (pass-through) quoting — allows SQL injection via identifier names.
|
|
def _quote_identifier(identifier: str, dialect: str) -> str:
|
|
"""Quote a SQL identifier per dialect rules.
|
|
|
|
PostgreSQL/Redshift uses double quotes; ClickHouse/MySQL uses backticks.
|
|
Unknown dialects default to ANSI double quotes as a safe fallback.
|
|
"""
|
|
if not identifier:
|
|
return identifier
|
|
# Remove any existing quotes to avoid double-quoting
|
|
cleaned = identifier.strip().strip('"').strip('`').strip('[]')
|
|
if dialect in POSTGRESQL_DIALECTS:
|
|
return f'"{cleaned}"'
|
|
elif dialect in BACKTICK_DIALECTS:
|
|
return f"`{cleaned}`"
|
|
else:
|
|
# Generic ANSI double-quote
|
|
return f'"{cleaned}"'
|
|
# #endregion Plugin.SqlGenerator.QuoteIdentifier
|
|
|
|
|
|
# #region Plugin.SqlGenerator.EncodeSqlValue [C:2] [TYPE Function] [SEMANTICS translate,sql,encoding]
|
|
# @BRIEF Encode a Python value into a SQL-safe literal with optional target-column type metadata.
|
|
def _encode_sql_value(
|
|
value: Any,
|
|
dialect: str | None = None,
|
|
target_type: str | None = None,
|
|
) -> str:
|
|
"""Encode a Python value into a SQL-safe literal."""
|
|
if value is None:
|
|
return "NULL"
|
|
if isinstance(value, bool):
|
|
return "TRUE" if value else "FALSE"
|
|
if isinstance(value, float) and not math.isfinite(value):
|
|
raise ValueError("NaN and infinity are not supported in SQL INSERT values")
|
|
if isinstance(value, Decimal):
|
|
return str(value)
|
|
|
|
if dialect in CLICKHOUSE_DIALECTS:
|
|
normalized = _normalize_timestamp_value(value, target_type)
|
|
if normalized:
|
|
return f"'{normalized}'"
|
|
|
|
if isinstance(value, (int, float)):
|
|
return str(value)
|
|
escaped = str(value).replace("'", "''")
|
|
return f"'{escaped}'"
|
|
# #endregion Plugin.SqlGenerator.EncodeSqlValue
|
|
|
|
|
|
# #region Plugin.SqlGenerator.BuildValuesClause [C:2] [TYPE Function] [SEMANTICS translate,sql,values]
|
|
# @BRIEF Build a VALUES clause for multiple rows with per-column value encoding and dialect-aware quoting.
|
|
def _build_values_clause(
|
|
columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
dialect: str | None = None,
|
|
column_types: dict[str, str] | None = None,
|
|
) -> str:
|
|
"""Build VALUES (...) clause for multiple rows."""
|
|
value_groups = []
|
|
for row in rows:
|
|
values = []
|
|
for col in columns:
|
|
lookup_key = col.strip('"').strip('`').strip('[]')
|
|
val = row.get(lookup_key)
|
|
target_type = (column_types or {}).get(lookup_key)
|
|
values.append(_encode_sql_value(val, dialect=dialect, target_type=target_type))
|
|
value_groups.append(f"({', '.join(values)})")
|
|
return ",\n".join(value_groups)
|
|
# #endregion Plugin.SqlGenerator.BuildValuesClause
|
|
|
|
|
|
# #region Plugin.SqlGenerator.GenerateInsertSql [C:3] [TYPE Function] [SEMANTICS translate,sql,insert]
|
|
# @ingroup Translate
|
|
# @BRIEF Generate a dialect-aware plain INSERT SQL statement for the given table, columns, and rows.
|
|
# @RELATION DEPENDS_ON -> [Plugin.SqlGenerator.QuoteIdentifier]
|
|
# @RELATION DEPENDS_ON -> [Plugin.SqlGenerator.BuildValuesClause]
|
|
def generate_insert_sql(
|
|
target_schema: str | None,
|
|
target_table: str,
|
|
columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
dialect: str | None = None,
|
|
column_types: dict[str, str] | None = None,
|
|
) -> str:
|
|
"""Generate a plain INSERT SQL."""
|
|
with belief_scope("generate_insert_sql"):
|
|
if not target_table:
|
|
raise ValueError("target_table is required for INSERT SQL generation")
|
|
if not columns:
|
|
raise ValueError("At least one column is required for INSERT SQL generation")
|
|
if not rows:
|
|
raise ValueError("At least one row is required for INSERT SQL generation")
|
|
|
|
col_list = ", ".join(columns)
|
|
values = _build_values_clause(columns, rows, dialect=dialect, column_types=column_types)
|
|
|
|
table_ref = target_table
|
|
if target_schema:
|
|
table_ref = f"{target_schema}.{target_table}"
|
|
|
|
sql = f"INSERT INTO {table_ref} ({col_list})\nVALUES\n{values};"
|
|
return sql
|
|
# #endregion Plugin.SqlGenerator.GenerateInsertSql
|
|
|
|
|
|
# #region Plugin.SqlGenerator.GenerateUpsertSql [C:4] [TYPE Function]
|
|
# @ingroup Translate
|
|
# @BRIEF Generate PostgreSQL dialect UPSERT SQL with ON CONFLICT DO UPDATE.
|
|
# @PRE dialect is postgresql-compatible. target_table, columns, key_columns are non-empty.
|
|
# @POST Returns UPSERT SQL string or raises ValueError.
|
|
# @RATIONALE ON CONFLICT DO UPDATE is the standard PostgreSQL upsert syntax.
|
|
# The DO NOTHING variant covers the pure-insert-without-overwrite case
|
|
# when all columns are conflict keys. EXCLUDED references the proposed row.
|
|
# @REJECTED MERGE INTO (SQL standard) — PostgreSQL only supports it via ON CONFLICT.
|
|
# Custom UPDATE+INSERT in a transaction — ON CONFLICT is atomic and faster.
|
|
|
|
def generate_upsert_sql(
|
|
target_schema: str | None,
|
|
target_table: str,
|
|
columns: list[str],
|
|
key_columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
column_types: dict[str, str] | None = None,
|
|
) -> str:
|
|
"""Generate INSERT ... ON CONFLICT (key_cols) DO UPDATE SET ... SQL."""
|
|
with belief_scope("generate_upsert_sql"):
|
|
if not target_table:
|
|
raise ValueError("target_table is required for UPSERT SQL generation")
|
|
if not columns:
|
|
raise ValueError("At least one column is required for UPSERT SQL generation")
|
|
if not key_columns:
|
|
raise ValueError("key_columns are required for UPSERT SQL generation")
|
|
if not rows:
|
|
raise ValueError("At least one row is required for UPSERT SQL generation")
|
|
|
|
col_list = ", ".join(columns)
|
|
key_list = ", ".join(key_columns)
|
|
values = _build_values_clause(columns, rows, column_types=column_types)
|
|
|
|
table_ref = target_table
|
|
if target_schema:
|
|
table_ref = f"{target_schema}.{target_table}"
|
|
|
|
# Build SET clause: exclude key columns from update
|
|
update_cols = [c for c in columns if c not in key_columns]
|
|
if not update_cols:
|
|
# If only key columns, use DO NOTHING
|
|
conflict_action = "DO NOTHING"
|
|
else:
|
|
set_parts = [f"{col} = EXCLUDED.{col}" for col in update_cols]
|
|
conflict_action = "DO UPDATE SET\n" + ",\n".join(set_parts)
|
|
|
|
sql = (
|
|
f"INSERT INTO {table_ref} ({col_list})\n"
|
|
f"VALUES\n"
|
|
f"{values}\n"
|
|
f"ON CONFLICT ({key_list}) {conflict_action};"
|
|
)
|
|
return sql
|
|
# #endregion Plugin.SqlGenerator.GenerateUpsertSql
|
|
|
|
|
|
# #region Plugin.SqlGenerator.SQLGenerator.Class [C:3] [TYPE Class]
|
|
# @defgroup Translate Module group.
|
|
# @BRIEF Generate safe, dialect-appropriate SQL INSERT/UPSERT statements.
|
|
# @PRE Job has target_schema, target_table, key columns configured.
|
|
# @POST Returns generated SQL string for the target dialect.
|
|
class SQLGenerator:
|
|
|
|
# region SQLGenerator.generate [TYPE Function]
|
|
# @BRIEF: Generate SQL for a set of rows, detecting dialect from the job configuration.
|
|
# @PRE dialect is a supported database dialect. columns list is non-empty. rows is non-empty.
|
|
# @POST Returns tuple of (sql_string, statement_count).
|
|
# @SIDE_EFFECT None — pure SQL generation.
|
|
@staticmethod
|
|
def generate(
|
|
dialect: str,
|
|
target_schema: str | None,
|
|
target_table: str,
|
|
columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
key_columns: list[str] | None = None,
|
|
upsert_strategy: str = "MERGE",
|
|
column_types: dict[str, str] | None = None,
|
|
) -> tuple[str, int]:
|
|
"""Generate dialect-appropriate INSERT/UPSERT SQL.
|
|
|
|
Args:
|
|
dialect: Target database dialect (e.g. 'postgresql', 'clickhouse').
|
|
target_schema: Optional schema name.
|
|
target_table: Target table name.
|
|
columns: List of column names to insert.
|
|
rows: List of row dicts with column values.
|
|
key_columns: Key columns for conflict resolution (UPSERT).
|
|
upsert_strategy: 'MERGE' (UPSERT), 'INSERT' (plain INSERT).
|
|
|
|
Returns:
|
|
Tuple of (sql_string, row_count).
|
|
"""
|
|
with belief_scope("SQLGenerator.generate"):
|
|
logger.reason("Generating SQL", {
|
|
"dialect": dialect,
|
|
"schema": target_schema,
|
|
"table": target_table,
|
|
"columns": len(columns),
|
|
"rows": len(rows),
|
|
"strategy": upsert_strategy,
|
|
})
|
|
|
|
# Validate inputs
|
|
if not target_table:
|
|
raise ValueError("target_table is required")
|
|
if not columns:
|
|
raise ValueError("At least one column is required")
|
|
if not rows:
|
|
raise ValueError("At least one row is required")
|
|
|
|
# Build fully qualified table reference
|
|
table_ref = target_table
|
|
if target_schema:
|
|
quoted_schema = _quote_identifier(target_schema, dialect)
|
|
quoted_table = _quote_identifier(target_table, dialect)
|
|
table_ref = f"{quoted_schema}.{quoted_table}"
|
|
else:
|
|
table_ref = _quote_identifier(target_table, dialect)
|
|
|
|
# Quote columns per dialect
|
|
quoted_columns = [_quote_identifier(c, dialect) for c in columns]
|
|
quoted_key_columns = (
|
|
[_quote_identifier(k, dialect) for k in key_columns]
|
|
if key_columns
|
|
else []
|
|
)
|
|
|
|
# Generate SQL per dialect and strategy
|
|
use_upsert = upsert_strategy.upper() == "MERGE" and key_columns
|
|
|
|
if dialect in UPSERT_SUPPORTED_DIALECTS:
|
|
# PostgreSQL/Redshift: support UPSERT via ON CONFLICT
|
|
if use_upsert:
|
|
sql = generate_upsert_sql(
|
|
target_schema=None,
|
|
target_table=table_ref,
|
|
columns=quoted_columns,
|
|
key_columns=quoted_key_columns,
|
|
rows=rows,
|
|
column_types=column_types,
|
|
)
|
|
else:
|
|
sql = generate_insert_sql(
|
|
target_schema=None,
|
|
target_table=table_ref,
|
|
columns=quoted_columns,
|
|
rows=rows,
|
|
dialect=dialect,
|
|
column_types=column_types,
|
|
)
|
|
elif dialect in CLICKHOUSE_DIALECTS:
|
|
# ClickHouse: plain INSERT, no ON CONFLICT support
|
|
sql = generate_insert_sql(
|
|
target_schema=None,
|
|
target_table=table_ref,
|
|
columns=quoted_columns,
|
|
rows=rows,
|
|
dialect=dialect,
|
|
column_types=column_types,
|
|
)
|
|
if use_upsert:
|
|
logger.reason("ClickHouse UPSERT not supported; using plain INSERT", {
|
|
"note": "ClickHouse does not support ON CONFLICT. Use ReplacingMergeTree for dedup.",
|
|
})
|
|
else:
|
|
# Other dialects (MySQL, MSSQL, Snowflake, etc.): plain INSERT
|
|
# Do NOT generate ON CONFLICT — most dialects don't support it
|
|
sql = generate_insert_sql(
|
|
target_schema=None,
|
|
target_table=table_ref,
|
|
columns=quoted_columns,
|
|
rows=rows,
|
|
dialect=dialect,
|
|
column_types=column_types,
|
|
)
|
|
if use_upsert:
|
|
logger.reason(
|
|
f"UPSERT not supported for dialect '{dialect}'; using plain INSERT", {
|
|
"note": "Use INSERT strategy explicitly to silence this warning.",
|
|
})
|
|
|
|
logger.reflect("SQL generated", {
|
|
"dialect": dialect,
|
|
"row_count": len(rows),
|
|
"sql_length": len(sql),
|
|
})
|
|
return sql, len(rows)
|
|
# endregion SQLGenerator.generate
|
|
|
|
# region SQLGenerator.generate_batch [TYPE Function]
|
|
# @BRIEF: Generate separate INSERT statements for each row (batch-safe version).
|
|
# @PRE Same as generate().
|
|
# @POST Returns list of (sql_string, row_index) tuples.
|
|
@staticmethod
|
|
def generate_batch(
|
|
dialect: str,
|
|
target_schema: str | None,
|
|
target_table: str,
|
|
columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
key_columns: list[str] | None = None,
|
|
upsert_strategy: str = "MERGE",
|
|
max_rows_per_statement: int = 500,
|
|
column_types: dict[str, str] | None = None,
|
|
) -> list[tuple[str, int]]:
|
|
"""Generate SQL in batches, splitting large row sets into multiple statements.
|
|
|
|
Returns:
|
|
List of (sql_string, row_count) tuples.
|
|
"""
|
|
with belief_scope("SQLGenerator.generate_batch"):
|
|
if not rows:
|
|
return []
|
|
|
|
statements = []
|
|
for i in range(0, len(rows), max_rows_per_statement):
|
|
chunk = rows[i:i + max_rows_per_statement]
|
|
sql, count = SQLGenerator.generate(
|
|
dialect=dialect,
|
|
target_schema=target_schema,
|
|
target_table=target_table,
|
|
columns=columns,
|
|
rows=chunk,
|
|
key_columns=key_columns,
|
|
upsert_strategy=upsert_strategy,
|
|
column_types=column_types,
|
|
)
|
|
statements.append((sql, count))
|
|
|
|
return statements
|
|
# endregion SQLGenerator.generate_batch
|
|
|
|
|
|
# #endregion Plugin.SqlGenerator.SQLGenerator.Class
|
|
# #endregion Plugin.SqlGenerator.SQLGenerator
|