DataOps-platform/tests/integration/test_rule_publication_lifecycle.py

from __future__ import annotations

import hashlib
import json
import os
import time
from datetime import UTC, datetime, timedelta

import polars as pl
import pytest
from minio import Minio
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.contracts import rule_spec_hash
from app.core.data_rules.execution_contracts import canonical_schema_hash
from app.core.data_rules.publication import (
    GenerationReceiptSigner,
    LogicalRuleCompiler,
    RulePublicationService,
    ServerOwnedLogicalDryRunRunner,
    ServerOwnedPhysicalPreflightRunner,
    generation_receipt_claims,
)
from app.core.data_rules.release import ProductionLineReleaseService
from app.core.data_rules.repository import DataRuleRepository
from app.runner.artifacts import ArtifactStore
from tests.core.data_rules.test_contracts import (
    valid_dataflow_spec,
    valid_rule_spec,
)
from tests.integration.test_data_rule_polars_execution import _compose_value

pytestmark = pytest.mark.integration


def _hash(value):
    return hashlib.sha256(
        json.dumps(
            value,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
        ).encode("utf-8")
    ).hexdigest()


@pytest.fixture()
def database_url():
    value = os.environ.get("TEST_DATABASE_URL")
    if not value:
        pytest.skip("TEST_DATABASE_URL is not configured")
    return value


def test_real_postgres_receipt_to_logical_compile_test_publish(database_url):
    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]+)")
    store = ArtifactStore(
        Minio(
            f"127.0.0.1:{minio_port}",
            access_key=minio_user,
            secret_key=minio_password,
            secure=False,
        ),
        bucket=bucket,
        max_artifact_bytes=32 * 1024 * 1024,
        max_rows=100_000,
        memory_limit_bytes=256 * 1024 * 1024,
        max_ttl_seconds=3600,
    )
    input_fields = [
        {"name": "name", "type": "string", "nullable": True},
        {"name": "mobile", "type": "string", "nullable": True},
    ]
    output_fields = [
        *input_fields,
        {"name": "name_copy", "type": "string", "nullable": True},
    ]
    sample = store.write(
        pl.DataFrame(
            {
                "name": [" Alice ", " Bob ", " Carol "],
                "mobile": ["13800138000", "invalid", "13900139000"],
            }
        ),
        new_governance_uid(),
        600,
        schema_fields=input_fields,
    )
    golden = store.write(
        pl.DataFrame(
            {
                "name": ["Alice", "Carol"],
                "mobile": ["13800138000", "13900139000"],
                "name_copy": ["Alice", "Carol"],
            }
        ),
        new_governance_uid(),
        600,
        schema_fields=output_fields,
    )
    engine = create_engine(database_url)
    with engine.connect() as connection:
        transaction = connection.begin()
        try:
            actor = connection.execute(
                text(
                    "SELECT id::text FROM public.users "
                    "WHERE status = 'active' ORDER BY created_at LIMIT 1"
                )
            ).scalar_one_or_none()
            if actor is None:
                actor = new_governance_uid()
                connection.execute(
                    text(
                        "INSERT INTO public.users "
                        "(id, username, display_name, password_hash, status) "
                        "VALUES (CAST(:id AS uuid), :username, "
                        "'Task7 Integration', 'not-a-login-secret', 'active')"
                    ),
                    {
                        "id": actor,
                        "username": f"task7-{actor[:8]}",
                    },
                )
            session = Session(bind=connection)
            repository = DataRuleRepository(session)
            input_snapshot_value = {
                    "schema_ref": "bd:rule-publication:input",
                    "source_revision": "integration:1",
                    "fields": input_fields,
                }
            input_snapshot_value["schema_hash"] = canonical_schema_hash(
                input_snapshot_value["fields"]
            )
            input_snapshot = repository.persist_schema_snapshot(
                snapshot=input_snapshot_value
            )
            output_snapshot_value = {
                "schema_ref": "bd:rule-publication:output",
                "source_revision": "integration:1",
                "fields": output_fields,
            }
            output_snapshot_value["schema_hash"] = canonical_schema_hash(
                output_snapshot_value["fields"]
            )
            output_snapshot = repository.persist_schema_snapshot(
                snapshot=output_snapshot_value
            )
            validation_context = {
                "input_schema_snapshot_id": input_snapshot["id"],
                "input_schema_hash": input_snapshot["schema_hash"],
                "input_fields": input_snapshot["fields"],
                "output_schema_snapshot_id": output_snapshot["id"],
                "output_schema_hash": output_snapshot["schema_hash"],
                "output_fields": output_snapshot["fields"],
                "input_sample_artifact_ref": sample["artifact_ref"],
                "input_sample_artifact_digest": sample["digest"],
                "golden_output_artifact_ref": golden["artifact_ref"],
                "golden_output_artifact_digest": golden["digest"],
            }
            spec = valid_rule_spec()
            spec["input_schema_ref"] = input_snapshot["schema_ref"]
            spec["output_schema_ref"] = output_snapshot["schema_ref"]
            spec["steps"][1]["on_failure"] = "quarantine"
            spec["steps"].append(
                {
                    "id": "copy_name",
                    "op": "derive",
                    "target": "name_copy",
                    "expression": "name",
                }
            )
            candidate = {
                "schema_version": "1.0",
                "candidate_type": "rule",
                "rule_spec": spec,
                "standard_spec": None,
                "assumptions": [],
                "ambiguities": [],
                "confidence": 0.99,
                "explanation": "integration candidate",
            }
            evidence = {
                "status": "ready",
                "source_text": "手机号去空格后必须为11位数字",
                "authoring_surface": "data_standard",
                "candidate": candidate,
                "model_provider": "integration",
                "model_name": "closed-fixture",
                "prompt_version": "integration-v1",
                "schema_version": "1.0",
                "context_hash": _hash(validation_context),
                "candidate_hash": repository.candidate_hash(candidate),
                "model_hash": "a" * 64,
                "prompt_hash": "b" * 64,
                "repair_attempts": 0,
                "generation_attempts": [],
            }
            generation = repository.record_generation_run(
                evidence=evidence,
                created_by=actor,
                validation_context=validation_context,
            )
            signer = GenerationReceiptSigner(
                "integration-receipt-secret-with-entropy"
            )
            claims = generation_receipt_claims(
                generation_run_id=generation["id"],
                actor_uid=actor,
                source_text=evidence["source_text"],
                candidate_hash=evidence["candidate_hash"],
                rule_spec=spec,
                model_hash=evidence["model_hash"],
                prompt_hash=evidence["prompt_hash"],
                context_hash=evidence["context_hash"],
                expires_at=datetime.now(UTC) + timedelta(seconds=2),
            )
            receipt = signer.issue(claims)
            service = RulePublicationService(
                repository,
                receipt_signer=signer,
                compiler=LogicalRuleCompiler(),
                test_runner=ServerOwnedLogicalDryRunRunner(store),
            )
            draft = service.create_draft(
                rule_spec=spec,
                source_text=evidence["source_text"],
                actor_uid=actor,
                generation_receipt=receipt,
                category="standard_clause",
                source_language="zh-CN",
                generated_kind="rulespec",
            )
            assert draft["status"] == "draft"
            assert draft["spec_hash"] == rule_spec_hash(spec)
            assert service.create_draft(
                rule_spec=spec,
                source_text=evidence["source_text"],
                actor_uid=actor,
                generation_receipt=receipt,
                category="standard_clause",
                source_language="zh-CN",
                generated_kind="rulespec",
            ) == draft
            while int(datetime.now(UTC).timestamp()) < claims["expires_at"]:
                time.sleep(0.05)
            assert service.create_draft(
                rule_spec=spec,
                source_text=evidence["source_text"],
                actor_uid=actor,
                generation_receipt=receipt,
                category="standard_clause",
                source_language="zh-CN",
                generated_kind="rulespec",
            ) == draft
            next_generation = repository.record_generation_run(
                evidence=evidence,
                created_by=actor,
                validation_context=validation_context,
            )
            next_claims = generation_receipt_claims(
                generation_run_id=next_generation["id"],
                actor_uid=actor,
                source_text=evidence["source_text"],
                candidate_hash=evidence["candidate_hash"],
                rule_spec=spec,
                model_hash=evidence["model_hash"],
                prompt_hash=evidence["prompt_hash"],
                context_hash=evidence["context_hash"],
                expires_at=datetime.now(UTC) + timedelta(seconds=2),
            )
            next_receipt = signer.issue(next_claims)
            while (
                int(datetime.now(UTC).timestamp())
                < next_claims["expires_at"]
            ):
                time.sleep(0.05)
            with pytest.raises(ValueError, match="expired"):
                service.create_draft(
                    rule_spec=spec,
                    source_text=evidence["source_text"],
                    actor_uid=actor,
                    generation_receipt=next_receipt,
                    category="standard_clause",
                    source_language="zh-CN",
                    generated_kind="rulespec",
                )

            compiled = service.validate(draft["id"], actor)
            assert compiled["plan_status"] == "compiled"
            assert service.validate(draft["id"], actor) == compiled
            logical_plan = (
                connection.execute(
                    text(
                        "SELECT plan, plan_hash, schema_hashes "
                        "FROM public.rule_logical_plans "
                        "WHERE id = CAST(:id AS uuid)"
                    ),
                    {"id": compiled["plan_id"]},
                )
                .mappings()
                .one()
            )
            physical_result = ServerOwnedPhysicalPreflightRunner(store).run(
                {
                    "backend": "polars_batch",
                    "plan": logical_plan["plan"],
                    "plan_hash": logical_plan["plan_hash"],
                    "schema_hashes": logical_plan["schema_hashes"],
                    "binding_hashes": {},
                    "sample_artifact": {
                        "artifact_ref": sample["artifact_ref"],
                        "digest": sample["digest"],
                        "schema_fields": input_fields,
                    },
                }
            )
            assert physical_result["counts"]["rows_quarantined"] == 1
            tampered_sample = repository.load_logical_test_context(
                version_id=draft["id"],
                plan_id=compiled["plan_id"],
            )
            tampered_sample["input_sample_artifact_digest"] = "f" * 64
            with pytest.raises(ValueError, match="drifted"):
                ServerOwnedLogicalDryRunRunner(store).run(tampered_sample)
            tested = service.test(
                draft["id"], actor, plan_id=compiled["plan_id"]
            )
            assert tested["version_status"] == "validated"
            assert tested["test_evidence"]["counts"]["rows_quarantined"] == 1
            assert tested["test_evidence"]["counts"]["rows_rejected"] == 0
            connection.execute(
                text(
                    "UPDATE public.data_schema_snapshots "
                    "SET schema_hash = :drifted_hash "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {
                    "id": output_snapshot["id"],
                    "drifted_hash": "f" * 64,
                },
            )
            with pytest.raises(ValueError, match="drifted"):
                service.publish(draft["id"], actor)
            connection.execute(
                text(
                    "UPDATE public.data_schema_snapshots "
                    "SET schema_hash = :schema_hash "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {
                    "id": output_snapshot["id"],
                    "schema_hash": output_snapshot["schema_hash"],
                },
            )
            published = service.publish(draft["id"], actor)
            assert published["status"] == "published"
            assert published["plan_status"] == "published"
            assert service.publish(draft["id"], actor) == published
            assert service.catalog(query=spec["name"], limit=10)[0][
                "id"
            ] == draft["id"]

            flow = valid_dataflow_spec(rule_version_id=draft["id"])
            flow["input_schema_refs"] = [spec["input_schema_ref"]]
            flow["output_schema_ref"] = spec["output_schema_ref"]
            flow["components"] = [
                component
                for component in flow["components"]
                if component["type"] == "rule.apply"
            ]

            class PinnedResolver:
                def resolve(self, schema_ref):
                    if schema_ref == input_snapshot["schema_ref"]:
                        return input_snapshot
                    if schema_ref == output_snapshot["schema_ref"]:
                        return output_snapshot
                    raise ValueError("unexpected schema ref")

            release_service = ProductionLineReleaseService(
                repository,
                schema_resolver=PinnedResolver(),
            )
            logical_test_id = tested["test_evidence_id"]
            connection.execute(
                text(
                    "UPDATE public.rule_logical_test_evidence "
                    "SET status = 'failed' "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": logical_test_id},
            )
            with pytest.raises(ValueError, match="published rule"):
                release_service.release(
                    dataflow_uid=flow["dataflow_uid"],
                    dataflow_spec=flow,
                    source_text="跨模式数据生产线",
                    created_by=actor,
                )
            connection.execute(
                text(
                    "UPDATE public.rule_logical_test_evidence "
                    "SET status = 'success', schema_hashes = '{}'::jsonb "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": logical_test_id},
            )
            with pytest.raises(ValueError, match="published rule"):
                release_service.release(
                    dataflow_uid=flow["dataflow_uid"],
                    dataflow_spec=flow,
                    source_text="跨模式数据生产线",
                    created_by=actor,
                )
            connection.execute(
                text(
                    "UPDATE public.rule_logical_test_evidence te "
                    "SET schema_hashes = lp.schema_hashes "
                    "FROM public.rule_logical_plans lp "
                    "WHERE te.logical_plan_id = lp.id "
                    "AND te.id = CAST(:id AS uuid)"
                ),
                {"id": logical_test_id},
            )
            released = release_service.release(
                dataflow_uid=flow["dataflow_uid"],
                dataflow_spec=flow,
                source_text="跨模式数据生产线",
                created_by=actor,
            )
            assert released["status"] == "released"

            rows = connection.execute(
                text(
                    "SELECT rv.status, lp.status, "
                    "(SELECT COUNT(*) FROM public.rule_logical_compile_evidence "
                    "WHERE logical_plan_id = lp.id) AS compile_count, "
                    "(SELECT COUNT(*) FROM public.rule_logical_test_evidence "
                    "WHERE logical_plan_id = lp.id) AS test_count "
                    "FROM public.data_rule_versions rv "
                    "JOIN public.rule_logical_plans lp "
                    "ON lp.rule_version_id = rv.id "
                    "WHERE rv.id = CAST(:id AS uuid)"
                ),
                {"id": draft["id"]},
            ).one()
            assert tuple(rows) == ("published", "published", 1, 1)
        finally:
            transaction.rollback()
            engine.dispose()
            store.delete(sample["artifact_ref"])
            store.delete(golden["artifact_ref"])