zelos.storage

  1"""
  2Pluggable Storage Backends — Event persistence + State storage.
  3
  4Phase 2: InMemory, Redis, PostgreSQL, MySQL backends.
  5All share a common interface: connect / disconnect / append / read / state / snapshot.
  6"""
  7
  8import json
  9import threading
 10import time
 11from abc import ABC, abstractmethod
 12
 13# ═══════════════════ Common Interface ═══════════════════
 14
 15
 16class StorageBackend(ABC):
 17    """Abstract storage backend for events and state."""
 18
 19    def __init__(self, config: dict | None = None):
 20        self.config = config or {}
 21        self._connected = False
 22
 23    @abstractmethod
 24    def connect(self) -> bool: ...
 25
 26    @abstractmethod
 27    def disconnect(self) -> None: ...
 28
 29    @abstractmethod
 30    def health(self) -> bool: ...
 31
 32    @abstractmethod
 33    def append(self, stream: str, events: list[dict]) -> int: ...
 34
 35    @abstractmethod
 36    def read(self, stream: str, from_position: int, count: int) -> list[dict]: ...
 37
 38    @abstractmethod
 39    def set_state(self, key: str, value: dict) -> None: ...
 40
 41    @abstractmethod
 42    def get_state(self, key: str) -> dict | None: ...
 43
 44    @abstractmethod
 45    def delete_state(self, key: str) -> None: ...
 46
 47    def create_snapshot(self, key: str, events_position: int, state: dict) -> None:
 48        self.set_state(
 49            f"snapshot:{key}", {"events_position": events_position, "state": state, "timestamp": time.time()}
 50        )
 51
 52    def get_snapshot(self, key: str) -> dict | None:
 53        return self.get_state(f"snapshot:{key}")
 54
 55    @property
 56    def is_connected(self) -> bool:
 57        return self._connected
 58
 59
 60# ═══════════════════ InMemory ═══════════════════
 61
 62
 63class InMemoryStorageBackend(StorageBackend):
 64    """Phase 1 compatible — stores everything in memory."""
 65
 66    def __init__(self, config: dict | None = None):
 67        super().__init__(config)
 68        self._streams: dict[str, list[dict]] = {}
 69        self._state: dict[str, dict] = {}
 70        self._lock = threading.Lock()
 71
 72    def connect(self) -> bool:
 73        self._connected = True
 74        return True
 75
 76    def disconnect(self) -> None:
 77        self._connected = False
 78
 79    def health(self) -> bool:
 80        return self._connected
 81
 82    def append(self, stream: str, events: list[dict]) -> int:
 83        with self._lock:
 84            if stream not in self._streams:
 85                self._streams[stream] = []
 86            self._streams[stream].extend(events)
 87            return len(self._streams[stream])
 88
 89    def read(self, stream: str, from_position: int, count: int) -> list[dict]:
 90        with self._lock:
 91            events = self._streams.get(stream, [])
 92            return events[from_position : from_position + count]
 93
 94    def set_state(self, key: str, value: dict) -> None:
 95        with self._lock:
 96            self._state[key] = value
 97
 98    def get_state(self, key: str) -> dict | None:
 99        return self._state.get(key)
100
101    def delete_state(self, key: str) -> None:
102        self._state.pop(key, None)
103
104
105# ═══════════════════ Redis ═══════════════════
106
107
108class RedisStorageBackend(StorageBackend):
109    """Redis-backed storage. Events stored as lists, state as hash.
110
111    Config:
112      url: redis://localhost:6379/0
113      prefix: "zelos" (key namespace)
114    """
115
116    def __init__(self, config: dict | None = None):
117        super().__init__(config)
118        self._url = (config or {}).get("url", "redis://localhost:6379/0")
119        self._prefix = (config or {}).get("prefix", "zelos")
120        self._client = None
121
122    def connect(self) -> bool:
123        try:
124            import redis
125
126            self._client = redis.Redis.from_url(self._url, decode_responses=True)
127            self._client.ping()
128            self._connected = True
129            return True
130        except Exception:
131            self._connected = False
132            return False
133
134    def disconnect(self) -> None:
135        if self._client:
136            self._client.close()
137        self._connected = False
138
139    def health(self) -> bool:
140        if not self._client:
141            return False
142        try:
143            self._client.ping()
144            return True
145        except Exception:
146            return False
147
148    def _stream_key(self, stream: str) -> str:
149        return f"{self._prefix}:stream:{stream}"
150
151    def _state_key(self, key: str) -> str:
152        return f"{self._prefix}:state:{key}"
153
154    def append(self, stream: str, events: list[dict]) -> int:
155        if not self._client:
156            return -1
157        pipe = self._client.pipeline()
158        for e in events:
159            pipe.rpush(self._stream_key(stream), json.dumps(e))
160        pipe.execute()
161        return self._client.llen(self._stream_key(stream))
162
163    def read(self, stream: str, from_position: int, count: int) -> list[dict]:
164        if not self._client:
165            return []
166        raw = self._client.lrange(self._stream_key(stream), from_position, from_position + count - 1)
167        return [json.loads(r) for r in raw]
168
169    def set_state(self, key: str, value: dict) -> None:
170        if self._client:
171            self._client.set(self._state_key(key), json.dumps(value))
172
173    def get_state(self, key: str) -> dict | None:
174        if not self._client:
175            return None
176        raw = self._client.get(self._state_key(key))
177        return json.loads(raw) if raw else None
178
179    def delete_state(self, key: str) -> None:
180        if self._client:
181            self._client.delete(self._state_key(key))
182
183
184# ═══════════════════ PostgreSQL ═══════════════════
185
186
187class PostgreSQLStorageBackend(StorageBackend):
188    """PostgreSQL-backed storage. Events table + state table.
189
190    Config:
191      url: postgresql://user:pass@localhost:5432/zelos
192    """
193
194    def __init__(self, config: dict | None = None):
195        super().__init__(config)
196        self._url = (config or {}).get("url", "postgresql://localhost:5432/zelos")
197        self._conn = None
198
199    def connect(self) -> bool:
200        try:
201            import psycopg2
202
203            self._conn = psycopg2.connect(self._url)
204            self._conn.autocommit = True
205            self._create_tables()
206            self._connected = True
207            return True
208        except Exception:
209            self._connected = False
210            return False
211
212    def disconnect(self) -> None:
213        if self._conn:
214            self._conn.close()
215        self._connected = False
216
217    def health(self) -> bool:
218        if not self._conn:
219            return False
220        try:
221            cur = self._conn.cursor()
222            cur.execute("SELECT 1")
223            cur.close()
224            return True
225        except Exception:
226            return False
227
228    def _create_tables(self) -> None:
229        cur = self._conn.cursor()
230        cur.execute("""
231            CREATE TABLE IF NOT EXISTS zelos_events (
232                id SERIAL PRIMARY KEY,
233                stream VARCHAR(255) NOT NULL,
234                position INTEGER NOT NULL,
235                event_data JSONB NOT NULL,
236                created_at TIMESTAMP DEFAULT NOW()
237            );
238            CREATE INDEX IF NOT EXISTS idx_zelos_events_stream_pos
239                ON zelos_events(stream, position);
240            CREATE TABLE IF NOT EXISTS zelos_state (
241                key VARCHAR(255) PRIMARY KEY,
242                value JSONB NOT NULL,
243                updated_at TIMESTAMP DEFAULT NOW()
244            );
245        """)
246        cur.close()
247
248    def append(self, stream: str, events: list[dict]) -> int:
249        if not self._conn:
250            return -1
251        cur = self._conn.cursor()
252        # Get current max position
253        cur.execute("SELECT COALESCE(MAX(position), -1) FROM zelos_events WHERE stream = %s", (stream,))
254        pos = cur.fetchone()[0]
255        for e in events:
256            pos += 1
257            cur.execute(
258                "INSERT INTO zelos_events (stream, position, event_data) VALUES (%s, %s, %s)",
259                (stream, pos, json.dumps(e)),
260            )
261        cur.close()
262        return pos + 1
263
264    def read(self, stream: str, from_position: int, count: int) -> list[dict]:
265        if not self._conn:
266            return []
267        cur = self._conn.cursor()
268        cur.execute(
269            "SELECT event_data FROM zelos_events WHERE stream = %s AND position >= %s ORDER BY position LIMIT %s",
270            (stream, from_position, count),
271        )
272        rows = cur.fetchall()
273        cur.close()
274        return [r[0] for r in rows]
275
276    def set_state(self, key: str, value: dict) -> None:
277        if not self._conn:
278            return
279        cur = self._conn.cursor()
280        cur.execute(
281            "INSERT INTO zelos_state (key, value, updated_at) VALUES (%s, %s, NOW()) "
282            "ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = NOW()",
283            (key, json.dumps(value)),
284        )
285        cur.close()
286
287    def get_state(self, key: str) -> dict | None:
288        if not self._conn:
289            return None
290        cur = self._conn.cursor()
291        cur.execute("SELECT value FROM zelos_state WHERE key = %s", (key,))
292        row = cur.fetchone()
293        cur.close()
294        return row[0] if row else None
295
296    def delete_state(self, key: str) -> None:
297        if not self._conn:
298            return
299        cur = self._conn.cursor()
300        cur.execute("DELETE FROM zelos_state WHERE key = %s", (key,))
301        cur.close()
302
303
304# ═══════════════════ MySQL ═══════════════════
305
306
307class MySQLStorageBackend(StorageBackend):
308    """MySQL-backed storage. Same schema as PostgreSQL.
309
310    Config:
311      url: mysql://user:pass@localhost:3306/zelos
312    """
313
314    def __init__(self, config: dict | None = None):
315        super().__init__(config)
316        self._url = (config or {}).get("url", "mysql://localhost:3306/zelos")
317        self._conn = None
318
319    def connect(self) -> bool:
320        try:
321            # Parse URL
322            from urllib.parse import urlparse
323
324            import mysql.connector
325
326            parsed = urlparse(self._url)
327            self._conn = mysql.connector.connect(
328                host=parsed.hostname or "localhost",
329                port=parsed.port or 3306,
330                user=parsed.username or "root",
331                password=parsed.password or "",
332                database=parsed.path.lstrip("/") or "zelos",
333                autocommit=True,
334            )
335            self._create_tables()
336            self._connected = True
337            return True
338        except Exception:
339            self._connected = False
340            return False
341
342    def disconnect(self) -> None:
343        if self._conn:
344            self._conn.close()
345        self._connected = False
346
347    def health(self) -> bool:
348        if not self._conn:
349            return False
350        try:
351            cur = self._conn.cursor()
352            cur.execute("SELECT 1")
353            cur.close()
354            return True
355        except Exception:
356            return False
357
358    def _create_tables(self) -> None:
359        cur = self._conn.cursor()
360        cur.execute("""
361            CREATE TABLE IF NOT EXISTS zelos_events (
362                id INT AUTO_INCREMENT PRIMARY KEY,
363                stream VARCHAR(255) NOT NULL,
364                position INT NOT NULL,
365                event_data JSON NOT NULL,
366                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
367                INDEX idx_stream_pos (stream, position)
368            )
369        """)
370        cur.execute("""
371            CREATE TABLE IF NOT EXISTS zelos_state (
372                `key` VARCHAR(255) PRIMARY KEY,
373                value JSON NOT NULL,
374                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
375            )
376        """)
377        cur.close()
378
379    def append(self, stream: str, events: list[dict]) -> int:
380        if not self._conn:
381            return -1
382        cur = self._conn.cursor()
383        cur.execute("SELECT COALESCE(MAX(position), -1) FROM zelos_events WHERE stream = %s", (stream,))
384        pos = cur.fetchone()[0]
385        for e in events:
386            pos += 1
387            cur.execute(
388                "INSERT INTO zelos_events (stream, position, event_data) VALUES (%s, %s, %s)",
389                (stream, pos, json.dumps(e)),
390            )
391        cur.close()
392        return pos + 1
393
394    def read(self, stream: str, from_position: int, count: int) -> list[dict]:
395        if not self._conn:
396            return []
397        cur = self._conn.cursor()
398        cur.execute(
399            "SELECT event_data FROM zelos_events WHERE stream = %s AND position >= %s ORDER BY position LIMIT %s",
400            (stream, from_position, count),
401        )
402        rows = cur.fetchall()
403        cur.close()
404        # JSON type returns str in MySQL connector
405        return [json.loads(r[0]) if isinstance(r[0], str) else r[0] for r in rows]
406
407    def set_state(self, key: str, value: dict) -> None:
408        if not self._conn:
409            return
410        cur = self._conn.cursor()
411        cur.execute(
412            "INSERT INTO zelos_state (`key`, value) VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = VALUES(value)",
413            (key, json.dumps(value)),
414        )
415        cur.close()
416
417    def get_state(self, key: str) -> dict | None:
418        if not self._conn:
419            return None
420        cur = self._conn.cursor()
421        cur.execute("SELECT value FROM zelos_state WHERE `key` = %s", (key,))
422        row = cur.fetchone()
423        cur.close()
424        if row:
425            return json.loads(row[0]) if isinstance(row[0], str) else row[0]
426        return None
427
428    def delete_state(self, key: str) -> None:
429        if not self._conn:
430            return
431        cur = self._conn.cursor()
432        cur.execute("DELETE FROM zelos_state WHERE `key` = %s", (key,))
433        cur.close()
434
435
436# ═══════════════════ Factory ═══════════════════
437
438BACKENDS = {
439    "memory": InMemoryStorageBackend,
440    "redis": RedisStorageBackend,
441    "postgresql": PostgreSQLStorageBackend,
442    "postgres": PostgreSQLStorageBackend,
443    "pgsql": PostgreSQLStorageBackend,
444    "mysql": MySQLStorageBackend,
445}
446
447
448def create_storage_backend(config: dict) -> StorageBackend:
449    """Factory: create a storage backend from configuration."""
450    backend_type = config.get("type", "memory").lower()
451    cls = BACKENDS.get(backend_type)
452    if cls is None:
453        raise ValueError(f"Unsupported storage backend: '{backend_type}'. Supported: {', '.join(BACKENDS.keys())}")
454    return cls(config)