From 5f7a5ed9b1e0a05fe13726a8aa3c697034bf61cc Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Thu, 19 Feb 2026 16:29:23 +0000 Subject: [PATCH] Added database snapshots, cleanup, better error messages. --- paskia/db/background.py | 8 +- paskia/db/jsonl.py | 194 ++++++++++++++++++++++------------------ paskia/db/lifecycle.py | 17 ++-- paskia/db/operations.py | 6 -- paskia/db/snapshot.py | 88 ++++++++++++++++++ 5 files changed, 209 insertions(+), 104 deletions(-) create mode 100644 paskia/db/snapshot.py diff --git a/paskia/db/background.py b/paskia/db/background.py index d1f8cd8..ef03c6c 100644 --- a/paskia/db/background.py +++ b/paskia/db/background.py @@ -21,7 +21,7 @@ _background_task: asyncio.Task | None = None async def flush() -> None: """Write all pending database changes to disk.""" - store = _ops._store + store = _ops._db._store if store is None: _logger.warning("flush() called but _store is None") return @@ -48,6 +48,10 @@ async def _background_loop(): cleanup_expired() await flush() # Flush cleanup changes last_cleanup = now + + # Conditionally write a snapshot to speed up future startups + if _ops._db._store is not None: + _ops._db._store.maybe_snapshot() except asyncio.CancelledError: # Final flush before exit await flush() @@ -99,7 +103,7 @@ async def stop_background(): except asyncio.CancelledError: pass _background_task = None - _ops._store.close() + _ops._db._store.close() # Aliases for backwards compatibility diff --git a/paskia/db/jsonl.py b/paskia/db/jsonl.py index f4f99f2..1f6193a 100644 --- a/paskia/db/jsonl.py +++ b/paskia/db/jsonl.py @@ -25,12 +25,56 @@ from paskia.db.migrations import ( apply_all_migrations, apply_migrations_readonly, ) +from paskia.db.snapshot import SnapshotState from paskia.db.structs import DB, Config, SessionContext _logger = logging.getLogger(__name__) -# Default database path -DB_PATH_DEFAULT = "paskia.jsonl" + +class ReplayResult(msgspec.Struct, frozen=False): + """Return value of _replay_from_data""" + + state: dict + v: int = 0 + ts: datetime | None = None + snapts: datetime | None = None + changes: int = 0 + + +class DatabaseError(Exception): + """Exception raised for database loading errors.""" + + pass + + +def _replay_from_data(data: bytes, db_path: str) -> ReplayResult: + """Replay database state from file data, using the last snapshot if available.""" + resolved_path = str(Path(db_path).resolve()) + result = ReplayResult(state={}) + + # Find and apply the last snapshot + snap, start_offset = SnapshotState.load(data) + if snap: + result.state = snap.state + result.v = snap.v + result.snapts = snap.ts + + # Replay change records after the snapshot + lines = data[start_offset:].split(b"\n") + for line_num, raw in enumerate(lines, start=1): # 1-based line numbering + line = raw.strip() + if not line: + continue + try: + change = msgspec.json.decode(line, type=ChangeRecord) + except msgspec.DecodeError as e: + raise DatabaseError(f"{resolved_path}:{line_num}: {e}") + result.state = jsondiff.patch(result.state, change.diff, marshal=True) + result.v = change.v + result.ts = change.ts + result.changes += 1 + + return result def load_readonly(db_path: str, *, rp_id: str = "localhost") -> DB: @@ -43,25 +87,20 @@ def load_readonly(db_path: str, *, rp_id: str = "localhost") -> DB: if not path.exists(): return DB(config=Config(rp_id=rp_id)) - data_dict: dict = {} - version = 0 try: with open(path, "rb") as f: content = f.read() - for line_num, line in enumerate(content.split(b"\n"), 1): - line = line.strip() - if not line: - continue - try: - change = msgspec.json.decode(line) - data_dict = jsondiff.patch(data_dict, change["diff"], marshal=True) - version = change.get("v", 0) - except Exception as e: - raise ValueError(f"Error parsing line {line_num}: {e}") + r = _replay_from_data(content, str(path.resolve())) + data_dict = r.state + version = r.v except OSError as e: - raise SystemExit(f"Failed to load database: {e}") - except (ValueError, msgspec.DecodeError) as e: - raise SystemExit(f"Failed to load database: {e}") + _logger.exception("Failed to load database") + raise SystemExit(f"{e}") + except (ValueError, msgspec.DecodeError, DatabaseError) as e: + raise SystemExit(f"{e}") + except Exception as e: + _logger.exception("Unexpected error loading database") + raise SystemExit(f"{e}") if not data_dict: return DB(config=Config(rp_id=rp_id)) @@ -74,45 +113,16 @@ def load_readonly(db_path: str, *, rp_id: str = "localhost") -> DB: return db -class _ChangeRecord(msgspec.Struct, omit_defaults=True): - """A single change record in the JSONL file.""" - - ts: datetime +class ChangeRecord(msgspec.Struct, omit_defaults=True, kw_only=True): + ts: datetime = msgspec.field(default_factory=lambda: datetime.now(UTC)) a: str # action - describes the operation (e.g., "migrate", "login", "create_user") - v: int # schema version after this change + v: int = 0 # schema version after this change u: str | None = None # user UUID who performed the action (None for system) - diff: dict = {} - - -# msgspec encoder for change records -_change_encoder = msgspec.json.Encoder() + diff: dict def compute_diff(previous: dict, current: dict) -> dict | None: - """Compute JSON diff between two states. - - Args: - previous: Previous state (JSON-compatible dict) - current: Current state (JSON-compatible dict) - - Returns: - The diff, or None if no changes - """ - diff = jsondiff.diff(previous, current, marshal=True) - return diff if diff else None - - -def create_change_record( - action: str, version: int, diff: dict, user: str | None = None -) -> _ChangeRecord: - """Create a change record for persistence.""" - return _ChangeRecord( - ts=datetime.now(UTC), - a=action, - v=version, - u=user, - diff=diff, - ) + return jsondiff.diff(previous, current, marshal=True) or None # Actions that are allowed to create a new database file @@ -122,18 +132,19 @@ _BOOTSTRAP_ACTIONS = frozenset({"bootstrap"}) class JsonlStore: """JSONL persistence layer for a DB instance.""" - def __init__(self, db: DB, db_path: str = DB_PATH_DEFAULT): + def __init__(self, db: DB, db_path: str): self.db: DB = db self.db_path = Path(db_path) self._file = LockedFile() self._flush_failed = False - self._previous_builtins: dict[str, Any] = {} - self._pending_changes: deque[_ChangeRecord] = deque() + self._statedict: dict[str, Any] = {} + self._pending_changes: deque[ChangeRecord] = deque() self._current_action: str = "system" self._current_user: str | None = None self._in_transaction: bool = False self._transaction_snapshot: dict[str, Any] | None = None - self._current_version: int = DBVER # Schema version for new databases + self._v: int = DBVER # Schema version for new databases + self._snapshot = SnapshotState() async def load( self, db_path: str | None = None, *, rp_id: str = "localhost" @@ -148,56 +159,49 @@ class JsonlStore: # Open with exclusive write lock and read contents — single threadpool call content = await asyncio.to_thread(self._file.open_and_read, self.db_path) - # Replay change log to reconstruct state - data_dict: dict = {} + # Replay change log to reconstruct state (snapshot-accelerated) try: - for line_num, line in enumerate(content.split(b"\n"), 1): - line = line.strip() - if not line: - continue - try: - change = msgspec.json.decode(line) - data_dict = jsondiff.patch(data_dict, change["diff"], marshal=True) - self._current_version = change.get("v", 0) - except Exception as e: - raise ValueError(f"Error parsing line {line_num}: {e}") - except OSError as e: - raise SystemExit(f"Failed to load database: {e}") - except (ValueError, msgspec.DecodeError) as e: - raise SystemExit(f"Failed to load database: {e}") + r = _replay_from_data(content, str(self.db_path.resolve())) + statedict = r.state + self._v = r.v + self._snapshot.ts = r.snapts + self._snapshot.changes = r.changes + except (OSError, ValueError, msgspec.DecodeError, DatabaseError) as e: + raise SystemExit(f"{e}") + except Exception as e: + _logger.exception("Unexpected error loading database") + raise SystemExit(f"{e}") - if not data_dict: + if not statedict: return # Set previous state for diffing (will be updated by _queue_change) - self._previous_builtins = copy.deepcopy(data_dict) + self._statedict = copy.deepcopy(statedict) # Callback to persist each migration async def persist_migration( action: str, new_version: int, current: dict ) -> None: - self._current_version = new_version + self._v = new_version self._queue_change(action, new_version, current) # Apply schema migrations one at a time await apply_all_migrations( - data_dict, - self._current_version, + statedict, + self._v, persist_migration, MigrationCtx(rp_id=rp_id), ) # Decode to msgspec struct decoder = msgspec.json.Decoder(DB) - self.db = decoder.decode(msgspec.json.encode(data_dict)) + self.db = decoder.decode(msgspec.json.encode(statedict)) self.db._store = self # Normalize via msgspec round-trip (handles omit_defaults etc.) # This ensures _previous_builtins matches what msgspec would produce normalized_dict = msgspec.to_builtins(self.db) - await persist_migration( - "migrate:msgspec", self._current_version, normalized_dict - ) + await persist_migration("migrate:msgspec", self._v, normalized_dict) def _queue_change( self, action: str, version: int, current: dict, user: str | None = None @@ -210,10 +214,17 @@ class JsonlStore: current: The current state as a plain dict user: Optional user UUID who performed the action """ - diff = compute_diff(self._previous_builtins, current) + diff = compute_diff(self._statedict, current) if not diff: return - self._pending_changes.append(create_change_record(action, version, diff, user)) + self._pending_changes.append( + ChangeRecord( + a=action, + v=version, + u=user, + diff=diff, + ) + ) # Log the change with user display name if available user_display = None @@ -225,8 +236,8 @@ class JsonlStore: except (ValueError, KeyError): user_display = user - log_change(action, diff, user_display, self._previous_builtins, self.db) - self._previous_builtins = copy.deepcopy(current) + log_change(action, diff, user_display, self._statedict, self.db) + self._statedict = copy.deepcopy(current) @contextmanager def transaction( @@ -248,13 +259,13 @@ class JsonlStore: # Check for out-of-transaction modifications current_state = msgspec.to_builtins(self.db) - if current_state != self._previous_builtins: + if current_state != self._statedict: # Allow bootstrap to create a new database from empty state is_bootstrap = action in _BOOTSTRAP_ACTIONS - if is_bootstrap and not self._previous_builtins: + if is_bootstrap and not self._statedict: pass # Expected: creating database from scratch else: - diff = compute_diff(self._previous_builtins, current_state) + diff = compute_diff(self._statedict, current_state) diff_json = msgspec.json.encode(diff).decode() _logger.critical( "Database state modified outside of transaction! " @@ -275,7 +286,7 @@ class JsonlStore: yield current = msgspec.to_builtins(self.db) self._queue_change( - self._current_action, self._current_version, current, self._current_user + self._current_action, self._v, current, self._current_user ) except Exception: # Rollback on error: restore from snapshot @@ -318,18 +329,23 @@ class JsonlStore: changes_to_write = list(self._pending_changes) try: - lines = [_change_encoder.encode(change) for change in changes_to_write] + lines = [msgspec.json.encode(change) for change in changes_to_write] if not lines: self._pending_changes.clear() return await asyncio.to_thread(self._file.write, b"\n".join(lines) + b"\n") + self._snapshot.record_lines(len(lines)) self._pending_changes.clear() except OSError as e: _logger.error("Failed to flush database: %s", e) self._flush_failed = True os.kill(os.getpid(), signal.SIGTERM) + def maybe_snapshot(self) -> None: + """Write a snapshot if conditions are met.""" + self._snapshot.maybe_write(self._file, self._v, self._statedict) + def close(self) -> None: """Release the file lock and close the file.""" self._file.close() diff --git a/paskia/db/lifecycle.py b/paskia/db/lifecycle.py index ce77eea..81cb489 100644 --- a/paskia/db/lifecycle.py +++ b/paskia/db/lifecycle.py @@ -9,20 +9,23 @@ from datetime import UTC, datetime import paskia.db.operations as _ops from paskia import oidc_notify from paskia.authsession import EXPIRES +from paskia.db.jsonl import JsonlStore _logger = logging.getLogger(__name__) -async def init(rp_id: str = "localhost", *args, **kwargs): +async def init(rp_id: str, *args, **kwargs): """Load database from JSONL file.""" - if _ops._initialized: + if _ops._db._store: _logger.debug("Database already initialized, skipping reload") return - default_path = f"{rp_id}.paskiadb" - db_path = os.environ.get("PASKIA_DB", default_path) - await _ops._store.load(db_path, rp_id=rp_id) - _ops._db = _ops._store.db - _ops._initialized = True + db_path = os.environ.get("PASKIA_DB", f"{rp_id}.paskiadb") + store = JsonlStore(_ops._db, db_path) + await store.load(db_path, rp_id=rp_id) + _ops._db = store.db + _ops._db._store = store + # Request a snapshot after successful startup + store._snapshot.request_force() def cleanup_expired() -> int: diff --git a/paskia/db/operations.py b/paskia/db/operations.py index d8dcaef..ccd8eae 100644 --- a/paskia/db/operations.py +++ b/paskia/db/operations.py @@ -15,9 +15,6 @@ import uuid7 from paskia import oidc_notify from paskia.config import SESSION_LIFETIME -from paskia.db.jsonl import ( - JsonlStore, -) from paskia.db.structs import ( DB, Client, @@ -41,9 +38,6 @@ _UNSET = object() # Global database instance (empty until init() loads data) _db = DB(config=Config(rp_id="uninitialized.invalid")) -_store = JsonlStore(_db) -_db._store = _store -_initialized = False def is_username_taken(username: str, exclude_uuid: UUID | None = None) -> bool: diff --git a/paskia/db/snapshot.py b/paskia/db/snapshot.py new file mode 100644 index 0000000..8a14f16 --- /dev/null +++ b/paskia/db/snapshot.py @@ -0,0 +1,88 @@ +""" +Snapshot handling for JSONL database persistence. +""" + +import logging +from datetime import UTC, datetime +from typing import Any + +import msgspec + +_logger = logging.getLogger(__name__) + +LINEPREFIX = b"SNAPSHOT " +MINDIFFS = 100 + + +class Snapshot(msgspec.Struct): + """Snapshot data structure for database persistence.""" + + ts: datetime + v: int + state: dict[str, Any] + + +class SnapshotState: + """Tracks snapshot timing and line counts for a database file.""" + + def __init__(self) -> None: + self.ts: datetime | None = None + self.changes: int = 0 + self._force_pending: bool = False + + def request_force(self) -> None: + """Request a forced snapshot on the next maybe_write call.""" + self._force_pending = True + + def record_lines(self, count: int) -> None: + self.changes += count + + def maybe_write(self, file, version: int, state: dict) -> None: + """Write a snapshot if conditions are met (enough changes, and Sunday UTC or forced).""" + if self.changes < MINDIFFS: + return + force = self._force_pending + now = datetime.now(UTC) + if not force and now.weekday() != 6: # 6 = Sunday + return + sunday_midnight = now.replace(hour=0, minute=0, second=0, microsecond=0) + if not force and self.ts is not None and self.ts >= sunday_midnight: + return + if not file.is_open: + return + try: + self._write(file, version, state, now) + self._force_pending = False + except Exception as exc: + _logger.error("snapshot: failed to write snapshot: %r", exc) + + def _write(self, file, version: int, state: dict, now: datetime) -> None: + """Write a snapshot and update internal state.""" + data = msgspec.json.encode(Snapshot(ts=now, v=version, state=state)) + file.write(LINEPREFIX + data + b"\n") + self.changes = 0 + self.ts = now + + @staticmethod + def load(data: bytes) -> tuple[Snapshot | None, int]: + """Find and parse the last snapshot in file data. + + Returns (snapshot, replay_offset) where replay_offset is the byte + position to start replaying change records from. If no valid snapshot + is found, returns (None, 0). + """ + marker = b"\n" + LINEPREFIX + pos = data.rfind(marker) + if pos != -1: + pos += 1 # skip the newline + elif data.startswith(LINEPREFIX): + pos = 0 + else: + return None, 0 + + end = data.find(b"\n", pos) + if end == -1: + raise ValueError("Incomplete snapshot line at end of file") + + snap = msgspec.json.decode(data[pos + len(LINEPREFIX) : end], type=Snapshot) + return snap, end + 1