DataOps-platform/tests/integration/test_data_rule_sql_execution.py

from __future__ import annotations

from contextlib import contextmanager

import pytest
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec
from app.core.data_rules.execution_contracts import canonical_schema_hash
from app.runner.nodes import NodeExecutionError

CASES = [
    (
        "postgresql",
        "postgresql+psycopg2://source_reader:source-test-password@127.0.0.1:25432/acceptance",
        "public",
        "C",
        "posix",
    ),
    (
        "mysql",
        "mysql+pymysql://source_reader:source-test-password@127.0.0.1:23306/acceptance",
        "acceptance",
        "utf8mb4_0900_bin",
        "icu",
    ),
]


class Definition:
    def __init__(self, dialect, capabilities):
        self.database_type = dialect
        self.extra_properties = {"sql_rule_capabilities": capabilities}


class Definitions:
    def __init__(self, definition):
        self.definition = definition

    def get(self, _uid):
        return self.definition


class DirectManager:
    def __init__(self, engine, definition):
        self.engine = engine
        self.definitions = Definitions(definition)

    @contextmanager
    def connect(self, _uid, purpose):
        assert purpose == "dataflow_write"
        with self.engine.connect() as connection:
            transaction = connection.begin()
            try:
                yield connection
                transaction.commit()
            except Exception:
                transaction.rollback()
                raise


class PlanRepository:
    def __init__(self, idempotency, context):
        self.idempotency = idempotency
        self.context = context
        self.record = None

    def load_bound_compile_context(self, **_ids):
        return self.context

    def persist_bound_component_plan(self, **kwargs):
        compiled = kwargs["compiled"]
        plan = compiled["plan"]
        self.record = {
            "component_binding_id": kwargs["component_binding_id"],
            "rule_version_id": kwargs["rule_version_id"],
            "backend": compiled["backend"],
            "compiler_version": compiled["compiler_version"],
            "plan": compiled["plan"],
            "plan_hash": compiled["plan_hash"],
            "schema_hashes": {
                "rule_spec_hash": plan["rule_spec_hash"],
                "input_schema_snapshot_id": plan["input_schema_snapshot_id"],
                "input_schema_hash": plan["input_schema_hash"],
                "output_schema_snapshot_id": plan["output_schema_snapshot_id"],
                "output_schema_hash": plan["output_schema_hash"],
            },
            "canonical_rule_spec_hash": plan["rule_spec_hash"],
            "canonical_input_schema_snapshot_id": plan[
                "input_schema_snapshot_id"
            ],
            "canonical_input_schema_hash": plan["input_schema_hash"],
            "canonical_output_schema_snapshot_id": plan[
                "output_schema_snapshot_id"
            ],
            "canonical_output_schema_hash": plan["output_schema_hash"],
            "plan_status": kwargs["status"],
            "rule_status": "published",
            "component_kind": "rule.apply",
            "binding_idempotency": self.idempotency,
        }
        return {
            "id": new_governance_uid(),
            "status": kwargs["status"],
            "plan_hash": compiled["plan_hash"],
        }

    def publish_with_evidence(self, plan_hash, evidence):
        assert self.record is not None
        assert self.record["plan_status"] == "compiled"
        assert self.record["plan_hash"] == plan_hash
        assert evidence["commit_outcome"] == "committed"
        assert evidence["rows_in"] >= evidence["rows_out"]
        self.record["plan_status"] = "tested"
        self.record["plan_status"] = "published"

    def load(self, **_kwargs):
        return dict(self.record)


def _snapshot(schema_ref):
    fields = [
        {"name": "customer_id", "type": "integer", "nullable": False},
        {"name": "name", "type": "string", "nullable": True},
        {"name": "mobile", "type": "string", "nullable": True},
    ]
    return {
        "id": new_governance_uid(),
        "schema_ref": schema_ref,
        "schema_hash": canonical_schema_hash(fields),
        "fields": fields,
        "source_revision": "task4:integration",
    }


@pytest.mark.parametrize(
    ("dialect", "url", "schema_name", "collation", "regex_engine"), CASES
)
def test_bound_rule_compiles_publishes_executes_and_rejects_tampering(
    dialect, url, schema_name, collation, regex_engine
):
    from app.core.data_rules.compilers import CompilerRegistry
    from app.core.data_rules.compilers.sql import SqlGlotRuleCompiler
    from app.core.data_rules.release import BoundSqlPlanService
    from app.runner.rule_sql import SqlGlotRulePlanAdapter
    from app.runner.rules import RulePlanExecutor

    engine = create_engine(url, pool_pre_ping=True)
    source_name = "task4_rule_source"
    target_name = "task4_rule_target"
    source_ref = f"{schema_name}.{source_name}"
    target_ref = f"{schema_name}.{target_name}"
    capabilities = {
        "dialect": dialect,
        "timezone": "Asia/Shanghai",
        "collation": collation,
        "rounding_mode": "half_away_from_zero",
        "regex_engine": regex_engine,
    }
    datasource_uid = new_governance_uid()
    input_schema = _snapshot("bd:task4:raw")
    output_schema = _snapshot("bd:task4:clean")
    input_binding = {
        "id": new_governance_uid(),
        "data_source_uid": datasource_uid,
        "object_kind": "table",
        "object_ref": source_ref,
        "schema_snapshot_id": input_schema["id"],
        "access_mode": "read",
        "dialect": dialect,
        "write_mode": "append",
    }
    output_binding = {
        "id": new_governance_uid(),
        "data_source_uid": datasource_uid,
        "object_kind": "table",
        "object_ref": target_ref,
        "schema_snapshot_id": output_schema["id"],
        "access_mode": "write",
        "dialect": dialect,
        "write_mode": "append",
    }
    spec = validate_rule_spec(
        {
            "schema_version": "2.0",
            "rule_uid": new_governance_uid(),
            "name": "task4_real_sql",
            "input_schema_ref": input_schema["schema_ref"],
            "output_schema_ref": output_schema["schema_ref"],
            "steps": [
                {
                    "id": "trim_name",
                    "op": "normalize_text",
                    "column": "name",
                    "trim": True,
                },
                {
                    "id": "mobile_format",
                    "op": "assert",
                    "expression": "matches(mobile, '^[0-9]{11}$')",
                    "on_failure": "reject",
                    "severity": "error",
                },
            ],
            "null_policy": "explicit",
            "timezone": "Asia/Shanghai",
        }
    )
    rule = {
        "id": new_governance_uid(),
        "status": "published",
        "rule_spec": spec,
        "spec_hash": rule_spec_hash(spec),
    }

    try:
        with engine.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {target_name}"))
            connection.execute(text(f"DROP TABLE IF EXISTS {source_name}"))
            connection.execute(
                text(
                    f"CREATE TABLE {source_name} ("
                    "customer_id BIGINT PRIMARY KEY, "
                    "name VARCHAR(100), mobile VARCHAR(30))"
                )
            )
            connection.execute(
                text(
                    f"CREATE TABLE {target_name} ("
                    "customer_id BIGINT PRIMARY KEY, "
                    "name VARCHAR(100), mobile VARCHAR(30))"
                )
            )
            connection.execute(
                text(
                    f"INSERT INTO {source_name} "
                    "(customer_id, name, mobile) VALUES "
                    "(1, ' Alice ', '13800138000'), "
                    "(2, ' Bad ', 'not-a-mobile')"
                )
            )

        component_binding_id = new_governance_uid()
        idempotency = {
            "strategy": "upsert",
            "key": "customer_id",
        }
        repository = PlanRepository(
            idempotency,
            {
                "component_binding": {
                    "id": component_binding_id,
                    "rule_version_id": rule["id"],
                },
                "rule_version": rule,
                "input_schema": input_schema,
                "output_schema": output_schema,
                "input_binding": input_binding,
                "output_binding": output_binding,
                "backend": capabilities,
            },
        )
        BoundSqlPlanService(
            repository,
            CompilerRegistry(
                {dialect: SqlGlotRuleCompiler(dialect)}
            ),
        ).compile_and_persist(
            component_binding_id=component_binding_id,
            rule_version_id=rule["id"],
            input_schema_snapshot_id=input_schema["id"],
            output_schema_snapshot_id=output_schema["id"],
            input_binding_id=input_binding["id"],
            output_binding_id=output_binding["id"],
        )
        record = repository.record
        assert record["plan_status"] == "compiled"
        compiled = {
            "plan": record["plan"],
            "plan_hash": record["plan_hash"],
        }
        adapter = SqlGlotRulePlanAdapter(
            DirectManager(
                engine,
                Definition(dialect, capabilities),
            )
        )
        node = {
            "id": "task4_real_rule",
            "type": "rule.apply",
            "purpose": "write",
            "idempotency": idempotency,
            "config": {
                "component_binding_id": component_binding_id,
                "rule_version_id": rule["id"],
                "execution_plan_hash": compiled["plan_hash"],
            },
        }
        preflight_evidence = adapter.execute(
            plan=compiled["plan"],
            node=node,
            parameters={},
            write_authorized=True,
        )
        with engine.begin() as connection:
            connection.execute(text(f"DELETE FROM {target_name}"))
        repository.publish_with_evidence(
            compiled["plan_hash"],
            preflight_evidence,
        )
        executor = RulePlanExecutor(
            repository,
            adapters={"sql_pushdown": adapter},
        )

        result = executor.execute(node, {}, write_authorized=True)

        assert result["rows_in"] == 2
        assert result["rows_out"] == 1
        assert result["rows_rejected"] == 1
        with engine.connect() as connection:
            rows = connection.execute(
                text(
                    f"SELECT customer_id, name, mobile "
                    f"FROM {target_name} ORDER BY customer_id"
                )
            ).tuples().all()
        assert rows == [(1, "Alice", "13800138000")]

        repeated = executor.execute(node, {}, write_authorized=True)
        assert repeated["rows_out"] == 1
        assert repeated["rows_rejected"] == 1
        with engine.connect() as connection:
            assert connection.execute(
                text(f"SELECT COUNT(*) FROM {target_name}")
            ).scalar_one() == 1

        repository.record["plan"] = {
            **repository.record["plan"],
            "result_contract": {
                **repository.record["plan"]["result_contract"],
                "rows_rejected": "unknown",
            },
        }
        with pytest.raises(NodeExecutionError, match="not executable"):
            executor.execute(node, {}, write_authorized=True)
    finally:
        with engine.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {target_name}"))
            connection.execute(text(f"DROP TABLE IF EXISTS {source_name}"))
        engine.dispose()


def test_mysql_upsert_rejects_real_nonunique_and_alternate_unique_targets():
    from app.core.data_rules.compilers.sql import SqlGlotRuleCompiler
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    dialect, url, schema_name, collation, regex_engine = CASES[1]
    engine = create_engine(url, pool_pre_ping=True)
    source_name = "task4_unique_source"
    nonunique_name = "task4_nonunique_target"
    alternate_name = "task4_alternate_target"
    capabilities = {
        "dialect": dialect,
        "timezone": "Asia/Shanghai",
        "collation": collation,
        "rounding_mode": "half_away_from_zero",
        "regex_engine": regex_engine,
    }
    datasource_uid = new_governance_uid()
    input_schema = _snapshot("bd:task4:unique:raw")
    output_schema = _snapshot("bd:task4:unique:clean")
    spec = validate_rule_spec(
        {
            "schema_version": "2.0",
            "rule_uid": new_governance_uid(),
            "name": "task4_unique_attestation",
            "input_schema_ref": input_schema["schema_ref"],
            "output_schema_ref": output_schema["schema_ref"],
            "steps": [
                {
                    "id": "trim_name",
                    "op": "normalize_text",
                    "column": "name",
                    "trim": True,
                }
            ],
            "null_policy": "explicit",
            "timezone": "Asia/Shanghai",
        }
    )
    rule = {
        "id": new_governance_uid(),
        "status": "published",
        "rule_spec": spec,
        "spec_hash": rule_spec_hash(spec),
    }
    input_binding = {
        "id": new_governance_uid(),
        "data_source_uid": datasource_uid,
        "object_kind": "table",
        "object_ref": f"{schema_name}.{source_name}",
        "schema_snapshot_id": input_schema["id"],
        "access_mode": "read",
        "dialect": dialect,
        "write_mode": "append",
    }

    try:
        with engine.begin() as connection:
            for name in (alternate_name, nonunique_name, source_name):
                connection.execute(text(f"DROP TABLE IF EXISTS {name}"))
            connection.execute(
                text(
                    f"CREATE TABLE {source_name} ("
                    "customer_id BIGINT PRIMARY KEY, "
                    "name VARCHAR(100), mobile VARCHAR(30))"
                )
            )
            connection.execute(
                text(
                    f"CREATE TABLE {nonunique_name} ("
                    "customer_id BIGINT, name VARCHAR(100), mobile VARCHAR(30))"
                )
            )
            connection.execute(
                text(
                    f"CREATE TABLE {alternate_name} ("
                    "customer_id BIGINT PRIMARY KEY, "
                    "name VARCHAR(100), mobile VARCHAR(30) UNIQUE)"
                )
            )
            connection.execute(
                text(
                    f"INSERT INTO {source_name} "
                    "(customer_id, name, mobile) "
                    "VALUES (1, ' Alice ', '13800138000')"
                )
            )

        adapter = SqlGlotRulePlanAdapter(
            DirectManager(engine, Definition(dialect, capabilities))
        )
        for target_name, error in (
            (nonunique_name, "exact unique key"),
            (alternate_name, "alternate unique"),
        ):
            output_binding = {
                "id": new_governance_uid(),
                "data_source_uid": datasource_uid,
                "object_kind": "table",
                "object_ref": f"{schema_name}.{target_name}",
                "schema_snapshot_id": output_schema["id"],
                "access_mode": "write",
                "dialect": dialect,
                "write_mode": "append",
            }
            compiled = SqlGlotRuleCompiler(dialect).compile(
                rule_version=rule,
                input_schema=input_schema,
                output_schema=output_schema,
                input_binding=input_binding,
                output_binding=output_binding,
                backend=capabilities,
            )
            node = {
                "id": "task4_unique_rule",
                "type": "rule.apply",
                "purpose": "write",
                "idempotency": {
                    "strategy": "upsert",
                    "key": "customer_id",
                },
                "config": {
                    "component_binding_id": new_governance_uid(),
                    "rule_version_id": rule["id"],
                    "execution_plan_hash": compiled["plan_hash"],
                },
            }
            with pytest.raises(NodeExecutionError, match=error):
                adapter.execute(
                    plan=compiled["plan"],
                    node=node,
                    parameters={},
                    write_authorized=True,
                )
            with engine.connect() as connection:
                assert connection.execute(
                    text(f"SELECT COUNT(*) FROM {target_name}")
                ).scalar_one() == 0
    finally:
        with engine.begin() as connection:
            for name in (alternate_name, nonunique_name, source_name):
                connection.execute(text(f"DROP TABLE IF EXISTS {name}"))
        engine.dispose()