Add is_read_only property to Store. Guard all write operations in repositories

This commit is contained in:
Yiorgis Gozadinos 2025-12-18 15:55:24 +02:00
parent 1dbbaf6dd8
commit 1a8f72f99a
No known key found for this signature in database
7 changed files with 287 additions and 8 deletions

View file

@ -1,4 +1,5 @@
from .engine import Store from .engine import Store
from .exceptions import ReadOnlyError
from .models import Chunk, Document from .models import Chunk, Document
__all__ = ["Store", "Chunk", "Document"] __all__ = ["Store", "Chunk", "Document", "ReadOnlyError"]

View file

@ -12,6 +12,7 @@ from pydantic import Field
from haiku.rag.config import AppConfig, Config from haiku.rag.config import AppConfig, Config
from haiku.rag.embeddings import get_embedder from haiku.rag.embeddings import get_embedder
from haiku.rag.store.exceptions import ReadOnlyError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -57,9 +58,11 @@ class Store:
config: AppConfig = Config, config: AppConfig = Config,
skip_validation: bool = False, skip_validation: bool = False,
create: bool = False, create: bool = False,
read_only: bool = False,
): ):
self.db_path: Path = db_path self.db_path: Path = db_path
self._config = config self._config = config
self._read_only = read_only
self.embedder = get_embedder(config=self._config) self.embedder = get_embedder(config=self._config)
self._vacuum_lock = asyncio.Lock() self._vacuum_lock = asyncio.Lock()
@ -87,15 +90,27 @@ class Store:
self._init_tables() self._init_tables()
# Run upgrades only on existing databases, set version for new ones # Run upgrades only on existing databases, set version for new ones
if is_new_db: # Skip upgrades in read-only mode (they would fail anyway)
self._set_initial_version() if not read_only:
else: if is_new_db:
self._run_upgrades() self._set_initial_version()
else:
self._run_upgrades()
# Validate config compatibility after connection is established # Validate config compatibility after connection is established
if not skip_validation: if not skip_validation:
self._validate_configuration() 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: async def vacuum(self, retention_seconds: int | None = None) -> None:
"""Optimize and clean up old versions across all tables to reduce disk usage. """Optimize and clean up old versions across all tables to reduce disk usage.
@ -106,7 +121,12 @@ class Store:
Note: Note:
If vacuum is already running, this method returns immediately without blocking. If vacuum is already running, this method returns immediately without blocking.
Use asyncio.create_task(store.vacuum()) for non-blocking background execution. 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( if self._has_cloud_config() and str(self._config.lancedb.uri).startswith(
"db://" "db://"
): ):
@ -317,7 +337,12 @@ class Store:
return "0.0.0" return "0.0.0"
def set_haiku_version(self, version: str) -> None: 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( settings_records = list(
self.settings_table.search().limit(1).to_pydantic(SettingsRecord) self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
) )
@ -343,7 +368,12 @@ class Store:
) )
def recreate_embeddings_table(self) -> None: 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 # Drop and recreate chunks table
try: try:
self.db.drop_table("chunks") self.db.drop_table("chunks")
@ -373,7 +403,12 @@ class Store:
} }
def restore_table_versions(self, versions: dict[str, int]) -> bool: 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.documents_table.restore(int(versions["documents"]))
self.chunks_table.restore(int(versions["chunks"])) self.chunks_table.restore(int(versions["chunks"]))
self.settings_table.restore(int(versions["settings"])) self.settings_table.restore(int(versions["settings"]))

View file

@ -0,0 +1,4 @@
class ReadOnlyError(Exception):
"""Raised when a write operation is attempted on a read-only store."""
pass

View file

@ -42,6 +42,7 @@ class ChunkRepository:
Chunks must have embeddings set before calling this method. Chunks must have embeddings set before calling this method.
Use client._ensure_chunks_embedded() to embed chunks if needed. Use client._ensure_chunks_embedded() to embed chunks if needed.
""" """
self.store._assert_writable()
# Handle single chunk # Handle single chunk
if isinstance(entity, Chunk): if isinstance(entity, Chunk):
assert entity.document_id, "Chunk must have a document_id to be created" 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. Chunk must have embedding set before calling this method.
""" """
self.store._assert_writable()
assert entity.id, "Chunk ID is required for update" assert entity.id, "Chunk ID is required for update"
assert entity.embedding is not None, "Chunk must have an embedding" 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: async def delete(self, entity_id: str) -> bool:
"""Delete a chunk by its ID.""" """Delete a chunk by its ID."""
self.store._assert_writable()
chunk = await self.get_by_id(entity_id) chunk = await self.get_by_id(entity_id)
if chunk is None: if chunk is None:
return False return False
@ -181,6 +184,7 @@ class ChunkRepository:
async def delete_all(self) -> None: async def delete_all(self) -> None:
"""Delete all chunks from the database.""" """Delete all chunks from the database."""
self.store._assert_writable()
# Drop and recreate table to clear all data # Drop and recreate table to clear all data
self.store.db.drop_table("chunks") self.store.db.drop_table("chunks")
self.store.chunks_table = self.store.db.create_table( 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: async def delete_by_document_id(self, document_id: str) -> bool:
"""Delete all chunks for a document.""" """Delete all chunks for a document."""
self.store._assert_writable()
chunks = await self.get_by_document_id(document_id) chunks = await self.get_by_document_id(document_id)
if not chunks: if not chunks:

View file

@ -47,6 +47,7 @@ class DocumentRepository:
async def create(self, entity: Document) -> Document: async def create(self, entity: Document) -> Document:
"""Create a document in the database.""" """Create a document in the database."""
self.store._assert_writable()
# Generate new UUID # Generate new UUID
doc_id = str(uuid4()) doc_id = str(uuid4())
@ -90,6 +91,7 @@ class DocumentRepository:
async def update(self, entity: Document) -> Document: async def update(self, entity: Document) -> Document:
"""Update an existing document.""" """Update an existing document."""
self.store._assert_writable()
from haiku.rag.store.models.document import invalidate_docling_document_cache from haiku.rag.store.models.document import invalidate_docling_document_cache
assert entity.id, "Document ID is required for update" assert entity.id, "Document ID is required for update"
@ -119,6 +121,7 @@ class DocumentRepository:
async def delete(self, entity_id: str) -> bool: async def delete(self, entity_id: str) -> bool:
"""Delete a document by its ID.""" """Delete a document by its ID."""
self.store._assert_writable()
from haiku.rag.store.models.document import invalidate_docling_document_cache from haiku.rag.store.models.document import invalidate_docling_document_cache
# Check if document exists # Check if document exists
@ -181,6 +184,7 @@ class DocumentRepository:
async def delete_all(self) -> None: async def delete_all(self) -> None:
"""Delete all documents from the database.""" """Delete all documents from the database."""
self.store._assert_writable()
# Delete all chunks first # Delete all chunks first
await self.chunk_repository.delete_all() await self.chunk_repository.delete_all()

View file

@ -72,6 +72,7 @@ class SettingsRepository:
def save_current_settings(self) -> None: def save_current_settings(self) -> None:
"""Save the current configuration to the database.""" """Save the current configuration to the database."""
self.store._assert_writable()
current_config = self.store._config.model_dump(mode="json") current_config = self.store._config.model_dump(mode="json")
# Check if settings exist # Check if settings exist

View file

@ -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()