danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

deduplicate escape_sql_string into utils

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-15 14:11:11 +03:00
parent 39b9d0f683
commit e2bac1e887
No known key found for this signature in database
6 changed files with 27 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
haiku_rag_slim/haiku/rag/client.py
View file

@ -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 import Document
from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.models.document_item import extract_items
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import ( from haiku.rag.store.repositories.document import DocumentRepository
DocumentRepository,
_escape_sql_string,
)
from haiku.rag.store.repositories.document_item import DocumentItemRepository from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.store.repositories.settings import SettingsRepository
from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
@ -910,7 +908,7 @@ class HaikuRAG:
if doc: if doc:
return 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}'") docs = await self.list_documents(filter=f"title = '{safe_input}'")
if docs and docs[0].id: if docs and docs[0].id:
return await self.get_document_by_id(docs[0].id) return await self.get_document_by_id(docs[0].id)

18
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -4,11 +4,7 @@ from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
from haiku.rag.utils import escape_sql_string
def _escape_sql_string(value: str) -> str:
"""Escape single quotes in SQL string literals."""
return value.replace("'", "''")
class DocumentRepository: class DocumentRepository:
@ -91,7 +87,7 @@ class DocumentRepository:
async def get_by_id(self, entity_id: str) -> Document | None: async def get_by_id(self, entity_id: str) -> Document | None:
"""Get a document by its ID.""" """Get a document by its ID."""
safe_id = _escape_sql_string(entity_id) safe_id = escape_sql_string(entity_id)
results = list( results = list(
self.store.documents_table.search() self.store.documents_table.search()
.where(f"id = '{safe_id}'") .where(f"id = '{safe_id}'")
@ -108,7 +104,7 @@ class DocumentRepository:
async def get_docling_data(self, entity_id: str) -> Document | None: async def get_docling_data(self, entity_id: str) -> Document | None:
"""Get a document with only docling data loaded (skips content blob).""" """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( results = list(
self.store.documents_table.search() self.store.documents_table.search()
.select(self._DOCLING_COLUMNS) .select(self._DOCLING_COLUMNS)
@ -130,7 +126,7 @@ class DocumentRepository:
async def get_pages_data(self, entity_id: str) -> Document | None: async def get_pages_data(self, entity_id: str) -> Document | None:
"""Get a document with only page image data loaded.""" """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( results = list(
self.store.documents_table.search() self.store.documents_table.search()
.select(["id", "docling_pages"]) .select(["id", "docling_pages"])
@ -160,7 +156,7 @@ class DocumentRepository:
entity.updated_at = datetime.fromisoformat(now) entity.updated_at = datetime.fromisoformat(now)
# Update the record # Update the record
safe_id = _escape_sql_string(entity.id) safe_id = escape_sql_string(entity.id)
self.store.documents_table.update( self.store.documents_table.update(
where=f"id = '{safe_id}'", where=f"id = '{safe_id}'",
values={ values={
@ -191,7 +187,7 @@ class DocumentRepository:
await self.document_item_repository.delete_by_document_id(entity_id) await self.document_item_repository.delete_by_document_id(entity_id)
# Delete the document # 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}'") self.store.documents_table.delete(f"id = '{safe_id}'")
return True return True
@ -261,7 +257,7 @@ class DocumentRepository:
async def get_by_uri(self, uri: str) -> Document | None: async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI.""" """Get a document by its URI."""
escaped_uri = _escape_sql_string(uri) escaped_uri = escape_sql_string(uri)
results = list( results = list(
self.store.documents_table.search() self.store.documents_table.search()
.where(f"uri = '{escaped_uri}'") .where(f"uri = '{escaped_uri}'")

16
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

@ -2,11 +2,7 @@ import json
from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.utils import escape_sql_string
def _escape_sql_string(value: str) -> str:
"""Escape single quotes in SQL string literals."""
return value.replace("'", "''")
class DocumentItemRepository: class DocumentItemRepository:
@ -48,7 +44,7 @@ class DocumentItemRepository:
self, document_id: str, start: int, end: int self, document_id: str, start: int, end: int
) -> list[DocumentItem]: ) -> list[DocumentItem]:
"""Get items for a document within a position range (inclusive).""" """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 = ( rows = (
self.store.document_items_table.search() self.store.document_items_table.search()
.where( .where(
@ -66,8 +62,8 @@ class DocumentItemRepository:
if not refs: if not refs:
return {} return {}
safe_id = _escape_sql_string(document_id) safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{_escape_sql_string(r)}'" for r in refs) refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
rows = ( rows = (
self.store.document_items_table.search() self.store.document_items_table.search()
.select(["self_ref", "position"]) .select(["self_ref", "position"])
@ -78,7 +74,7 @@ class DocumentItemRepository:
async def get_item_count(self, document_id: str) -> int: async def get_item_count(self, document_id: str) -> int:
"""Count items for a document.""" """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( return self.store.document_items_table.count_rows(
filter=f"document_id = '{safe_id}'" filter=f"document_id = '{safe_id}'"
) )
@ -86,5 +82,5 @@ class DocumentItemRepository:
async def delete_by_document_id(self, document_id: str) -> None: async def delete_by_document_id(self, document_id: str) -> None:
"""Delete all items for a document.""" """Delete all items for a document."""
self.store._assert_writable() 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}'") self.store.document_items_table.delete(f"document_id = '{safe_id}'")

3
haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py
View file

@ -3,6 +3,7 @@ import logging
from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.upgrades import Upgrade from haiku.rag.store.upgrades import Upgrade
from haiku.rag.utils import escape_sql_string
logger = logging.getLogger(__name__) 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): for idx, doc_id in enumerate(ids, 1):
# Load only docling data # Load only docling data
safe_id = doc_id.replace("'", "''") safe_id = escape_sql_string(doc_id)
rows = ( rows = (
store.documents_table.search() store.documents_table.search()
.select(["id", "docling_document"]) .select(["id", "docling_document"])

5
haiku_rag_slim/haiku/rag/utils.py
View file

@ -452,6 +452,11 @@ def build_prompt(base_prompt: str, config: "AppConfig") -> str:
return base_prompt 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]: def get_package_versions() -> dict[str, str]:
"""Get versions of haiku.rag and its dependencies. """Get versions of haiku.rag and its dependencies.

8
tests/test_client.py
View file

@ -1502,12 +1502,12 @@ async def test_client_convert_with_html_format(temp_db_path):
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @pytest.mark.vcr()
async def test_sql_injection_is_blocked_with_escaping(temp_db_path): 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. 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: async with HaikuRAG(temp_db_path, create=True) as client:
# Create documents # 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 # With proper escaping, single quotes become double quotes
# so the filter becomes: title = 'x'' OR title LIKE ''%' # so the filter becomes: title = 'x'' OR title LIKE ''%'
# which searches for a literal title containing the injection string # 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}'") docs = await client.list_documents(filter=f"title = '{safe_payload}'")
# Should find 0 documents (injection is escaped, searching for literal string) # Should find 0 documents (injection is escaped, searching for literal string)