deduplicate escape_sql_string into utils

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)

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}'")

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}'")

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"])

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.
 

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)