DataOps-platform/tests/integration/test_rule_publication_lifecycle.py

from __future__ import annotations

import hashlib
import json
import os
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.repository import DataRuleRepository
from app.runner.artifacts import ArtifactStore
from tests.core.data_rules.test_contracts import 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,
    )
    fields = [
        {"name": "name", "type": "string", "nullable": True},
        {"name": "mobile", "type": "string", "nullable": True},
    ]
    sample = store.write(
        pl.DataFrame(
            {
                "name": [" Alice ", " Bob ", " Carol "],
                "mobile": ["13800138000", "invalid", "13900139000"],
            }
        ),
        new_governance_uid(),
        600,
        schema_fields=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()
            session = Session(bind=connection)
            repository = DataRuleRepository(session)
            snapshot_value = {
                    "schema_ref": "bd:rule-publication:integration",
                    "source_revision": "integration:1",
                    "fields": fields,
                }
            snapshot_value["schema_hash"] = canonical_schema_hash(
                snapshot_value["fields"]
            )
            snapshot = repository.persist_schema_snapshot(
                snapshot=snapshot_value
            )
            validation_context = {
                "schema_snapshot_id": snapshot["id"],
                "schema_hash": snapshot["schema_hash"],
                "fields": snapshot["fields"],
                "sample_artifact_ref": sample["artifact_ref"],
                "sample_artifact_digest": sample["digest"],
            }
            spec = valid_rule_spec()
            spec["input_schema_ref"] = snapshot["schema_ref"]
            spec["output_schema_ref"] = snapshot["schema_ref"]
            spec["steps"][1]["on_failure"] = "quarantine"
            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(minutes=5),
            )
            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)
            with pytest.raises(ValueError, match="consumable"):
                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",
                )

            compiled = service.validate(draft["id"], actor)
            assert compiled["plan_status"] == "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": 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["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
            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"]

            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"])