DataOps-platform/app/runner/rule_sql.py

"""Runner adapter for exact, published SQLGlot rule plans."""

from __future__ import annotations

import copy
import re

from sqlalchemy import text
from sqlglot import exp, parse_one

from app.core.data_rules.compilers.sql import (
    bound_sql_plan_hash,
    validate_bound_sql_plan,
)
from app.core.data_source.errors import DataSourceWriteOutcomeUnknown
from app.runner.nodes import NodeExecutionError


def _dialect(value):
    normalized = str(value or "").strip().lower()
    return "postgresql" if normalized == "postgres" else normalized


def _source_count_statement(plan):
    dialect = "postgres" if plan["dialect"] == "postgresql" else plan["dialect"]
    insert = parse_one(plan["statements"][0]["sql"], read=dialect)
    source_tables = list(insert.expression.find_all(exp.Table))
    if len(source_tables) != 1:
        raise NodeExecutionError("published SQL rule plan has an invalid source")
    count = exp.Select(
        expressions=[exp.Count(this=exp.Star())]
    ).from_(copy.deepcopy(source_tables[0]))
    return count.sql(dialect=dialect)


def _accepted_count_statement(plan):
    dialect = "postgres" if plan["dialect"] == "postgresql" else plan["dialect"]
    insert = parse_one(plan["statements"][0]["sql"], read=dialect)
    accepted = exp.Select(
        expressions=[exp.Count(this=exp.Star())]
    ).from_(
        exp.Subquery(
            this=copy.deepcopy(insert.expression),
            alias=exp.TableAlias(
                this=exp.Identifier(this="_dataops_accepted")
            ),
        )
    )
    statement = accepted.sql(dialect=dialect)
    if plan["dialect"] == "postgresql":
        statement = re.sub(
            r"%\(([A-Za-z_][A-Za-z0-9_]*)\)s",
            r":\1",
            statement,
        )
    return statement


def _upsert_statement(plan, key):
    if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,127}", str(key or "")) is None:
        raise NodeExecutionError("governed SQL rule idempotency key is invalid")
    dialect = "postgres" if plan["dialect"] == "postgresql" else plan["dialect"]
    insert = parse_one(plan["statements"][0]["sql"], read=dialect)
    target = insert.this
    if not isinstance(target, exp.Schema):
        raise NodeExecutionError("published SQL rule target schema is invalid")
    columns = [column.name for column in target.expressions]
    if key not in columns:
        raise NodeExecutionError(
            "governed SQL rule idempotency key is not an output field"
        )
    assignments = []
    if plan["dialect"] == "postgresql":
        for name in columns:
            if name == key:
                continue
            assignments.append(
                exp.EQ(
                    this=exp.Column(
                        this=exp.Identifier(this=name, quoted=True)
                    ),
                    expression=exp.Column(
                        this=exp.Identifier(this=name, quoted=True),
                        table=exp.Identifier(this="EXCLUDED"),
                    ),
                )
            )
        insert.set(
            "conflict",
            exp.OnConflict(
                duplicate=False,
                expressions=assignments,
                action=exp.Var(
                    this="DO UPDATE" if assignments else "DO NOTHING"
                ),
                conflict_keys=[
                    exp.Ordered(
                        this=exp.Column(
                            this=exp.Identifier(this=key, quoted=True)
                        )
                    )
                ],
            ),
        )
    else:
        for name in columns:
            assignments.append(
                exp.EQ(
                    this=exp.Column(
                        this=exp.Identifier(this=name, quoted=True)
                    ),
                    expression=exp.Anonymous(
                        this="VALUES",
                        expressions=[
                            exp.Identifier(this=name, quoted=True)
                        ],
                    ),
                )
            )
        insert.set(
            "conflict",
            exp.OnConflict(
                duplicate=True,
                expressions=assignments,
                action=exp.Var(this="UPDATE"),
            ),
        )
    statement = insert.sql(dialect=dialect)
    if plan["dialect"] == "postgresql":
        statement = re.sub(
            r"%\(([A-Za-z_][A-Za-z0-9_]*)\)s",
            r":\1",
            statement,
        )
    return statement


def _target_relation(plan):
    dialect = "postgres" if plan["dialect"] == "postgresql" else plan["dialect"]
    insert = parse_one(plan["statements"][0]["sql"], read=dialect)
    target = insert.this
    if not isinstance(target, exp.Schema) or not isinstance(
        target.this, exp.Table
    ):
        raise NodeExecutionError("published SQL rule target schema is invalid")
    table = target.this
    if not table.db or not table.name:
        raise NodeExecutionError(
            "published SQL rule target must be schema-qualified"
        )
    return table.db, table.name


def _attest_upsert_key(connection, plan, key):
    """Prove the exact server-side unique-key contract before any write."""

    schema_name, table_name = _target_relation(plan)
    if plan["dialect"] == "postgresql":
        query = text(
            "SELECT tc.constraint_name, tc.constraint_type, "
            "string_agg(kcu.column_name, ',' ORDER BY kcu.ordinal_position) "
            "AS columns "
            "FROM information_schema.table_constraints tc "
            "JOIN information_schema.key_column_usage kcu "
            "ON kcu.constraint_catalog = tc.constraint_catalog "
            "AND kcu.constraint_schema = tc.constraint_schema "
            "AND kcu.constraint_name = tc.constraint_name "
            "WHERE tc.table_schema = :schema_name "
            "AND tc.table_name = :table_name "
            "AND tc.constraint_type IN ('PRIMARY KEY', 'UNIQUE') "
            "GROUP BY tc.constraint_name, tc.constraint_type"
        )
    else:
        query = text(
            "SELECT index_name AS constraint_name, "
            "CASE WHEN index_name = 'PRIMARY' THEN 'PRIMARY KEY' ELSE 'UNIQUE' END "
            "AS constraint_type, "
            "GROUP_CONCAT(column_name ORDER BY seq_in_index SEPARATOR ',') "
            "AS columns "
            "FROM information_schema.statistics "
            "WHERE table_schema = :schema_name "
            "AND table_name = :table_name AND non_unique = 0 "
            "GROUP BY index_name"
        )
    rows = (
        connection.execute(
            query,
            {"schema_name": schema_name, "table_name": table_name},
        )
        .mappings()
        .all()
    )

    def columns(row):
        value = row["columns"]
        if isinstance(value, str):
            return value.split(",") if value else []
        return list(value or [])

    unique_keys = [columns(row) for row in rows]
    if [key] not in unique_keys:
        raise NodeExecutionError(
            "governed SQL rule idempotency key is not an exact unique key"
        )
    if plan["dialect"] == "mysql" and any(
        unique_key != [key] for unique_key in unique_keys
    ):
        raise NodeExecutionError(
            "MySQL upsert target has an alternate unique collision path"
        )


class SqlGlotRulePlanAdapter:
    """Execute a validated INSERT plan and count its source in one transaction."""

    def __init__(self, manager):
        self.manager = manager

    def execute(
        self,
        *,
        plan,
        node,
        parameters,
        write_authorized,
    ):
        try:
            normalized = validate_bound_sql_plan(plan)
        except ValueError as exc:
            raise NodeExecutionError("published SQL rule plan is invalid") from exc
        config = node.get("config") or {}
        if config.get("execution_plan_hash") != bound_sql_plan_hash(normalized):
            raise NodeExecutionError("published SQL rule plan hash does not match")
        if config.get("rule_version_id") != normalized["rule_version_id"]:
            raise NodeExecutionError("published SQL rule plan rule id does not match")
        idempotency = node.get("idempotency")
        if (
            node.get("type") != "rule.apply"
            or node.get("purpose") != "write"
            or not write_authorized
            or not isinstance(idempotency, dict)
            or idempotency.get("strategy") != "upsert"
            or not str(idempotency.get("key") or "").strip()
        ):
            raise NodeExecutionError(
                "governed write authorization and idempotency are required"
            )
        if parameters not in ({}, None):
            raise NodeExecutionError(
                "bound SQL rule plans do not accept unbound runtime parameters"
            )
        definition = self.manager.definitions.get(
            normalized["data_source_uid"]
        )
        if definition is None:
            raise NodeExecutionError("published SQL rule datasource was not found")
        if _dialect(getattr(definition, "database_type", None)) != normalized[
            "dialect"
        ]:
            raise NodeExecutionError(
                "published SQL rule datasource dialect does not match"
            )
        datasource_capabilities = dict(
            getattr(definition, "extra_properties", {}) or {}
        ).get("sql_rule_capabilities")
        if datasource_capabilities != normalized["capabilities"]:
            raise NodeExecutionError(
                "published SQL rule datasource capabilities do not match"
            )

        statement = normalized["statements"][0]
        try:
            with self.manager.connect(
                normalized["data_source_uid"],
                purpose="dataflow_write",
            ) as connection:
                rows_in = int(
                    connection.execute(
                        text(_source_count_statement(normalized)),
                        {},
                    ).scalar_one()
                    or 0
                )
                rows_out = int(
                    connection.execute(
                        text(_accepted_count_statement(normalized)),
                        statement["parameters"],
                    ).scalar_one()
                    or 0
                )
                _attest_upsert_key(
                    connection,
                    normalized,
                    idempotency["key"],
                )
                result = connection.execute(
                    text(
                        _upsert_statement(
                            normalized,
                            idempotency["key"],
                        )
                    ),
                    statement["parameters"],
                )
                if result.rowcount is not None and int(result.rowcount) < 0:
                    raise NodeExecutionError(
                        "governed SQL rule returned an invalid write count"
                    )
                metrics = {
                    "rows_in": rows_in,
                    "rows_out": rows_out,
                    "rows_rejected": max(0, rows_in - rows_out),
                    "commit_outcome": "committed",
                }
        except DataSourceWriteOutcomeUnknown as exc:
            raise NodeExecutionError(
                "governed SQL rule commit outcome is unknown",
                commit_outcome="unknown",
            ) from exc
        except NodeExecutionError:
            raise
        except Exception as exc:
            raise NodeExecutionError(
                "governed SQL rule write failed",
                commit_outcome="not_committed",
            ) from exc
        return metrics


class SqlGlotQualityPlanAdapter:
    """Evaluate one attested SQLGlot plan without executing its write."""

    def __init__(self, manager):
        self.manager = manager

    def execute(
        self,
        *,
        plan,
        node,
        parameters,
        write_authorized,
    ):
        try:
            normalized = validate_bound_sql_plan(plan)
        except ValueError as exc:
            raise NodeExecutionError(
                "published SQL quality plan is invalid"
            ) from exc
        config = node.get("config") or {}
        if config.get("execution_plan_hash") != bound_sql_plan_hash(normalized):
            raise NodeExecutionError(
                "published SQL quality plan hash does not match"
            )
        if config.get("rule_version_id") != normalized["rule_version_id"]:
            raise NodeExecutionError(
                "published SQL quality plan rule id does not match"
            )
        if (
            node.get("type") != "quality.check"
            or node.get("purpose") != "read"
            or write_authorized
            or node.get("idempotency") is not None
        ):
            raise NodeExecutionError("quality check must be read only")
        if parameters not in ({}, None):
            raise NodeExecutionError(
                "bound SQL quality plans do not accept runtime parameters"
            )
        definition = self.manager.definitions.get(
            normalized["data_source_uid"]
        )
        if definition is None:
            raise NodeExecutionError(
                "published SQL quality datasource was not found"
            )
        if _dialect(getattr(definition, "database_type", None)) != normalized[
            "dialect"
        ]:
            raise NodeExecutionError(
                "published SQL quality datasource dialect does not match"
            )
        datasource_capabilities = dict(
            getattr(definition, "extra_properties", {}) or {}
        ).get("sql_rule_capabilities")
        if datasource_capabilities != normalized["capabilities"]:
            raise NodeExecutionError(
                "published SQL quality datasource capabilities do not match"
            )

        statement = normalized["statements"][0]
        try:
            with self.manager.connect(
                normalized["data_source_uid"],
                purpose="dataflow_read",
            ) as connection:
                rows_in = int(
                    connection.execute(
                        text(_source_count_statement(normalized)),
                        {},
                    ).scalar_one()
                    or 0
                )
                rows_out = int(
                    connection.execute(
                        text(_accepted_count_statement(normalized)),
                        statement["parameters"],
                    ).scalar_one()
                    or 0
                )
        except NodeExecutionError:
            raise
        except Exception as exc:
            raise NodeExecutionError(
                "governed SQL quality check failed",
                commit_outcome="not_applicable",
            ) from exc
        rows_rejected = max(0, rows_in - rows_out)
        return {
            "rows_in": rows_in,
            "rows_out": rows_out,
            "rows_rejected": rows_rejected,
            "rows_quarantined": 0,
            "violation_count": rows_rejected,
            "violations": (
                [
                    {
                        "step_id": "quality_check",
                        "count": rows_rejected,
                    }
                ]
                if rows_rejected
                else []
            ),
            "commit_outcome": "not_applicable",
        }


__all__ = ["SqlGlotQualityPlanAdapter", "SqlGlotRulePlanAdapter"]