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


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
                )
                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:
    """Explicit fail-closed placeholder until a read-only quality plan exists."""

    def execute(self, **_kwargs):
        raise NodeExecutionError(
            "quality_check requires an explicit read-only compiled plan"
        )


__all__ = ["SqlGlotQualityPlanAdapter", "SqlGlotRulePlanAdapter"]