DataOps-platform/tests/integration/test_data_rule_polars_execution.py

from __future__ import annotations

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",
    }


class Resolver:
    def __init__(self, artifacts):
        self.artifacts = artifacts

    def resolve(self, *, binding_id, correlation_id):
        artifact = self.artifacts[binding_id]
        assert f"/rules/{correlation_id}/" in artifact["artifact_ref"]
        return artifact


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
    from app.runner.rule_polars import PolarsRulePlanAdapter

    source_user = _compose_value(
        r"source-postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    source_password = _compose_value(
        r"source-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"')
    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,
    )
    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"

    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()
            ]
        customer_artifact = store.write(
            pl.DataFrame(customer_rows).lazy(), correlation_id, 900
        )
        segment_artifact = store.write(
            pl.DataFrame(segment_rows).lazy(), correlation_id, 900
        )

        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": new_governance_uid(),
            "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]
        resolver = Resolver(
            {
                input_binding["id"]: {
                    **customer_artifact,
                    "binding_hash": compiled["plan"][
                        "input_binding_hash"
                    ],
                },
                lookup_binding["id"]: {
                    **segment_artifact,
                    "binding_hash": lookup_operation[
                        "lookup_binding_hash"
                    ],
                },
            }
        )
        node = {
            "id": "task5_real_polars",
            "type": "rule.apply",
            "purpose": "write",
            "idempotency": {
                "strategy": "deduplication_key",
                "key": "customer_id",
            },
            "config": {
                "component_binding_id": new_governance_uid(),
                "rule_version_id": rule["id"],
                "execution_plan_hash": compiled["plan_hash"],
            },
        }
        result = PolarsRulePlanAdapter(
            artifact_store=store,
            artifact_resolver=resolver,
            artifact_ttl_seconds=900,
        ).execute(
            plan=compiled["plan"],
            node=node,
            parameters={},
            write_authorized=True,
            correlation_id=correlation_id,
        )

        assert result["rows_in"] == 4
        assert result["rows_out"] == 2
        assert result["rows_rejected"] == 2
        assert result["violation_count"] == 1
        assert result["violations"] == [
            {"step_id": "valid_mobile", "count": 1}
        ]
        output = store.read(
            result["artifact_ref"], result["digest"]
        ).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
            )
        )
    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}"))
        postgres.dispose()
        mysql.dispose()