DataOps-platform/tests/runner/test_rule_polars.py

from __future__ import annotations

import copy

import polars as pl
import pytest

from app.core.common.identifiers import new_governance_uid
from app.runner.nodes import NodeExecutionError
from tests.core.data_rules.test_polars_compiler import (
    _backend,
    _binding,
    _published_rule,
    _schema,
)
from tests.runner.test_artifacts import FakeMinio, _store


class Resolver:
    def __init__(self, values):
        self.values = values
        self.calls = []

    def resolve(self, *, binding_id, correlation_id):
        self.calls.append((binding_id, correlation_id))
        return self.values[binding_id]


def _compiled_plan(steps):
    from app.core.data_rules.compilers.polars import PolarsRuleCompiler

    fields = [
        ("customer_id", "integer", False),
        ("name", "string", True),
        ("mobile", "string", True),
    ]
    input_schema = _schema("bd:customer:raw", fields)
    output_schema = _schema("bd:customer:clean", fields)
    input_binding = _binding(input_schema, access_mode="read")
    output_binding = _binding(output_schema, access_mode="write")
    rule = _published_rule(input_schema, output_schema, steps)
    compiled = PolarsRuleCompiler().compile(
        rule_version=rule,
        input_schema=input_schema,
        output_schema=output_schema,
        input_binding=input_binding,
        output_binding=output_binding,
        backend=_backend(),
    )
    return compiled, input_binding


def _node(compiled):
    return {
        "id": "task5_polars",
        "type": "rule.apply",
        "purpose": "write",
        "idempotency": {
            "strategy": "deduplication_key",
            "key": "customer_id",
        },
        "config": {
            "component_binding_id": new_governance_uid(),
            "rule_version_id": compiled["plan"]["rule_version_id"],
            "execution_plan_hash": compiled["plan_hash"],
        },
    }


def test_polars_adapter_reconstructs_assert_and_deduplicate_and_writes_artifact():
    from app.runner.rule_polars import PolarsRulePlanAdapter

    compiled, input_binding = _compiled_plan(
        [
            {
                "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",
            },
            {
                "id": "one_customer",
                "op": "deduplicate",
                "keys": ["customer_id"],
                "order_by": ["name"],
                "keep": "first",
            },
        ]
    )
    store = _store(FakeMinio())
    correlation_id = new_governance_uid()
    source = store.write(
        pl.DataFrame(
            {
                "customer_id": [1, 1, 2],
                "name": [" Alice ", "Alice B", " Bad "],
                "mobile": ["13800138000", "13800138000", "invalid"],
            }
        ).lazy(),
        correlation_id,
        600,
    )
    resolver = Resolver(
        {
            input_binding["id"]: {
                **source,
                "binding_hash": compiled["plan"]["input_binding_hash"],
            }
        }
    )
    adapter = PolarsRulePlanAdapter(
        artifact_store=store,
        artifact_resolver=resolver,
        masking_policies={
            "customer_mobile_last4": "preserve_last_4"
        },
        artifact_ttl_seconds=300,
    )

    result = adapter.execute(
        plan=compiled["plan"],
        node=_node(compiled),
        parameters={},
        write_authorized=True,
        correlation_id=correlation_id,
    )

    assert result["rows_in"] == 3
    assert result["rows_out"] == 1
    assert result["rows_rejected"] == 2
    assert result["violation_count"] == 1
    assert result["violations"] == [
        {"step_id": "mobile_format", "count": 1}
    ]
    assert result["commit_outcome"] == "committed"
    assert store.read(
        result["artifact_ref"], result["digest"]
    ).collect().to_dicts() == [
        {
            "customer_id": 1,
            "name": "Alice",
            "mobile": "13800138000",
        }
    ]


def test_polars_adapter_fails_closed_for_plan_hash_binding_and_authorization():
    from app.runner.rule_polars import PolarsRulePlanAdapter

    compiled, input_binding = _compiled_plan(
        [
            {
                "id": "trim_name",
                "op": "normalize_text",
                "column": "name",
                "trim": True,
            }
        ]
    )
    store = _store(FakeMinio())
    correlation_id = new_governance_uid()
    source = store.write(
        pl.DataFrame(
            {"customer_id": [1], "name": [" A "], "mobile": ["1"]}
        ).lazy(),
        correlation_id,
        600,
    )
    resolver = Resolver(
        {
            input_binding["id"]: {
                **source,
                "binding_hash": "0" * 64,
            }
        }
    )
    adapter = PolarsRulePlanAdapter(
        artifact_store=store,
        artifact_resolver=resolver,
    )
    node = _node(compiled)

    with pytest.raises(NodeExecutionError, match="authorization"):
        adapter.execute(
            plan=compiled["plan"],
            node=node,
            parameters={},
            write_authorized=False,
            correlation_id=correlation_id,
        )
    with pytest.raises(NodeExecutionError, match="binding"):
        adapter.execute(
            plan=compiled["plan"],
            node=node,
            parameters={},
            write_authorized=True,
            correlation_id=correlation_id,
        )
    resolver.values[input_binding["id"]]["binding_hash"] = compiled["plan"][
        "input_binding_hash"
    ]
    node["config"]["execution_plan_hash"] = "0" * 64
    with pytest.raises(NodeExecutionError, match="hash"):
        adapter.execute(
            plan=compiled["plan"],
            node=node,
            parameters={},
            write_authorized=True,
            correlation_id=correlation_id,
        )

    tampered = copy.deepcopy(compiled["plan"])
    tampered["operations"][0]["callable"] = "unsafe"
    node["config"]["execution_plan_hash"] = compiled["plan_hash"]
    with pytest.raises(NodeExecutionError, match="invalid"):
        adapter.execute(
            plan=tampered,
            node=node,
            parameters={},
            write_authorized=True,
            correlation_id=correlation_id,
        )


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

    compiled, _input_binding = _compiled_plan(
        [
            {
                "id": "trim_name",
                "op": "normalize_text",
                "column": "name",
                "trim": True,
            }
        ]
    )
    node = _node(compiled)
    correlation_id = new_governance_uid()
    plan = compiled["plan"]
    record = {
        "component_binding_id": node["config"]["component_binding_id"],
        "rule_version_id": plan["rule_version_id"],
        "backend": "polars_batch",
        "compiler_version": compiled["compiler_version"],
        "plan": 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": "published",
        "rule_status": "published",
        "component_kind": "rule.apply",
        "binding_idempotency": node["idempotency"],
    }

    class Repository:
        def load(self, **_kwargs):
            return record

    class Adapter:
        def __init__(self):
            self.kwargs = None

        def execute(self, **kwargs):
            self.kwargs = kwargs
            return {"rows_in": 1, "rows_out": 1, "rows_rejected": 0}

    adapter = Adapter()
    result = RulePlanExecutor(
        Repository(), adapters={"polars_batch": adapter}
    ).execute(
        node,
        {},
        write_authorized=True,
        correlation_id=correlation_id,
    )

    assert result["rows_out"] == 1
    assert adapter.kwargs["correlation_id"] == correlation_id

    record["canonical_input_schema_hash"] = "0" * 64
    with pytest.raises(NodeExecutionError, match="attestation"):
        RulePlanExecutor(
            Repository(), adapters={"polars_batch": adapter}
        ).execute(
            node,
            {},
            write_authorized=True,
            correlation_id=correlation_id,
        )


def test_node_registry_forwards_trusted_correlation_context():
    from app.runner.nodes import NodeRegistry

    class Executor:
        def __init__(self):
            self.correlation_id = None

        def execute(self, _node, _parameters, **kwargs):
            self.correlation_id = kwargs["correlation_id"]
            return {"ok": True}

    executor = Executor()
    correlation_id = new_governance_uid()
    assert NodeRegistry({"rule.apply": executor}).execute(
        {"type": "rule.apply"},
        {},
        write_authorized=True,
        correlation_id=correlation_id,
    ) == {"ok": True}
    assert executor.correlation_id == correlation_id