orchestrata_ultimate.postman.message_store

Reference › orchestrata_ultimate › postman › message_store

The durable store for message envelopes and payloads, SPEC-002 story US-006 (REQ-002-US-006-001 to -003, -008, -014). MessageStore implements IMessageStore from interfaces and keeps messages in three layers: SQLite as the primary record, a Redis cache, and an audit event queue for the graph. The replay_engine reads from it, and receipt_store is its companion for receipts. Value types are in types. The curated postman page describes Redis with an SQLite fallback; in this module SQLite is always written first and Redis is the cache.

initialize creates the messages table in the SQLite file (./data/messages.db by default). SQL runs in a worker thread through _execute_sql and _execute_sql_write. store(message_id, envelope, payload) inserts the row, and when a Redis client exists it also caches the message under msg:<id> for an hour and pushes a message_stored event onto an audit list. retrieve tries the Redis cache first, falls back to SQLite and refills the cache. delete removes the message everywhere. query filters by sender, recipient, topic, time range and a limit (100 by default). cleanup(older_than) deletes rows before a cutoff, records a messages_cleaned_up audit event and returns the count. get_stats reports storage statistics.

At a glance

Class diagram

Dependency diagram

API

Classes

MessageStore

Bases: IMessageStore · Decorators: traces(spec='002', story='US-006', reqs=['REQ-002-US-006-001', 'REQ-002-US-006-002', 'REQ-002-US-006-003']) · Line 63

Message store using SQLite (primary), Redis (cache), graph (audit).

Functions

Source

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

# mypy: ignore-errors
"""Message Store module for Orchestrata."""

from __future__ import annotations

import asyncio
import json
import pickle
import sqlite3
from collections.abc import Awaitable
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.postman.constants import (
    CONTENT_TYPE_COMPRESSED,
    DEFAULT_MESSAGE_QUERY_LIMIT,
)
from orchestrata_ultimate.postman.interfaces import IMessageStore


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``). ``SQLiteMessageStore`` 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.types import MessageEnvelope, MessagePayload

"""
SQLite message store with Redis cache and graph audit.
Implements IMessageStore interface.
"""

AUDIT_QUEUE_KEY = "orchestrata:graph:audit:queue"


try:
    from redis.asyncio import Redis as AsyncRedis  # noqa: F401

    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False


@traces(
    spec="002",
    story="US-006",
    reqs=[
        "REQ-002-US-006-001",
        "REQ-002-US-006-002",
        "REQ-002-US-006-003",
    ],
)
class MessageStore(IMessageStore):
    """Message store using SQLite (primary), Redis (cache), graph (audit)."""

    def __init__(
        self,
        sqlite_path: str = "./data/messages.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.sqlite_path = sqlite_path

        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

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

        Call once after construction before first use, e.g.::

            store = MessageStore(...)
            await store.initialize()
        """
        await asyncio.to_thread(self._init_sqlite)

    def _init_sqlite(self) -> None:
        """Create messages table if not exists."""
        conn = sqlite3.connect(self.sqlite_path)
        cur = conn.cursor()
        cur.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                message_id TEXT PRIMARY KEY,
                sender TEXT,
                recipient TEXT,
                timestamp DATETIME,
                content_type TEXT,
                message_type TEXT,
                envelope_json TEXT,
                payload_data BLOB
            )
        """)
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_sender ON messages(sender)
        """)
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_recipient ON messages(recipient)
        """)
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_timestamp ON messages(timestamp)
        """)
        conn.commit()
        conn.close()

    async def _execute_sql(self, query: str, args: tuple = ()) -> list[tuple]:
        """Run SQL query in a thread."""

        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:
        """Run write SQL, return last row id."""

        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-001", "REQ-002-US-006-002", "REQ-002-US-006-003"],
    )
    async def store(
        self,
        message_id: str,
        envelope: MessageEnvelope,
        payload: MessagePayload,
    ) -> None:
        """Store message in SQLite, cache in Redis, audit in graph."""
        # SQLite
        envelope_json = json.dumps(envelope.to_dict())
        await self._execute_sql_write(
            """
            INSERT INTO messages (message_id, sender, recipient, timestamp,
                                  content_type, message_type, envelope_json, payload_data)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                message_id,
                envelope.from_,
                envelope.to,
                envelope.created_at.isoformat(),
                payload.content_type,
                envelope.message_type.value,
                envelope_json,
                (
                    payload.content.encode("utf-8")
                    if isinstance(payload.content, str)
                    else payload.content
                ),
            ),
        )

        # Redis cache
        if self.redis_client:
            payload_bytes = (
                payload.content.encode("utf-8")
                if isinstance(payload.content, str)
                else payload.content
            )
            cache_data = pickle.dumps(
                (
                    envelope.to_dict(),
                    payload_bytes,
                    payload.content_type,
                    payload.compressed,
                    payload.size_bytes,
                )
            )
            await self.redis_client.setex(f"msg:{message_id}", 3600, cache_data)

        # Graph audit via Redis queue
        if self.redis_client:
            audit_event = {
                "event_type": "message_stored",
                "message_id": message_id,
                "sender": envelope.from_,
                "recipient": envelope.to,
                "timestamp": envelope.created_at.isoformat(),
                "action": "stored",
            }
            await _redis_result(
                self.redis_client.lpush(AUDIT_QUEUE_KEY, json.dumps(audit_event))
            )

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-014"])
    async def retrieve(
        self, message_id: str
    ) -> tuple[MessageEnvelope, MessagePayload] | None:
        """Retrieve message by ID, trying Redis first."""
        from orchestrata_ultimate.postman.types import MessageEnvelope, MessagePayload

        # Redis cache
        if self.redis_client:
            cached = await self.redis_client.get(f"msg:{message_id}")
            if cached:
                envelope_dict, data, content_type, compressed, size_bytes = (
                    pickle.loads(cached)
                )
                envelope = MessageEnvelope.from_dict(envelope_dict)
                payload = MessagePayload(
                    content=data.decode("utf-8") if isinstance(data, bytes) else data,
                    content_type=content_type,
                    compressed=compressed,
                    size_bytes=size_bytes,
                )
                return envelope, payload

        # SQLite
        rows = await self._execute_sql(
            "SELECT envelope_json, payload_data, content_type FROM messages WHERE message_id = ?",
            (message_id,),
        )
        if not rows:
            return None
        envelope_json, payload_data, content_type = rows[0]
        envelope_dict = json.loads(envelope_json)
        envelope = MessageEnvelope.from_dict(envelope_dict)
        # The SQLite schema has no compressed/size_bytes columns (no migration
        # for this restoration, see MessagePayload's docstring). compressed is
        # derived from content_type instead: PayloadHandler.create_payload
        # always sets content_type=CONTENT_TYPE_COMPRESSED when it compresses
        # (payload.py), so the two are equivalent in this codebase today.
        # size_bytes is left unset (0) so MessagePayload.__post_init__
        # recomputes it from the reconstructed content, which is correct
        # since content itself came straight from payload_data.
        payload = MessagePayload(
            content=(
                payload_data.decode("utf-8")
                if isinstance(payload_data, bytes)
                else payload_data
            ),
            content_type=content_type,
            compressed=(content_type == CONTENT_TYPE_COMPRESSED),
        )
        # Populate cache
        if self.redis_client:
            cache_data = pickle.dumps(
                (
                    envelope_dict,
                    payload_data,
                    content_type,
                    payload.compressed,
                    payload.size_bytes,
                )
            )
            await self.redis_client.setex(f"msg:{message_id}", 3600, cache_data)

        return envelope, payload

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-009"])
    async def delete(self, message_id: str) -> bool:
        """Delete message from all stores."""
        rows = await self._execute_sql(
            "SELECT 1 FROM messages WHERE message_id = ?", (message_id,)
        )
        if not rows:
            return False

        await self._execute_sql_write(
            "DELETE FROM messages WHERE message_id = ?", (message_id,)
        )

        if self.redis_client:
            await self.redis_client.delete(f"msg:{message_id}")

        # Graph audit via Redis queue
        if self.redis_client:
            audit_event = {
                "event_type": "message_deleted",
                "message_id": message_id,
                "timestamp": datetime.now().isoformat(),
                "action": "deleted",
            }
            await _redis_result(
                self.redis_client.lpush(AUDIT_QUEUE_KEY, json.dumps(audit_event))
            )

        return True

    @traces(
        spec="002", story="US-006", reqs=["REQ-002-US-006-011", "REQ-002-US-006-012"]
    )
    async def query(
        self,
        sender: str | None = None,
        recipient: str | None = None,
        topic: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        limit: int = DEFAULT_MESSAGE_QUERY_LIMIT,
    ) -> list[str]:
        """Query message IDs matching criteria."""
        query_parts = ["SELECT message_id FROM messages WHERE 1=1"]
        params = []
        if sender:
            query_parts.append("AND sender = ?")
            params.append(sender)
        # Use recipient for both direct recipient and topic
        if recipient:
            query_parts.append("AND recipient = ?")
            params.append(recipient)
        if topic:
            # treat topic as a LIKE pattern on recipient
            query_parts.append("AND recipient LIKE ?")
            params.append(topic.replace("*", "%"))
        if start_time:
            query_parts.append("AND timestamp >= ?")
            params.append(start_time.isoformat())
        if end_time:
            query_parts.append("AND timestamp <= ?")
            params.append(end_time.isoformat())

        query_parts.append(f"ORDER BY timestamp DESC LIMIT {limit}")
        sql = " ".join(query_parts)
        rows = await self._execute_sql(sql, tuple(params))
        return [row[0] for row in rows]

    @traces(spec="002", story="US-006", reqs=["REQ-002-US-006-008"])
    async def cleanup(self, older_than: datetime) -> int:
        """Remove messages older than cutoff."""
        rows_before = await self._execute_sql(
            "SELECT COUNT(*) FROM messages WHERE timestamp < ?",
            (older_than.isoformat(),),
        )
        count = rows_before[0][0] if rows_before else 0
        await self._execute_sql_write(
            "DELETE FROM messages WHERE timestamp < ?",
            (older_than.isoformat(),),
        )
        # Graph audit via Redis queue
        if self.redis_client:
            audit_event = {
                "event_type": "messages_cleaned_up",
                "count": count,
                "cutoff": older_than.isoformat(),
                "timestamp": datetime.now().isoformat(),
                "action": "cleaned_up",
            }
            await _redis_result(
                self.redis_client.lpush(AUDIT_QUEUE_KEY, json.dumps(audit_event))
            )
        return count

    async def get_stats(self) -> dict[str, Any]:
        """Get storage statistics."""
        rows = await self._execute_sql(
            "SELECT COUNT(*), MIN(timestamp), MAX(timestamp) FROM messages"
        )
        total, oldest, newest = rows[0] if rows else (0, None, None)
        return {
            "total_messages": total,
            "oldest_timestamp": oldest,
            "newest_timestamp": newest,
        }