from __future__ import annotations import json import time from concurrent.futures import ThreadPoolExecutor import pytest from sqlalchemy import create_engine, text from sqlalchemy.orm import Session from app.core.common.identifiers import new_governance_uid from app.core.data_rules.compilers.sql import SqlGlotRuleCompiler from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec from app.core.data_rules.publication import ( PhysicalPlanPublicationService, ServerOwnedPhysicalPreflightRunner, ) from app.core.data_rules.repository import DataRuleRepository from tests.integration.test_data_rule_polars_execution import _compose_value from tests.integration.test_data_rule_sql_execution import ( CASES, ReadOnlyPreflightManager, _snapshot, ) pytestmark = pytest.mark.integration def _platform_url() -> str: user = _compose_value(r"POSTGRES_USER:\s*([^\s]+)") password = _compose_value(r"POSTGRES_PASSWORD:\s*([^\s]+)") port = _compose_value(r'"(15432):5432"') return ( f"postgresql+psycopg2://{user}:{password}" f"@127.0.0.1:{port}/dataops" ) def _insert_logical_trust( connection, *, rule_id: str, actor_id: str, input_schema: dict, output_schema: dict, compiled: dict, ) -> None: profile_id = new_governance_uid() logical_id = new_governance_uid() logical_schema_hashes = { "input": input_schema["schema_hash"], "output": output_schema["schema_hash"], } capabilities = compiled["plan"]["capabilities"] connection.execute( text( """ INSERT INTO public.rule_validation_profiles (id, rule_version_id, input_schema_snapshot_id, input_schema_hash, input_fields, output_schema_snapshot_id, output_schema_hash, output_fields, input_sample_artifact_ref, input_sample_artifact_digest, context_hash) VALUES (CAST(:id AS uuid), CAST(:rule_id AS uuid), CAST(:input_id AS uuid), :input_hash, CAST(:input_fields AS jsonb), CAST(:output_id AS uuid), :output_hash, CAST(:output_fields AS jsonb), 'test://physical-publication-input', :digest, :digest) """ ), { "id": profile_id, "rule_id": rule_id, "input_id": input_schema["id"], "input_hash": input_schema["schema_hash"], "input_fields": json.dumps(input_schema["fields"]), "output_id": output_schema["id"], "output_hash": output_schema["schema_hash"], "output_fields": json.dumps(output_schema["fields"]), "digest": "a" * 64, }, ) connection.execute( text( """ INSERT INTO public.rule_logical_plans (id, rule_version_id, validation_profile_id, compiler_version, backend, plan, plan_hash, schema_hashes, capabilities, status) VALUES (CAST(:id AS uuid), CAST(:rule_id AS uuid), CAST(:profile_id AS uuid), :compiler_version, 'sql_pushdown', CAST(:plan AS jsonb), :plan_hash, CAST(:schema_hashes AS jsonb), CAST(:capabilities AS jsonb), 'published') """ ), { "id": logical_id, "rule_id": rule_id, "profile_id": profile_id, "compiler_version": compiled["compiler_version"], "plan": json.dumps(compiled["plan"]), "plan_hash": compiled["plan_hash"], "schema_hashes": json.dumps(logical_schema_hashes), "capabilities": json.dumps(capabilities), }, ) connection.execute( text( """ INSERT INTO public.rule_logical_compile_evidence (id, logical_plan_id, compiler_version, compiler_digest, plan_hash, schema_hashes, capabilities, status, created_by) VALUES (CAST(:id AS uuid), CAST(:logical_id AS uuid), :compiler_version, :digest, :plan_hash, CAST(:schema_hashes AS jsonb), CAST(:capabilities AS jsonb), 'success', CAST(:actor_id AS uuid)) """ ), { "id": new_governance_uid(), "logical_id": logical_id, "compiler_version": compiled["compiler_version"], "digest": "b" * 64, "plan_hash": compiled["plan_hash"], "schema_hashes": json.dumps(logical_schema_hashes), "capabilities": json.dumps(capabilities), "actor_id": actor_id, }, ) connection.execute( text( """ INSERT INTO public.rule_logical_test_evidence (id, logical_plan_id, test_kind, evidence_hash, run_id, plan_hash, schema_hashes, evidence, status, created_by) VALUES (CAST(:id AS uuid), CAST(:logical_id AS uuid), 'dry_run', :digest, CAST(:run_id AS uuid), :plan_hash, CAST(:schema_hashes AS jsonb), '{}'::jsonb, 'success', CAST(:actor_id AS uuid)) """ ), { "id": new_governance_uid(), "logical_id": logical_id, "digest": "c" * 64, "run_id": new_governance_uid(), "plan_hash": compiled["plan_hash"], "schema_hashes": json.dumps(logical_schema_hashes), "actor_id": actor_id, }, ) @pytest.mark.parametrize( ("dialect", "url", "schema_name", "collation", "regex_engine"), CASES ) def test_physical_service_rejects_post_test_canonical_drift_and_replays( dialect, url, schema_name, collation, regex_engine ): source_engine = create_engine(url, pool_pre_ping=True) platform_engine = create_engine(_platform_url(), pool_pre_ping=True) suffix = dialect.replace("postgresql", "pg") source_table = f"task7_publish_source_{suffix}" target_table = f"task7_publish_target_{suffix}" source_ref = f"{schema_name}.{source_table}" target_ref = f"{schema_name}.{target_table}" input_schema = _snapshot(f"bd:task7:publish:{dialect}:input") output_schema = _snapshot(f"bd:task7:publish:{dialect}:output") data_source_uid = new_governance_uid() input_binding = { "id": new_governance_uid(), "data_source_uid": data_source_uid, "object_kind": "table", "object_ref": source_ref, "schema_snapshot_id": input_schema["id"], "access_mode": "read", "dialect": dialect, "write_mode": "append", } output_binding = { "id": new_governance_uid(), "data_source_uid": data_source_uid, "object_kind": "table", "object_ref": target_ref, "schema_snapshot_id": output_schema["id"], "access_mode": "write", "dialect": dialect, "write_mode": "append", } spec = validate_rule_spec( { "schema_version": "2.0", "rule_uid": new_governance_uid(), "name": f"task7_physical_publication_{dialect}", "input_schema_ref": input_schema["schema_ref"], "output_schema_ref": output_schema["schema_ref"], "steps": [ { "id": "trim_name", "op": "normalize_text", "column": "name", "trim": True, } ], "null_policy": "explicit", "timezone": "Asia/Shanghai", } ) rule_id = new_governance_uid() compiled = SqlGlotRuleCompiler(dialect).compile( rule_version={ "id": rule_id, "status": "published", "rule_spec": spec, "spec_hash": rule_spec_hash(spec), }, input_schema=input_schema, output_schema=output_schema, input_binding=input_binding, output_binding=output_binding, backend={ "dialect": dialect, "timezone": "Asia/Shanghai", "collation": collation, "rounding_mode": "half_away_from_zero", "regex_engine": regex_engine, }, ) try: with source_engine.begin() as source: source.execute(text(f"DROP TABLE IF EXISTS {target_table}")) source.execute(text(f"DROP TABLE IF EXISTS {source_table}")) source.execute( text( f"CREATE TABLE {source_table} (" "customer_id BIGINT PRIMARY KEY, name VARCHAR(100), " "mobile VARCHAR(30))" ) ) source.execute( text( f"CREATE TABLE {target_table} (" "customer_id BIGINT PRIMARY KEY, name VARCHAR(100), " "mobile VARCHAR(30))" ) ) source.execute( text( f"INSERT INTO {source_table} " "(customer_id, name, mobile) VALUES " "(1, ' Alice ', '13800138000')" ) ) with platform_engine.connect() as connection: transaction = connection.begin() try: actor_id = new_governance_uid() rule_uid = spec["rule_uid"] dataflow_version_id = new_governance_uid() deployment_id = new_governance_uid() component_id = new_governance_uid() connection.execute( text( "INSERT INTO public.users " "(id, username, display_name, password_hash, status) " "VALUES (CAST(:id AS uuid), :username, 'Task7', " "'not-a-login-secret', 'active')" ), { "id": actor_id, "username": f"task7-{actor_id[:8]}", }, ) for snapshot in (input_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)" ), { **snapshot, "fields": json.dumps(snapshot["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, " "'Task7 physical publication', 'en', " "CAST(:rule_spec AS jsonb), :spec_hash, 'sql', " "'published', CURRENT_TIMESTAMP)" ), { "id": rule_id, "rule_uid": rule_uid, "rule_spec": json.dumps(spec), "spec_hash": rule_spec_hash(spec), }, ) _insert_logical_trust( connection, rule_id=rule_id, actor_id=actor_id, input_schema=input_schema, output_schema=output_schema, compiled=compiled, ) 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(:uid AS uuid), 1, " "'Task7 physical publication', '{}'::jsonb, " "CAST(:inputs AS jsonb), :output, 'released', " "CURRENT_TIMESTAMP)" ), { "id": dataflow_version_id, "uid": new_governance_uid(), "inputs": json.dumps([input_schema["schema_hash"]]), "output": output_schema["schema_hash"], }, ) connection.execute( text( "INSERT INTO public.dataflow_deployments " "(id, dataflow_version_id, environment, " "deployment_config, status) VALUES " "(CAST(:id AS uuid), CAST(:version_id AS uuid), " "'test', '{}'::jsonb, 'disabled')" ), { "id": deployment_id, "version_id": dataflow_version_id, }, ) 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(:version_id AS uuid), " "'task7_publish', 'rule.apply', " "CAST(:rule_id AS uuid), 'transform', 0, " "'{\"strategy\":\"upsert\",\"key\":\"customer_id\"}'" "::jsonb, '{}'::jsonb)" ), { "id": component_id, "version_id": dataflow_version_id, "rule_id": rule_id, }, ) input_binding_hash = "d" * 64 output_binding_hash = "e" * 64 for logical_ref, binding, binding_hash in ( ( "input", input_binding, input_binding_hash, ), ( "output", output_binding, 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), 'table', " ":object_ref, CAST(:snapshot_id AS uuid), " ":dialect, :access_mode, 'append', :binding_hash)" ), { "id": binding["id"], "deployment_id": deployment_id, "logical_ref": logical_ref, "source_uid": data_source_uid, "object_ref": binding["object_ref"], "snapshot_id": binding["schema_snapshot_id"], "dialect": dialect, "access_mode": binding["access_mode"], "binding_hash": binding_hash, }, ) session = Session(bind=connection) repository = DataRuleRepository(session) persisted = repository.persist_bound_component_plan( component_binding_id=component_id, rule_version_id=rule_id, input_binding_id=input_binding["id"], output_binding_id=output_binding["id"], compiled=compiled, ) service = PhysicalPlanPublicationService( repository, test_runner=ServerOwnedPhysicalPreflightRunner( artifact_store=None, datasource_manager=ReadOnlyPreflightManager( source_engine ), ), ) validated = service.validate(persisted["id"], actor_id) assert service.validate(persisted["id"], actor_id) == validated tested = service.test(persisted["id"], actor_id) assert service.test(persisted["id"], actor_id) == tested drift_cases = ( ( "dataflow_dataset_bindings", input_binding["id"], "binding_hash", "f" * 64, input_binding_hash, ), ( "dataflow_dataset_bindings", input_binding["id"], "data_source_uid", new_governance_uid(), data_source_uid, ), ( "dataflow_dataset_bindings", input_binding["id"], "object_ref", f"{schema_name}.drifted_source", source_ref, ), ( "dataflow_dataset_bindings", input_binding["id"], "dialect", "mysql" if dialect == "postgresql" else "postgresql", dialect, ), ( "data_schema_snapshots", output_schema["id"], "schema_hash", "f" * 64, output_schema["schema_hash"], ), ( "data_rule_versions", rule_id, "status", "revoked", "published", ), ( "rule_execution_plans", persisted["id"], "compiler_version", "drifted-compiler", compiled["compiler_version"], ), ) for table_name, row_id, column, drifted, canonical in drift_cases: connection.execute( text( f"UPDATE public.{table_name} SET {column} = :value " "WHERE id = CAST(:id AS uuid)" ), {"id": row_id, "value": drifted}, ) with pytest.raises(ValueError, match="drifted|not found"): service.publish(persisted["id"], actor_id) connection.execute( text( f"UPDATE public.{table_name} SET {column} = :value " "WHERE id = CAST(:id AS uuid)" ), {"id": row_id, "value": canonical}, ) connection.execute( text( "UPDATE public.rule_test_evidence " "SET schema_hashes = '{}'::jsonb " "WHERE rule_execution_plan_id = CAST(:id AS uuid)" ), {"id": persisted["id"]}, ) with pytest.raises(ValueError, match="drifted"): service.publish(persisted["id"], actor_id) connection.execute( text( "UPDATE public.rule_test_evidence te " "SET schema_hashes = p.schema_hashes " "FROM public.rule_execution_plans p " "WHERE te.rule_execution_plan_id = p.id " "AND p.id = CAST(:id AS uuid)" ), {"id": persisted["id"]}, ) if dialect == "postgresql": second_actor_id = new_governance_uid() second_component_id = new_governance_uid() connection.execute( text( "INSERT INTO public.users " "(id, username, display_name, password_hash, " "status) VALUES (CAST(:id AS uuid), :username, " "'Task7 Concurrent', 'not-a-login-secret', " "'active')" ), { "id": second_actor_id, "username": f"task7-{second_actor_id}", }, ) 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(:version_id AS uuid), " "'task7_publish_concurrent', 'rule.apply', " "CAST(:rule_id AS uuid), 'transform', 1, " "'{\"strategy\":\"upsert\"," "\"key\":\"customer_id\"}'::jsonb, '{}'::jsonb)" ), { "id": second_component_id, "version_id": dataflow_version_id, "rule_id": rule_id, }, ) second_plan = repository.persist_bound_component_plan( component_binding_id=second_component_id, rule_version_id=rule_id, input_binding_id=input_binding["id"], output_binding_id=output_binding["id"], compiled=compiled, ) service.validate(second_plan["id"], actor_id) service.test(second_plan["id"], actor_id) transaction.commit() def publish_in_session( plan_id: str, publishing_actor: str, delay: float = 0.0, ): if delay: time.sleep(delay) with platform_engine.begin() as worker_connection: worker_service = PhysicalPlanPublicationService( DataRuleRepository( Session(bind=worker_connection) ), test_runner=None, ) try: return ( "ok", worker_service.publish( plan_id, publishing_actor ), ) except ValueError as exc: return ("rejected", str(exc)) with ThreadPoolExecutor(max_workers=2) as executor: same_actor = [ executor.submit( publish_in_session, persisted["id"], actor_id, ) for _ in range(2) ] same_actor_results = [ future.result() for future in same_actor ] assert [item[0] for item in same_actor_results] == [ "ok", "ok", ] assert ( same_actor_results[0][1] == same_actor_results[1][1] ) with ThreadPoolExecutor(max_workers=2) as executor: owner_future = executor.submit( publish_in_session, second_plan["id"], actor_id, ) other_future = executor.submit( publish_in_session, second_plan["id"], second_actor_id, 0.05, ) assert owner_future.result()[0] == "ok" assert other_future.result()[0] == "rejected" with platform_engine.begin() as cleanup: plan_ids = [ persisted["id"], second_plan["id"], ] cleanup.execute( text( "DELETE FROM public.rule_publication_audits " "WHERE rule_execution_plan_id = " "ANY(CAST(:ids AS uuid[]))" ), {"ids": plan_ids}, ) for evidence_table in ( "rule_test_evidence", "rule_compile_evidence", ): cleanup.execute( text( f"DELETE FROM public.{evidence_table} " "WHERE rule_execution_plan_id = " "ANY(CAST(:ids AS uuid[]))" ), {"ids": plan_ids}, ) cleanup.execute( text( "DELETE FROM public.rule_execution_plans " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), {"ids": plan_ids}, ) cleanup.execute( text( "DELETE FROM public.dataflow_dataset_bindings " "WHERE dataflow_deployment_id = " "CAST(:id AS uuid)" ), {"id": deployment_id}, ) cleanup.execute( text( "DELETE FROM " "public.dataflow_component_bindings " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), { "ids": [ component_id, second_component_id, ] }, ) cleanup.execute( text( "DELETE FROM public.dataflow_deployments " "WHERE id = CAST(:id AS uuid)" ), {"id": deployment_id}, ) cleanup.execute( text( "DELETE FROM public.dataflow_versions " "WHERE id = CAST(:id AS uuid)" ), {"id": dataflow_version_id}, ) cleanup.execute( text( "DELETE FROM " "public.rule_logical_test_evidence " "WHERE logical_plan_id IN (SELECT id FROM " "public.rule_logical_plans WHERE " "rule_version_id = CAST(:id AS uuid))" ), {"id": rule_id}, ) cleanup.execute( text( "DELETE FROM " "public.rule_logical_compile_evidence " "WHERE logical_plan_id IN (SELECT id FROM " "public.rule_logical_plans WHERE " "rule_version_id = CAST(:id AS uuid))" ), {"id": rule_id}, ) cleanup.execute( text( "DELETE FROM public.rule_logical_plans " "WHERE rule_version_id = CAST(:id AS uuid)" ), {"id": rule_id}, ) cleanup.execute( text( "DELETE FROM " "public.rule_validation_profiles " "WHERE rule_version_id = CAST(:id AS uuid)" ), {"id": rule_id}, ) cleanup.execute( text( "DELETE FROM public.data_rule_versions " "WHERE id = CAST(:id AS uuid)" ), {"id": rule_id}, ) cleanup.execute( text( "DELETE FROM public.data_rules " "WHERE rule_uid = CAST(:id AS uuid)" ), {"id": rule_uid}, ) cleanup.execute( text( "DELETE FROM public.data_schema_snapshots " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), { "ids": [ input_schema["id"], output_schema["id"], ] }, ) cleanup.execute( text( "DELETE FROM public.users " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), {"ids": [actor_id, second_actor_id]}, ) else: published = service.publish(persisted["id"], actor_id) assert ( service.publish(persisted["id"], actor_id) == published ) audit_count = connection.execute( text( "SELECT COUNT(*) FROM " "public.rule_publication_audits " "WHERE rule_execution_plan_id = " "CAST(:id AS uuid) AND action = 'published'" ), {"id": persisted["id"]}, ).scalar_one() assert audit_count == 1 finally: if transaction.is_active: transaction.rollback() finally: with source_engine.begin() as source: source.execute(text(f"DROP TABLE IF EXISTS {target_table}")) source.execute(text(f"DROP TABLE IF EXISTS {source_table}")) source_engine.dispose() platform_engine.dispose()