Coverage for src/local_deep_research/database/alembic_runner.py: 93%
181 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +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, text
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 # The catalog can contain an embedded quote even with this prefix.
260 # Use the dialect's identifier quoting, which doubles that quote,
261 # before interpolating a name into DDL (identifiers cannot be bound).
262 quoted_name = conn.dialect.identifier_preparer.quote_identifier(name)
263 # bearer:disable python_lang_sql_injection
264 # nosemgrep: semgrep.rules.sql-string-concatenation, reason: dialect quote_identifier escapes the catalog name before SQL interpolation
265 conn.exec_driver_sql(f"DROP TABLE IF EXISTS {quoted_name}") # noqa: S608
268def _disable_fk_for_migration(conn: Connection) -> None:
269 """Disable FK enforcement on the migration connection BEFORE any
270 transaction opens (issue #3990).
272 ``apply_performance_pragmas`` set ``PRAGMA foreign_keys = ON`` at
273 connect. SQLite then *silently ignores* further toggles of
274 ``foreign_keys`` once any transaction (explicit or driver-implicit)
275 is active. ``run_migrations`` starts the physical transaction at its
276 ``conn.begin()`` boundary, before the first revision runs, so this
277 helper must run before that boundary for the toggle to take effect.
279 Before this helper, the first DML of a multi-migration upgrade
280 (revision 0001 → 0009) auto-began the driver transaction and froze
281 FK in the connect-time ON state for the rest of the upgrade.
282 That defeated migration 0007's defensive PRAGMA OFF and made its
283 orphan-scrub DELETE fail with ``foreign key mismatch`` on tables
284 whose FK target lacks a UNIQUE backing — exactly the broken-schema
285 state migration 0007 is meant to repair.
287 The caller is responsible for re-enabling FK after the migration
288 transaction commits, BEFORE returning the connection to the pool —
289 see ``run_migrations``.
290 """
291 # Clear any DBAPI transaction a custom engine/event configuration may
292 # already have opened. SQLite silently ignores the toggle in a transaction.
293 conn.rollback()
294 conn.exec_driver_sql("PRAGMA foreign_keys = OFF")
295 fk_enabled = conn.exec_driver_sql("PRAGMA foreign_keys").scalar_one()
296 # ``exec_driver_sql`` triggers SQLAlchemy autobegin even though no
297 # sqlite-level transaction was opened (PRAGMA isn't DML, so the
298 # driver doesn't auto-begin). Roll back the no-op SQLAlchemy
299 # transaction so the caller's ``conn.begin()`` is allowed to start
300 # a fresh one. ``PRAGMA foreign_keys`` is connection-level state and
301 # survives ROLLBACK at the SQLite level — see
302 # https://www.sqlite.org/pragma.html#pragma_foreign_keys.
303 conn.rollback()
304 if fk_enabled != 0:
305 raise RuntimeError(
306 "SQLite foreign-key enforcement remained enabled; refusing to "
307 "run migrations with an unsafe connection transaction state"
308 )
311def _restore_fk_after_migration(conn: Connection) -> None:
312 """Restore and verify FK enforcement before returning ``conn`` to its pool.
314 SQLite silently ignores ``PRAGMA foreign_keys = ON`` while a DBAPI
315 transaction is active. Merely executing the statement is therefore not
316 enough: clear any transaction, enable the pragma, and read it back. If any
317 part fails, invalidate the physical connection so an FK-disabled handle
318 can never be reused.
319 """
320 try:
321 conn.rollback()
322 conn.exec_driver_sql("PRAGMA foreign_keys = ON")
323 fk_enabled = conn.exec_driver_sql("PRAGMA foreign_keys").scalar_one()
324 conn.rollback()
325 except BaseException:
326 logger.exception(
327 "Failed to restore and verify foreign-key enforcement after "
328 "migration"
329 )
330 try:
331 conn.invalidate()
332 except Exception:
333 logger.exception(
334 "Failed to invalidate migration connection after "
335 "foreign-key restoration failure"
336 )
337 raise
339 if fk_enabled != 1:
340 logger.error(
341 "SQLite foreign-key enforcement remained disabled after "
342 "migration cleanup"
343 )
344 try:
345 conn.invalidate()
346 except Exception:
347 logger.exception(
348 "Failed to invalidate migration connection after "
349 "foreign-key verification failure"
350 )
351 raise RuntimeError(
352 "SQLite foreign-key enforcement remained disabled after "
353 "migration cleanup"
354 )
357# Revisions a database can be at before 0004 runs. A pre-Alembic database is
358# stamped at 0001, and one whose earlier upgrade failed at 0004 stays below it.
359_REVISIONS_BEFORE_0004 = frozenset({"0001", "0002", "0003"})
362def _add_settings_env_var_before_0004(
363 conn: Connection, current: Optional[str]
364) -> None:
365 """Give a pre-0004 ``settings`` table the ``env_var`` column 0004 copies.
367 Revision 0004 names ``settings.env_var`` when it re-scopes legacy
368 ``app.*`` keys, and some pre-Alembic settings tables predate that column,
369 so the upgrade failed there. 0004 has shipped and cannot change, and 0032
370 adds the column only after 0004 runs, so add it here first. The column
371 is nullable, matching the model; databases that have it are untouched.
372 """
373 if current not in _REVISIONS_BEFORE_0004:
374 return
375 inspector = inspect(conn)
376 if not inspector.has_table("settings"):
377 return
378 if "env_var" in {c["name"] for c in inspector.get_columns("settings")}:
379 return
380 conn.execute(text("ALTER TABLE settings ADD COLUMN env_var VARCHAR(255)"))
381 logger.info("Added settings.env_var before revision 0004")
384def run_migrations(engine: Engine, target: str = "head") -> None:
385 """
386 Run pending migrations on a database.
388 The initial migration is idempotent (only creates tables that don't exist),
389 so this function runs migrations rather than just stamping existing
390 databases. This ensures any missing tables are created.
392 When ``target == "head"`` and the database is already at head, the call
393 short-circuits before the migration transaction opens. On SQLite that
394 transaction starts with ``BEGIN IMMEDIATE`` (unless a begin-event hook
395 already started one), so running it anyway would take the write lock
396 just to discover there's nothing to apply, making other writers wait
397 behind a no-op on every cold engine reopen. Any other ``target``, even
398 the head revision's own ID, is not resolved up front, so the
399 transaction opens and ``command.upgrade()`` runs even when the database
400 is already at that revision.
402 Security validations performed before running migrations:
403 - Migration directory path is within expected package boundary
404 - Migration files are not world-writable
406 Pre-upgrade hygiene (run outside the migration transaction):
407 - Drop orphan ``_alembic_tmp_*`` tables from prior failed
408 ``batch_alter_table`` runs (issue #3817).
409 - Disable ``PRAGMA foreign_keys`` so 0007's orphan scrub can run
410 (issue #3990). Re-enabled after a successful upgrade, before the
411 connection returns to the pool.
413 Inside the migration transaction, before Alembic runs, a database at
414 revision 0001, 0002 or 0003 gains ``settings.env_var`` if its settings
415 table lacks it (see ``_add_settings_env_var_before_0004``), so a failed
416 upgrade rolls that column back too.
418 On failure inside the migration transaction, the inner
419 ``conn.begin()`` rolls back automatically — the database stays at
420 its previous revision. On SQLite the physical transaction starts
421 before the first revision runs, so the rollback covers schema changes
422 as well as data. The original exception is re-raised so callers can
423 decide how to handle it.
425 Args:
426 engine: SQLAlchemy engine to migrate
427 target: Target revision (default "head" for latest)
429 Raises:
430 Exception: If migration fails (database is safely rolled back)
431 """
432 migration_start = time.perf_counter()
434 # Security: Validate migrations directory and file permissions
435 migrations_dir = get_migrations_dir()
436 _validate_migrations_permissions(migrations_dir)
438 head = get_head_revision()
440 if head is None:
441 # No migrations exist yet - nothing to do
442 logger.debug("No migrations found, skipping")
443 return
445 current = get_current_revision(engine)
447 # BUG-3747: Pre-Alembic baseline detection.
448 #
449 # A database that has schema tables but no alembic_version row was
450 # created before commit 4fde036df (v1.4.0, 2026-03-21) via
451 # Base.metadata.create_all(). Without stamping, command.upgrade() runs
452 # 0001 (no-op for existing tables) followed by 0002+ against a legacy
453 # column shape. Migration 0007's index backfill silently fails on
454 # missing columns (e.g. settings.category), leaving the DB in a
455 # corrupted state. Stamping at "0001" bypasses the broken path.
456 if current is None:
457 inspector = inspect(engine)
458 existing_tables = set(inspector.get_table_names())
460 # Defensive guard: refuse what looks like an auth database. The
461 # auth DB has its own initialization path (`init_auth_database()`
462 # in `auth_db.py`) and contains ONLY the `users` table. Pre-
463 # Alembic user DBs ALSO contain `users` (created by the old
464 # `Base.metadata.create_all()` path before migration 0001 added
465 # the explicit skip), so we cannot just check "users present".
466 # Instead we check the auth-DB *shape*: only `users`, optionally
467 # alongside `alembic_version`. A real user DB always has 50+
468 # other tables. If the auth engine is ever accidentally routed
469 # through this function, this guard will refuse loudly rather
470 # than silently pollute the auth DB with user-DB tables.
471 non_metadata_tables = existing_tables - {"alembic_version"}
472 if non_metadata_tables == {"users"}:
473 raise RuntimeError(
474 "Refusing to run migrations on what looks like an auth "
475 f"database (only 'users' table present; tables: "
476 f"{sorted(existing_tables)}). Auth DB is initialized via "
477 "init_auth_database()."
478 )
480 # User-DB sentinels: both tables date to project inception
481 # (2025-06-29) and have never been renamed. We require BOTH —
482 # any single one could be present on a partial-init test DB
483 # (e.g. one that ran `Setting.__table__.create()` directly)
484 # where we'd want 0001's `create_all()` to add the missing
485 # tables, not be skipped by stamping. A real pre-Alembic
486 # production DB has 60+ tables and definitely has both sentinels.
487 PRE_ALEMBIC_SENTINELS = {"settings", "research_history"}
488 if PRE_ALEMBIC_SENTINELS.issubset(existing_tables):
489 logger.warning(
490 "BUG-3747: pre-Alembic database detected "
491 f"({len(existing_tables)} tables, no alembic_version). "
492 "Stamping at revision 0001 before applying migrations."
493 )
494 stamp_database(engine, "0001")
495 current = get_current_revision(engine)
496 logger.info(
497 f"BUG-3747: pre-Alembic DB stamped at {current}; "
498 "proceeding with upgrade to head"
499 )
501 # Short-circuit when the database is already at head. The migration
502 # transaction below starts with BEGIN IMMEDIATE, which takes SQLite's
503 # write lock even when there is nothing to apply, making other writers
504 # wait behind a no-op on every cold engine reopen. The fresh-DB path
505 # (current is None) still runs the upgrade so tables and the
506 # alembic_version row get created.
507 if current is not None and current == head and target == "head":
508 logger.info(f"Database already at revision {head}; skipping upgrade")
509 return
511 if current is None:
512 logger.warning(
513 "Database has no migration history — applying migrations "
514 f"(target={target})"
515 )
516 elif current != head and target == "head":
517 logger.warning(
518 f"Database schema outdated (revision {current}, "
519 f"head is {head}) — applying migrations"
520 )
522 config = get_alembic_config(engine)
524 try:
525 with engine.connect() as conn:
526 _drop_orphan_alembic_temp_tables(conn)
527 try:
528 # Keep disabling inside the cleanup boundary: the helper can
529 # raise after SQLite has already accepted PRAGMA ... = OFF.
530 _disable_fk_for_migration(conn)
531 with conn.begin():
532 # SQLite's legacy driver begins only on DML. Start its
533 # physical transaction before DDL so schema changes roll
534 # back too. IMMEDIATE takes the write lock now, so
535 # busy_timeout covers waiting for another writer. With a
536 # DEFERRED BEGIN, Alembic's first reads would pin a WAL
537 # snapshot, and a commit from another connection before
538 # the first write would fail that write at once with
539 # SQLITE_BUSY_SNAPSHOT, which busy_timeout never retries.
540 # A begin-event hook may already have started the
541 # transaction; its mode is kept.
542 if (
543 conn.dialect.name == "sqlite"
544 and not conn.connection.driver_connection.in_transaction
545 ):
546 conn.exec_driver_sql("BEGIN IMMEDIATE")
547 # After the BEGIN, so a failed upgrade also rolls back
548 # this ADD COLUMN instead of leaving it committed.
549 _add_settings_env_var_before_0004(conn, current)
550 config.attributes["connection"] = conn
551 command.upgrade(config, target)
552 except BaseException:
553 # Restore connection-local state on ordinary failures and
554 # cancellation alike. Never hide the original migration error
555 # if cleanup itself fails.
556 try:
557 _restore_fk_after_migration(conn)
558 except Exception:
559 logger.debug(
560 "Foreign-key cleanup failure was already logged; "
561 "preserving the original migration exception"
562 )
563 raise
564 else:
565 # We cannot rely on ``engine.dispose()`` below because
566 # creator-backed engines have ``url.database is None``.
567 _restore_fk_after_migration(conn)
568 except Exception:
569 logger.exception(
570 "Database migration failed — transaction rolled back where "
571 "supported"
572 )
573 raise
575 # Belt-and-suspenders: dispose pooled connections after a successful
576 # upgrade. With FK explicitly re-enabled above this is no longer
577 # load-bearing for FK state, but it forces the next checkout through
578 # ``apply_performance_pragmas`` which also resets temp_store, cache_size,
579 # journal_mode, etc. — protecting against any future migration that
580 # touches connection-level PRAGMAs not handled by the FK fix-up.
581 # Skip for ``:memory:`` engines — those use a single shared connection
582 # and disposing it would destroy the just-migrated database.
583 db_name = engine.url.database
584 if db_name and db_name != ":memory:":
585 engine.dispose()
587 new_revision = get_current_revision(engine)
588 elapsed_ms = (time.perf_counter() - migration_start) * 1000
589 if current != new_revision:
590 logger.warning(
591 f"Database migrated: {current} -> {new_revision} "
592 f"({elapsed_ms:.0f}ms)"
593 )
594 elif elapsed_ms > 100: 594 ↛ 595line 594 didn't jump to line 595 because the condition on line 594 was never true
595 logger.info(
596 f"Database already at revision {new_revision} "
597 f"(no-op upgrade took {elapsed_ms:.0f}ms)"
598 )
599 else:
600 logger.info(
601 f"Database already at revision {new_revision} ({elapsed_ms:.0f}ms)"
602 )