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


def _plan_store(client):
    from app.runner.artifacts import ArtifactStore

    return ArtifactStore(
        client,
        bucket="dataops-rules",
        max_artifact_bytes=8 * 1024 * 1024,
        max_rows=10_000,
        memory_limit_bytes=32 * 1024 * 1024,
        max_ttl_seconds=3600,
    )


class Resolver:
    def __init__(self, values):
        self.values = values
        self.calls = []
        self.events = []
        self.attest_error = None
        self.registrations = []

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

    def attest_binding(self, *, binding_id, binding_hash, access_mode):
        self.events.append(("attest", binding_id, binding_hash, access_mode))
        if self.attest_error is not None:
            raise self.attest_error
        return {"binding_hash": binding_hash}

    def register(
        self,
        *,
        binding_id,
        correlation_id,
        artifact,
        kind,
        binding_hash,
    ):
        self.events.append(("register", binding_id, kind))
        self.registrations.append(
            {
                "binding_id": binding_id,
                "correlation_id": correlation_id,
                "artifact": artifact,
                "kind": kind,
                "binding_hash": binding_hash,
            }
        )


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 = _plan_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,
        schema_fields=compiled["plan"]["input_fields"],
        limits=compiled["plan"]["resource_limits"],
    )
    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"] == 1
    assert result["rows_filtered"] == 0
    assert result["rows_deduplicated"] == 1
    assert result["rows_join_dropped"] == 0
    assert result["rows_aggregated"] == 0
    assert result["violation_count"] == 1
    assert result["violations"] == [
        {"step_id": "mobile_format", "count": 1}
    ]
    assert result["commit_outcome"] == "committed"
    assert resolver.events[0] == (
        "attest",
        compiled["plan"]["output_binding_id"],
        compiled["plan"]["output_binding_hash"],
        "write",
    )
    assert resolver.registrations[0]["kind"] == "output"
    assert store.read(
        result["artifact_ref"],
        result["digest"],
        expected_schema_fields=compiled["plan"]["output_fields"],
        limits=compiled["plan"]["resource_limits"],
    ).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 = _plan_store(FakeMinio())
    correlation_id = new_governance_uid()
    source = store.write(
        pl.DataFrame(
            {"customer_id": [1], "name": [" A "], "mobile": ["1"]}
        ).lazy(),
        correlation_id,
        600,
        schema_fields=compiled["plan"]["input_fields"],
        limits=compiled["plan"]["resource_limits"],
    )
    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_polars_adapter_attests_current_output_binding_before_reading_input():
    from app.runner.rule_polars import PolarsRulePlanAdapter

    compiled, input_binding = _compiled_plan(
        [
            {
                "id": "trim_name",
                "op": "normalize_text",
                "column": "name",
                "trim": True,
            }
        ]
    )
    store = _plan_store(FakeMinio())
    correlation_id = new_governance_uid()
    source = store.write(
        pl.DataFrame(
            {"customer_id": [1], "name": [" A "], "mobile": ["1"]}
        ),
        correlation_id,
        600,
        schema_fields=compiled["plan"]["input_fields"],
        limits=compiled["plan"]["resource_limits"],
    )
    resolver = Resolver(
        {
            input_binding["id"]: {
                **source,
                "binding_hash": compiled["plan"]["input_binding_hash"],
            }
        }
    )
    resolver.attest_error = ValueError("binding changed")
    reads_before_execute = list(store.client.get_calls)

    with pytest.raises(NodeExecutionError, match="output binding"):
        PolarsRulePlanAdapter(
            artifact_store=store,
            artifact_resolver=resolver,
        ).execute(
            plan=compiled["plan"],
            node=_node(compiled),
            parameters={},
            write_authorized=True,
            correlation_id=correlation_id,
        )

    assert resolver.events == [
        (
            "attest",
            compiled["plan"]["output_binding_id"],
            compiled["plan"]["output_binding_hash"],
            "write",
        )
    ]
    assert store.client.get_calls == reads_before_execute


def test_polars_adapter_uses_exact_decimal_and_timestamptz_output_contracts():
    from app.core.data_rules.compilers.polars import PolarsRuleCompiler
    from app.core.data_rules.execution_contracts import canonical_schema_hash
    from app.runner.rule_polars import PolarsRulePlanAdapter

    input_schema = _schema(
        "bd:payment:raw",
        [
            ("amount", "string", False),
            ("occurred_at", "string", False),
        ],
    )
    output_schema = _schema(
        "bd:payment:clean",
        [
            ("amount", "decimal", False),
            ("occurred_at", "timestamptz", False),
        ],
    )
    output_schema["fields"][0].update({"precision": 12, "scale": 2})
    output_schema["fields"][1]["timezone"] = "Asia/Shanghai"
    output_schema["schema_hash"] = canonical_schema_hash(
        output_schema["fields"]
    )
    source_binding = _binding(input_schema, access_mode="read")
    output_binding = _binding(output_schema, access_mode="write")
    rule = _published_rule(
        input_schema,
        output_schema,
        [
            {
                "id": "cast_amount",
                "op": "cast",
                "column": "amount",
                "to": "decimal",
                "on_error": "fail",
            },
            {
                "id": "cast_time",
                "op": "cast",
                "column": "occurred_at",
                "to": "timestamptz",
                "on_error": "fail",
            },
        ],
    )
    compiled = PolarsRuleCompiler().compile(
        rule_version=rule,
        input_schema=input_schema,
        output_schema=output_schema,
        input_binding=source_binding,
        output_binding=output_binding,
        backend=_backend(),
    )
    store = _plan_store(FakeMinio())
    correlation_id = new_governance_uid()
    source = store.write(
        pl.DataFrame(
            {
                "amount": ["12.34"],
                "occurred_at": ["2026-07-23T12:30:00+08:00"],
            }
        ),
        correlation_id,
        600,
        schema_fields=compiled["plan"]["input_fields"],
        limits=compiled["plan"]["resource_limits"],
    )
    resolver = Resolver(
        {
            source_binding["id"]: {
                **source,
                "binding_hash": compiled["plan"]["input_binding_hash"],
            }
        }
    )

    result = PolarsRulePlanAdapter(
        artifact_store=store,
        artifact_resolver=resolver,
    ).execute(
        plan=compiled["plan"],
        node=_node(compiled),
        parameters={},
        write_authorized=True,
        correlation_id=correlation_id,
    )

    output = store.read(
        result["artifact_ref"],
        result["digest"],
        expected_schema_fields=compiled["plan"]["output_fields"],
        limits=compiled["plan"]["resource_limits"],
    ).collect()
    assert output.schema["amount"] == pl.Decimal(precision=12, scale=2)
    assert output.schema["occurred_at"] == pl.Datetime(
        time_zone="Asia/Shanghai"
    )


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