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()