Coverage for src/local_deep_research/database/alembic_runner.py: 91%
166 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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 # Clear any DBAPI transaction a custom engine/event configuration may
291 # already have opened. SQLite silently ignores the toggle in a transaction.
292 conn.rollback()
293 conn.exec_driver_sql("PRAGMA foreign_keys = OFF")
294 fk_enabled = conn.exec_driver_sql("PRAGMA foreign_keys").scalar_one()
295 # ``exec_driver_sql`` triggers SQLAlchemy autobegin even though no
296 # sqlite-level transaction was opened (PRAGMA isn't DML, so the
297 # driver doesn't auto-begin). Roll back the no-op SQLAlchemy
298 # transaction so the caller's ``conn.begin()`` is allowed to start
299 # a fresh one. ``PRAGMA foreign_keys`` is connection-level state and
300 # survives ROLLBACK at the SQLite level — see
301 # https://www.sqlite.org/pragma.html#pragma_foreign_keys.
302 conn.rollback()
303 if fk_enabled != 0:
304 raise RuntimeError(
305 "SQLite foreign-key enforcement remained enabled; refusing to "
306 "run migrations with an unsafe connection transaction state"
307 )
310def _restore_fk_after_migration(conn: Connection) -> None:
311 """Restore and verify FK enforcement before returning ``conn`` to its pool.
313 SQLite silently ignores ``PRAGMA foreign_keys = ON`` while a DBAPI
314 transaction is active. Merely executing the statement is therefore not
315 enough: clear any transaction, enable the pragma, and read it back. If any
316 part fails, invalidate the physical connection so an FK-disabled handle
317 can never be reused.
318 """
319 try:
320 conn.rollback()
321 conn.exec_driver_sql("PRAGMA foreign_keys = ON")
322 fk_enabled = conn.exec_driver_sql("PRAGMA foreign_keys").scalar_one()
323 conn.rollback()
324 except BaseException:
325 logger.exception(
326 "Failed to restore and verify foreign-key enforcement after "
327 "migration"
328 )
329 try:
330 conn.invalidate()
331 except Exception:
332 logger.exception(
333 "Failed to invalidate migration connection after "
334 "foreign-key restoration failure"
335 )
336 raise
338 if fk_enabled != 1:
339 logger.error(
340 "SQLite foreign-key enforcement remained disabled after "
341 "migration cleanup"
342 )
343 try:
344 conn.invalidate()
345 except Exception:
346 logger.exception(
347 "Failed to invalidate migration connection after "
348 "foreign-key verification failure"
349 )
350 raise RuntimeError(
351 "SQLite foreign-key enforcement remained disabled after "
352 "migration cleanup"
353 )
356def run_migrations(engine: Engine, target: str = "head") -> None:
357 """
358 Run pending migrations on a database.
360 The initial migration is idempotent (only creates tables that don't exist),
361 so this function runs migrations rather than just stamping existing
362 databases. This ensures any missing tables are created.
364 When ``target == "head"`` and the database is already at head, the call
365 short-circuits without opening a write transaction — calling
366 ``command.upgrade()`` unconditionally would still open a write transaction
367 via ``engine.begin()`` (taking a RESERVED lock on the SQLite file) just to
368 discover there's nothing to apply, serialising concurrent readers behind
369 a no-op on every cold engine reopen.
371 Security validations performed before running migrations:
372 - Migration directory path is within expected package boundary
373 - Migration files are not world-writable
375 Pre-upgrade hygiene (run outside the migration transaction):
376 - Drop orphan ``_alembic_tmp_*`` tables from prior failed
377 ``batch_alter_table`` runs (issue #3817).
378 - Disable ``PRAGMA foreign_keys`` so 0007's orphan scrub can run
379 (issue #3990). Re-enabled after a successful upgrade, before the
380 connection returns to the pool.
382 On failure inside the migration transaction, the inner
383 ``conn.begin()`` rolls back automatically — the database stays at
384 its previous revision. The original exception is re-raised so
385 callers can decide how to handle it.
387 Args:
388 engine: SQLAlchemy engine to migrate
389 target: Target revision (default "head" for latest)
391 Raises:
392 Exception: If migration fails (database is safely rolled back)
393 """
394 migration_start = time.perf_counter()
396 # Security: Validate migrations directory and file permissions
397 migrations_dir = get_migrations_dir()
398 _validate_migrations_permissions(migrations_dir)
400 head = get_head_revision()
402 if head is None:
403 # No migrations exist yet - nothing to do
404 logger.debug("No migrations found, skipping")
405 return
407 current = get_current_revision(engine)
409 # BUG-3747: Pre-Alembic baseline detection.
410 #
411 # A database that has schema tables but no alembic_version row was
412 # created before commit 4fde036df (v1.4.0, 2026-03-21) via
413 # Base.metadata.create_all(). Without stamping, command.upgrade() runs
414 # 0001 (no-op for existing tables) followed by 0002+ against a legacy
415 # column shape. Migration 0007's index backfill silently fails on
416 # missing columns (e.g. settings.category), leaving the DB in a
417 # corrupted state. Stamping at "0001" bypasses the broken path.
418 if current is None:
419 inspector = inspect(engine)
420 existing_tables = set(inspector.get_table_names())
422 # Defensive guard: refuse what looks like an auth database. The
423 # auth DB has its own initialization path (`init_auth_database()`
424 # in `auth_db.py`) and contains ONLY the `users` table. Pre-
425 # Alembic user DBs ALSO contain `users` (created by the old
426 # `Base.metadata.create_all()` path before migration 0001 added
427 # the explicit skip), so we cannot just check "users present".
428 # Instead we check the auth-DB *shape*: only `users`, optionally
429 # alongside `alembic_version`. A real user DB always has 50+
430 # other tables. If the auth engine is ever accidentally routed
431 # through this function, this guard will refuse loudly rather
432 # than silently pollute the auth DB with user-DB tables.
433 non_metadata_tables = existing_tables - {"alembic_version"}
434 if non_metadata_tables == {"users"}:
435 raise RuntimeError(
436 "Refusing to run migrations on what looks like an auth "
437 f"database (only 'users' table present; tables: "
438 f"{sorted(existing_tables)}). Auth DB is initialized via "
439 "init_auth_database()."
440 )
442 # User-DB sentinels: both tables date to project inception
443 # (2025-06-29) and have never been renamed. We require BOTH —
444 # any single one could be present on a partial-init test DB
445 # (e.g. one that ran `Setting.__table__.create()` directly)
446 # where we'd want 0001's `create_all()` to add the missing
447 # tables, not be skipped by stamping. A real pre-Alembic
448 # production DB has 60+ tables and definitely has both sentinels.
449 PRE_ALEMBIC_SENTINELS = {"settings", "research_history"}
450 if PRE_ALEMBIC_SENTINELS.issubset(existing_tables):
451 logger.warning(
452 "BUG-3747: pre-Alembic database detected "
453 f"({len(existing_tables)} tables, no alembic_version). "
454 "Stamping at revision 0001 before applying migrations."
455 )
456 stamp_database(engine, "0001")
457 current = get_current_revision(engine)
458 logger.info(
459 f"BUG-3747: pre-Alembic DB stamped at {current}; "
460 "proceeding with upgrade to head"
461 )
463 # Short-circuit when the database is already at head. Calling
464 # command.upgrade() unconditionally opens a write transaction via
465 # engine.begin() even when there is nothing to apply — SQLite takes
466 # a RESERVED lock on the file as soon as the first DML lands inside
467 # that transaction, serialising concurrent readers behind a no-op on
468 # every cold engine reopen. The fresh-DB path (current is None) still
469 # runs the upgrade so tables and the alembic_version row get created.
470 if current is not None and current == head and target == "head":
471 logger.info(f"Database already at revision {head}; skipping upgrade")
472 return
474 if current is None:
475 logger.warning(
476 "Database has no migration history — applying migrations "
477 f"(target={target})"
478 )
479 elif current != head and target == "head":
480 logger.warning(
481 f"Database schema outdated (revision {current}, "
482 f"head is {head}) — applying migrations"
483 )
485 config = get_alembic_config(engine)
487 try:
488 with engine.connect() as conn:
489 _drop_orphan_alembic_temp_tables(conn)
490 try:
491 # Keep disabling inside the cleanup boundary: the helper can
492 # raise after SQLite has already accepted PRAGMA ... = OFF.
493 _disable_fk_for_migration(conn)
494 with conn.begin():
495 config.attributes["connection"] = conn
496 command.upgrade(config, target)
497 except BaseException:
498 # Restore connection-local state on ordinary failures and
499 # cancellation alike. Never hide the original migration error
500 # if cleanup itself fails.
501 try:
502 _restore_fk_after_migration(conn)
503 except Exception:
504 logger.debug(
505 "Foreign-key cleanup failure was already logged; "
506 "preserving the original migration exception"
507 )
508 raise
509 else:
510 # We cannot rely on ``engine.dispose()`` below because
511 # creator-backed engines have ``url.database is None``.
512 _restore_fk_after_migration(conn)
513 except Exception:
514 logger.exception(
515 "Database migration failed — transaction rolled back where "
516 "supported; SQLite DDL from a failed revision may require repair"
517 )
518 raise
520 # Belt-and-suspenders: dispose pooled connections after a successful
521 # upgrade. With FK explicitly re-enabled above this is no longer
522 # load-bearing for FK state, but it forces the next checkout through
523 # ``apply_performance_pragmas`` which also resets temp_store, cache_size,
524 # journal_mode, etc. — protecting against any future migration that
525 # touches connection-level PRAGMAs not handled by the FK fix-up.
526 # Skip for ``:memory:`` engines — those use a single shared connection
527 # and disposing it would destroy the just-migrated database.
528 db_name = engine.url.database
529 if db_name and db_name != ":memory:":
530 engine.dispose()
532 new_revision = get_current_revision(engine)
533 elapsed_ms = (time.perf_counter() - migration_start) * 1000
534 if current != new_revision: 534 ↛ 539line 534 didn't jump to line 539 because the condition on line 534 was always true
535 logger.warning(
536 f"Database migrated: {current} -> {new_revision} "
537 f"({elapsed_ms:.0f}ms)"
538 )
539 elif elapsed_ms > 100:
540 logger.info(
541 f"Database already at revision {new_revision} "
542 f"(no-op upgrade took {elapsed_ms:.0f}ms)"
543 )
544 else:
545 logger.info(
546 f"Database already at revision {new_revision} ({elapsed_ms:.0f}ms)"
547 )