diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index d06a7856..2ff5b509 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -22,12 +22,10 @@ from haiku.rag.store.models.chunk import Chunk, SearchResult from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.chunk import ChunkRepository -from haiku.rag.store.repositories.document import ( - DocumentRepository, - _escape_sql_string, -) +from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document_item import DocumentItemRepository from haiku.rag.store.repositories.settings import SettingsRepository +from haiku.rag.utils import escape_sql_string if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -910,7 +908,7 @@ class HaikuRAG: if doc: return doc - safe_input = _escape_sql_string(id_or_title) + safe_input = escape_sql_string(id_or_title) docs = await self.list_documents(filter=f"title = '{safe_input}'") if docs and docs[0].id: return await self.get_document_by_id(docs[0].id) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 42a24c2e..0c9de830 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -4,11 +4,7 @@ from uuid import uuid4 from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema from haiku.rag.store.models.document import Document - - -def _escape_sql_string(value: str) -> str: - """Escape single quotes in SQL string literals.""" - return value.replace("'", "''") +from haiku.rag.utils import escape_sql_string class DocumentRepository: @@ -91,7 +87,7 @@ class DocumentRepository: async def get_by_id(self, entity_id: str) -> Document | None: """Get a document by its ID.""" - safe_id = _escape_sql_string(entity_id) + safe_id = escape_sql_string(entity_id) results = list( self.store.documents_table.search() .where(f"id = '{safe_id}'") @@ -108,7 +104,7 @@ class DocumentRepository: async def get_docling_data(self, entity_id: str) -> Document | None: """Get a document with only docling data loaded (skips content blob).""" - safe_id = _escape_sql_string(entity_id) + safe_id = escape_sql_string(entity_id) results = list( self.store.documents_table.search() .select(self._DOCLING_COLUMNS) @@ -130,7 +126,7 @@ class DocumentRepository: async def get_pages_data(self, entity_id: str) -> Document | None: """Get a document with only page image data loaded.""" - safe_id = _escape_sql_string(entity_id) + safe_id = escape_sql_string(entity_id) results = list( self.store.documents_table.search() .select(["id", "docling_pages"]) @@ -160,7 +156,7 @@ class DocumentRepository: entity.updated_at = datetime.fromisoformat(now) # Update the record - safe_id = _escape_sql_string(entity.id) + safe_id = escape_sql_string(entity.id) self.store.documents_table.update( where=f"id = '{safe_id}'", values={ @@ -191,7 +187,7 @@ class DocumentRepository: await self.document_item_repository.delete_by_document_id(entity_id) # Delete the document - safe_id = _escape_sql_string(entity_id) + safe_id = escape_sql_string(entity_id) self.store.documents_table.delete(f"id = '{safe_id}'") return True @@ -261,7 +257,7 @@ class DocumentRepository: async def get_by_uri(self, uri: str) -> Document | None: """Get a document by its URI.""" - escaped_uri = _escape_sql_string(uri) + escaped_uri = escape_sql_string(uri) results = list( self.store.documents_table.search() .where(f"uri = '{escaped_uri}'") diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 972a4532..9b30a271 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -2,11 +2,7 @@ import json from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.models.document_item import DocumentItem - - -def _escape_sql_string(value: str) -> str: - """Escape single quotes in SQL string literals.""" - return value.replace("'", "''") +from haiku.rag.utils import escape_sql_string class DocumentItemRepository: @@ -48,7 +44,7 @@ class DocumentItemRepository: self, document_id: str, start: int, end: int ) -> list[DocumentItem]: """Get items for a document within a position range (inclusive).""" - safe_id = _escape_sql_string(document_id) + safe_id = escape_sql_string(document_id) rows = ( self.store.document_items_table.search() .where( @@ -66,8 +62,8 @@ class DocumentItemRepository: if not refs: return {} - safe_id = _escape_sql_string(document_id) - refs_sql = ", ".join(f"'{_escape_sql_string(r)}'" for r in refs) + safe_id = escape_sql_string(document_id) + refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) rows = ( self.store.document_items_table.search() .select(["self_ref", "position"]) @@ -78,7 +74,7 @@ class DocumentItemRepository: async def get_item_count(self, document_id: str) -> int: """Count items for a document.""" - safe_id = _escape_sql_string(document_id) + safe_id = escape_sql_string(document_id) return self.store.document_items_table.count_rows( filter=f"document_id = '{safe_id}'" ) @@ -86,5 +82,5 @@ class DocumentItemRepository: async def delete_by_document_id(self, document_id: str) -> None: """Delete all items for a document.""" self.store._assert_writable() - safe_id = _escape_sql_string(document_id) + safe_id = escape_sql_string(document_id) self.store.document_items_table.delete(f"document_id = '{safe_id}'") diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py index 33f5dcea..907069ff 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py @@ -3,6 +3,7 @@ import logging from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.upgrades import Upgrade +from haiku.rag.utils import escape_sql_string logger = logging.getLogger(__name__) @@ -31,7 +32,7 @@ def _apply_populate_document_items(store: Store) -> None: # pragma: no cover for idx, doc_id in enumerate(ids, 1): # Load only docling data - safe_id = doc_id.replace("'", "''") + safe_id = escape_sql_string(doc_id) rows = ( store.documents_table.search() .select(["id", "docling_document"]) diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index c7d554ec..90c67c14 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -452,6 +452,11 @@ def build_prompt(base_prompt: str, config: "AppConfig") -> str: return base_prompt +def escape_sql_string(value: str) -> str: + """Escape single quotes in SQL string literals.""" + return value.replace("'", "''") + + def get_package_versions() -> dict[str, str]: """Get versions of haiku.rag and its dependencies. diff --git a/tests/test_client.py b/tests/test_client.py index 40a1ee73..7c61a9b8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1502,12 +1502,12 @@ async def test_client_convert_with_html_format(temp_db_path): @pytest.mark.asyncio @pytest.mark.vcr() async def test_sql_injection_is_blocked_with_escaping(temp_db_path): - """SQL injection is blocked when using _escape_sql_string. + """SQL injection is blocked when using escape_sql_string. - This test verifies that _escape_sql_string properly prevents SQL injection + This test verifies that escape_sql_string properly prevents SQL injection by escaping single quotes in user input. """ - from haiku.rag.store.repositories.document import _escape_sql_string + from haiku.rag.utils import escape_sql_string async with HaikuRAG(temp_db_path, create=True) as client: # Create documents @@ -1529,7 +1529,7 @@ async def test_sql_injection_is_blocked_with_escaping(temp_db_path): # With proper escaping, single quotes become double quotes # so the filter becomes: title = 'x'' OR title LIKE ''%' # which searches for a literal title containing the injection string - safe_payload = _escape_sql_string(injection_payload) + safe_payload = escape_sql_string(injection_payload) docs = await client.list_documents(filter=f"title = '{safe_payload}'") # Should find 0 documents (injection is escaped, searching for literal string)