"""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"]