agentflow-runtime 1.1.0__py3-none-any.whl

src/orchestration/dags/daily_batch.py

@@ -0,0 +1,201 @@
+"""Daily batch DAG: compaction, aggregation, and quality reports.
+
+Runs daily at 02:00 UTC when streaming traffic is lowest.
+Handles Iceberg maintenance and pre-computes aggregates for fast serving.
+
+Local mode uses DuckDB; production uses Trino/Iceberg catalog.
+"""
+
+import os
+
+import duckdb
+from dagster import (
+    AssetExecutionContext,
+    Definitions,
+    ScheduleDefinition,
+    asset,
+    define_asset_job,
+)
+
+DB_PATH = os.getenv("DUCKDB_PATH", ":memory:")
+
+
+def _get_conn():
+    return duckdb.connect(DB_PATH)
+
+
+@asset(group_name="maintenance")
+def iceberg_snapshot_expiry(context: AssetExecutionContext):
+    """Expire old Iceberg snapshots to prevent metadata bloat.
+
+    Keeps last 30 days of snapshots for time-travel debugging.
+    In local mode: no-op (DuckDB has no snapshots).
+    In production: calls Iceberg REST catalog expire_snapshots.
+    """
+    if DB_PATH == ":memory:":
+        context.log.info("Local mode: snapshot expiry is a no-op")
+        return {"mode": "local", "tables_processed": 0}
+
+    context.log.info("Expiring Iceberg snapshots older than 30 days")
+    tables = ["orders_v2", "sessions_aggregated", "products_current"]
+    # Production: catalog.expire_snapshots(t, older_than_days=30)
+    return {"tables_processed": len(tables)}
+
+
+@asset(group_name="maintenance", deps=[iceberg_snapshot_expiry])
+def iceberg_compaction(context: AssetExecutionContext):
+    """Compact small Iceberg data files into larger ones.
+
+    Target: 128-512 MB per file for optimal read performance.
+    In local mode: runs DuckDB VACUUM/CHECKPOINT.
+    """
+    if DB_PATH == ":memory:":
+        context.log.info("Local mode: compaction is a no-op")
+        return {"mode": "local", "tables_compacted": 0}
+
+    context.log.info("Compacting data files")
+    conn = _get_conn()
+    conn.execute("CHECKPOINT")
+    conn.close()
+    return {"tables_compacted": 3}
+
+
+@asset(group_name="aggregation")
+def daily_user_profiles(context: AssetExecutionContext):
+    """Pre-compute user profile aggregates for fast entity lookups.
+
+    Materializes users_enriched from orders_v2.
+    """
+    conn = _get_conn()
+
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS users_enriched (
+            user_id VARCHAR PRIMARY KEY,
+            total_orders INTEGER DEFAULT 0,
+            total_spent DECIMAL(10,2) DEFAULT 0,
+            first_order_at TIMESTAMP,
+            last_order_at TIMESTAMP,
+            preferred_category VARCHAR
+        )
+    """)
+
+    result = conn.execute("""
+        SELECT
+            user_id,
+            COUNT(*) as total_orders,
+            SUM(total_amount) as total_spent,
+            MIN(created_at) as first_order_at,
+            MAX(created_at) as last_order_at
+        FROM orders_v2
+        WHERE status != 'cancelled'
+        GROUP BY user_id
+    """).fetchall()
+
+    for row in result:
+        conn.execute(
+            """
+            INSERT OR REPLACE INTO users_enriched
+            (user_id, total_orders, total_spent,
+             first_order_at, last_order_at)
+            VALUES (?, ?, ?, ?, ?)
+        """,
+            list(row),
+        )
+
+    conn.close()
+    context.log.info("User profiles updated: %d", len(result))
+    return {"users_updated": len(result)}
+
+
+@asset(group_name="aggregation")
+def daily_product_metrics(context: AssetExecutionContext):
+    """Pre-compute product-level metrics from pipeline events."""
+    conn = _get_conn()
+
+    count = conn.execute("SELECT COUNT(*) FROM products_current").fetchone()[0]
+    conn.close()
+
+    context.log.info("Product metrics refreshed: %d", count)
+    return {"products_updated": count}
+
+
+@asset(
+    group_name="quality",
+    deps=[daily_user_profiles, daily_product_metrics],
+)
+def daily_quality_report(context: AssetExecutionContext):
+    """Generate daily data quality report.
+
+    Checks:
+    - Row counts per table (alert on >20% daily deviation)
+    - Null rates for required fields
+    - Dead letter ratio
+    """
+    conn = _get_conn()
+
+    checks = {}
+    for table in [
+        "orders_v2",
+        "users_enriched",
+        "products_current",
+        "sessions_aggregated",
+    ]:
+        try:
+            row = conn.execute(
+                f"SELECT COUNT(*) FROM {table}"  # nosec B608 - table comes from the fixed health-check allowlist
+            ).fetchone()
+            checks[table] = {"rows": row[0], "status": "ok"}
+        except duckdb.Error as e:
+            checks[table] = {"rows": 0, "status": f"error: {e}"}
+
+    # Dead letter ratio
+    try:
+        total = conn.execute("SELECT COUNT(*) FROM pipeline_events").fetchone()[0]
+        dead = conn.execute(
+            "SELECT COUNT(*) FROM pipeline_events WHERE topic = 'events.deadletter'"
+        ).fetchone()[0]
+        dl_ratio = dead / total if total > 0 else 0.0
+        checks["dead_letter_ratio"] = {"value": round(dl_ratio, 4)}
+    except duckdb.Error:
+        checks["dead_letter_ratio"] = {"value": None}
+
+    conn.close()
+
+    failed = sum(
+        1
+        for v in checks.values()
+        if isinstance(v, dict) and v.get("status", "").startswith("error")
+    )
+    context.log.info("Quality report: %s", checks)
+    return {"checks": checks, "checks_failed": failed}
+
+
+# ── Job & Schedule ──────────────────────────────────────────────
+
+daily_maintenance_job = define_asset_job(
+    name="daily_maintenance",
+    selection=[
+        iceberg_snapshot_expiry,
+        iceberg_compaction,
+        daily_user_profiles,
+        daily_product_metrics,
+        daily_quality_report,
+    ],
+)
+
+daily_schedule = ScheduleDefinition(
+    job=daily_maintenance_job,
+    cron_schedule="0 2 * * *",  # 02:00 UTC daily
+)
+
+defs = Definitions(
+    assets=[
+        iceberg_snapshot_expiry,
+        iceberg_compaction,
+        daily_user_profiles,
+        daily_product_metrics,
+        daily_quality_report,
+    ],
+    jobs=[daily_maintenance_job],
+    schedules=[daily_schedule],
+)

src/processing/event_replayer.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Callable
+from contextlib import nullcontext
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from uuid import uuid4
+
+import structlog
+from opentelemetry import trace
+
+from src.processing.outbox import OutboxProcessor, ensure_outbox_table
+from src.processing.tracing import inject_trace_to_kafka_headers, telemetry_disabled
+from src.quality.validators.schema_validator import validate_event
+from src.quality.validators.semantic_validator import validate_semantics
+
+DEFAULT_KAFKA_BOOTSTRAP = os.getenv("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
+DEFAULT_REPLAY_TOPIC = "events.raw"
+tracer = trace.get_tracer("agentflow.event_replayer")
+
+
+class DeadLetterEventNotFoundError(LookupError):
+    pass
+
+
+class ReplayValidationError(ValueError):
+    pass
+
+
+@dataclass
+class ReplayResult:
+    event_id: str
+    status: str
+    retry_count: int
+    last_retried_at: datetime
+    payload: dict
+
+
+def ensure_dead_letter_table(conn) -> None:
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS dead_letter_events (
+            event_id TEXT PRIMARY KEY,
+            tenant_id TEXT DEFAULT 'default',
+            event_type TEXT,
+            payload JSON,
+            failure_reason TEXT,
+            failure_detail TEXT,
+            received_at TIMESTAMP,
+            retry_count INTEGER DEFAULT 0,
+            last_retried_at TIMESTAMP,
+            status TEXT DEFAULT 'failed'
+        )
+        """
+    )
+    conn.execute(
+        "ALTER TABLE dead_letter_events ADD COLUMN IF NOT EXISTS tenant_id TEXT DEFAULT 'default'"
+    )
+
+
+class EventReplayer:
+    def __init__(
+        self,
+        conn,
+        producer: Callable[[str, dict], None] | None = None,
+        bootstrap_servers: str | None = None,
+    ) -> None:
+        self._conn = conn
+        self._producer = producer or self._produce_to_kafka
+        self._bootstrap_servers = bootstrap_servers or DEFAULT_KAFKA_BOOTSTRAP
+        ensure_dead_letter_table(self._conn)
+        ensure_outbox_table(self._conn)
+
+    def replay(
+        self,
+        event_id: str,
+        corrected_payload: dict | None = None,
+    ) -> ReplayResult:
+        row = self._load_row(event_id)
+        payload = self._decoded_payload(row["payload"])
+        candidate = dict(payload)
+        if corrected_payload:
+            candidate.update(corrected_payload)
+        self._validate(candidate)
+
+        replayed_at = datetime.now(UTC)
+        retry_count = int(row["retry_count"] or 0) + 1
+        outbox_id = str(uuid4())
+        self._conn.execute("BEGIN TRANSACTION")
+        try:
+            self._conn.execute(
+                """
+                UPDATE dead_letter_events
+                SET payload = ?, status = 'replay_pending', retry_count = ?, last_retried_at = ?
+                WHERE event_id = ?
+                """,
+                [
+                    json.dumps(candidate),
+                    retry_count,
+                    replayed_at,
+                    event_id,
+                ],
+            )
+            self._conn.execute(
+                """
+                INSERT INTO outbox (
+                    id,
+                    event_id,
+                    payload,
+                    topic,
+                    created_at,
+                    sent_at,
+                    status,
+                    retry_count,
+                    next_attempt_at,
+                    last_error
+                )
+                VALUES (?, ?, ?, ?, ?, NULL, 'pending', 0, ?, NULL)
+                """,
+                [
+                    outbox_id,
+                    event_id,
+                    json.dumps(candidate),
+                    DEFAULT_REPLAY_TOPIC,
+                    replayed_at,
+                    replayed_at,
+                ],
+            )
+            self._conn.execute("COMMIT")
+        except Exception:  # nosec B110 - rollback must preserve the original replay failure
+            # Transaction rollback must happen before unexpected errors propagate.
+            self._conn.execute("ROLLBACK")
+            raise
+        status = "replay_pending"
+        processor = OutboxProcessor(
+            conn=self._conn,
+            producer=self._producer,
+            bootstrap_servers=self._bootstrap_servers,
+        )
+        if processor.process_entry(outbox_id):
+            status = "replayed"
+        return ReplayResult(
+            event_id=event_id,
+            status=status,
+            retry_count=retry_count,
+            last_retried_at=replayed_at,
+            payload=candidate,
+        )
+
+    def dismiss(self, event_id: str) -> None:
+        self._load_row(event_id)
+        self._conn.execute(
+            "UPDATE dead_letter_events SET status = 'dismissed' WHERE event_id = ?",
+            [event_id],
+        )
+
+    def _load_row(self, event_id: str) -> dict:
+        row = self._conn.execute(
+            """
+            SELECT
+                event_id,
+                payload,
+                retry_count
+            FROM dead_letter_events
+            WHERE event_id = ?
+            """,
+            [event_id],
+        ).fetchone()
+        if row is None:
+            raise DeadLetterEventNotFoundError(event_id)
+        return {
+            "event_id": row[0],
+            "payload": row[1],
+            "retry_count": row[2],
+        }
+
+    def _decoded_payload(self, payload) -> dict:
+        if isinstance(payload, dict):
+            return payload
+        if isinstance(payload, str):
+            decoded = json.loads(payload)
+            if isinstance(decoded, dict):
+                return decoded
+        raise ReplayValidationError("Dead-letter payload must be a JSON object.")
+
+    def _validate(self, payload: dict) -> None:
+        schema_result = validate_event(payload)
+        if not schema_result.is_valid:
+            first_error = schema_result.errors[0] if schema_result.errors else {}
+            raise ReplayValidationError(first_error.get("msg", "Schema validation failed."))
+
+        semantic_result = validate_semantics(payload)
+        semantic_errors = [
+            issue.message for issue in semantic_result.issues if issue.severity == "error"
+        ]
+        if semantic_errors:
+            raise ReplayValidationError(semantic_errors[0])
+
+    def _produce_to_kafka(self, topic: str, payload: dict) -> None:
+        from confluent_kafka import Producer
+
+        delivery_errors: list[str] = []
+
+        def on_delivery(err, msg) -> None:
+            del msg
+            if err is not None:
+                delivery_errors.append(str(err))
+
+        producer = Producer({"bootstrap.servers": self._bootstrap_servers})
+        produce_span = (
+            tracer.start_as_current_span("kafka.produce")
+            if not telemetry_disabled()
+            else nullcontext()
+        )
+        with produce_span as span:
+            if span is not None and span.is_recording():
+                span.set_attribute("topic", topic)
+                event_type = payload.get("event_type")
+                if event_type is not None:
+                    span.set_attribute("event_type", str(event_type))
+                tenant_id = payload.get("tenant_id") or structlog.contextvars.get_contextvars().get(
+                    "tenant_id"
+                )
+                if tenant_id is not None:
+                    span.set_attribute("tenant_id", str(tenant_id))
+            headers = inject_trace_to_kafka_headers({})
+            try:
+                producer.produce(
+                    topic,
+                    key=str(payload.get("event_id", "")),
+                    value=json.dumps(payload).encode("utf-8"),
+                    headers=list(headers.items()) or None,
+                    on_delivery=on_delivery,
+                )
+            except TypeError as exc:
+                if "on_delivery" not in str(exc):
+                    raise
+                producer.produce(
+                    topic,
+                    key=str(payload.get("event_id", "")),
+                    value=json.dumps(payload).encode("utf-8"),
+                    headers=list(headers.items()) or None,
+                )
+            remaining = producer.flush(10)
+            if delivery_errors:
+                raise RuntimeError(delivery_errors[0])
+            if remaining != 0:
+                raise RuntimeError(f"{remaining} Kafka message(s) were not delivered")

src/processing/flink_jobs/Dockerfile

@@ -0,0 +1,55 @@
+FROM python:3.11-slim-bookworm
+
+ARG FLINK_VERSION=1.19.1
+ARG SCALA_VERSION=2.12
+ARG PYFLINK_KAFKA_JAR_VERSION=3.3.0-1.19
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends bash ca-certificates curl openjdk-17-jre-headless \
+    && curl -fsSL "https://archive.apache.org/dist/flink/flink-${FLINK_VERSION}/flink-${FLINK_VERSION}-bin-scala_${SCALA_VERSION}.tgz" -o /tmp/flink.tgz \
+    && tar -xzf /tmp/flink.tgz -C /opt \
+    && ln -s "/opt/flink-${FLINK_VERSION}" /opt/flink \
+    && curl -fsSL "https://repo.maven.apache.org/maven2/org/apache/flink/flink-sql-connector-kafka/${PYFLINK_KAFKA_JAR_VERSION}/flink-sql-connector-kafka-${PYFLINK_KAFKA_JAR_VERSION}.jar" -o "/opt/flink/lib/flink-sql-connector-kafka-${PYFLINK_KAFKA_JAR_VERSION}.jar" \
+    && mkdir -p /opt/flink/plugins/s3-fs-hadoop \
+    && ln -s "/opt/flink/opt/flink-s3-fs-hadoop-${FLINK_VERSION}.jar" "/opt/flink/plugins/s3-fs-hadoop/flink-s3-fs-hadoop-${FLINK_VERSION}.jar" \
+    && python -m pip install --no-cache-dir --upgrade pip \
+    && python -m pip install --no-cache-dir "apache-flink==${FLINK_VERSION}" confluent-kafka pydantic structlog \
+    && rm -rf /var/lib/apt/lists/* /tmp/flink.tgz
+
+WORKDIR /opt/agentflow
+ENV FLINK_HOME=/opt/flink
+ENV JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64
+ENV PATH="/opt/flink/bin:${PATH}"
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+ENV PYFLINK_CLIENT_EXECUTABLE=/usr/local/bin/python
+ENV PYTHONPATH=/opt/agentflow
+
+COPY . /opt/agentflow/src
+
+RUN cat <<'PY' > /opt/agentflow/sitecustomize.py
+from datetime import timedelta as _timedelta
+
+try:
+    from pyflink.common import WatermarkStrategy
+    from pyflink.common.time import Duration, Time
+    from pyflink.datastream.state import StateTtlConfig
+except Exception:
+    pass
+else:
+    _original_watermark = WatermarkStrategy.for_bounded_out_of_orderness
+    _original_ttl_builder = StateTtlConfig.new_builder
+
+    def _patched_watermark(max_out_of_orderness):
+        if isinstance(max_out_of_orderness, _timedelta):
+            max_out_of_orderness = Duration.of_millis(int(max_out_of_orderness.total_seconds() * 1000))
+        return _original_watermark(max_out_of_orderness)
+
+    def _patched_ttl_builder(ttl):
+        if isinstance(ttl, _timedelta):
+            ttl = Time.milliseconds(int(ttl.total_seconds() * 1000))
+        return _original_ttl_builder(ttl)
+
+    WatermarkStrategy.for_bounded_out_of_orderness = staticmethod(_patched_watermark)
+    StateTtlConfig.new_builder = staticmethod(_patched_ttl_builder)
+PY

src/processing/flink_jobs/checkpointing.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+
+def _checkpoint_constants() -> tuple[object, object]:
+    try:
+        from pyflink.datastream import CheckpointingMode
+        from pyflink.datastream.checkpoint_config import ExternalizedCheckpointCleanup
+    except ModuleNotFoundError:
+        return "EXACTLY_ONCE", "RETAIN_ON_CANCELLATION"
+
+    return (
+        CheckpointingMode.EXACTLY_ONCE,
+        ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION,
+    )
+
+
+def configure_checkpointing(env: Any) -> None:
+    checkpoint_mode, cleanup_mode = _checkpoint_constants()
+
+    env.enable_checkpointing(60_000)
+    config = env.get_checkpoint_config()
+    config.set_checkpointing_mode(checkpoint_mode)
+    config.set_min_pause_between_checkpoints(30_000)
+    config.set_checkpoint_timeout(120_000)
+    config.set_max_concurrent_checkpoints(1)
+    config.enable_externalized_checkpoints(cleanup_mode)
+    config.set_checkpoint_storage(
+        os.getenv("FLINK_CHECKPOINT_DIR", "file:///tmp/flink-checkpoints")
+    )