From 4522fbdf1b3117217b9311a2333b669b3a7806b0 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 15 Jul 2026 17:07:58 +0300 Subject: [PATCH] Add tag restore CLI --- haiku_rag_slim/haiku/rag/app.py | 32 ++++++++-- haiku_rag_slim/haiku/rag/cli.py | 38 +++++++++++- tests/test_cli.py | 100 ++++++++++++++++++++++++++++++++ 3 files changed, 163 insertions(+), 7 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 302664ca..5a9dcb91 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -363,8 +363,7 @@ class HaikuRAGApp: # pragma: no cover async def create_tag(self, name: str): """Tag the current version of every table.""" if self._is_local and not self.db_path.exists(): - self.console.print("[red]Database path does not exist.[/red]") - return + raise ValueError(f"Database path does not exist: {self.db_path}") async with self._tag_write_store() as store: await store.create_tag(name) self.console.print(f"[green]Created tag '{escape(name)}'[/green]") @@ -372,8 +371,7 @@ class HaikuRAGApp: # pragma: no cover async def list_tags(self): """List database tags, flagging partial ones.""" if self._is_local and not self.db_path.exists(): - self.console.print("[red]Database path does not exist.[/red]") - return + raise ValueError(f"Database path does not exist: {self.db_path}") async with self._tag_read_store() as store: tags = await store.list_tags() @@ -394,12 +392,34 @@ class HaikuRAGApp: # pragma: no cover async def delete_tag(self, name: str): """Delete a tag from every table that has it.""" if self._is_local and not self.db_path.exists(): - self.console.print("[red]Database path does not exist.[/red]") - return + raise ValueError(f"Database path does not exist: {self.db_path}") async with self._tag_write_store() as store: await store.delete_tag(name) self.console.print(f"[green]Deleted tag '{escape(name)}'[/green]") + async def restore_tag(self, name: str): + """Restore the database to a tagged state and report the outcome. + + The Store context exits before anything is printed; no high-level + database access happens after the restore. + + Raises: + ValueError: If the database path does not exist. + """ + if self._is_local and not self.db_path.exists(): + raise ValueError(f"Database path does not exist: {self.db_path}") + async with self._tag_write_store() as store: + safety_tag = await store.restore_tag(name) + self.console.print(f"[green]Restored database to tag '{escape(name)}'.[/green]") + self.console.print( + f"The previous state is preserved as '{escape(safety_tag)}'." + ) + self.console.print( + "The restored state is now live. Later historical versions remain " + "until eligible for vacuum. Run [cyan]haiku-rag migrate[/cyan] if " + "migration is required." + ) + async def list_documents(self, filter: str | None = None): async with HaikuRAG( db_path=self.db_path, diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index bb1f713b..3dddbfe1 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -646,7 +646,11 @@ def tag_list( # pragma: no cover ), ): app = create_app(db) - asyncio.run(app.list_tags()) + try: + asyncio.run(app.list_tags()) + except (ValueError, RuntimeError) as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) @tag_cli.command("delete", help="Delete a tag") @@ -666,6 +670,38 @@ def tag_delete( # pragma: no cover raise typer.Exit(1) +@tag_cli.command("restore", help="Restore the database to a tagged state") +def tag_restore( # pragma: no cover + name: str = typer.Argument(help="Name of the tag to restore"), + yes: bool = typer.Option( + False, + "--yes", + help="Skip the confirmation prompt. Provides no locking or " + "concurrent-writer protection.", + ), + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + app = create_app(db) + if not yes: + typer.echo(f"Database: {app.db_path}") + typer.echo(f"Tag: {name}") + typer.echo("This changes the live database state across all tables.") + typer.echo("Stop all ingestion and other writers before continuing.") + typer.echo("The operation is coordinated but not transactionally atomic.") + typer.echo("A safety tag will preserve the current state.") + if not typer.confirm("Continue?", default=False): + raise typer.Exit(1) + try: + asyncio.run(app.restore_tag(name)) + except (ValueError, RuntimeError) as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) + + @_cli.command("download-models", help="Download Docling and Ollama models per config") def download_models_cmd(): # pragma: no cover app = HaikuRAGApp(db_path=Path(), config=get_config(), read_only=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4d770cb5..de404775 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -153,6 +153,17 @@ class TestTagCommands: assert "document_meta" not in asyncio.run(_table_names()) + def test_tag_commands_missing_database_exit_nonzero(self, tmp_path): + missing = str(tmp_path / "does_not_exist.lancedb") + for args in ( + ["tag", "create", "r1", "--db", missing], + ["tag", "delete", "r1", "--db", missing], + ["tag", "list", "--db", missing], + ): + result = runner.invoke(cli, args) + assert result.exit_code == 1, args + assert "does not exist" in result.output, args + def test_tag_create_invalid_name_fails_cleanly(self, temp_db_path): """lance restricts ref names to alphanumeric, '.', '-', '_'; the CLI surfaces that as a clean error instead of a traceback.""" @@ -164,3 +175,92 @@ class TestTagCommands: assert result.exit_code == 1 assert "Error:" in result.output assert "Ref characters" in result.output + + +class TestTagRestore: + def test_restore_requires_confirmation_and_decline_changes_nothing( + self, temp_db_path + ): + db = str(temp_db_path) + assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0 + assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0 + + result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db], input="n\n") + assert result.exit_code == 1 + assert "live database state" in result.output + assert "Stop all ingestion" in result.output + assert "not transactionally atomic" in result.output + assert "safety tag" in result.output + + result = runner.invoke(cli, ["tag", "list", "--db", db]) + assert "before-restore" not in result.output + + def test_restore_non_interactive_without_yes_fails(self, temp_db_path): + db = str(temp_db_path) + assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0 + assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0 + + result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db]) + assert result.exit_code == 1 + + result = runner.invoke(cli, ["tag", "list", "--db", db]) + assert "before-restore" not in result.output + + def test_restore_with_yes(self, temp_db_path): + db = str(temp_db_path) + assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0 + assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0 + + result = runner.invoke(cli, ["tag", "restore", "r1", "--yes", "--db", db]) + assert result.exit_code == 0 + assert "Restored database to tag 'r1'" in result.output + assert "before-restore-" in result.output + assert "now live" in result.output + assert "migrate" in result.output + + result = runner.invoke(cli, ["tag", "list", "--db", db]) + assert "before-restore-" in result.output + + def test_restore_missing_tag_errors(self, temp_db_path): + db = str(temp_db_path) + assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0 + + result = runner.invoke(cli, ["tag", "restore", "nope", "--yes", "--db", db]) + assert result.exit_code == 1 + assert "does not exist" in result.output + + def test_restore_partial_tag_errors(self, temp_db_path): + import asyncio + + from haiku.rag.store.engine import Store + + async def _partial_tag(): + async with Store(temp_db_path, create=True) as store: + version = await store.chunks_table.version() + await store.chunks_table.tags.create("stale", version) + + asyncio.run(_partial_tag()) + + result = runner.invoke( + cli, ["tag", "restore", "stale", "--yes", "--db", str(temp_db_path)] + ) + assert result.exit_code == 1 + assert "partial" in result.output + assert "documents" in result.output + + def test_restore_missing_database_exits_nonzero(self, tmp_path): + missing = tmp_path / "does_not_exist.lancedb" + result = runner.invoke( + cli, ["tag", "restore", "r1", "--yes", "--db", str(missing)] + ) + assert result.exit_code == 1 + assert "does not exist" in result.output + + def test_tag_help_includes_restore(self): + result = runner.invoke(cli, ["tag", "--help"]) + assert result.exit_code == 0 + assert "restore" in result.output + + result = runner.invoke(cli, ["--help"]) + assert "--before" not in result.output + assert "--at" not in result.output