Merge pull request #242 from ggozad/chore/out-out-migration
Require explicit migrate command for database migrations
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
2f4354d100
15 changed files with 547 additions and 73 deletions
|
|
@ -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
|
||||
|
|
|
|||
27
docs/cli.md
27
docs/cli.md
|
|
@ -305,6 +305,33 @@ haiku-rag init [--db /path/to/your.lancedb]
|
|||
|
||||
This creates the database with the configured settings. **All other commands require an existing database** - they will fail with an informative error if the database doesn't exist.
|
||||
|
||||
### Migrate Database
|
||||
|
||||
Apply pending database migrations:
|
||||
|
||||
```bash
|
||||
haiku-rag migrate [--db /path/to/your.lancedb]
|
||||
```
|
||||
|
||||
When you upgrade haiku.rag to a new version that includes schema changes, the database requires migration. Opening a database with pending migrations will display an error:
|
||||
|
||||
```
|
||||
Error: Database requires migration from 0.19.0 to 0.26.5. 3 migration(s) pending. Run 'haiku-rag migrate' to upgrade.
|
||||
```
|
||||
|
||||
Run `haiku-rag migrate` to apply the pending migrations. The command shows which migrations were applied:
|
||||
|
||||
```
|
||||
Applied 3 migration(s):
|
||||
- 0.20.0: Add 'docling_document_json' and 'docling_version' columns
|
||||
- 0.23.1: Add content_fts column for contextualized FTS search
|
||||
- 0.25.0: Compress docling_document with gzip
|
||||
Migration completed successfully.
|
||||
```
|
||||
|
||||
!!! tip
|
||||
Back up your database before running migrations. While migrations are designed to be safe, having a backup provides peace of mind for production databases.
|
||||
|
||||
### Info
|
||||
|
||||
Display database metadata:
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ async with HaikuRAG("path/to/database.lancedb") as client:
|
|||
# Your code here
|
||||
pass
|
||||
|
||||
# Open in read-only mode (blocks writes, skips upgrades)
|
||||
# Open in read-only mode (blocks writes)
|
||||
async with HaikuRAG("path/to/database.lancedb", read_only=True) as client:
|
||||
results = await client.search("query") # Read operations work
|
||||
# await client.create_document(...) # Would raise ReadOnlyError
|
||||
|
|
@ -28,7 +28,10 @@ async with HaikuRAG("path/to/database.lancedb", read_only=True) as client:
|
|||
Databases must be explicitly created with `create=True` or via `haiku-rag init` before use. Operations on non-existent databases will raise `FileNotFoundError`.
|
||||
|
||||
!!! note
|
||||
Read-only mode is useful for safely accessing databases without risk of modification. It blocks all write operations, skips database upgrades on open, and prevents settings from being saved.
|
||||
Read-only mode is useful for safely accessing databases without risk of modification. It blocks all write operations and prevents settings from being saved.
|
||||
|
||||
!!! warning "Database Migrations"
|
||||
When upgrading haiku.rag to a version with schema changes, opening an existing database will raise `MigrationRequiredError`. Run `haiku-rag migrate` to apply pending migrations before using the database. See [CLI Database Management](cli.md#migrate-database) for details.
|
||||
|
||||
## Document Management
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from haiku.rag.cli import cli
|
||||
from haiku.rag.cli import _cli as cli
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
|
|
|||
191
tests/store/test_migrations.py
Normal file
191
tests/store/test_migrations.py
Normal 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 == []
|
||||
|
|
@ -602,3 +602,74 @@ async def test_rebuild_empty_database(tmp_path, monkeypatch):
|
|||
|
||||
calls = [str(c) for c in mock_print.call_args_list]
|
||||
assert any("No documents found" in c for c in calls)
|
||||
|
||||
|
||||
def test_migrate_with_pending_migrations(tmp_path):
|
||||
"""Test migrate method when migrations are applied."""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
db_path = tmp_path / "test.lancedb"
|
||||
store = Store(db_path, create=True)
|
||||
store.close()
|
||||
|
||||
app = HaikuRAGApp(db_path=db_path)
|
||||
|
||||
with patch("haiku.rag.store.engine.Store") as mock_store_class:
|
||||
mock_store = MagicMock()
|
||||
mock_store.migrate.return_value = ["Migration 1", "Migration 2"]
|
||||
mock_store_class.return_value = mock_store
|
||||
|
||||
result = app.migrate()
|
||||
|
||||
mock_store_class.assert_called_once_with(
|
||||
db_path,
|
||||
config=app.config,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
)
|
||||
mock_store.migrate.assert_called_once()
|
||||
mock_store.close.assert_called_once()
|
||||
assert result == ["Migration 1", "Migration 2"]
|
||||
|
||||
|
||||
def test_migrate_no_pending_migrations(tmp_path):
|
||||
"""Test migrate method when no migrations are pending."""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
db_path = tmp_path / "test.lancedb"
|
||||
store = Store(db_path, create=True)
|
||||
store.close()
|
||||
|
||||
app = HaikuRAGApp(db_path=db_path)
|
||||
|
||||
with patch("haiku.rag.store.engine.Store") as mock_store_class:
|
||||
mock_store = MagicMock()
|
||||
mock_store.migrate.return_value = []
|
||||
mock_store_class.return_value = mock_store
|
||||
|
||||
result = app.migrate()
|
||||
|
||||
mock_store.migrate.assert_called_once()
|
||||
mock_store.close.assert_called_once()
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_migrate_closes_store_on_exception(tmp_path):
|
||||
"""Test migrate method closes store even if migration fails."""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
db_path = tmp_path / "test.lancedb"
|
||||
store = Store(db_path, create=True)
|
||||
store.close()
|
||||
|
||||
app = HaikuRAGApp(db_path=db_path)
|
||||
|
||||
with patch("haiku.rag.store.engine.Store") as mock_store_class:
|
||||
mock_store = MagicMock()
|
||||
mock_store.migrate.side_effect = Exception("Migration error")
|
||||
mock_store_class.return_value = mock_store
|
||||
|
||||
with pytest.raises(Exception, match="Migration error"):
|
||||
app.migrate()
|
||||
|
||||
mock_store.close.assert_called_once()
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from haiku.rag.cli import cli
|
||||
from haiku.rag.cli import _cli as cli
|
||||
from haiku.rag.cli import cli as cli_wrapper
|
||||
from haiku.rag.store.exceptions import MigrationRequiredError
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
|
@ -386,3 +388,63 @@ def test_add_document_src_directory(tmp_path):
|
|||
mock_app_instance.add_document_from_source.assert_called_once()
|
||||
call_args = mock_app_instance.add_document_from_source.call_args
|
||||
assert call_args[1]["source"] == str(test_dir)
|
||||
|
||||
|
||||
def test_migrate_with_applied_migrations():
|
||||
"""Test migrate command when migrations are applied."""
|
||||
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
|
||||
mock_app_instance = MagicMock()
|
||||
mock_app_instance.migrate.return_value = [
|
||||
"Add full-text search index",
|
||||
"Add metadata column",
|
||||
]
|
||||
mock_app.return_value = mock_app_instance
|
||||
|
||||
result = runner.invoke(cli, ["migrate"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_app_instance.migrate.assert_called_once()
|
||||
assert "Applied 2 migration(s)" in result.output
|
||||
assert "Add full-text search index" in result.output
|
||||
assert "Add metadata column" in result.output
|
||||
assert "Migration completed successfully" in result.output
|
||||
|
||||
|
||||
def test_migrate_no_pending_migrations():
|
||||
"""Test migrate command when no migrations are pending."""
|
||||
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
|
||||
mock_app_instance = MagicMock()
|
||||
mock_app_instance.migrate.return_value = []
|
||||
mock_app.return_value = mock_app_instance
|
||||
|
||||
result = runner.invoke(cli, ["migrate"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_app_instance.migrate.assert_called_once()
|
||||
assert "No migrations pending" in result.output
|
||||
assert "Database is up to date" in result.output
|
||||
|
||||
|
||||
def test_migrate_failure():
|
||||
"""Test migrate command when migration fails."""
|
||||
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
|
||||
mock_app_instance = MagicMock()
|
||||
mock_app_instance.migrate.side_effect = Exception("Migration failed")
|
||||
mock_app.return_value = mock_app_instance
|
||||
|
||||
result = runner.invoke(cli, ["migrate"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "Migration failed" in result.output
|
||||
|
||||
|
||||
def test_cli_wrapper_catches_migration_required_error():
|
||||
"""Test that cli() wrapper catches MigrationRequiredError and exits with code 1."""
|
||||
with patch("haiku.rag.cli._cli") as mock_cli:
|
||||
mock_cli.side_effect = MigrationRequiredError(
|
||||
"Database requires migration. Run 'haiku-rag migrate' to upgrade."
|
||||
)
|
||||
|
||||
with patch("sys.exit") as mock_exit:
|
||||
cli_wrapper()
|
||||
mock_exit.assert_called_once_with(1)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, patch
|
|||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from haiku.rag.cli import cli
|
||||
from haiku.rag.cli import _cli as cli
|
||||
from haiku.rag.store.models import Document
|
||||
|
||||
runner = CliRunner()
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Reference in a new issue