orchestrata_ultimate.postman.replay_engine

Reference › orchestrata_ultimate › postman › replay_engine

The replay engine of the broker, SPEC-002 stories US-006 and FR-010 (REQ-002-US-006-014, -015, -018, REQ-002-FR10-010). ReplayEngine implements IReplayEngine from interfaces. It is built from a MessageStore and a ReceiptStore and keeps its own progress table in a SQLite file (./data/replays.db by default). Nothing else in the package imports it.

replay(message_ids, options) creates a replay id, writes a replays row with status running, sets a Redis progress counter, pushes a replay_started audit event and starts _process_replay as a background task, returning a ReplayResult at once. _process_replay walks the ids: it retrieves each message from the message store, counts it as replayed and writes a message_replayed audit event if found, and counts it as failed if it is missing or raises. After each message it updates the SQLite row and the Redis progress counter, and it waits delay_between_ms between messages. The final status is completed or completed_with_failures, followed by a replay_completed event. In this module a replayed message is read from the store and audited; the module does not call a router or the postman to send it again. replay_by_query selects ids by sender, recipient, topic and time range (up to 1000 by default) and replays them, and cancel_replay cancels a running task.

At a glance

Class diagram

Dependency diagram

API

Classes

ReplayEngine

Bases: IReplayEngine · Decorators: traces(spec='002', story='FR-010', reqs=['REQ-002-FR10-010', 'REQ-002-US-006-014', 'REQ-002-US-006-015']) · Line 67

Engine for replaying stored messages.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/postman/replay_engine.py

# mypy: ignore-errors
"""Replay Engine module for Orchestrata."""

from __future__ import annotations

import asyncio
import json
import logging
import sqlite3
import uuid
from collections.abc import Awaitable
from datetime import UTC, datetime
from typing import TYPE_CHECKING, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.postman.constants import DEFAULT_REPLAY_QUERY_LIMIT
from orchestrata_ultimate.postman.interfaces import IReplayEngine


def _redis_result[T](value: Awaitable[T] | T) -> Awaitable[T]:
    """Narrow a redis-py command result back to its true async type.

    redis-py types command methods (e.g. ``lpush``) as ``Awaitable[T] | T``
    because the same command mixin backs both the sync and async clients
    (see ``redis.typing.ResponseT``). ``SQLiteReplayEngine`` only ever holds
    an async client (``create_async_redis``), which always returns an
    awaitable at runtime -- this documents that fact for pyright instead of
    suppressing the check.
    """
    return cast(Awaitable[T], value)


if TYPE_CHECKING:
    from orchestrata_ultimate.postman.interfaces import IMessageStore, IReceiptStore
    from orchestrata_ultimate.postman.types import ReplayOptions, ReplayResult

# Spec: 002 - Postman Message Broker


"""
Replay engine using SQLite for state, Redis for progress, graph for audit.
Implements IReplayEngine interface.
"""


logger = logging.getLogger(__name__)

try:
    # The probe must actually import redis: with the old bare `pass` here,
    # REDIS_AVAILABLE was hardcoded True even when redis is not installed.
    import redis.asyncio  # noqa: F401

    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False


@traces(
    spec="002",
    story="FR-010",
    reqs=[
        "REQ-002-FR10-010",
        "REQ-002-US-006-014",
        "REQ-002-US-006-015",
    ],
)
class ReplayEngine(IReplayEngine):
    """Engine for replaying stored messages."""

    def __init__(
        self,
        message_store: IMessageStore,  # IMessageStore instance
        receipt_store: IReceiptStore,  # IReceiptStore instance (optional)
        sqlite_path: str = "./data/replays.db",
        redis_host: str | None = None,
        # No literal port default (RUL-003/RUL-011): "6379" would silently
        # aim the async client at nothing when REDIS_PORT isn't the target,
        # reading as an outage instead of the config gap it is. Resolved
        # lazily below via resolve_redis_port() so importing this module
        # never raises.
        redis_port: int | None = None,
        graph_uri: str | None = None,
        graph_user: str = "",
        graph_password: str | None = None,
    ) -> None:
        self.message_store = message_store
        self.receipt_store = receipt_store
        self.sqlite_path = sqlite_path
        self._init_sqlite()

        if REDIS_AVAILABLE:
            from orchestrata_ultimate.core.redis_factory import (
                create_async_redis,
                resolve_redis_host,
                resolve_redis_port,
            )

            self.redis_client = create_async_redis(
                host=redis_host if redis_host is not None else resolve_redis_host(),
                port=redis_port if redis_port is not None else resolve_redis_port(),
                decode_responses=False,
            )
        else:
            self.redis_client = None

        self.active_tasks: dict[str, asyncio.Task] = {}

    async def initialize(self) -> None:
        """Complete async initialisation.

        No-op: the SQLite schema is created synchronously in __init__.
        Kept for API compatibility with callers that await
        ``engine.initialize()`` after construction, e.g.::

            engine = ReplayEngine(...)
            await engine.initialize()
        """
        return None

    def _init_sqlite(self) -> None:
        """Create replays table."""

        def sync_init() -> None:
            conn = sqlite3.connect(self.sqlite_path)
            cur = conn.cursor()
            cur.execute("""
                CREATE TABLE IF NOT EXISTS replays (
                    replay_id TEXT PRIMARY KEY,
                    status TEXT,
                    total_messages INTEGER,
                    replayed INTEGER,
                    failed INTEGER,
                    created_at DATETIME,
                    updated_at DATETIME
                )
            """)
            conn.commit()
            conn.close()

        sync_init()

    async def _execute_sql(self, query: str, args: tuple = ()) -> list[tuple]:
        def sync_query() -> list[tuple]:
            conn = sqlite3.connect(self.sqlite_path)
            cur = conn.cursor()
            cur.execute(query, args)
            result = cur.fetchall()
            conn.commit()
            conn.close()
            return result

        return await asyncio.to_thread(sync_query)

    async def _execute_sql_write(self, query: str, args: tuple = ()) -> int:
        def sync_write() -> int:
            conn = sqlite3.connect(self.sqlite_path)
            cur = conn.cursor()
            cur.execute(query, args)
            rowcount = cur.rowcount
            conn.commit()
            conn.close()
            return rowcount

        return await asyncio.to_thread(sync_write)

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-018"])
    async def _update_replay(self, replay_id: str, **kwargs) -> None:
        """Update replay row with given fields."""
        set_clause = ", ".join(f"{k}=?" for k in kwargs)
        values = list(kwargs.values())
        await self._execute_sql_write(
            f"UPDATE replays SET {set_clause}, updated_at=? WHERE replay_id=?",
            (*values, datetime.now(UTC).isoformat(), replay_id),
        )

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-014"])
    async def replay(
        self,
        message_ids: list[str],
        options: ReplayOptions,
    ) -> ReplayResult:
        """Start replay of given message IDs."""
        from orchestrata_ultimate.postman.types import ReplayResult

        replay_id = str(uuid.uuid4())
        result = ReplayResult(
            replay_id=replay_id,
            total_messages=len(message_ids),
            replayed=0,
            failed=0,
            status="running",
        )
        # Store initial state
        await self._execute_sql_write(
            """
            INSERT INTO replays (replay_id, status, total_messages, replayed, failed, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                replay_id,
                "running",
                len(message_ids),
                0,
                0,
                datetime.now(UTC).isoformat(),
                datetime.now(UTC).isoformat(),
            ),
        )
        # Redis progress
        if self.redis_client:
            await self.redis_client.set(f"replay:{replay_id}:progress", 0)

        # Graph audit start (via queue)
        if self.redis_client:
            audit_event = {
                "event_type": "replay_started",
                "replay_id": replay_id,
                "total_messages": len(message_ids),
                "timestamp": datetime.now(UTC).isoformat(),
            }
            await _redis_result(
                self.redis_client.lpush(
                    "orchestrata:graph:audit:queue", json.dumps(audit_event)
                )
            )

        # Start background task
        task = asyncio.create_task(
            self._process_replay(replay_id, message_ids, options)
        )
        self.active_tasks[replay_id] = task

        return result

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-018"])
    async def _process_replay(
        self, replay_id: str, message_ids: list[str], options: ReplayOptions
    ) -> None:
        """Background processing of a replay."""

        replayed = 0
        failed = 0
        for mid in message_ids:
            try:
                # Retrieve message
                msg = await self.message_store.retrieve(mid)
                if msg:
                    envelope, payload = msg
                    # Simulate replay: log to audit queue
                    if self.redis_client:
                        audit_event = {
                            "event_type": "message_replayed",
                            "replay_id": replay_id,
                            "message_id": mid,
                            "sender": envelope.from_,
                            "recipient": envelope.to,
                            "timestamp": envelope.created_at.isoformat(),
                        }
                        await _redis_result(
                            self.redis_client.lpush(
                                "orchestrata:graph:audit:queue",
                                json.dumps(audit_event),
                            )
                        )
                    replayed += 1
                else:
                    failed += 1
            except Exception as exc:
                logger.warning("Replay failed for message_id=%s: %s", mid, exc)
                failed += 1

            # Update progress
            await self._update_replay(
                replay_id,
                replayed=replayed,
                failed=failed,
                status="running",
            )
            if self.redis_client:
                await self.redis_client.set(f"replay:{replay_id}:progress", replayed)

            # Optional delay
            if options.delay_between_ms > 0:
                await asyncio.sleep(options.delay_between_ms / 1000.0)

        # Mark finished
        final_status = "completed" if failed == 0 else "completed_with_failures"
        await self._update_replay(replay_id, status=final_status)
        if self.redis_client:
            await self.redis_client.delete(f"replay:{replay_id}:progress")
            # Audit completion
            audit_event = {
                "event_type": "replay_completed",
                "replay_id": replay_id,
                "status": final_status,
                "replayed": replayed,
                "failed": failed,
                "timestamp": datetime.now(UTC).isoformat(),
            }
            await _redis_result(
                self.redis_client.lpush(
                    "orchestrata:graph:audit:queue", json.dumps(audit_event)
                )
            )

        # Clean up task
        self.active_tasks.pop(replay_id, None)

    @traces(
        spec="002",
        story="US-006",
        reqs=[
            "REQ-002-US-006-011",
            "REQ-002-US-006-012",
            "REQ-002-US-006-013",
            "REQ-002-US-006-015",
        ],
    )
    async def replay_by_query(
        self,
        sender: str | None = None,
        recipient: str | None = None,
        topic: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        options: ReplayOptions | None = None,
    ) -> ReplayResult:
        """Replay messages matching query."""
        from orchestrata_ultimate.postman.types import ReplayOptions

        # Use message_store query to get IDs
        ids = await self.message_store.query(
            sender=sender,
            recipient=recipient,
            topic=topic,
            start_time=start_time,
            end_time=end_time,
            limit=DEFAULT_REPLAY_QUERY_LIMIT,
        )
        opts = options or ReplayOptions()
        return await self.replay(ids, opts)

    async def cancel_replay(self, replay_id: str) -> bool:
        """Cancel in‑progress replay."""
        task = self.active_tasks.get(replay_id)
        if task and not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                logger.debug("replay task %s cancelled", replay_id)
            await self._update_replay(replay_id, status="cancelled")
            if self.redis_client:
                await self.redis_client.delete(f"replay:{replay_id}:progress")
                # Audit cancellation
                audit_event = {
                    "event_type": "replay_cancelled",
                    "replay_id": replay_id,
                    "timestamp": datetime.now(UTC).isoformat(),
                }
                await _redis_result(
                    self.redis_client.lpush(
                        "orchestrata:graph:audit:queue", json.dumps(audit_event)
                    )
                )
            self.active_tasks.pop(replay_id, None)
            return True
        return False