require explicit migrate command for database migrations

This commit is contained in:
Yiorgis Gozadinos 2026-01-19 13:22:09 +02:00
parent 5f2f180960
commit 5464d24483
No known key found for this signature in database
9 changed files with 379 additions and 68 deletions

View file

@ -1,6 +1,15 @@
# Changelog
## [Unreleased]
### Changed
- **Explicit Database Migrations**: Database migrations are no longer applied automatically on open
- Opening a database with pending migrations now raises `MigrationRequiredError` with a clear message
- New `haiku-rag migrate` command to explicitly apply pending migrations
- Version-only updates (no schema changes) are applied silently in writable mode
- New `skip_migration_check` parameter on `Store` for tools that need to bypass the check
- `Store.migrate()` method returns list of applied migration descriptions
## [0.26.5] - 2026-01-16
### Added

View file

@ -578,6 +578,26 @@ class HaikuRAGApp:
await client.vacuum()
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
def migrate(self) -> list[str]:
"""Run pending database migrations.
Returns:
List of descriptions of applied migrations.
"""
from haiku.rag.store.engine import Store
store = Store(
self.db_path,
config=self.config,
skip_validation=True,
skip_migration_check=True,
)
try:
applied = store.migrate()
return applied
finally:
store.close()
async def create_index(self):
"""Create vector index on the chunks table."""
async with HaikuRAG(

View file

@ -1,5 +1,6 @@
import asyncio
import json
import sys
import warnings
from datetime import datetime
from importlib.metadata import version
@ -18,15 +19,25 @@ from haiku.rag.config import (
set_config,
)
from haiku.rag.logging import configure_cli_logging
from haiku.rag.store.exceptions import MigrationRequiredError
from haiku.rag.utils import is_up_to_date
# Load environment variables from .env file for API keys and service URLs
load_dotenv()
cli = typer.Typer(
_cli = typer.Typer(
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
)
def cli():
try:
_cli()
except MigrationRequiredError as e:
typer.echo(f"Error: {e}", err=True)
sys.exit(1)
# Module-level flags set by callback
_read_only: bool = False
_before: datetime | None = None
@ -65,7 +76,7 @@ def version_callback(value: bool):
raise typer.Exit()
@cli.callback()
@_cli.callback()
def main(
_version: bool = typer.Option(
False,
@ -141,7 +152,7 @@ def main(
pass
@cli.command("list", help="List all stored documents")
@_cli.command("list", help="List all stored documents")
def list_documents(
db: Path | None = typer.Option(
None,
@ -183,7 +194,7 @@ def _parse_meta_options(meta: list[str] | None) -> dict[str, Any]:
return result
@cli.command("add", help="Add a document from text input")
@_cli.command("add", help="Add a document from text input")
def add_document_text(
text: str = typer.Argument(
help="The text content of the document to add",
@ -205,7 +216,7 @@ def add_document_text(
asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None))
@cli.command("add-src", help="Add a document from a file path, directory, or URL")
@_cli.command("add-src", help="Add a document from a file path, directory, or URL")
def add_document_src(
source: str = typer.Argument(
help="The file path, directory, or URL of the document(s) to add",
@ -236,7 +247,7 @@ def add_document_src(
)
@cli.command("get", help="Get and display a document by its ID")
@_cli.command("get", help="Get and display a document by its ID")
def get_document(
doc_id: str = typer.Argument(
help="The ID of the document to get",
@ -251,7 +262,7 @@ def get_document(
asyncio.run(app.get_document(doc_id=doc_id))
@cli.command("delete", help="Delete a document by its ID")
@_cli.command("delete", help="Delete a document by its ID")
def delete_document(
doc_id: str = typer.Argument(
help="The ID of the document to delete",
@ -267,10 +278,12 @@ def delete_document(
# Add alias `rm` for delete
cli.command("rm", help="Alias for delete: remove a document by its ID")(delete_document)
_cli.command("rm", help="Alias for delete: remove a document by its ID")(
delete_document
)
@cli.command("search", help="Search for documents by a query")
@_cli.command("search", help="Search for documents by a query")
def search(
query: str = typer.Argument(
help="The search query to use",
@ -297,7 +310,7 @@ def search(
asyncio.run(app.search(query=query, limit=limit, filter=filter))
@cli.command("visualize", help="Show visual grounding for a chunk")
@_cli.command("visualize", help="Show visual grounding for a chunk")
def visualize(
chunk_id: str = typer.Argument(
help="The ID of the chunk to visualize",
@ -312,7 +325,7 @@ def visualize(
asyncio.run(app.visualize_chunk(chunk_id=chunk_id))
@cli.command("ask", help="Ask a question using the QA agent")
@_cli.command("ask", help="Ask a question using the QA agent")
def ask(
question: str = typer.Argument(
help="The question to ask",
@ -368,7 +381,7 @@ def ask(
)
@cli.command("research", help="Run multi-agent research and output a concise report")
@_cli.command("research", help="Run multi-agent research and output a concise report")
def research(
question: str = typer.Argument(..., help="The research question to investigate"),
db: Path | None = typer.Option(
@ -408,14 +421,14 @@ def research(
)
@cli.command("settings", help="Display current configuration settings")
@_cli.command("settings", help="Display current configuration settings")
def settings():
config = get_config()
app = HaikuRAGApp(db_path=Path(), config=config)
app.show_settings()
@cli.command("init-config", help="Generate a YAML configuration file")
@_cli.command("init-config", help="Generate a YAML configuration file")
def init_config(
output: Path = typer.Argument(
Path("haiku.rag.yaml"),
@ -447,7 +460,7 @@ def init_config(
typer.echo("Edit the file to customize your settings.")
@cli.command(
@_cli.command(
"rebuild",
help="Rebuild the database by deleting all chunks and re-indexing all documents",
)
@ -485,7 +498,7 @@ def rebuild(
asyncio.run(app.rebuild(mode=mode))
@cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
@_cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
def vacuum(
db: Path | None = typer.Option(
None,
@ -497,7 +510,32 @@ def vacuum(
asyncio.run(app.vacuum())
@cli.command("create-index", help="Create vector index for efficient similarity search")
@_cli.command("migrate", help="Run pending database migrations")
def migrate(
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
):
app = create_app(db)
try:
applied = app.migrate()
if applied:
typer.echo(f"Applied {len(applied)} migration(s):")
for desc in applied:
typer.echo(f" - {desc}")
typer.echo("Migration completed successfully.")
else:
typer.echo("No migrations pending. Database is up to date.")
except Exception as e:
typer.echo(f"Migration failed: {e}")
raise typer.Exit(1)
@_cli.command(
"create-index", help="Create vector index for efficient similarity search"
)
def create_index(
db: Path | None = typer.Option(
None,
@ -509,7 +547,7 @@ def create_index(
asyncio.run(app.create_index())
@cli.command("init", help="Initialize a new database")
@_cli.command("init", help="Initialize a new database")
def init_db(
db: Path | None = typer.Option(
None,
@ -521,7 +559,7 @@ def init_db(
asyncio.run(app.init())
@cli.command("info", help="Show database info")
@_cli.command("info", help="Show database info")
def info(
db: Path | None = typer.Option(
None,
@ -533,7 +571,7 @@ def info(
asyncio.run(app.info())
@cli.command("history", help="Show version history for database tables")
@_cli.command("history", help="Show version history for database tables")
def history(
db: Path | None = typer.Option(
None,
@ -557,7 +595,7 @@ def history(
asyncio.run(app.history(table=table, limit=limit))
@cli.command("download-models", help="Download Docling and Ollama models per config")
@_cli.command("download-models", help="Download Docling and Ollama models per config")
def download_models_cmd():
app = HaikuRAGApp(db_path=Path(), config=get_config())
try:
@ -567,7 +605,7 @@ def download_models_cmd():
raise typer.Exit(1)
@cli.command("inspect", help="Launch interactive TUI to inspect database contents")
@_cli.command("inspect", help="Launch interactive TUI to inspect database contents")
def inspect(
db: Path | None = typer.Option(
None,
@ -586,7 +624,7 @@ def inspect(
run_inspector(db_path, read_only=_read_only, before=_before)
@cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
@_cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
def chat(
db: Path | None = typer.Option(
None,
@ -624,7 +662,7 @@ def chat(
)
@cli.command(
@_cli.command(
"serve",
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
)

View file

@ -1,5 +1,5 @@
from .engine import Store
from .exceptions import ReadOnlyError
from .exceptions import MigrationRequiredError, ReadOnlyError
from .models import Chunk, Document
__all__ = ["Store", "Chunk", "Document", "ReadOnlyError"]
__all__ = ["Store", "Chunk", "Document", "MigrationRequiredError", "ReadOnlyError"]

View file

@ -14,7 +14,7 @@ from pydantic import Field
from haiku.rag.config import AppConfig, Config
from haiku.rag.embeddings import get_embedder
from haiku.rag.store.exceptions import ReadOnlyError
from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
logger = logging.getLogger(__name__)
@ -84,6 +84,7 @@ class Store:
create: bool = False,
read_only: bool = False,
before: datetime | None = None,
skip_migration_check: bool = False,
):
self.db_path: Path = db_path
self._config = config
@ -129,13 +130,12 @@ class Store:
if before is not None:
self._checkout_tables_before(before)
# Run upgrades only on existing databases, set version for new ones
# Skip upgrades in read-only mode (they would fail anyway)
if not self._read_only:
if is_new_db:
# Set version for new databases, check migrations for existing ones
if is_new_db:
if not self._read_only:
self._set_initial_version()
else:
self._run_upgrades()
elif not skip_migration_check:
self._check_migrations()
# Validate config compatibility after connection is established
if not skip_validation:
@ -371,26 +371,55 @@ class Store:
"""Set the initial version for a new database."""
self.set_haiku_version(metadata.version("haiku.rag-slim"))
def _run_upgrades(self):
"""Run pending database upgrades."""
try:
from haiku.rag.store.upgrades import run_pending_upgrades
def _check_migrations(self) -> None:
"""Check if migrations are pending and error or update version accordingly.
current_version = metadata.version("haiku.rag-slim")
db_version = self.get_haiku_version()
Raises:
MigrationRequiredError: If migrations are pending.
"""
from haiku.rag.store.upgrades import get_pending_upgrades
run_pending_upgrades(self, db_version, current_version)
current_version = metadata.version("haiku.rag-slim")
db_version = self.get_haiku_version()
self.set_haiku_version(current_version)
except Exception as e:
# Avoid hard failure on initial connection; log and continue so CLI remains usable.
logger.warning(
"Skipping upgrade due to error (db=%s -> pkg=%s): %s",
self.get_haiku_version(),
metadata.version("haiku.rag-slim"),
e,
pending = get_pending_upgrades(db_version)
if pending:
# Migrations are pending - require explicit migrate command
raise MigrationRequiredError(
f"Database requires migration from {db_version} to {current_version}. "
f"{len(pending)} migration(s) pending. "
"Run 'haiku-rag migrate' to upgrade."
)
# No pending migrations - update version silently if needed (writable only)
if not self._read_only and db_version != current_version:
self.set_haiku_version(current_version)
def migrate(self) -> list[str]:
"""Run pending database migrations.
Returns:
List of descriptions of applied upgrades.
Raises:
ReadOnlyError: If the store is in read-only mode.
"""
self._assert_writable()
from haiku.rag.store.upgrades import run_pending_upgrades
db_version = self.get_haiku_version()
current_version = metadata.version("haiku.rag-slim")
applied = run_pending_upgrades(self, db_version)
# Update version after successful migration
if applied or db_version != current_version:
self.set_haiku_version(current_version)
return applied
def get_haiku_version(self) -> str:
"""Returns the user version stored in settings."""
settings_records = list(

View file

@ -2,3 +2,9 @@ class ReadOnlyError(Exception):
"""Raised when a write operation is attempted on a read-only store."""
pass
class MigrationRequiredError(Exception):
"""Database requires migration. Run 'haiku-rag migrate' to upgrade."""
pass

View file

@ -1,10 +1,12 @@
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING
from packaging.version import Version, parse
from haiku.rag.store.engine import Store
if TYPE_CHECKING:
from haiku.rag.store.engine import Store
logger = logging.getLogger(__name__)
@ -14,7 +16,7 @@ class Upgrade:
"""Represents a database upgrade step."""
version: str
apply: Callable[[Store], None]
apply: Callable[["Store"], None]
description: str = ""
@ -22,24 +24,31 @@ class Upgrade:
upgrades: list[Upgrade] = []
def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> None:
"""Run upgrades where from_version < step.version <= to_version."""
def get_pending_upgrades(from_version: str) -> list[Upgrade]:
"""Get pending upgrades from the given version.
Returns:
List of Upgrade objects where from_version < upgrade.version,
sorted by version in ascending order.
"""
v_from: Version = parse(from_version)
v_to: Version = parse(to_version)
# Ensure that tests/development run available code upgrades even if the
# installed package version hasn't been bumped to include them yet.
if upgrades:
highest_step_version: Version = max(parse(u.version) for u in upgrades)
if highest_step_version > v_to:
v_to = highest_step_version
# Determine applicable steps
sorted_steps = sorted(upgrades, key=lambda u: parse(u.version))
applicable = [s for s in sorted_steps if v_from < parse(s.version) <= v_to]
return [s for s in sorted_steps if v_from < parse(s.version)]
def run_pending_upgrades(store: "Store", from_version: str) -> list[str]:
"""Run upgrades where from_version < step.version.
Returns:
List of descriptions of applied upgrades.
"""
applicable = get_pending_upgrades(from_version)
if applicable:
logger.info("%d upgrade step(s) pending", len(applicable))
applied: list[str] = []
# Apply in ascending order
for idx, step in enumerate(applicable, start=1):
logger.info(
@ -51,6 +60,11 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No
)
step.apply(store)
logger.info("Completed upgrade %s", step.version)
applied.append(
f"{step.version}: {step.description}" if step.description else step.version
)
return applied
# Import upgrade modules AFTER Upgrade class is defined to avoid circular imports

View file

@ -0,0 +1,191 @@
from importlib import metadata
import pytest
from haiku.rag.store import Store
from haiku.rag.store.exceptions import MigrationRequiredError
class TestMigrationRequiredError:
def test_migration_required_error_is_exception(self):
"""MigrationRequiredError should be a subclass of Exception."""
assert issubclass(MigrationRequiredError, Exception)
def test_migration_required_error_can_be_raised(self):
"""MigrationRequiredError can be raised and caught."""
with pytest.raises(MigrationRequiredError) as exc_info:
raise MigrationRequiredError("Run 'haiku-rag migrate' to upgrade")
assert "migrate" in str(exc_info.value)
class TestMigrationCheck:
def test_new_database_sets_version(self, temp_db_path):
"""New database should set the current package version."""
store = Store(temp_db_path, create=True)
version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
store.close()
def test_existing_database_same_version_no_error(self, temp_db_path):
"""Opening a database with the same version should not error."""
store = Store(temp_db_path, create=True)
store.close()
# Re-open - should work without error
store = Store(temp_db_path)
store.close()
def test_version_bump_without_pending_migrations_updates_silently(
self, temp_db_path
):
"""When version is outdated but no migrations pending, update version silently."""
store = Store(temp_db_path, create=True)
# Set an older version that has no pending migrations
# (newer than all current upgrade steps)
store.set_haiku_version("100.0.0")
store.close()
# Re-open - should update version silently, no error
store = Store(temp_db_path)
# Version should now be current
version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
store.close()
def test_pending_migrations_raises_error(self, temp_db_path):
"""When actual migrations are pending, should raise MigrationRequiredError."""
store = Store(temp_db_path, create=True)
# Set version to before the first upgrade step
store.set_haiku_version("0.19.0")
store.close()
# Re-open should raise
with pytest.raises(MigrationRequiredError) as exc_info:
Store(temp_db_path)
assert "migrate" in str(exc_info.value).lower()
def test_pending_migrations_read_only_raises_error(self, temp_db_path):
"""Read-only mode with pending migrations should still raise."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
with pytest.raises(MigrationRequiredError):
Store(temp_db_path, read_only=True)
def test_read_only_version_bump_without_migrations_ok(self, temp_db_path):
"""Read-only mode with version bump but no migrations should work."""
store = Store(temp_db_path, create=True)
# Set a version newer than all upgrade steps
store.set_haiku_version("100.0.0")
store.close()
# Read-only open should work (version not updated, but no error)
store = Store(temp_db_path, read_only=True)
# Version should stay at the old value (can't update in read-only)
assert store.get_haiku_version() == "100.0.0"
store.close()
def test_skip_migration_check_bypasses_error(self, temp_db_path):
"""skip_migration_check=True should bypass migration error."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
# Open with skip_migration_check should work
store = Store(temp_db_path, skip_migration_check=True)
# Version should remain old (no auto-migration)
assert store.get_haiku_version() == "0.19.0"
store.close()
class TestMigrateMethod:
def test_migrate_applies_pending_upgrades(self, temp_db_path):
"""Store.migrate() should apply pending upgrades and update version."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
# Open with skip_migration_check to avoid error
store = Store(temp_db_path, skip_migration_check=True)
old_version = store.get_haiku_version()
assert old_version == "0.19.0"
# Run migration
applied = store.migrate()
# Should have applied migrations
assert len(applied) > 0
# Version should be updated
new_version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert new_version == expected
store.close()
def test_migrate_returns_applied_upgrades(self, temp_db_path):
"""Store.migrate() should return list of applied upgrade descriptions."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
# Should return descriptions of applied upgrades
assert isinstance(applied, list)
for item in applied:
assert isinstance(item, str)
store.close()
def test_migrate_with_no_pending_returns_empty(self, temp_db_path):
"""Store.migrate() with no pending migrations returns empty list."""
store = Store(temp_db_path, create=True)
# Already at current version
store.close()
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
assert applied == []
store.close()
def test_migrate_raises_read_only_error(self, temp_db_path):
"""Store.migrate() should raise ReadOnlyError in read-only mode."""
from haiku.rag.store.exceptions import ReadOnlyError
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
store = Store(temp_db_path, skip_migration_check=True, read_only=True)
with pytest.raises(ReadOnlyError):
store.migrate()
store.close()
class TestGetPendingUpgrades:
def test_get_pending_upgrades_returns_list(self):
"""get_pending_upgrades() should return a list of Upgrade objects."""
from haiku.rag.store.upgrades import get_pending_upgrades
pending = get_pending_upgrades("0.19.0")
assert isinstance(pending, list)
# Should have at least the v0.20.0 upgrade
assert len(pending) > 0
def test_get_pending_upgrades_from_current_version_is_empty(self):
"""get_pending_upgrades() from current version should be empty."""
from haiku.rag.store.upgrades import get_pending_upgrades
current = metadata.version("haiku.rag-slim")
pending = get_pending_upgrades(current)
assert pending == []
def test_get_pending_upgrades_from_future_version_is_empty(self):
"""get_pending_upgrades() from a future version should be empty."""
from haiku.rag.store.upgrades import get_pending_upgrades
pending = get_pending_upgrades("100.0.0")
assert pending == []

View file

@ -72,20 +72,24 @@ def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
Store(temp_db_path, create=True)
def test_existing_database_runs_upgrades(monkeypatch, temp_db_path):
def test_existing_database_checks_migrations(monkeypatch, temp_db_path):
Store(temp_db_path, create=True)
called = {"value": False}
from haiku.rag.store import upgrades
def mark_called(*_args, **_kwargs):
called = {"value": False}
original_get_pending = upgrades.get_pending_upgrades
def mark_called(*args, **kwargs):
called["value"] = True
return original_get_pending(*args, **kwargs)
monkeypatch.setattr(
"haiku.rag.store.upgrades.run_pending_upgrades",
"haiku.rag.store.upgrades.get_pending_upgrades",
mark_called,
)
# Opening an existing database should trigger upgrades
# Opening an existing database should check for pending migrations
Store(temp_db_path)
assert called["value"]