Coverage for src/local_deep_research/database/alembic_runner.py: 91%
136 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Programmatic Alembic migration runner for per-user encrypted databases.
4This module provides functions to run Alembic migrations against SQLCipher
5encrypted databases without using the Alembic CLI. Each user database
6tracks its own migration version via the alembic_version table.
7"""
9import os
10import time
11from pathlib import Path
12from typing import Optional
14from alembic import command
15from alembic.config import Config
16from alembic.runtime.migration import MigrationContext
17from alembic.script import ScriptDirectory
18from loguru import logger
19from sqlalchemy import Connection, Engine, inspect
20from sqlalchemy.exc import IntegrityError, OperationalError
23def get_migrations_dir() -> Path:
24 """
25 Get the path to the migrations directory with security validation.
27 Validates that the migrations directory is within the expected package
28 boundary to prevent symlink attacks that could redirect migration loading
29 to arbitrary locations.
31 Returns:
32 Path to the migrations directory
34 Raises:
35 ValueError: If the migrations path is outside expected boundaries
36 """
37 migrations_dir = Path(__file__).parent / "migrations"
38 real_path = migrations_dir.resolve()
39 expected_parent = Path(__file__).parent.resolve()
41 # Security: Ensure migrations directory is within expected package boundary
42 # This prevents symlink attacks that could load arbitrary Python code
43 if not real_path.is_relative_to(expected_parent): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true
44 raise ValueError(
45 "Invalid migrations path (possible symlink attack): "
46 "migrations dir resolves outside expected package boundary"
47 )
49 return migrations_dir
52def _validate_migrations_permissions(migrations_dir: Path) -> None:
53 """
54 Validate migration files are not world-writable.
56 World-writable migration files could be replaced with malicious code
57 that would execute during database migrations with the application's
58 privileges.
60 Args:
61 migrations_dir: Path to the migrations directory
63 Raises:
64 ValueError: If any migration file is world-writable
66 Note:
67 This check is skipped on Windows where file permissions work differently.
68 """
69 if os.name == "nt": # Skip permission checks on Windows
70 return
72 versions_dir = migrations_dir / "versions"
73 if not versions_dir.exists(): 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true
74 return
76 # Check the versions directory itself
77 st = versions_dir.stat()
78 if st.st_mode & 0o002:
79 raise ValueError(
80 f"Migrations directory has insecure permissions (world-writable): "
81 f"{versions_dir}. Fix with: chmod o-w {versions_dir}"
82 )
84 for migration_file in versions_dir.glob("*.py"):
85 st = migration_file.stat()
86 if st.st_mode & 0o002: # World-writable bit
87 raise ValueError(
88 f"Migration file has insecure permissions (world-writable): "
89 f"{migration_file.name}. "
90 f"Fix with: chmod o-w {migration_file}"
91 )
94def get_alembic_config(engine: Engine) -> Config:
95 """
96 Create an Alembic Config object for programmatic usage.
98 Args:
99 engine: SQLAlchemy engine to run migrations against
101 Returns:
102 Configured Alembic Config object
103 """
104 migrations_dir = get_migrations_dir()
106 # Create config object without ini file
107 config = Config()
109 # Set script location
110 config.set_main_option("script_location", str(migrations_dir))
112 # Set SQLAlchemy URL (not actually used since we pass connection directly)
113 # But Alembic requires it to be set
114 config.set_main_option("sqlalchemy.url", "sqlite:///:memory:")
116 return config
119def get_current_revision(engine: Engine) -> Optional[str]:
120 """
121 Get the current migration revision for a database.
123 Args:
124 engine: SQLAlchemy engine
126 Returns:
127 Current revision string or None if no migrations have run
128 """
129 with engine.connect() as conn:
130 context = MigrationContext.configure(conn)
131 return context.get_current_revision()
134def get_head_revision() -> Optional[str]:
135 """
136 Get the latest migration revision.
138 Returns:
139 Head revision string, or None if no migrations exist
140 """
141 migrations_dir = get_migrations_dir()
142 config = Config()
143 config.set_main_option("script_location", str(migrations_dir))
145 script = ScriptDirectory.from_config(config)
146 return script.get_current_head()
149def needs_migration(engine: Engine) -> bool:
150 """
151 Check if a database needs migrations.
153 Args:
154 engine: SQLAlchemy engine
156 Returns:
157 True if migrations are pending
158 """
159 head = get_head_revision()
161 if head is None: 161 ↛ 163line 161 didn't jump to line 163 because the condition on line 161 was never true
162 # No migrations exist yet
163 return False
165 current = get_current_revision(engine)
167 if current is None:
168 # Check if this is a fresh database or existing without migrations
169 inspector = inspect(engine)
170 tables = inspector.get_table_names()
172 if not tables:
173 # Fresh database, needs initial migration
174 return True
175 if "alembic_version" not in tables: 175 ↛ 179line 175 didn't jump to line 179 because the condition on line 175 was always true
176 # Existing database without Alembic - needs stamping then check
177 return True
179 return current != head
182def stamp_database(engine: Engine, revision: str = "head") -> None:
183 """
184 Stamp a database with a revision without running migrations.
185 Used for baselining existing databases.
187 Concurrency: If two callers race to stamp a fresh database, one will hit
188 "table alembic_version already exists" (OperationalError) or a duplicate
189 PK on version_num (IntegrityError). Both outcomes are benign — the DB
190 ends up stamped — so we swallow them after verifying the table+row are
191 in place. A genuine failure (no row appeared) is re-raised.
193 Args:
194 engine: SQLAlchemy engine
195 revision: Revision to stamp (default "head")
196 """
197 config = get_alembic_config(engine)
199 try:
200 with engine.begin() as conn:
201 config.attributes["connection"] = conn
202 command.stamp(config, revision)
203 except (IntegrityError, OperationalError) as exc:
204 # Only swallow errors that look like a benign concurrent-stamp
205 # race on the alembic_version table itself. A genuine failure
206 # (disk full, SQLITE_BUSY on an unrelated table, corruption,
207 # etc.) must propagate so callers see the real error.
208 msg = str(exc).lower()
209 looks_like_race = (
210 "alembic_version" in msg # IntegrityError or table-exists race
211 or "already exists" in msg # CREATE TABLE race
212 )
213 if not looks_like_race or get_current_revision(engine) is None: 213 ↛ 217line 213 didn't jump to line 217 because the condition on line 213 was always true
214 raise
215 # Race-loss path: another caller stamped first. Don't claim we
216 # stamped it ourselves — log at debug only.
217 logger.debug(
218 f"stamp_database({revision}) lost race to concurrent caller "
219 f"({type(exc).__name__}); database is stamped, continuing"
220 )
221 return
223 logger.info(f"Stamped database at revision: {revision}")
226def _drop_orphan_alembic_temp_tables(conn: Connection) -> None:
227 """Drop leftover ``_alembic_tmp_<table>`` tables from prior failed
228 batch_alter_table runs (issue #3817).
230 ``op.batch_alter_table`` rebuilds a table by creating
231 ``_alembic_tmp_<table>``, copying data, dropping the original, and
232 renaming. On a clean run alembic drops the temp table automatically.
233 If a previous attempt failed in a way that bypassed transaction
234 rollback (e.g., an older migration runner that auto-committed each
235 migration, or a process killed mid-DDL on a non-transactional
236 sqlite build), the temp table persists. The next attempt then fails
237 at ``op.batch_alter_table`` with ``table _alembic_tmp_* already exists``.
239 This runs in autocommit mode at the SQLite level — each ``DROP TABLE``
240 briefly takes the file write lock and releases it. If a concurrent
241 migration is mid-batch_alter_table on the same table, our DROP blocks
242 on the SQLite write lock (busy_timeout=10000); by the time we
243 acquire it, the concurrent migration's rename has consumed the temp
244 table and our DROP IF EXISTS is a no-op. The race is benign.
245 """
246 inspector = inspect(conn)
247 temp_tables = [
248 name
249 for name in inspector.get_table_names()
250 if name.startswith("_alembic_tmp_")
251 ]
252 if not temp_tables:
253 return
254 logger.warning(
255 f"Found {len(temp_tables)} orphan alembic temp table(s) from a "
256 f"prior failed migration: {sorted(temp_tables)}. Dropping before retry."
257 )
258 for name in temp_tables:
259 # Identifier is constrained to the ``_alembic_tmp_`` prefix + a
260 # parent table name from ``inspector.get_table_names()``; both
261 # come from the database's own catalog and cannot contain
262 # injection vectors.
263 # bearer:disable python_lang_sql_injection
264 conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{name}"') # noqa: S608
267def _disable_fk_for_migration(conn: Connection) -> None:
268 """Disable FK enforcement on the migration connection BEFORE any
269 transaction opens (issue #3990).
271 ``apply_performance_pragmas`` set ``PRAGMA foreign_keys = ON`` at
272 connect. SQLite then *silently ignores* further toggles of
273 ``foreign_keys`` once any transaction (explicit or driver-implicit)
274 is active. The sqlite3/sqlcipher3 driver auto-begins on the first
275 DML; PRAGMA itself isn't DML, so issuing the PRAGMA before any DML
276 is the only window where it actually takes effect.
278 With multi-migration upgrades (revision 0001 → 0009), the first
279 migration to issue DML auto-begins the driver transaction and
280 freezes FK in the connect-time ON state for the rest of the upgrade.
281 That defeats migration 0007's defensive PRAGMA OFF and makes its
282 orphan-scrub DELETE fail with ``foreign key mismatch`` on tables
283 whose FK target lacks a UNIQUE backing — exactly the broken-schema
284 state migration 0007 is meant to repair.
286 The caller is responsible for re-enabling FK after the migration
287 transaction commits, BEFORE returning the connection to the pool —
288 see ``run_migrations``.
289 """
290 conn.exec_driver_sql("PRAGMA foreign_keys = OFF")
291 # ``exec_driver_sql`` triggers SQLAlchemy autobegin even though no
292 # sqlite-level transaction was opened (PRAGMA isn't DML, so the
293 # driver doesn't auto-begin). Roll back the no-op SQLAlchemy
294 # transaction so the caller's ``conn.begin()`` is allowed to start
295 # a fresh one. ``PRAGMA foreign_keys`` is connection-level state and
296 # survives ROLLBACK at the SQLite level — see
297 # https://www.sqlite.org/pragma.html#pragma_foreign_keys.
298 conn.rollback()
301def run_migrations(engine: Engine, target: str = "head") -> None:
302 """
303 Run pending migrations on a database.
305 The initial migration is idempotent (only creates tables that don't exist),
306 so this function runs migrations rather than just stamping existing
307 databases. This ensures any missing tables are created.
309 When ``target == "head"`` and the database is already at head, the call
310 short-circuits without opening a write transaction — calling
311 ``command.upgrade()`` unconditionally would still open a write transaction
312 via ``engine.begin()`` (taking a RESERVED lock on the SQLite file) just to
313 discover there's nothing to apply, serialising concurrent readers behind
314 a no-op on every cold engine reopen.
316 Security validations performed before running migrations:
317 - Migration directory path is within expected package boundary
318 - Migration files are not world-writable
320 Pre-upgrade hygiene (run outside the migration transaction):
321 - Drop orphan ``_alembic_tmp_*`` tables from prior failed
322 ``batch_alter_table`` runs (issue #3817).
323 - Disable ``PRAGMA foreign_keys`` so 0007's orphan scrub can run
324 (issue #3990). Re-enabled after a successful upgrade, before the
325 connection returns to the pool.
327 On failure inside the migration transaction, the inner
328 ``conn.begin()`` rolls back automatically — the database stays at
329 its previous revision. The original exception is re-raised so
330 callers can decide how to handle it.
332 Args:
333 engine: SQLAlchemy engine to migrate
334 target: Target revision (default "head" for latest)
336 Raises:
337 Exception: If migration fails (database is safely rolled back)
338 """
339 migration_start = time.perf_counter()
341 # Security: Validate migrations directory and file permissions
342 migrations_dir = get_migrations_dir()
343 _validate_migrations_permissions(migrations_dir)
345 head = get_head_revision()
347 if head is None:
348 # No migrations exist yet - nothing to do
349 logger.debug("No migrations found, skipping")
350 return
352 current = get_current_revision(engine)
354 # BUG-3747: Pre-Alembic baseline detection.
355 #
356 # A database that has schema tables but no alembic_version row was
357 # created before commit 4fde036df (v1.4.0, 2026-03-21) via
358 # Base.metadata.create_all(). Without stamping, command.upgrade() runs
359 # 0001 (no-op for existing tables) followed by 0002+ against a legacy
360 # column shape. Migration 0007's index backfill silently fails on
361 # missing columns (e.g. settings.category), leaving the DB in a
362 # corrupted state. Stamping at "0001" bypasses the broken path.
363 if current is None:
364 inspector = inspect(engine)
365 existing_tables = set(inspector.get_table_names())
367 # Defensive guard: refuse what looks like an auth database. The
368 # auth DB has its own initialization path (`init_auth_database()`
369 # in `auth_db.py`) and contains ONLY the `users` table. Pre-
370 # Alembic user DBs ALSO contain `users` (created by the old
371 # `Base.metadata.create_all()` path before migration 0001 added
372 # the explicit skip), so we cannot just check "users present".
373 # Instead we check the auth-DB *shape*: only `users`, optionally
374 # alongside `alembic_version`. A real user DB always has 50+
375 # other tables. If the auth engine is ever accidentally routed
376 # through this function, this guard will refuse loudly rather
377 # than silently pollute the auth DB with user-DB tables.
378 non_metadata_tables = existing_tables - {"alembic_version"}
379 if non_metadata_tables == {"users"}:
380 raise RuntimeError(
381 "Refusing to run migrations on what looks like an auth "
382 f"database (only 'users' table present; tables: "
383 f"{sorted(existing_tables)}). Auth DB is initialized via "
384 "init_auth_database()."
385 )
387 # User-DB sentinels: both tables date to project inception
388 # (2025-06-29) and have never been renamed. We require BOTH —
389 # any single one could be present on a partial-init test DB
390 # (e.g. one that ran `Setting.__table__.create()` directly)
391 # where we'd want 0001's `create_all()` to add the missing
392 # tables, not be skipped by stamping. A real pre-Alembic
393 # production DB has 60+ tables and definitely has both sentinels.
394 PRE_ALEMBIC_SENTINELS = {"settings", "research_history"}
395 if PRE_ALEMBIC_SENTINELS.issubset(existing_tables):
396 logger.warning(
397 "BUG-3747: pre-Alembic database detected "
398 f"({len(existing_tables)} tables, no alembic_version). "
399 "Stamping at revision 0001 before applying migrations."
400 )
401 stamp_database(engine, "0001")
402 current = get_current_revision(engine)
403 logger.info(
404 f"BUG-3747: pre-Alembic DB stamped at {current}; "
405 "proceeding with upgrade to head"
406 )
408 # Short-circuit when the database is already at head. Calling
409 # command.upgrade() unconditionally opens a write transaction via
410 # engine.begin() even when there is nothing to apply — SQLite takes
411 # a RESERVED lock on the file as soon as the first DML lands inside
412 # that transaction, serialising concurrent readers behind a no-op on
413 # every cold engine reopen. The fresh-DB path (current is None) still
414 # runs the upgrade so tables and the alembic_version row get created.
415 if current is not None and current == head and target == "head":
416 logger.info(f"Database already at revision {head}; skipping upgrade")
417 return
419 if current is None:
420 logger.warning(
421 "Database has no migration history — applying migrations "
422 f"(target={target})"
423 )
424 elif current != head and target == "head":
425 logger.warning(
426 f"Database schema outdated (revision {current}, "
427 f"head is {head}) — applying migrations"
428 )
430 config = get_alembic_config(engine)
432 try:
433 with engine.connect() as conn:
434 _drop_orphan_alembic_temp_tables(conn)
435 _disable_fk_for_migration(conn)
436 with conn.begin():
437 config.attributes["connection"] = conn
438 command.upgrade(config, target)
439 # Re-enable FK on this connection BEFORE it returns to the
440 # pool so subsequent checkouts see the production-default ON
441 # state. The migration transaction has just committed, so we
442 # are back outside any active transaction — PRAGMA toggles
443 # work again. We can't rely on ``engine.dispose()`` to force
444 # a fresh connection because engines built with ``creator=``
445 # have ``url.database is None``, which fails the dispose
446 # guard below and leaves FK=OFF leaking into the pool.
447 conn.exec_driver_sql("PRAGMA foreign_keys = ON")
448 conn.rollback()
449 except Exception:
450 logger.exception(
451 "Database migration failed — database remains at previous "
452 "revision (auto-rollback by transaction manager)"
453 )
454 raise
456 # Belt-and-suspenders: dispose pooled connections after a successful
457 # upgrade. With FK explicitly re-enabled above this is no longer
458 # load-bearing for FK state, but it forces the next checkout through
459 # ``apply_performance_pragmas`` which also resets temp_store, cache_size,
460 # journal_mode, etc. — protecting against any future migration that
461 # touches connection-level PRAGMAs not handled by the FK fix-up.
462 # Skip for ``:memory:`` engines — those use a single shared connection
463 # and disposing it would destroy the just-migrated database.
464 db_name = engine.url.database
465 if db_name and db_name != ":memory:":
466 engine.dispose()
468 new_revision = get_current_revision(engine)
469 elapsed_ms = (time.perf_counter() - migration_start) * 1000
470 if current != new_revision: 470 ↛ 475line 470 didn't jump to line 475 because the condition on line 470 was always true
471 logger.warning(
472 f"Database migrated: {current} -> {new_revision} "
473 f"({elapsed_ms:.0f}ms)"
474 )
475 elif elapsed_ms > 100:
476 logger.info(
477 f"Database already at revision {new_revision} "
478 f"(no-op upgrade took {elapsed_ms:.0f}ms)"
479 )
480 else:
481 logger.info(
482 f"Database already at revision {new_revision} ({elapsed_ms:.0f}ms)"
483 )