From 1a8f72f99a64fc54ea591871e6f9a81540d3564f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 18 Dec 2025 15:55:24 +0200 Subject: [PATCH] Add is_read_only property to Store. Guard all write operations in repositories --- haiku_rag_slim/haiku/rag/store/__init__.py | 3 +- haiku_rag_slim/haiku/rag/store/engine.py | 49 +++- haiku_rag_slim/haiku/rag/store/exceptions.py | 4 + .../haiku/rag/store/repositories/chunk.py | 5 + .../haiku/rag/store/repositories/document.py | 4 + .../haiku/rag/store/repositories/settings.py | 1 + tests/store/test_read_only.py | 229 ++++++++++++++++++ 7 files changed, 287 insertions(+), 8 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/store/exceptions.py create mode 100644 tests/store/test_read_only.py diff --git a/haiku_rag_slim/haiku/rag/store/__init__.py b/haiku_rag_slim/haiku/rag/store/__init__.py index 362c9f66..df965f8f 100644 --- a/haiku_rag_slim/haiku/rag/store/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/__init__.py @@ -1,4 +1,5 @@ from .engine import Store +from .exceptions import ReadOnlyError from .models import Chunk, Document -__all__ = ["Store", "Chunk", "Document"] +__all__ = ["Store", "Chunk", "Document", "ReadOnlyError"] diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 5a7cb79d..554b86bd 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -12,6 +12,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 logger = logging.getLogger(__name__) @@ -57,9 +58,11 @@ class Store: config: AppConfig = Config, skip_validation: bool = False, create: bool = False, + read_only: bool = False, ): self.db_path: Path = db_path self._config = config + self._read_only = read_only self.embedder = get_embedder(config=self._config) self._vacuum_lock = asyncio.Lock() @@ -87,15 +90,27 @@ class Store: self._init_tables() # Run upgrades only on existing databases, set version for new ones - if is_new_db: - self._set_initial_version() - else: - self._run_upgrades() + # Skip upgrades in read-only mode (they would fail anyway) + if not read_only: + if is_new_db: + self._set_initial_version() + else: + self._run_upgrades() # Validate config compatibility after connection is established if not skip_validation: self._validate_configuration() + @property + def is_read_only(self) -> bool: + """Whether the store is in read-only mode.""" + return self._read_only + + def _assert_writable(self) -> None: + """Raise ReadOnlyError if the store is in read-only mode.""" + if self._read_only: + raise ReadOnlyError("Cannot modify database in read-only mode") + async def vacuum(self, retention_seconds: int | None = None) -> None: """Optimize and clean up old versions across all tables to reduce disk usage. @@ -106,7 +121,12 @@ class Store: Note: If vacuum is already running, this method returns immediately without blocking. Use asyncio.create_task(store.vacuum()) for non-blocking background execution. + + Raises: + ReadOnlyError: If the store is in read-only mode. """ + self._assert_writable() + if self._has_cloud_config() and str(self._config.lancedb.uri).startswith( "db://" ): @@ -317,7 +337,12 @@ class Store: return "0.0.0" def set_haiku_version(self, version: str) -> None: - """Updates the user version in settings.""" + """Updates the user version in settings. + + Raises: + ReadOnlyError: If the store is in read-only mode. + """ + self._assert_writable() settings_records = list( self.settings_table.search().limit(1).to_pydantic(SettingsRecord) ) @@ -343,7 +368,12 @@ class Store: ) def recreate_embeddings_table(self) -> None: - """Recreate the chunks table with current vector dimensions.""" + """Recreate the chunks table with current vector dimensions. + + Raises: + ReadOnlyError: If the store is in read-only mode. + """ + self._assert_writable() # Drop and recreate chunks table try: self.db.drop_table("chunks") @@ -373,7 +403,12 @@ class Store: } def restore_table_versions(self, versions: dict[str, int]) -> bool: - """Restore tables to the provided versions using LanceDB's API.""" + """Restore tables to the provided versions using LanceDB's API. + + Raises: + ReadOnlyError: If the store is in read-only mode. + """ + self._assert_writable() self.documents_table.restore(int(versions["documents"])) self.chunks_table.restore(int(versions["chunks"])) self.settings_table.restore(int(versions["settings"])) diff --git a/haiku_rag_slim/haiku/rag/store/exceptions.py b/haiku_rag_slim/haiku/rag/store/exceptions.py new file mode 100644 index 00000000..4d5c7e7c --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/exceptions.py @@ -0,0 +1,4 @@ +class ReadOnlyError(Exception): + """Raised when a write operation is attempted on a read-only store.""" + + pass diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 1577c6e0..40cdadae 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -42,6 +42,7 @@ class ChunkRepository: Chunks must have embeddings set before calling this method. Use client._ensure_chunks_embedded() to embed chunks if needed. """ + self.store._assert_writable() # Handle single chunk if isinstance(entity, Chunk): assert entity.document_id, "Chunk must have a document_id to be created" @@ -126,6 +127,7 @@ class ChunkRepository: Chunk must have embedding set before calling this method. """ + self.store._assert_writable() assert entity.id, "Chunk ID is required for update" assert entity.embedding is not None, "Chunk must have an embedding" @@ -145,6 +147,7 @@ class ChunkRepository: async def delete(self, entity_id: str) -> bool: """Delete a chunk by its ID.""" + self.store._assert_writable() chunk = await self.get_by_id(entity_id) if chunk is None: return False @@ -181,6 +184,7 @@ class ChunkRepository: async def delete_all(self) -> None: """Delete all chunks from the database.""" + self.store._assert_writable() # Drop and recreate table to clear all data self.store.db.drop_table("chunks") self.store.chunks_table = self.store.db.create_table( @@ -193,6 +197,7 @@ class ChunkRepository: async def delete_by_document_id(self, document_id: str) -> bool: """Delete all chunks for a document.""" + self.store._assert_writable() chunks = await self.get_by_document_id(document_id) if not chunks: diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index f9a2c037..3e97042f 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -47,6 +47,7 @@ class DocumentRepository: async def create(self, entity: Document) -> Document: """Create a document in the database.""" + self.store._assert_writable() # Generate new UUID doc_id = str(uuid4()) @@ -90,6 +91,7 @@ class DocumentRepository: async def update(self, entity: Document) -> Document: """Update an existing document.""" + self.store._assert_writable() from haiku.rag.store.models.document import invalidate_docling_document_cache assert entity.id, "Document ID is required for update" @@ -119,6 +121,7 @@ class DocumentRepository: async def delete(self, entity_id: str) -> bool: """Delete a document by its ID.""" + self.store._assert_writable() from haiku.rag.store.models.document import invalidate_docling_document_cache # Check if document exists @@ -181,6 +184,7 @@ class DocumentRepository: async def delete_all(self) -> None: """Delete all documents from the database.""" + self.store._assert_writable() # Delete all chunks first await self.chunk_repository.delete_all() diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index 50de331d..4e8b8dc2 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -72,6 +72,7 @@ class SettingsRepository: def save_current_settings(self) -> None: """Save the current configuration to the database.""" + self.store._assert_writable() current_config = self.store._config.model_dump(mode="json") # Check if settings exist diff --git a/tests/store/test_read_only.py b/tests/store/test_read_only.py new file mode 100644 index 00000000..26f3ec95 --- /dev/null +++ b/tests/store/test_read_only.py @@ -0,0 +1,229 @@ +import pytest + +from haiku.rag.store import ReadOnlyError, Store +from haiku.rag.store.models import Chunk, Document +from haiku.rag.store.repositories.chunk import ChunkRepository +from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.store.repositories.settings import SettingsRepository + + +class TestReadOnlyError: + def test_read_only_error_is_exception(self): + """ReadOnlyError should be a subclass of Exception.""" + assert issubclass(ReadOnlyError, Exception) + + def test_read_only_error_can_be_raised(self): + """ReadOnlyError can be raised and caught.""" + with pytest.raises(ReadOnlyError) as exc_info: + raise ReadOnlyError("Cannot modify database in read-only mode") + assert "read-only" in str(exc_info.value) + + +class TestStoreReadOnly: + def test_store_default_is_not_read_only(self, temp_db_path): + """Store defaults to not read-only.""" + store = Store(temp_db_path, create=True) + assert store.is_read_only is False + store.close() + + def test_store_can_be_created_read_only(self, temp_db_path): + """Store can be created with read_only=True.""" + # First create a normal store to initialize the database + store = Store(temp_db_path, create=True) + store.close() + + # Now open in read-only mode + store = Store(temp_db_path, read_only=True) + assert store.is_read_only is True + store.close() + + def test_assert_writable_raises_when_read_only(self, temp_db_path): + """_assert_writable() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + with pytest.raises(ReadOnlyError): + store._assert_writable() + store.close() + + def test_assert_writable_passes_when_not_read_only(self, temp_db_path): + """_assert_writable() does not raise when read_only=False.""" + store = Store(temp_db_path, create=True) + store._assert_writable() # Should not raise + store.close() + + def test_vacuum_raises_when_read_only(self, temp_db_path): + """vacuum() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + with pytest.raises(ReadOnlyError): + import asyncio + + asyncio.get_event_loop().run_until_complete(store.vacuum()) + store.close() + + def test_set_haiku_version_raises_when_read_only(self, temp_db_path): + """set_haiku_version() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + with pytest.raises(ReadOnlyError): + store.set_haiku_version("1.0.0") + store.close() + + def test_recreate_embeddings_table_raises_when_read_only(self, temp_db_path): + """recreate_embeddings_table() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + with pytest.raises(ReadOnlyError): + store.recreate_embeddings_table() + store.close() + + def test_restore_table_versions_raises_when_read_only(self, temp_db_path): + """restore_table_versions() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + versions = store.current_table_versions() + store.close() + + store = Store(temp_db_path, read_only=True) + with pytest.raises(ReadOnlyError): + store.restore_table_versions(versions) + store.close() + + +class TestDocumentRepositoryReadOnly: + @pytest.mark.asyncio + async def test_create_raises_when_read_only(self, temp_db_path): + """DocumentRepository.create() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = DocumentRepository(store) + doc = Document(content="test content") + + with pytest.raises(ReadOnlyError): + await repo.create(doc) + store.close() + + @pytest.mark.asyncio + async def test_update_raises_when_read_only(self, temp_db_path): + """DocumentRepository.update() raises ReadOnlyError when read_only=True.""" + # First create a document + store = Store(temp_db_path, create=True) + repo = DocumentRepository(store) + doc = Document(content="test content") + created_doc = await repo.create(doc) + store.close() + + # Try to update in read-only mode + store = Store(temp_db_path, read_only=True) + repo = DocumentRepository(store) + created_doc.content = "updated content" + + with pytest.raises(ReadOnlyError): + await repo.update(created_doc) + store.close() + + @pytest.mark.asyncio + async def test_delete_raises_when_read_only(self, temp_db_path): + """DocumentRepository.delete() raises ReadOnlyError when read_only=True.""" + # First create a document + store = Store(temp_db_path, create=True) + repo = DocumentRepository(store) + doc = Document(content="test content") + created_doc = await repo.create(doc) + assert created_doc.id is not None + doc_id = created_doc.id + store.close() + + # Try to delete in read-only mode + store = Store(temp_db_path, read_only=True) + repo = DocumentRepository(store) + + with pytest.raises(ReadOnlyError): + await repo.delete(doc_id) + store.close() + + @pytest.mark.asyncio + async def test_delete_all_raises_when_read_only(self, temp_db_path): + """DocumentRepository.delete_all() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = DocumentRepository(store) + + with pytest.raises(ReadOnlyError): + await repo.delete_all() + store.close() + + +class TestChunkRepositoryReadOnly: + @pytest.mark.asyncio + async def test_create_raises_when_read_only(self, temp_db_path): + """ChunkRepository.create() raises ReadOnlyError when read_only=True.""" + # First create a document to have a valid document_id + store = Store(temp_db_path, create=True) + doc_repo = DocumentRepository(store) + doc = Document(content="test content") + created_doc = await doc_repo.create(doc) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = ChunkRepository(store) + chunk = Chunk( + content="test chunk", + document_id=created_doc.id, + embedding=[0.0] * store.embedder._vector_dim, + ) + + with pytest.raises(ReadOnlyError): + await repo.create(chunk) + store.close() + + @pytest.mark.asyncio + async def test_delete_by_document_id_raises_when_read_only(self, temp_db_path): + """ChunkRepository.delete_by_document_id() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = ChunkRepository(store) + + with pytest.raises(ReadOnlyError): + await repo.delete_by_document_id("some-id") + store.close() + + @pytest.mark.asyncio + async def test_delete_all_raises_when_read_only(self, temp_db_path): + """ChunkRepository.delete_all() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = ChunkRepository(store) + + with pytest.raises(ReadOnlyError): + await repo.delete_all() + store.close() + + +class TestSettingsRepositoryReadOnly: + def test_save_current_settings_raises_when_read_only(self, temp_db_path): + """SettingsRepository.save_current_settings() raises ReadOnlyError when read_only=True.""" + store = Store(temp_db_path, create=True) + store.close() + + store = Store(temp_db_path, read_only=True) + repo = SettingsRepository(store) + + with pytest.raises(ReadOnlyError): + repo.save_current_settings() + store.close()