DataOps-platform/tests/integration/test_data_rule_polars_execution.py

from __future__ import annotations

import json
import re
from pathlib import Path

import polars as pl
from minio import Minio
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec
from app.core.data_rules.execution_contracts import canonical_schema_hash

COMPOSE = (
    Path(__file__).resolve().parents[2]
    / "deploy"
    / "docker"
    / "docker-compose.yml"
)


def _compose_value(pattern):
    source = COMPOSE.read_text(encoding="utf-8")
    match = re.search(pattern, source, flags=re.DOTALL)
    assert match is not None
    return match.group(1)


def _schema(schema_ref, fields):
    normalized = [
        {"name": name, "type": field_type, "nullable": nullable}
        for name, field_type, nullable in fields
    ]
    return {
        "id": new_governance_uid(),
        "schema_ref": schema_ref,
        "schema_hash": canonical_schema_hash(normalized),
        "fields": normalized,
        "source_revision": "task5:real-cross-source",
    }


def _binding(schema, *, source_uid, access_mode, object_ref):
    return {
        "id": new_governance_uid(),
        "data_source_uid": source_uid,
        "object_kind": "parquet_artifact",
        "object_ref": object_ref,
        "schema_snapshot_id": schema["id"],
        "access_mode": access_mode,
        "dialect": "parquet",
        "write_mode": "append",
    }


def test_real_postgres_mysql_minio_polars_cross_source_execution():
    from app.core.data_rules.compilers.polars import PolarsRuleCompiler
    from app.runner.artifacts import ArtifactStore, PostgresArtifactResolver
    from app.runner.rule_polars import PolarsRulePlanAdapter
    from app.runner.rules import PostgresRulePlanRepository, RulePlanExecutor

    source_user = _compose_value(
        r"source-postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    source_password = _compose_value(
        r"source-postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
    )
    platform_user = _compose_value(
        r"\n  postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    platform_password = _compose_value(
        r"\n  postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
    )
    postgres_port = _compose_value(r'"(25432):5432"')
    mysql_port = _compose_value(r'"(23306):3306"')
    minio_user = _compose_value(r"MINIO_ROOT_USER:\s*([^\s]+)")
    minio_password = _compose_value(r"MINIO_ROOT_PASSWORD:\s*([^\s]+)")
    minio_port = _compose_value(r'"(19000):9000"')
    platform_port = _compose_value(r'"(15432):5432"')
    bucket = _compose_value(r"mc mb --ignore-existing local/([^\s]+)")
    postgres = create_engine(
        f"postgresql+psycopg2://{source_user}:{source_password}"
        f"@127.0.0.1:{postgres_port}/acceptance",
        pool_pre_ping=True,
    )
    mysql = create_engine(
        f"mysql+pymysql://{source_user}:{source_password}"
        f"@127.0.0.1:{mysql_port}/acceptance",
        pool_pre_ping=True,
    )
    platform = create_engine(
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/dataops",
        pool_pre_ping=True,
    )
    minio = Minio(
        f"127.0.0.1:{minio_port}",
        access_key=minio_user,
        secret_key=minio_password,
        secure=False,
    )
    store = ArtifactStore(
        minio,
        bucket=bucket,
        max_artifact_bytes=4 * 1024 * 1024,
        max_rows=1_000,
        memory_limit_bytes=16 * 1024 * 1024,
        max_ttl_seconds=3600,
    )
    correlation_id = new_governance_uid()
    prefix = f"rules/{correlation_id}/"
    customer_table = "task5_polars_customers"
    segment_table = "task5_polars_segments"
    rule_uid = new_governance_uid()
    rule_id = new_governance_uid()
    dataflow_uid = new_governance_uid()
    dataflow_version_id = new_governance_uid()
    deployment_id = new_governance_uid()
    component_binding_id = new_governance_uid()
    plan_id = new_governance_uid()

    try:
        with postgres.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
            connection.execute(
                text(
                    f"CREATE TABLE {customer_table} ("
                    "customer_id BIGINT NOT NULL, "
                    "name VARCHAR(100), mobile VARCHAR(30), "
                    "segment_code VARCHAR(20), version_no BIGINT NOT NULL)"
                )
            )
            connection.execute(
                text(
                    f"INSERT INTO {customer_table} "
                    "(customer_id, name, mobile, segment_code, version_no) "
                    "VALUES "
                    "(1, ' Alice ', '13800138000', 'A', 1), "
                    "(1, ' Alice Updated ', '13800138000', 'A', 2), "
                    "(2, ' Bad ', 'invalid', 'B', 1), "
                    "(3, ' Carol ', '13900139000', 'C', 1)"
                )
            )
        with mysql.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
            connection.execute(
                text(
                    f"CREATE TABLE {segment_table} ("
                    "code VARCHAR(20) PRIMARY KEY, "
                    "segment_name VARCHAR(100) NOT NULL)"
                )
            )
            connection.execute(
                text(
                    f"INSERT INTO {segment_table} (code, segment_name) "
                    "VALUES ('A', 'Gold'), ('B', 'Basic'), ('C', 'Silver')"
                )
            )
        with postgres.connect() as connection:
            customer_rows = [
                dict(row)
                for row in connection.execute(
                    text(
                        f"SELECT customer_id, name, mobile, "
                        f"segment_code, version_no FROM {customer_table}"
                    )
                ).mappings()
            ]
        with mysql.connect() as connection:
            segment_rows = [
                dict(row)
                for row in connection.execute(
                    text(
                        f"SELECT code, segment_name FROM {segment_table}"
                    )
                ).mappings()
            ]
        input_schema = _schema(
            "bd:task5:customer:raw",
            [
                ("customer_id", "integer", False),
                ("name", "string", True),
                ("mobile", "string", True),
                ("segment_code", "string", True),
                ("version_no", "integer", False),
            ],
        )
        lookup_schema = _schema(
            "bd:task5:segment:lookup",
            [
                ("code", "string", False),
                ("segment_name", "string", False),
            ],
        )
        output_schema = _schema(
            "bd:task5:customer:enriched",
            [
                ("customer_id", "integer", False),
                ("name", "string", True),
                ("mobile", "string", True),
                ("segment_code", "string", True),
                ("version_no", "integer", False),
                ("segment_name", "string", True),
            ],
        )
        input_binding = _binding(
            input_schema,
            source_uid=new_governance_uid(),
            access_mode="read",
            object_ref="postgres-customer-artifact",
        )
        lookup_binding = _binding(
            lookup_schema,
            source_uid=new_governance_uid(),
            access_mode="read",
            object_ref="mysql-segment-artifact",
        )
        output_binding = _binding(
            output_schema,
            source_uid=new_governance_uid(),
            access_mode="write",
            object_ref="polars-output-artifact",
        )
        spec = validate_rule_spec(
            {
                "schema_version": "2.0",
                "rule_uid": new_governance_uid(),
                "name": "task5_real_cross_source",
                "input_schema_ref": input_schema["schema_ref"],
                "output_schema_ref": output_schema["schema_ref"],
                "steps": [
                    {
                        "id": "normalize_name",
                        "op": "normalize_text",
                        "column": "name",
                        "trim": True,
                    },
                    {
                        "id": "join_segment",
                        "op": "lookup_join",
                        "lookup": {
                            "binding_id": lookup_binding["id"],
                            "left_on": ["segment_code"],
                            "right_on": ["code"],
                            "select": {
                                "segment_name": "segment_name"
                            },
                            "how": "left",
                        },
                    },
                    {
                        "id": "valid_mobile",
                        "op": "assert",
                        "expression": "matches(mobile, '^[0-9]{11}$')",
                        "on_failure": "reject",
                        "severity": "error",
                    },
                    {
                        "id": "latest_customer",
                        "op": "deduplicate",
                        "keys": ["customer_id"],
                        "order_by": ["version_no"],
                        "keep": "last",
                    },
                ],
                "null_policy": "explicit",
                "timezone": "Asia/Shanghai",
            }
        )
        rule = {
            "id": rule_id,
            "status": "published",
            "rule_spec": spec,
            "spec_hash": rule_spec_hash(spec),
        }
        compiled = PolarsRuleCompiler().compile(
            rule_version=rule,
            input_schema=input_schema,
            output_schema=output_schema,
            input_binding=input_binding,
            output_binding=output_binding,
            backend={
                "max_rows": 1_000,
                "max_artifact_bytes": 4 * 1024 * 1024,
                "memory_limit_bytes": 16 * 1024 * 1024,
                "masking_policies": {},
                "lookup_bindings": {
                    lookup_binding["id"]: {
                        "binding": lookup_binding,
                        "schema": lookup_schema,
                    }
                },
            },
        )
        lookup_operation = compiled["plan"]["operations"][1]
        schema_hashes = {
            "rule_spec_hash": compiled["plan"]["rule_spec_hash"],
            "input_schema_snapshot_id": input_schema["id"],
            "input_schema_hash": input_schema["schema_hash"],
            "output_schema_snapshot_id": output_schema["id"],
            "output_schema_hash": output_schema["schema_hash"],
        }
        with platform.begin() as connection:
            for schema in (input_schema, lookup_schema, output_schema):
                connection.execute(
                    text(
                        """
                        INSERT INTO public.data_schema_snapshots
                        (id, schema_ref, schema_hash, fields, source_revision)
                        VALUES (CAST(:id AS uuid), :schema_ref, :schema_hash,
                                CAST(:fields AS jsonb), :source_revision)
                        """
                    ),
                    {**schema, "fields": json.dumps(schema["fields"])},
                )
            connection.execute(
                text(
                    """
                    INSERT INTO public.data_rules
                    (id, rule_uid, name, category, status)
                    VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid),
                            :name, 'general', 'active')
                    """
                ),
                {
                    "id": new_governance_uid(),
                    "rule_uid": rule_uid,
                    "name": spec["name"],
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.data_rule_versions
                    (id, rule_uid, version_no, source_text, source_language,
                     rule_spec, spec_hash, generated_kind, status, published_at)
                    VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1,
                            :source_text, 'en', CAST(:rule_spec AS jsonb),
                            :spec_hash, 'polars', 'published',
                            CURRENT_TIMESTAMP)
                    """
                ),
                {
                    "id": rule_id,
                    "rule_uid": rule_uid,
                    "source_text": "Task 5 real cross-source integration",
                    "rule_spec": json.dumps(spec),
                    "spec_hash": rule["spec_hash"],
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_versions
                    (id, dataflow_uid, version_no, name, dataflow_spec,
                     input_schema_hashes, output_schema_hash, status,
                     released_at)
                    VALUES (CAST(:id AS uuid), CAST(:dataflow_uid AS uuid), 1,
                            :name, '{}'::jsonb,
                            CAST(:input_schema_hashes AS jsonb),
                            :output_schema_hash, 'released', CURRENT_TIMESTAMP)
                    """
                ),
                {
                    "id": dataflow_version_id,
                    "dataflow_uid": dataflow_uid,
                    "name": "Task 5 real cross-source integration",
                    "input_schema_hashes": json.dumps(
                        [
                            input_schema["schema_hash"],
                            lookup_schema["schema_hash"],
                        ]
                    ),
                    "output_schema_hash": output_schema["schema_hash"],
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_deployments
                    (id, dataflow_version_id, environment, deployment_config,
                     status, activated_at)
                    VALUES (CAST(:id AS uuid),
                            CAST(:dataflow_version_id AS uuid), 'test',
                            '{}'::jsonb, 'active', CURRENT_TIMESTAMP)
                    """
                ),
                {
                    "id": deployment_id,
                    "dataflow_version_id": dataflow_version_id,
                },
            )
            for logical_ref, binding, binding_hash in (
                (
                    "customers",
                    input_binding,
                    compiled["plan"]["input_binding_hash"],
                ),
                (
                    "segments",
                    lookup_binding,
                    lookup_operation["lookup_binding_hash"],
                ),
                (
                    "enriched",
                    output_binding,
                    compiled["plan"]["output_binding_hash"],
                ),
            ):
                connection.execute(
                    text(
                        """
                        INSERT INTO public.dataflow_dataset_bindings
                        (id, dataflow_deployment_id, logical_ref,
                         data_source_uid, object_kind, object_ref,
                         schema_snapshot_id, dialect, access_mode, write_mode,
                         binding_hash)
                        VALUES (CAST(:id AS uuid), CAST(:deployment_id AS uuid),
                                :logical_ref, CAST(:source_uid AS uuid),
                                'parquet_artifact', :object_ref,
                                CAST(:schema_snapshot_id AS uuid), 'parquet',
                                :access_mode, 'append', :binding_hash)
                        """
                    ),
                    {
                        "id": binding["id"],
                        "deployment_id": deployment_id,
                        "logical_ref": logical_ref,
                        "source_uid": binding["data_source_uid"],
                        "object_ref": binding["object_ref"],
                        "schema_snapshot_id": binding["schema_snapshot_id"],
                        "access_mode": binding["access_mode"],
                        "binding_hash": binding_hash,
                    },
                )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_component_bindings
                    (id, dataflow_version_id, component_id, component_kind,
                     rule_version_id, stage, order_no, idempotency, provenance)
                    VALUES (CAST(:id AS uuid),
                            CAST(:dataflow_version_id AS uuid),
                            'task5_real_polars', 'rule.apply',
                            CAST(:rule_version_id AS uuid), 'transform', 0,
                            CAST(:idempotency AS jsonb), '{}'::jsonb)
                    """
                ),
                {
                    "id": component_binding_id,
                    "dataflow_version_id": dataflow_version_id,
                    "rule_version_id": rule_id,
                    "idempotency": json.dumps(
                        {
                            "strategy": "deduplication_key",
                            "key": "customer_id",
                        }
                    ),
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_execution_plans
                    (id, component_binding_id, backend, compiler_version, plan,
                     plan_hash, schema_hashes, status)
                    VALUES (CAST(:id AS uuid),
                            CAST(:component_binding_id AS uuid),
                            'polars_batch', :compiler_version,
                            CAST(:plan AS jsonb), :plan_hash,
                            CAST(:schema_hashes AS jsonb), 'published')
                    """
                ),
                {
                    "id": plan_id,
                    "component_binding_id": component_binding_id,
                    "compiler_version": compiled["compiler_version"],
                    "plan": json.dumps(compiled["plan"]),
                    "plan_hash": compiled["plan_hash"],
                    "schema_hashes": json.dumps(schema_hashes),
                },
            )
        customer_artifact = store.write(
            pl.DataFrame(customer_rows),
            correlation_id,
            900,
            schema_fields=input_schema["fields"],
            limits=compiled["plan"]["resource_limits"],
        )
        segment_artifact = store.write(
            pl.DataFrame(segment_rows),
            correlation_id,
            900,
            schema_fields=lookup_schema["fields"],
            limits=compiled["plan"]["resource_limits"],
        )
        resolver = PostgresArtifactResolver(platform, store)
        resolver.register(
            binding_id=input_binding["id"],
            binding_hash=compiled["plan"]["input_binding_hash"],
            correlation_id=correlation_id,
            artifact=customer_artifact,
            kind="input",
        )
        resolver.register(
            binding_id=lookup_binding["id"],
            binding_hash=lookup_operation["lookup_binding_hash"],
            correlation_id=correlation_id,
            artifact=segment_artifact,
            kind="lookup",
        )
        node = {
            "id": "task5_real_polars",
            "type": "rule.apply",
            "purpose": "write",
            "idempotency": {
                "strategy": "deduplication_key",
                "key": "customer_id",
            },
            "config": {
                "component_binding_id": component_binding_id,
                "rule_version_id": rule["id"],
                "execution_plan_hash": compiled["plan_hash"],
            },
        }
        executor = RulePlanExecutor(
            PostgresRulePlanRepository(platform),
            adapters={
                "polars_batch": PolarsRulePlanAdapter(
                    artifact_store=store,
                    artifact_resolver=resolver,
                    artifact_ttl_seconds=900,
                )
            },
        )
        result = executor.execute(
            node,
            {},
            write_authorized=True,
            correlation_id=correlation_id,
        )
        repeated = executor.execute(
            node,
            {},
            write_authorized=True,
            correlation_id=correlation_id,
        )

        assert result["rows_in"] == 4
        assert result["rows_out"] == 2
        assert result["rows_rejected"] == 1
        assert result["rows_deduplicated"] == 1
        assert result["rows_filtered"] == 0
        assert result["rows_join_dropped"] == 0
        assert result["rows_aggregated"] == 0
        assert result["violation_count"] == 1
        assert result["violations"] == [
            {"step_id": "valid_mobile", "count": 1}
        ]
        output = store.read(
            result["artifact_ref"],
            result["digest"],
            expected_schema_fields=output_schema["fields"],
            limits=compiled["plan"]["resource_limits"],
        ).collect()
        assert output.sort("customer_id").to_dicts() == [
            {
                "customer_id": 1,
                "mobile": "13800138000",
                "name": "Alice Updated",
                "segment_code": "A",
                "segment_name": "Gold",
                "version_no": 2,
            },
            {
                "customer_id": 3,
                "mobile": "13900139000",
                "name": "Carol",
                "segment_code": "C",
                "segment_name": "Silver",
                "version_no": 1,
            },
        ]
        assert all(
            item.object_name.startswith(prefix)
            for item in minio.list_objects(
                bucket, prefix=prefix, recursive=True
            )
        )
        assert repeated["rows_out"] == 2
        with platform.connect() as connection:
            catalog_count = connection.execute(
                text(
                    """
                    SELECT COUNT(*)
                    FROM public.rule_run_artifacts
                    WHERE correlation_id = CAST(:correlation_id AS uuid)
                    """
                ),
                {"correlation_id": correlation_id},
            ).scalar_one()
        # The repeated deterministic output has the same digest and is
        # idempotently retained as one stable catalog handoff.
        assert catalog_count == 3
    finally:
        for item in list(
            minio.list_objects(bucket, prefix=prefix, recursive=True)
        ):
            minio.remove_object(bucket, item.object_name)
        assert list(
            minio.list_objects(bucket, prefix=prefix, recursive=True)
        ) == []
        with postgres.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
        with mysql.begin() as connection:
            connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
        with platform.begin() as connection:
            connection.execute(
                text(
                    "DELETE FROM public.rule_run_artifacts "
                    "WHERE correlation_id = CAST(:correlation_id AS uuid)"
                ),
                {"correlation_id": correlation_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.rule_execution_plans "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": plan_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_dataset_bindings "
                    "WHERE dataflow_deployment_id = CAST(:id AS uuid)"
                ),
                {"id": deployment_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_component_bindings "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": component_binding_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_deployments "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": deployment_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_versions "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": dataflow_version_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.data_rule_versions "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": rule_id},
            )
            connection.execute(
                text(
                    "DELETE FROM public.data_rules "
                    "WHERE rule_uid = CAST(:rule_uid AS uuid)"
                ),
                {"rule_uid": rule_uid},
            )
            for schema in (input_schema, lookup_schema, output_schema):
                connection.execute(
                    text(
                        "DELETE FROM public.data_schema_snapshots "
                        "WHERE id = CAST(:id AS uuid)"
                    ),
                    {"id": schema["id"]},
                )
        postgres.dispose()
        mysql.dispose()
        platform.dispose()