Add tag restore CLI

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,

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)

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