Centralize vacuuming

src/haiku/rag/client.py

@@ -617,13 +617,13 @@ class HaikuRAG:
 
         # Final maintenance: centralized vacuum to curb disk usage
         try:
-            self.store.vacuum()
+            await self.store.vacuum()
         except Exception:
             pass
 
     async def vacuum(self) -> None:
         """Optimize and clean up old versions across all tables."""
-        self.store.vacuum()
+        await self.store.vacuum()
 
     def close(self):
         """Close the underlying store connection."""

src/haiku/rag/migration.py

@@ -27,7 +27,7 @@ class SQLiteToLanceDBMigrator:
         self.lancedb_path = lancedb_path
         self.console = Console()
 
-    def migrate(self) -> bool:
+    async def migrate(self) -> bool:
         """Perform the migration."""
         try:
             self.console.print(
@@ -94,7 +94,7 @@ class SQLiteToLanceDBMigrator:
             # Optimize and cleanup using centralized vacuum
             self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
             try:
-                lance_store.vacuum()
+                await lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")
             except Exception as e:
                 self.console.print(
@@ -313,4 +313,4 @@ async def migrate_sqlite_to_lancedb(
         lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
 
     migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
-    return migrator.migrate()
+    return await migrator.migrate()

src/haiku/rag/store/engine.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 from datetime import timedelta
@@ -51,6 +52,7 @@ class Store:
     def __init__(self, db_path: Path, skip_validation: bool = False):
         self.db_path: Path = db_path
         self.embedder = get_embedder()
+        self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
@@ -78,20 +80,38 @@ class Store:
         if not skip_validation:
             self._validate_configuration()
 
-    def vacuum(self, retention_seconds: int = Config.VACUUM_RETENTION_SECONDS) -> None:
+    async def vacuum(
+        self, retention_seconds: int = Config.VACUUM_RETENTION_SECONDS
+    ) -> None:
         """Optimize and clean up old versions across all tables to reduce disk usage.
 
         Args:
             retention_seconds: Retention threshold in seconds. Only versions older
                               than this will be removed. Defaults to Config.VACUUM_RETENTION_SECONDS.
+
+        Note:
+            If vacuum is already running, this method returns immediately without blocking.
         """
         if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
             return
 
-        # Perform maintenance per table using optimize() with configurable retention
+        # Skip if already running (non-blocking)
-        retention = timedelta(seconds=retention_seconds)
+        if self._vacuum_lock.locked():
-        for table in [self.documents_table, self.chunks_table, self.settings_table]:
+            return
-            table.optimize(cleanup_older_than=retention)
+
+        async with self._vacuum_lock:
+            try:
+                # Perform maintenance per table using optimize() with configurable retention
+                retention = timedelta(seconds=retention_seconds)
+                for table in [
+                    self.documents_table,
+                    self.chunks_table,
+                    self.settings_table,
+                ]:
+                    table.optimize(cleanup_older_than=retention)
+            except (RuntimeError, OSError) as e:
+                # Handle resource errors gracefully
+                logger.debug(f"Vacuum skipped due to resource constraints: {e}")
 
     def _connect_to_lancedb(self, db_path: Path):
         """Establish connection to LanceDB (local, cloud, or object storage)."""

src/haiku/rag/store/repositories/chunk.py

@@ -1,4 +1,3 @@
-import asyncio
 import inspect
 import json
 import logging
@@ -23,7 +22,6 @@ class ChunkRepository:
     def __init__(self, store: Store) -> None:
         self.store = store
         self.embedder = get_embedder()
-        self._optimize_lock = asyncio.Lock()
 
     def _ensure_fts_index(self) -> None:
         """Ensure FTS index exists on the content column."""
@@ -35,21 +33,6 @@ class ChunkRepository:
             # Log the error but don't fail - FTS might already exist
             logger.debug(f"FTS index creation skipped: {e}")
 
-    async def _optimize(self) -> None:
-        """Optimize the chunks table to refresh indexes."""
-        # Skip optimization for LanceDB Cloud as it handles this automatically
-        if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
-            return
-
-        async with self._optimize_lock:
-            try:
-                self.store.chunks_table.optimize()
-            except (RuntimeError, OSError) as e:
-                # Handle "too many open files" and other resource errors gracefully
-                logger.debug(
-                    f"Table optimization skipped due to resource constraints: {e}"
-                )
-
     async def create(self, entity: Chunk) -> Chunk:
         """Create a chunk in the database."""
         assert entity.document_id, "Chunk must have a document_id to be created"
@@ -77,11 +60,6 @@ class ChunkRepository:
         self.store.chunks_table.add([chunk_record])
 
         entity.id = chunk_id
-
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def get_by_id(self, entity_id: str) -> Chunk | None:
@@ -125,10 +103,6 @@ class ChunkRepository:
                 "vector": embedding,
             },
         )
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def delete(self, entity_id: str) -> bool:
@@ -227,8 +201,6 @@ class ChunkRepository:
         if chunk_records:
             self.store.chunks_table.add(chunk_records)
 
-        # Force optimization once at the end for bulk operations
-        await self._optimize()
         return created_chunks
 
     async def delete_all(self) -> None:

src/haiku/rag/store/repositories/document.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from datetime import datetime
 from typing import TYPE_CHECKING
@@ -200,8 +201,9 @@ class DocumentRepository:
                     chunk.order = order
                     await self.chunk_repository.create(chunk)
 
-            # Vacuum old versions after successful creation
+            # Vacuum old versions in background (non-blocking)
-            self.store.vacuum()
+            asyncio.create_task(self.store.vacuum())
+
             return created_doc
         except Exception:
             # Roll back to the captured versions and re-raise
@@ -232,8 +234,9 @@ class DocumentRepository:
                 updated_doc.id, docling_document
             )
 
-            # Vacuum old versions after successful update
+            # Vacuum old versions in background (non-blocking)
-            self.store.vacuum()
+            asyncio.create_task(self.store.vacuum())
+
             return updated_doc
         except Exception:
             # Roll back to the captured versions and re-raise

tests/test_lancedb_connection.py

@@ -4,43 +4,24 @@ import pytest
 
 from haiku.rag.config import Config
 from haiku.rag.store.engine import Store
-from haiku.rag.store.models.chunk import Chunk
-from haiku.rag.store.models.document import Document
-from haiku.rag.store.repositories.chunk import ChunkRepository
-from haiku.rag.store.repositories.document import DocumentRepository
 
 
 @pytest.mark.asyncio
 async def test_lancedb_cloud_skips_optimization(temp_db_path):
-    """Test that optimization is skipped when using LanceDB Cloud (db:// URI)."""
+    """Test that vacuum is skipped when using LanceDB Cloud (db:// URI)."""
     # Create a store
     store = Store(temp_db_path)
-    chunk_repo = ChunkRepository(store)
-    doc_repo = DocumentRepository(store)
 
-    # Create a document
+    # Mock all cloud config to simulate LanceDB Cloud usage
-    document = Document(content="Test document content", metadata={})
+    with (
-    created_document = await doc_repo.create(document)
+        patch.object(Config, "LANCEDB_URI", "db://test-database"),
-    document_id = created_document.id
+        patch.object(Config, "LANCEDB_API_KEY", "test-api-key"),
-
+        patch.object(Config, "LANCEDB_REGION", "us-east-1"),
-    # Mock LANCEDB_URI to simulate LanceDB Cloud usage
+    ):
-    with patch.object(Config, "LANCEDB_URI", "db://test-database"):
         # Mock the optimize method to track if it's called
         with patch.object(store.chunks_table, "optimize") as mock_optimize:
-            # Create a chunk - this should trigger optimization logic
+            # Call vacuum - this should skip optimization for LanceDB Cloud
-            chunk = Chunk(
+            await store.vacuum()
-                document_id=document_id,
-                content="Test chunk content",
-                metadata={"test": "value"},
-            )
-
-            created_chunk = await chunk_repo.create(chunk)
-            assert created_chunk.id is not None
-
-            # Wait a moment to ensure any async optimization would complete
-            import asyncio
-
-            await asyncio.sleep(0.1)
 
             # The optimize method should NOT have been called for LanceDB Cloud
             mock_optimize.assert_not_called()
@@ -50,35 +31,16 @@ async def test_lancedb_cloud_skips_optimization(temp_db_path):
 
 @pytest.mark.asyncio
 async def test_local_storage_calls_optimization(temp_db_path):
-    """Test that optimization is called for local storage."""
+    """Test that vacuum calls optimization for local storage."""
     # Create a store
     store = Store(temp_db_path)
-    chunk_repo = ChunkRepository(store)
-    doc_repo = DocumentRepository(store)
-
-    # Create a document
-    document = Document(content="Test document content", metadata={})
-    created_document = await doc_repo.create(document)
-    document_id = created_document.id
 
     # Ensure LANCEDB_URI is empty (local storage)
     with patch.object(Config, "LANCEDB_URI", ""):
         # Mock the optimize method to track if it's called
         with patch.object(store.chunks_table, "optimize") as mock_optimize:
-            # Create a chunk - this should trigger optimization logic
+            # Call vacuum - this should optimize all tables for local storage
-            chunk = Chunk(
+            await store.vacuum()
-                document_id=document_id,
-                content="Test chunk content",
-                metadata={"test": "value"},
-            )
-
-            created_chunk = await chunk_repo.create(chunk)
-            assert created_chunk.id is not None
-
-            # Wait a moment to ensure async optimization completes
-            import asyncio
-
-            await asyncio.sleep(0.1)
 
             # The optimize method SHOULD have been called for local storage
             mock_optimize.assert_called()

tests/test_versioning.py

@@ -158,7 +158,7 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
 
     # Vacuum with default threshold (60 seconds) - should keep recent versions
     # Note: vacuum may create new versions even when not cleaning up old ones
-    store.vacuum()
+    await store.vacuum()
 
     after_default_doc_versions = len(list(store.documents_table.list_versions()))
     after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
@@ -173,7 +173,7 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
     )
 
     # Vacuum with 0 threshold - should significantly reduce versions
-    store.vacuum(retention_seconds=0)
+    await store.vacuum(retention_seconds=0)
 
     after_zero_doc_versions = len(list(store.documents_table.list_versions()))
     after_zero_chunk_versions = len(list(store.chunks_table.list_versions()))