From 094f0e0d4484b4e261da88de18242bf0a7df9c53 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Oct 2025 12:05:25 +0300 Subject: [PATCH] Centralize vacuuming --- src/haiku/rag/client.py | 4 +- src/haiku/rag/migration.py | 6 +- src/haiku/rag/store/engine.py | 30 ++++++++-- src/haiku/rag/store/repositories/chunk.py | 28 --------- src/haiku/rag/store/repositories/document.py | 11 ++-- tests/test_lancedb_connection.py | 62 ++++---------------- tests/test_versioning.py | 4 +- 7 files changed, 51 insertions(+), 94 deletions(-) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 1c4377ed..97e6f931 100644 --- a/src/haiku/rag/client.py +++ b/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.""" diff --git a/src/haiku/rag/migration.py b/src/haiku/rag/migration.py index a0e8fce6..da2d21e6 100644 --- a/src/haiku/rag/migration.py +++ b/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() diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 5c4da84c..9c6b2e88 100644 --- a/src/haiku/rag/store/engine.py +++ b/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 - retention = timedelta(seconds=retention_seconds) - for table in [self.documents_table, self.chunks_table, self.settings_table]: - table.optimize(cleanup_older_than=retention) + # Skip if already running (non-blocking) + if self._vacuum_lock.locked(): + return + + 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).""" diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index fe3bd540..c8fac655 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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: diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index dc31878d..ef951d6c 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/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 - self.store.vacuum() + # Vacuum old versions in background (non-blocking) + 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 - self.store.vacuum() + # Vacuum old versions in background (non-blocking) + asyncio.create_task(self.store.vacuum()) + return updated_doc except Exception: # Roll back to the captured versions and re-raise diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index ae821fe9..c61eae0d 100644 --- a/tests/test_lancedb_connection.py +++ b/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 - document = Document(content="Test document content", metadata={}) - created_document = await doc_repo.create(document) - document_id = created_document.id - - # Mock LANCEDB_URI to simulate LanceDB Cloud usage - with patch.object(Config, "LANCEDB_URI", "db://test-database"): + # Mock all cloud config to simulate LanceDB Cloud usage + with ( + patch.object(Config, "LANCEDB_URI", "db://test-database"), + patch.object(Config, "LANCEDB_API_KEY", "test-api-key"), + patch.object(Config, "LANCEDB_REGION", "us-east-1"), + ): # 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 - chunk = Chunk( - 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) + # Call vacuum - this should skip optimization for LanceDB Cloud + await store.vacuum() # 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 - chunk = Chunk( - 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) + # Call vacuum - this should optimize all tables for local storage + await store.vacuum() # The optimize method SHOULD have been called for local storage mock_optimize.assert_called() diff --git a/tests/test_versioning.py b/tests/test_versioning.py index 1a35336e..869aea8e 100644 --- a/tests/test_versioning.py +++ b/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()))