DataOps-platform/tests/runner/test_rule_sql.py

from __future__ import annotations

from contextlib import contextmanager

import pytest

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.errors import DataSourceWriteOutcomeUnknown
from app.runner.nodes import NodeExecutionError


class Definition:
    def __init__(self, dialect):
        self.database_type = dialect
        self.extra_properties = {
            "sql_rule_capabilities": {
                "dialect": dialect,
                "timezone": "Asia/Shanghai",
                "collation": "C" if dialect == "postgresql" else "utf8mb4_0900_bin",
                "rounding_mode": "half_away_from_zero",
                "regex_engine": "posix" if dialect == "postgresql" else "icu",
            }
        }


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

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


class Result:
    def __init__(self, *, scalar=None, rowcount=0, rows=None):
        self._scalar = scalar
        self.rowcount = rowcount
        self._rows = rows or []

    def scalar_one(self):
        return self._scalar

    def mappings(self):
        return self

    def all(self):
        return self._rows


class Connection:
    def __init__(self, unique_rows=None):
        self.calls = []
        self.unique_rows = unique_rows

    def execute(self, statement, parameters):
        self.calls.append((str(statement), parameters))
        if len(self.calls) == 1:
            return Result(scalar=3)
        if len(self.calls) == 2:
            return Result(scalar=2)
        if "information_schema" in str(statement):
            return Result(
                rows=self.unique_rows
                if self.unique_rows is not None
                else [
                    {
                        "constraint_name": "customer_pkey",
                        "constraint_type": "PRIMARY KEY",
                        "columns": ["id"],
                    }
                ]
            )
        return Result(rowcount=2)


class Manager:
    def __init__(
        self,
        dialect="postgresql",
        *,
        unknown_commit=False,
        unique_rows=None,
    ):
        self.definitions = Definitions(Definition(dialect))
        self.connection = Connection(unique_rows)
        self.unknown_commit = unknown_commit
        self.calls = []

    @contextmanager
    def connect(self, uid, purpose):
        self.calls.append((uid, purpose))
        yield self.connection
        if self.unknown_commit:
            raise DataSourceWriteOutcomeUnknown()


def sql_plan(dialect="postgresql"):
    from app.core.data_rules.compilers.sql import (
        COMPILER_VERSION,
        bound_sql_plan_hash,
    )

    uid = new_governance_uid()
    capabilities = Definition(dialect).extra_properties["sql_rule_capabilities"]
    quote = '"' if dialect == "postgresql" else "`"
    plan = {
        "schema_version": "1.0",
        "compiler_version": COMPILER_VERSION,
        "dialect": dialect,
        "capabilities": capabilities,
        "data_source_uid": uid,
        "rule_version_id": new_governance_uid(),
        "rule_spec_hash": "a" * 64,
        "input_schema_snapshot_id": new_governance_uid(),
        "input_schema_hash": "b" * 64,
        "output_schema_snapshot_id": new_governance_uid(),
        "output_schema_hash": "c" * 64,
        "input_binding_id": new_governance_uid(),
        "output_binding_id": new_governance_uid(),
        "statements": [
            {
                "purpose": "write",
                "sql": (
                    f"INSERT INTO {quote}clean{quote}.{quote}customer{quote} "
                    f"({quote}id{quote}) SELECT {quote}id{quote} "
                    f"FROM {quote}raw{quote}.{quote}customer{quote}"
                ),
                "parameters": {},
            }
        ],
        "result_contract": {
            "rows_in": "counted",
            "rows_out": "counted",
            "rows_rejected": "counted",
        },
    }
    return plan, bound_sql_plan_hash(plan)


def node_for(plan, plan_hash):
    return {
        "id": "task4_rule",
        "type": "rule.apply",
        "purpose": "write",
        "idempotency": {"strategy": "upsert", "key": "id"},
        "config": {
            "component_binding_id": new_governance_uid(),
            "rule_version_id": plan["rule_version_id"],
            "execution_plan_hash": plan_hash,
        },
    }


def test_sqlglot_rule_adapter_verifies_and_executes_one_transaction():
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, plan_hash = sql_plan()
    node = node_for(plan, plan_hash)
    manager = Manager()

    result = SqlGlotRulePlanAdapter(manager).execute(
        plan=plan,
        node=node,
        parameters={},
        write_authorized=True,
    )

    assert result == {
        "rows_in": 3,
        "rows_out": 2,
        "rows_rejected": 1,
        "commit_outcome": "committed",
    }
    assert manager.calls == [(plan["data_source_uid"], "dataflow_write")]
    assert len(manager.connection.calls) == 4


def test_sqlglot_rule_adapter_fails_closed_for_dialect_hash_and_authorization():
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, plan_hash = sql_plan()
    node = node_for(plan, plan_hash)

    with pytest.raises(NodeExecutionError, match="dialect"):
        SqlGlotRulePlanAdapter(Manager("mysql")).execute(
            plan=plan,
            node=node,
            parameters={},
            write_authorized=True,
        )

    node["config"]["execution_plan_hash"] = "0" * 64
    with pytest.raises(NodeExecutionError, match="hash"):
        SqlGlotRulePlanAdapter(Manager()).execute(
            plan=plan,
            node=node,
            parameters={},
            write_authorized=True,
        )

    node["config"]["execution_plan_hash"] = plan_hash
    with pytest.raises(NodeExecutionError, match="authorization"):
        SqlGlotRulePlanAdapter(Manager()).execute(
            plan=plan,
            node=node,
            parameters={},
            write_authorized=False,
        )


def test_sqlglot_rule_adapter_rejects_unimplemented_idempotency_strategy():
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, plan_hash = sql_plan()
    node = node_for(plan, plan_hash)
    node["idempotency"] = {
        "strategy": "partition_replace",
        "key": "task4",
    }

    with pytest.raises(NodeExecutionError, match="idempotency"):
        SqlGlotRulePlanAdapter(Manager()).execute(
            plan=plan,
            node=node,
            parameters={},
            write_authorized=True,
        )


def test_rule_executor_matches_node_idempotency_to_persisted_component():
    from app.runner.rules import RulePlanExecutor

    plan, plan_hash = sql_plan()
    node = node_for(plan, plan_hash)

    class Repository:
        def load(self, **_kwargs):
            return {
                "component_binding_id": node["config"]["component_binding_id"],
                "rule_version_id": plan["rule_version_id"],
                "backend": "sql_pushdown",
                "plan": plan,
                "plan_hash": plan_hash,
                "plan_status": "published",
                "rule_status": "published",
                "component_kind": "rule.apply",
                "binding_idempotency": {
                    "strategy": "upsert",
                    "key": "different_id",
                },
            }

    class Adapter:
        def execute(self, **_kwargs):
            return {"rows_in": 0, "rows_out": 0, "rows_rejected": 0}

    with pytest.raises(NodeExecutionError, match="idempotency"):
        RulePlanExecutor(
            Repository(),
            adapters={"sql_pushdown": Adapter()},
        ).execute(node, {}, write_authorized=True)


def test_rule_executor_rechecks_canonical_rule_schema_and_compiler_attestations():
    from app.core.data_rules.compilers.sql import COMPILER_VERSION
    from app.runner.rules import RulePlanExecutor

    plan, plan_hash = sql_plan()
    node = node_for(plan, plan_hash)

    def record(**overrides):
        value = {
            "component_binding_id": node["config"]["component_binding_id"],
            "rule_version_id": plan["rule_version_id"],
            "backend": "sql_pushdown",
            "compiler_version": COMPILER_VERSION,
            "plan": plan,
            "plan_hash": 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": "published",
            "rule_status": "published",
            "component_kind": "rule.apply",
            "binding_idempotency": node["idempotency"],
        }
        value.update(overrides)
        return value

    class Repository:
        def __init__(self, value):
            self.value = value

        def load(self, **_kwargs):
            return self.value

    class Adapter:
        def execute(self, **_kwargs):
            return {"rows_in": 0, "rows_out": 0, "rows_rejected": 0}

    for tampered in (
        {"compiler_version": "dataops-sqlglot-99.0.0"},
        {"canonical_rule_spec_hash": "f" * 64},
        {"canonical_input_schema_hash": "e" * 64},
        {"canonical_output_schema_snapshot_id": new_governance_uid()},
    ):
        with pytest.raises(NodeExecutionError, match="attestation"):
            RulePlanExecutor(
                Repository(record(**tampered)),
                adapters={"sql_pushdown": Adapter()},
            ).execute(node, {}, write_authorized=True)


def test_sqlglot_rule_adapter_reports_unknown_commit_outcome():
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, plan_hash = sql_plan()

    with pytest.raises(NodeExecutionError) as error:
        SqlGlotRulePlanAdapter(Manager(unknown_commit=True)).execute(
            plan=plan,
            node=node_for(plan, plan_hash),
            parameters={},
            write_authorized=True,
        )

    assert error.value.commit_outcome == "unknown"


def test_sqlglot_rule_adapter_rejects_function_injected_into_attested_plan():
    from app.core.data_rules.compilers.sql import bound_sql_plan_hash
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, _ = sql_plan()
    plan["statements"][0]["sql"] = plan["statements"][0]["sql"].replace(
        'SELECT "id"', 'SELECT pg_sleep(1)'
    )
    with pytest.raises(ValueError, match="unsupported"):
        bound_sql_plan_hash(plan)

    # A malicious publisher cannot bypass the AST allowlist by recomputing a
    # hash because the adapter independently validates the compiler subset.
    node = node_for(plan, "0" * 64)
    with pytest.raises(NodeExecutionError, match="invalid"):
        SqlGlotRulePlanAdapter(Manager()).execute(
            plan=plan,
            node=node,
            parameters={},
            write_authorized=True,
        )


@pytest.mark.parametrize(
    "unique_rows",
    [
        [],
        [
            {
                "constraint_name": "customer_id_idx",
                "constraint_type": "UNIQUE",
                "columns": "id",
            },
            {
                "constraint_name": "mobile_idx",
                "constraint_type": "UNIQUE",
                "columns": "mobile",
            },
        ],
    ],
)
def test_mysql_upsert_requires_exact_key_and_no_alternate_unique_path(
    unique_rows,
):
    from app.runner.rule_sql import SqlGlotRulePlanAdapter

    plan, plan_hash = sql_plan("mysql")
    with pytest.raises(NodeExecutionError, match="unique"):
        SqlGlotRulePlanAdapter(
            Manager("mysql", unique_rows=unique_rows)
        ).execute(
            plan=plan,
            node=node_for(plan, plan_hash),
            parameters={},
            write_authorized=True,
        )