353 lines
12 KiB
Python
353 lines
12 KiB
Python
"""
|
|
JSONL persistence layer for the database.
|
|
"""
|
|
|
|
import asyncio
|
|
import copy
|
|
import logging
|
|
import os
|
|
import signal
|
|
from collections import deque
|
|
from contextlib import contextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
import jsondiff
|
|
import msgspec
|
|
|
|
from paskia.db.filelock import LockedFile
|
|
from paskia.db.logging import log_change
|
|
from paskia.db.migrations import (
|
|
DBVER,
|
|
MigrationCtx,
|
|
apply_all_migrations,
|
|
apply_migrations_readonly,
|
|
)
|
|
from paskia.db.snapshot import SnapshotState
|
|
from paskia.db.structs import DB, Config, SessionContext
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
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(ValueError):
|
|
"""Exception raised for database loading errors."""
|
|
|
|
|
|
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()
|
|
|
|
# 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 raw in lines:
|
|
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}: {e}\n{line.decode(errors='replace')}"
|
|
)
|
|
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:
|
|
"""Replay JSONL and apply migrations to produce a DB, without writing anything.
|
|
|
|
This is suitable for reading settings before the server starts.
|
|
Migrations are applied in-memory only; nothing is queued or flushed.
|
|
"""
|
|
path = Path(db_path)
|
|
if not path.exists():
|
|
return DB(config=Config(rp_id=rp_id))
|
|
|
|
try:
|
|
content = path.read_bytes()
|
|
r = _replay_from_data(content, str(path.resolve()))
|
|
data_dict = r.state
|
|
version = r.v
|
|
|
|
if not data_dict:
|
|
return DB(config=Config(rp_id=rp_id))
|
|
|
|
# Apply migrations in-memory (no persistence)
|
|
apply_migrations_readonly(data_dict, version, MigrationCtx(rp_id=rp_id))
|
|
|
|
# Decode to msgspec struct
|
|
try:
|
|
return msgspec.json.decode(msgspec.json.encode(data_dict), type=DB)
|
|
except msgspec.ValidationError as e:
|
|
raise DatabaseError(f"{path.resolve()}: {e}") from None
|
|
except OSError as e:
|
|
_logger.exception("Failed to load database")
|
|
raise SystemExit(f"{e}")
|
|
except (ValueError, msgspec.DecodeError) as e:
|
|
raise SystemExit(f"{e}")
|
|
except Exception as e:
|
|
_logger.exception("Unexpected error loading database")
|
|
raise SystemExit(f"{e}")
|
|
|
|
|
|
class ChangeRecord(msgspec.Struct, omit_defaults=True, kw_only=True):
|
|
ts: datetime = msgspec.field(default_factory=lambda: datetime.now(UTC))
|
|
a: str = "" # action (e.g., "migrate", "login", "create_user")
|
|
v: int = 0 # schema version after this change
|
|
u: str | None = None # user UUID who performed the action (None for system)
|
|
diff: dict
|
|
|
|
|
|
def compute_diff(previous: dict, current: dict) -> dict | None:
|
|
return jsondiff.diff(previous, current, marshal=True) or None
|
|
|
|
|
|
# Actions that are allowed to create a new database file
|
|
_BOOTSTRAP_ACTIONS = frozenset({"bootstrap"})
|
|
|
|
|
|
class JsonlStore:
|
|
"""JSONL persistence layer for a DB instance."""
|
|
|
|
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._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._v: int = DBVER # Schema version for new databases
|
|
self._snapshot = SnapshotState()
|
|
|
|
async def load(
|
|
self, db_path: str | None = None, *, rp_id: str = "localhost"
|
|
) -> None:
|
|
"""Load data from JSONL change log."""
|
|
if db_path is not None:
|
|
self.db_path = Path(db_path)
|
|
self._rp_id = rp_id
|
|
if not self.db_path.exists():
|
|
return
|
|
|
|
# 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 (snapshot-accelerated)
|
|
try:
|
|
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 statedict:
|
|
return
|
|
|
|
# Set previous state for diffing (will be updated by _queue_change)
|
|
self._statedict = copy.deepcopy(statedict)
|
|
|
|
# Callback to persist each migration
|
|
async def persist_migration(
|
|
action: str, new_version: int, current: dict
|
|
) -> None:
|
|
self._v = new_version
|
|
self._queue_change(action, new_version, current)
|
|
|
|
# Apply schema migrations one at a time
|
|
await apply_all_migrations(
|
|
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(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._v, normalized_dict)
|
|
|
|
def _queue_change(
|
|
self, action: str, version: int, current: dict, user: str | None = None
|
|
) -> None:
|
|
"""Queue a change record and log it.
|
|
|
|
Args:
|
|
action: The action name for the change record
|
|
version: The schema version for the change record
|
|
current: The current state as a plain dict
|
|
user: Optional user UUID who performed the action
|
|
"""
|
|
diff = compute_diff(self._statedict, current)
|
|
if not diff:
|
|
return
|
|
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
|
|
if user:
|
|
try:
|
|
user_uuid = UUID(user)
|
|
if user_uuid in self.db.users:
|
|
user_display = self.db.users[user_uuid].display_name
|
|
except (ValueError, KeyError):
|
|
user_display = user
|
|
|
|
log_change(action, diff, user_display, self._statedict, self.db)
|
|
self._statedict = copy.deepcopy(current)
|
|
|
|
@contextmanager
|
|
def transaction(
|
|
self,
|
|
action: str,
|
|
ctx: SessionContext | None = None,
|
|
*,
|
|
user: str | None = None,
|
|
):
|
|
"""Wrap writes in transaction. Queues change on successful exit.
|
|
|
|
Args:
|
|
action: Describes the operation (e.g., "Created user", "Login")
|
|
ctx: Session context of user performing the action (None for system operations)
|
|
user: User UUID string (alternative to ctx when full context unavailable)
|
|
"""
|
|
if self._in_transaction:
|
|
raise RuntimeError("Nested transactions are not supported")
|
|
|
|
# Check for out-of-transaction modifications
|
|
current_state = msgspec.to_builtins(self.db)
|
|
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._statedict:
|
|
pass # Expected: creating database from scratch
|
|
else:
|
|
diff = compute_diff(self._statedict, current_state)
|
|
diff_json = msgspec.json.encode(diff).decode()
|
|
_logger.critical(
|
|
"Database state modified outside of transaction! "
|
|
"This indicates a bug where DB changes occurred without a transaction wrapper.\n"
|
|
f"Changes detected:\n{diff_json}"
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
old_action = self._current_action
|
|
old_user = self._current_user
|
|
self._current_action = action
|
|
# Prefer ctx.user.uuid if ctx provided, otherwise use user param
|
|
self._current_user = str(ctx.user.uuid) if ctx else user
|
|
self._in_transaction = True
|
|
self._transaction_snapshot = current_state
|
|
|
|
try:
|
|
yield
|
|
current = msgspec.to_builtins(self.db)
|
|
self._queue_change(
|
|
self._current_action, self._v, current, self._current_user
|
|
)
|
|
except Exception:
|
|
# Rollback on error: restore from snapshot
|
|
_logger.warning("Transaction '%s' failed, rolling back changes", action)
|
|
if self._transaction_snapshot is not None:
|
|
decoder = msgspec.json.Decoder(DB)
|
|
self.db = decoder.decode(
|
|
msgspec.json.encode(self._transaction_snapshot)
|
|
)
|
|
self.db._store = self
|
|
raise
|
|
finally:
|
|
self._current_action = old_action
|
|
self._current_user = old_user
|
|
self._in_transaction = False
|
|
self._transaction_snapshot = None
|
|
|
|
async def flush(self) -> None:
|
|
"""Write all pending changes to disk.
|
|
|
|
On failure, logs an error and sends SIGTERM to trigger graceful shutdown.
|
|
"""
|
|
if self._flush_failed or not self._pending_changes:
|
|
return
|
|
|
|
if not self._file.is_open:
|
|
first_action = self._pending_changes[0].a
|
|
if first_action not in _BOOTSTRAP_ACTIONS:
|
|
_logger.error(
|
|
"Refusing to create database file with action '%s' - "
|
|
"only bootstrap can create a new database",
|
|
first_action,
|
|
)
|
|
self._flush_failed = True
|
|
os.kill(os.getpid(), signal.SIGTERM)
|
|
return
|
|
# Bootstrap: create and open the file with lock
|
|
await asyncio.to_thread(self._file.open, self.db_path, create=True)
|
|
|
|
changes_to_write = list(self._pending_changes)
|
|
|
|
try:
|
|
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()
|