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):
        self._scalar = scalar
        self.rowcount = rowcount

    def scalar_one(self):
        return self._scalar


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

    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)
        return Result(rowcount=2)


class Manager:
    def __init__(self, dialect="postgresql", *, unknown_commit=False):
        self.definitions = Definitions(Definition(dialect))
        self.connection = Connection()
        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 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",
        "dialect": dialect,
        "capabilities": capabilities,
        "data_source_uid": uid,
        "rule_version_id": new_governance_uid(),
        "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) == 3


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