from __future__ import annotations import copy from datetime import date, datetime 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=256 * 1024 * 1024, max_ttl_seconds=3600, ) class Resolver: def __init__(self, values, artifact_store): self.values = values self.artifact_store = artifact_store self.calls = [] self.events = [] self.attest_error = None self.publish_error = None self.registrations = [] self.handoffs = {} def resolve(self, *, binding_id, correlation_id, kind): self.calls.append((binding_id, correlation_id)) self.events.append(("resolve", binding_id, kind)) 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 resolve_handoff(self, *, artifact_ref, correlation_id): self.events.append(("handoff", artifact_ref, correlation_id)) return self.handoffs[artifact_ref] def publish_path( self, path, *, binding_id, correlation_id, kind, binding_hash, ttl_seconds, schema_fields, limits, ): self.events.append(("publish", binding_id, kind)) if self.publish_error is not None: raise self.publish_error artifact = self.artifact_store.write_path( path, correlation_id, ttl_seconds, schema_fields=schema_fields, limits=limits, ) self.registrations.append( { "binding_id": binding_id, "correlation_id": correlation_id, "artifact": artifact, "kind": kind, "binding_hash": binding_hash, } ) return artifact 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(memory_limit_bytes=128 * 1024 * 1024), ) 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"], } }, store, ) 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 len(result["_violation_sample"]) == 1 assert result["_violation_sample"][0]["mobile"] == "[REDACTED]" assert result["commit_outcome"] == "committed" assert "schema_fields" not in result 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_executes_quality_node_without_business_write_authority(): from app.runner.rule_polars import PolarsRulePlanAdapter compiled, input_binding = _compiled_plan( [ { "id": "mobile_format", "op": "assert", "expression": "matches(mobile, '^[0-9]{11}$')", "on_failure": "quarantine", "severity": "error", } ] ) store = _plan_store(FakeMinio()) correlation_id = new_governance_uid() source = store.write( pl.DataFrame( { "customer_id": [1, 2], "name": ["Alice", "Bad"], "mobile": ["13800138000", "invalid"], } ), 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"], } }, store, ) node = _node(compiled) node["type"] = "quality.check" node["purpose"] = "read" del node["idempotency"] result = PolarsRulePlanAdapter( artifact_store=store, artifact_resolver=resolver, ).execute( plan=compiled["plan"], node=node, parameters={}, write_authorized=False, correlation_id=correlation_id, ) assert result["rows_in"] == 2 assert result["rows_out"] == 1 assert result["rows_quarantined"] == 1 assert result["violations"] == [ {"step_id": "mobile_format", "count": 1} ] assert result["commit_outcome"] == "not_applicable" assert "artifact_ref" not in result assert resolver.registrations == [] def test_polars_adapter_reports_unknown_catalog_commit_outcome(): from app.runner.artifacts import ArtifactCommitUnknown 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"], } }, store, ) resolver.publish_error = ArtifactCommitUnknown("lost acknowledgement") with pytest.raises(NodeExecutionError, match="commit outcome") as error: PolarsRulePlanAdapter( artifact_store=store, artifact_resolver=resolver, ).execute( plan=compiled["plan"], node=_node(compiled), parameters={}, write_authorized=True, correlation_id=correlation_id, ) assert error.value.commit_outcome == "unknown" def test_polars_adapter_consumes_only_attested_upstream_artifact_ref(): 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": ["Alice"], "mobile": ["13800138000"], } ), correlation_id, 600, schema_fields=compiled["plan"]["input_fields"], limits=compiled["plan"]["resource_limits"], ) resolver = Resolver({}, store) resolver.handoffs[source["artifact_ref"]] = source adapter = PolarsRulePlanAdapter( artifact_store=store, artifact_resolver=resolver, artifact_ttl_seconds=300, ) result = adapter.execute( plan=compiled["plan"], node=_node(compiled), parameters={"input_artifact": source["artifact_ref"]}, write_authorized=True, correlation_id=correlation_id, ) assert result["rows_out"] == 1 assert ( "handoff", source["artifact_ref"], correlation_id, ) in resolver.events with pytest.raises(NodeExecutionError, match="only one artifact"): adapter.execute( plan=compiled["plan"], node=_node(compiled), parameters={"rows": [{"customer_id": 1}]}, write_authorized=True, correlation_id=correlation_id, ) 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, } }, store, ) 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"], } }, store, ) 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(memory_limit_bytes=128 * 1024 * 1024), ) 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"], } }, store, ) 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_polars_expression_date_and_timestamp_use_plan_timezone(): from app.core.data_rules.compilers.polars import PolarsRuleCompiler from app.runner.rule_polars import PolarsRulePlanAdapter input_schema = _schema( "bd:event:raw", [ ("raw_time", "string", False), ("cast_time", "string", False), ], ) output_schema = _schema( "bd:event:clean", [ ("raw_time", "string", False), ("cast_time", "timestamp", False), ("local_date", "date", False), ("local_time", "timestamp", False), ], ) 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_timestamp", "op": "cast", "column": "cast_time", "to": "timestamp", "on_error": "fail", }, { "id": "derive_date", "op": "derive", "target": "local_date", "expression": "date(raw_time)", }, { "id": "derive_timestamp", "op": "derive", "target": "local_time", "expression": "timestamp(raw_time)", }, ], ) compiled = PolarsRuleCompiler().compile( rule_version=rule, input_schema=input_schema, output_schema=output_schema, input_binding=source_binding, output_binding=output_binding, backend=_backend(memory_limit_bytes=128 * 1024 * 1024), ) store = _plan_store(FakeMinio()) correlation_id = new_governance_uid() source = store.write( pl.DataFrame( { "raw_time": ["2026-07-22T16:30:00+00:00"], "cast_time": ["2026-07-22T16:30:00+00: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"], } }, store, ) 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["local_date"].item() == date(2026, 7, 23) assert output["local_time"].item() == datetime(2026, 7, 23, 0, 30) assert output["cast_time"].item() == datetime(2026, 7, 23, 0, 30) assert output.schema["local_time"] == pl.Datetime 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_input_binding_hash": plan["input_binding_hash"], "canonical_input_object_kind": "parquet_artifact", "canonical_output_schema_snapshot_id": plan[ "output_schema_snapshot_id" ], "canonical_output_schema_hash": plan["output_schema_hash"], "canonical_output_binding_hash": plan["output_binding_hash"], "canonical_output_object_kind": "parquet_artifact", "plan_status": "published", "rule_status": "published", "publication_audit_trusted": True, "logical_evidence_trusted": True, "physical_evidence_trusted": True, "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