125 lines
5.2 KiB
Python
125 lines
5.2 KiB
Python
"""Private SQLite setup and immutable ordered utility migrations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import sqlite3
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from stat import S_ISDIR, S_ISLNK
|
|
|
|
|
|
class MigrationError(RuntimeError):
|
|
"""The local migration history cannot safely be applied."""
|
|
|
|
|
|
def _private_directory(path: Path) -> None:
|
|
if path.is_symlink():
|
|
raise MigrationError("utility root must not be a symlink")
|
|
path.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
mode = path.lstat().st_mode
|
|
if S_ISLNK(mode) or not S_ISDIR(mode):
|
|
raise MigrationError("utility root must be a directory")
|
|
path.chmod(0o700)
|
|
|
|
|
|
class StorageDatabase:
|
|
def __init__(self, root: Path, migrations_dir: Path) -> None:
|
|
if not root.is_absolute() or not migrations_dir.is_absolute():
|
|
raise ValueError("storage paths must be absolute")
|
|
self.root = root
|
|
self.migrations_dir = migrations_dir
|
|
self.path = root / "capture.sqlite3"
|
|
|
|
def connect(self) -> sqlite3.Connection:
|
|
_private_directory(self.root)
|
|
if self.path.is_symlink():
|
|
raise MigrationError("database file must not be a symlink")
|
|
connection: sqlite3.Connection | None = None
|
|
try:
|
|
connection = sqlite3.connect(self.path, isolation_level=None, timeout=5.0)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
journal = connection.execute("PRAGMA journal_mode = WAL").fetchone()[0]
|
|
connection.execute("PRAGMA busy_timeout = 5000")
|
|
connection.execute("PRAGMA synchronous = FULL")
|
|
checks = (
|
|
connection.execute("PRAGMA foreign_keys").fetchone()[0] == 1,
|
|
journal == "wal",
|
|
connection.execute("PRAGMA busy_timeout").fetchone()[0] == 5000,
|
|
connection.execute("PRAGMA synchronous").fetchone()[0] == 2,
|
|
)
|
|
if not all(checks):
|
|
raise MigrationError("SQLite connection invariants were not enabled")
|
|
for candidate in (
|
|
self.path,
|
|
self.path.with_name(self.path.name + "-wal"),
|
|
self.path.with_name(self.path.name + "-shm"),
|
|
):
|
|
if candidate.exists() and not candidate.is_symlink():
|
|
candidate.chmod(0o600)
|
|
return connection
|
|
except Exception:
|
|
if connection is not None:
|
|
connection.close()
|
|
raise
|
|
|
|
def _migrations(self) -> list[tuple[int, Path, str]]:
|
|
found: list[tuple[int, Path, str]] = []
|
|
versions: set[int] = set()
|
|
for path in self.migrations_dir.glob("*.sql"):
|
|
if len(path.name) < 5 or not path.name[:3].isdigit() or path.name[3] != "_":
|
|
raise MigrationError("migration filename is invalid")
|
|
version = int(path.name[:3])
|
|
if version in versions:
|
|
raise MigrationError("migration versions must be unique")
|
|
versions.add(version)
|
|
checksum = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
found.append((version, path, checksum))
|
|
found.sort(key=lambda item: item[0])
|
|
if not found:
|
|
raise MigrationError("no storage migrations found")
|
|
return found
|
|
|
|
def migrate(self) -> None:
|
|
migrations = self._migrations()
|
|
with self.connect() as connection:
|
|
connection.execute(
|
|
"CREATE TABLE IF NOT EXISTS schema_migrations ("
|
|
"version INTEGER PRIMARY KEY, filename TEXT NOT NULL UNIQUE, "
|
|
"checksum TEXT NOT NULL, applied_at TEXT NOT NULL)"
|
|
)
|
|
applied = list(
|
|
connection.execute(
|
|
"SELECT version, filename, checksum FROM schema_migrations ORDER BY version"
|
|
)
|
|
)
|
|
if len(applied) > len(migrations):
|
|
raise MigrationError("applied migration is no longer present")
|
|
expected_prefix = migrations[: len(applied)]
|
|
for row, expected in zip(applied, expected_prefix, strict=True):
|
|
version, path, checksum = expected
|
|
if (row["version"], row["filename"], row["checksum"]) != (
|
|
version,
|
|
path.name,
|
|
checksum,
|
|
):
|
|
raise MigrationError("migration history is not an exact ordered prefix")
|
|
for version, path, checksum in migrations[len(applied) :]:
|
|
sql = path.read_text(encoding="utf-8")
|
|
try:
|
|
applied_at = datetime.now(UTC).isoformat().replace("'", "''")
|
|
connection.executescript(
|
|
"BEGIN IMMEDIATE;\n"
|
|
+ sql
|
|
+ "\nINSERT INTO schema_migrations VALUES "
|
|
+ f"({version}, '{path.name}', '{checksum}', '{applied_at}');"
|
|
+ "\nCOMMIT;"
|
|
)
|
|
except Exception:
|
|
if connection.in_transaction:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
os.chmod(self.path, 0o600)
|