ss-tools/backend/src/plugins/translate/sql_generator.py

# #region SQLGenerator [C:4] [TYPE Module] [SEMANTICS translate,sql,generator,dialect]
# @BRIEF Dialect-aware safe SQL generation for INSERT/UPSERT operations with identifier quoting and value encoding.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [TranslationJob]
# @RELATION DEPENDS_ON -> [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.

from typing import Any, Dict, List, Optional, Tuple
from datetime import datetime, timezone
from ...core.logger import logger, belief_scope

# PostgreSQL/Greenplum dialects that support ON CONFLICT
POSTGRESQL_DIALECTS = {"postgresql", "redshift", "greenplum"}
# Dialects that use backtick or no quoting
CLICKHOUSE_DIALECTS = {"clickhouse"}


# #region _normalize_timestamp_value [C:2] [TYPE Function] [SEMANTICS translate,sql,timestamp]
# @BRIEF Detect Unix timestamp strings (seconds or millis) and convert to 'YYYY-MM-DD' for Date columns.
def _normalize_timestamp_value(value: Any) -> Optional[str]:
    """Detect Unix timestamp values and convert to date string.

    Handles:
    - Integer/float Unix timestamps (seconds or milliseconds)
    - String representations of Unix timestamps (e.g. '1726358400000.0')

    Returns 'YYYY-MM-DD' if conversion succeeds, None if value is not a timestamp.
    """
    # Try numeric conversion first
    try:
        ts = float(value)
    except (ValueError, TypeError):
        return None

    # Heuristic: Unix timestamps in seconds are ~10 digits (1e9 range for 2001-2033)
    # Unix timestamps in milliseconds are ~13 digits (1e12 range for 2001-2033)
    if 1e9 <= ts < 1e12:
        # Already in seconds
        pass
    elif 1e12 <= ts < 1e15:
        # Milliseconds — convert to seconds
        ts = ts / 1000.0
    else:
        # Not a plausible Unix timestamp
        return None

    try:
        dt = datetime.fromtimestamp(ts, tz=timezone.utc)
        return dt.strftime("%Y-%m-%d")
    except (OSError, OverflowError, ValueError):
        return None
# #endregion _normalize_timestamp_value


# #region _quote_identifier [C:2] [TYPE Function] [SEMANTICS translate,sql,quoting]
# @BRIEF Quote a SQL identifier per dialect rules — double quotes for PostgreSQL, backticks for ClickHouse.
def _quote_identifier(identifier: str, dialect: str) -> str:
    """Quote a SQL identifier per dialect rules."""
    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 CLICKHOUSE_DIALECTS:
        return f"`{cleaned}`"
    else:
        # Generic ANSI double-quote
        return f'"{cleaned}"'
# #endregion _quote_identifier


# #region _encode_sql_value [C:2] [TYPE Function] [SEMANTICS translate,sql,encoding]
# @BRIEF Encode a Python value into a SQL-safe literal for INSERT VALUES, with ClickHouse timestamp normalization.
def _encode_sql_value(value: Any, dialect: Optional[str] = None) -> str:
    """Encode a Python value into a SQL-safe literal.

    For ClickHouse dialect, attempts to detect Unix timestamp strings
    and convert them to 'YYYY-MM-DD' format for Date column compatibility.
    """
    if value is None:
        return "NULL"
    if isinstance(value, bool):
        return "TRUE" if value else "FALSE"

    # For ClickHouse: try to normalize timestamp-like string values
    if dialect in CLICKHOUSE_DIALECTS and isinstance(value, str) and value:
        normalized = _normalize_timestamp_value(value)
        if normalized:
            return f"'{normalized}'"

    if isinstance(value, (int, float)):
        return str(value)
    # String — escape single quotes by doubling them
    escaped = str(value).replace("'", "''")
    return f"'{escaped}'"
# #endregion _encode_sql_value


# #region _build_values_clause [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: Optional[str] = None) -> str:
    """Build VALUES (...) clause for multiple rows.

    NOTE: columns may be quoted (e.g. '"col"' or '`col`') for the SQL column list,
    but row dicts have UNQUOTED keys. Strip quotes before value lookup.
    """
    value_groups = []
    for row in rows:
        values = []
        for col in columns:
            # Strip quoting for value lookup — row keys are unquoted
            lookup_key = col.strip('"').strip('`').strip('[]')
            val = row.get(lookup_key)
            values.append(_encode_sql_value(val, dialect=dialect))
        value_groups.append(f"({', '.join(values)})")
    return ",\n".join(value_groups)
# #endregion _build_values_clause


# #region generate_insert_sql [C:3] [TYPE Function] [SEMANTICS translate,sql,insert]
# @BRIEF Generate a dialect-aware plain INSERT SQL statement for the given table, columns, and rows.
# @RELATION DEPENDS_ON -> [_quote_identifier]
# @RELATION DEPENDS_ON -> [_build_values_clause]
def generate_insert_sql(
    target_schema: Optional[str],
    target_table: str,
    columns: List[str],
    rows: List[Dict[str, Any]],
    dialect: Optional[str] = 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)

        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 generate_insert_sql


# #region generate_upsert_sql [C:3] [TYPE Function] [SEMANTICS translate,sql,upsert]
# @BRIEF Generate PostgreSQL dialect UPSERT SQL with ON CONFLICT DO UPDATE / DO NOTHING.
# @RELATION DEPENDS_ON -> [_quote_identifier]
# @RELATION DEPENDS_ON -> [_build_values_clause]
def generate_upsert_sql(
    target_schema: Optional[str],
    target_table: str,
    columns: List[str],
    key_columns: List[str],
    rows: List[Dict[str, Any]],
    dialect: Optional[str] = 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, dialect=dialect)

        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 generate_upsert_sql


# #region SQLGenerator [C:4] [TYPE Class] [SEMANTICS translate,sql,generator]
# @BRIEF Generate safe, dialect-appropriate SQL INSERT/UPSERT statements with batching support.
# @PRE Job has target_schema, target_table, key columns configured.
# @POST Returns generated SQL string for the target dialect.
# @SIDE_EFFECT None — pure SQL generation.
# @RELATION DEPENDS_ON -> [generate_insert_sql]
# @RELATION DEPENDS_ON -> [generate_upsert_sql]
class SQLGenerator:

    # #region SQLGenerator.generate [C:4] [TYPE Function] [SEMANTICS translate,sql,generate]
    # @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: Optional[str],
        target_table: str,
        columns: List[str],
        rows: List[Dict[str, Any]],
        key_columns: Optional[List[str]] = None,
        upsert_strategy: str = "MERGE",
    ) -> 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 POSTGRESQL_DIALECTS or dialect not in CLICKHOUSE_DIALECTS:
                # PostgreSQL and other ANSI dialects: 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,
                        dialect=dialect,
                    )
                else:
                    sql = generate_insert_sql(
                        target_schema=None,
                        target_table=table_ref,
                        columns=quoted_columns,
                        rows=rows,
                        dialect=dialect,
                    )
            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,
                )
                if use_upsert:
                    logger.reason("ClickHouse UPSERT not supported; using plain INSERT", {
                        "note": "ClickHouse does not support ON CONFLICT. Use ReplacingMergeTree for dedup.",
                    })
            else:
                # Fallback: plain INSERT
                sql = generate_insert_sql(
                    target_schema=None,
                    target_table=table_ref,
                    columns=quoted_columns,
                    rows=rows,
                    dialect=dialect,
                )

            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 [C:3] [TYPE Function] [SEMANTICS translate,sql,batch]
    # @BRIEF Generate separate INSERT statements for each row chunk (batch-safe version).
    # @RELATION DEPENDS_ON -> [SQLGenerator.generate]
    @staticmethod
    def generate_batch(
        dialect: str,
        target_schema: Optional[str],
        target_table: str,
        columns: List[str],
        rows: List[Dict[str, Any]],
        key_columns: Optional[List[str]] = None,
        upsert_strategy: str = "MERGE",
        max_rows_per_statement: int = 500,
    ) -> 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,
                )
                statements.append((sql, count))

            return statements
    # #endregion SQLGenerator.generate_batch


# #endregion SQLGenerator
# #endregion SQLGenerator