Coverage for src/qdrant_loader/core/state/transitions.py: 82%
185 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1from __future__ import annotations
3import logging
4from collections.abc import Awaitable, Callable
5from datetime import UTC, datetime
6from typing import Any
8from sqlalchemy import select
10from qdrant_loader.core.document import Document
11from qdrant_loader.core.state.models import DocumentStateRecord, IngestionHistory
13AsyncSessionFactory = Callable[[], Awaitable[Any]]
14logger = logging.getLogger(__name__)
17async def update_last_ingestion(
18 session_factory: AsyncSessionFactory,
19 *,
20 source_type: str,
21 source: str,
22 status: str,
23 error_message: str | None,
24 document_count: int,
25 project_id: str | None,
26) -> None:
27 async with session_factory() as session: # type: ignore
28 now = datetime.now(UTC)
29 query = (
30 select(IngestionHistory)
31 .filter(IngestionHistory.source_type == source_type)
32 .filter(IngestionHistory.source == source)
33 )
34 if project_id is not None:
35 query = query.filter(IngestionHistory.project_id == project_id)
36 result = await session.execute(query)
37 ingestion = result.scalar_one_or_none()
38 if ingestion:
39 ingestion.last_successful_ingestion = (
40 now if status == "SUCCESS" else ingestion.last_successful_ingestion
41 ) # type: ignore
42 ingestion.status = status # type: ignore
43 ingestion.document_count = (
44 document_count if document_count else ingestion.document_count
45 ) # type: ignore
46 ingestion.updated_at = now # type: ignore
47 ingestion.error_message = error_message # type: ignore
48 else:
49 ingestion = IngestionHistory(
50 project_id=project_id,
51 source_type=source_type,
52 source=source,
53 last_successful_ingestion=now,
54 status=status,
55 document_count=document_count,
56 error_message=error_message,
57 created_at=now,
58 updated_at=now,
59 )
60 session.add(ingestion)
61 await session.commit()
64async def get_last_ingestion(
65 session_factory: AsyncSessionFactory,
66 *,
67 source_type: str,
68 source: str,
69 project_id: str | None,
70) -> IngestionHistory | None:
71 async with session_factory() as session: # type: ignore
72 query = (
73 select(IngestionHistory)
74 .filter(IngestionHistory.source_type == source_type)
75 .filter(IngestionHistory.source == source)
76 )
77 if project_id is not None:
78 query = query.filter(IngestionHistory.project_id == project_id)
79 result = await session.execute(query)
80 return result.scalar_one_or_none()
83async def mark_document_deleted(
84 session_factory: AsyncSessionFactory,
85 *,
86 source_type: str,
87 source: str,
88 document_id: str,
89 project_id: str | None,
90) -> None:
91 async with session_factory() as session: # type: ignore
92 now = datetime.now(UTC)
93 query = select(DocumentStateRecord).filter(
94 DocumentStateRecord.source_type == source_type,
95 DocumentStateRecord.source == source,
96 DocumentStateRecord.document_id == document_id,
97 )
98 if project_id is not None:
99 query = query.filter(DocumentStateRecord.project_id == project_id)
100 result = await session.execute(query)
101 state = result.scalar_one_or_none()
102 if state:
103 state.is_deleted = True # type: ignore
104 state.updated_at = now # type: ignore
105 await session.commit()
108async def get_document_state_record(
109 session_factory: AsyncSessionFactory,
110 *,
111 source_type: str,
112 source: str,
113 document_id: str,
114 project_id: str | None,
115) -> DocumentStateRecord | None:
116 async with session_factory() as session: # type: ignore
117 query = select(DocumentStateRecord).filter(
118 DocumentStateRecord.source_type == source_type,
119 DocumentStateRecord.source == source,
120 DocumentStateRecord.document_id == document_id,
121 )
122 if project_id is not None:
123 query = query.filter(DocumentStateRecord.project_id == project_id)
124 result = await session.execute(query)
125 return result.scalar_one_or_none()
128async def get_document_state_records(
129 session_factory: AsyncSessionFactory,
130 *,
131 source_type: str,
132 source: str,
133 since: datetime | None,
134) -> list[DocumentStateRecord]:
135 async with session_factory() as session: # type: ignore
136 # Select records for the given source_type and source; the previous
137 # duplicate assignment filtered on the wrong field and was overridden
138 # immediately — remove the erroneous assignment.
139 query = select(DocumentStateRecord).filter(
140 DocumentStateRecord.source_type == source_type,
141 DocumentStateRecord.source == source,
142 )
143 if since:
144 query = query.filter(DocumentStateRecord.updated_at >= since)
145 result = await session.execute(query)
146 return list(result.scalars().all())
149async def get_document_state_records_by_ids(
150 session_factory: AsyncSessionFactory,
151 *,
152 source_type: str,
153 source: str,
154 document_ids: list[str],
155 project_id: str | None = None,
156) -> list[DocumentStateRecord]:
157 """Fetch a set of DocumentStateRecord rows for a given source and list of document IDs in one query."""
158 if not document_ids:
159 return []
160 async with session_factory() as session: # type: ignore
161 query = select(DocumentStateRecord).filter(
162 DocumentStateRecord.source_type == source_type,
163 DocumentStateRecord.source == source,
164 DocumentStateRecord.document_id.in_(document_ids),
165 DocumentStateRecord.is_deleted.is_(False),
166 )
167 if project_id is not None:
168 query = query.filter(DocumentStateRecord.project_id == project_id)
169 result = await session.execute(query)
170 return list(result.scalars().all())
173async def update_document_state(
174 session_factory: AsyncSessionFactory,
175 *,
176 document: Document,
177 project_id: str | None,
178) -> DocumentStateRecord:
179 async with session_factory() as session: # type: ignore
180 document_state_record = await _apply_document_state_update(
181 session, document=document, project_id=project_id
182 )
183 await session.commit()
184 return document_state_record
187async def update_document_states_batch(
188 session_factory: AsyncSessionFactory,
189 *,
190 documents: list[Document],
191 project_id: str | None,
192) -> list[tuple[Document, DocumentStateRecord | None, Exception | None]]:
193 """Update state for multiple documents in a single session and commit.
195 Each document's read+write is wrapped in its own SAVEPOINT
196 (``session.begin_nested()``), so a failure on one document rolls back
197 only that document's changes -- the rest of the batch still lands in the
198 single, final ``commit()``. This turns N sequential fsync'd transactions
199 (the previous per-document behavior) into one, while preserving the
200 per-document error isolation callers rely on.
202 Returns a list of ``(document, record_or_None, exception_or_None)``
203 aligned with ``documents``, so callers can log per-document
204 success/failure exactly as before.
205 """
206 results: list[tuple[Document, DocumentStateRecord | None, Exception | None]] = []
207 async with session_factory() as session: # type: ignore
208 for document in documents:
209 try:
210 async with session.begin_nested():
211 record = await _apply_document_state_update(
212 session, document=document, project_id=project_id
213 )
214 results.append((document, record, None))
215 except Exception as e: # noqa: BLE001 - reported to caller, not swallowed
216 results.append((document, None, e))
218 try:
219 await session.commit()
220 except (
221 Exception
222 ) as commit_error: # noqa: BLE001 - return as per-document failures
223 try:
224 await session.rollback()
225 except Exception as rollback_error:
226 # Best-effort rollback; keep the original commit error as the reported cause.
227 logger.warning(
228 "Best-effort rollback failed after batch commit error",
229 exc_info=rollback_error,
230 )
232 # Commit failed, so no successful writes in this batch were persisted.
233 # Convert previously "successful" items into per-document failures.
234 results = [
235 (
236 doc,
237 None,
238 commit_error if error is None else error,
239 )
240 for doc, _record, error in results
241 ]
243 return results
246async def _apply_document_state_update(
247 session: Any,
248 *,
249 document: Document,
250 project_id: str | None,
251) -> DocumentStateRecord:
252 """Fetch (if present) and update/create a document's state record.
254 Operates within the caller's session/transaction and does not commit --
255 callers own the commit boundary so single- and batch-update paths can
256 share this logic while controlling transaction granularity themselves.
257 """
258 query = select(DocumentStateRecord).filter(
259 DocumentStateRecord.source_type == document.source_type,
260 DocumentStateRecord.source == document.source,
261 DocumentStateRecord.document_id == document.id,
262 )
263 if project_id is not None:
264 query = query.filter(DocumentStateRecord.project_id == project_id)
265 result = await session.execute(query)
266 document_state_record = result.scalar_one_or_none()
268 now = datetime.now(UTC)
270 metadata = document.metadata
271 conversion_method = metadata.get("conversion_method")
272 is_converted = conversion_method is not None
273 conversion_failed = metadata.get("conversion_failed", False)
275 is_attachment = metadata.get("is_attachment", False)
276 parent_document_id = metadata.get("parent_document_id")
277 attachment_id = metadata.get("attachment_id")
279 if document_state_record:
280 document_state_record.title = document.title # type: ignore
281 document_state_record.content_hash = document.content_hash # type: ignore
282 document_state_record.is_deleted = False # type: ignore
283 document_state_record.updated_at = now # type: ignore
285 document_state_record.is_converted = is_converted # type: ignore
286 document_state_record.conversion_method = conversion_method # type: ignore
287 document_state_record.original_file_type = metadata.get("original_file_type") # type: ignore
288 document_state_record.original_filename = metadata.get("original_filename") # type: ignore
289 document_state_record.file_size = metadata.get("file_size") # type: ignore
290 document_state_record.conversion_failed = conversion_failed # type: ignore
291 document_state_record.conversion_error = metadata.get("conversion_error") # type: ignore
292 document_state_record.conversion_time = metadata.get("conversion_time") # type: ignore
294 document_state_record.is_attachment = is_attachment # type: ignore
295 document_state_record.parent_document_id = parent_document_id # type: ignore
296 document_state_record.attachment_id = attachment_id # type: ignore
297 document_state_record.attachment_filename = metadata.get("attachment_filename") # type: ignore
298 document_state_record.attachment_mime_type = metadata.get("attachment_mime_type") # type: ignore
299 document_state_record.attachment_download_url = metadata.get("attachment_download_url") # type: ignore
300 document_state_record.attachment_author = metadata.get("attachment_author") # type: ignore
302 attachment_created_str = metadata.get("attachment_created_at")
303 if attachment_created_str:
304 try:
305 if isinstance(attachment_created_str, str):
306 document_state_record.attachment_created_at = (
307 datetime.fromisoformat(
308 attachment_created_str.replace("Z", "+00:00")
309 )
310 ) # type: ignore
311 elif isinstance(attachment_created_str, datetime):
312 document_state_record.attachment_created_at = attachment_created_str # type: ignore
313 except (ValueError, TypeError):
314 document_state_record.attachment_created_at = None # type: ignore
315 else:
316 attachment_created_at = None
317 attachment_created_str = metadata.get("attachment_created_at")
318 if attachment_created_str:
319 try:
320 if isinstance(attachment_created_str, str):
321 attachment_created_at = datetime.fromisoformat(
322 attachment_created_str.replace("Z", "+00:00")
323 )
324 elif isinstance(attachment_created_str, datetime):
325 attachment_created_at = attachment_created_str
326 except (ValueError, TypeError):
327 attachment_created_at = None
329 document_state_record = DocumentStateRecord(
330 project_id=project_id,
331 document_id=document.id,
332 source_type=document.source_type,
333 source=document.source,
334 url=document.url,
335 title=document.title,
336 content_hash=document.content_hash,
337 is_deleted=False,
338 created_at=now,
339 updated_at=now,
340 is_converted=is_converted,
341 conversion_method=conversion_method,
342 original_file_type=metadata.get("original_file_type"),
343 original_filename=metadata.get("original_filename"),
344 file_size=metadata.get("file_size"),
345 conversion_failed=conversion_failed,
346 conversion_error=metadata.get("conversion_error"),
347 conversion_time=metadata.get("conversion_time"),
348 is_attachment=is_attachment,
349 parent_document_id=parent_document_id,
350 attachment_id=attachment_id,
351 attachment_filename=metadata.get("attachment_filename"),
352 attachment_mime_type=metadata.get("attachment_mime_type"),
353 attachment_download_url=metadata.get("attachment_download_url"),
354 attachment_author=metadata.get("attachment_author"),
355 attachment_created_at=attachment_created_at,
356 )
357 session.add(document_state_record)
359 return document_state_record
362async def update_conversion_metrics(
363 session_factory: AsyncSessionFactory,
364 *,
365 source_type: str,
366 source: str,
367 converted_files_count: int,
368 conversion_failures_count: int,
369 attachments_processed_count: int,
370 total_conversion_time: float,
371) -> None:
372 async with session_factory() as session: # type: ignore
373 result = await session.execute(
374 select(IngestionHistory).filter_by(source_type=source_type, source=source)
375 )
376 ingestion = result.scalar_one_or_none()
377 if ingestion:
378 ingestion.converted_files_count = (
379 ingestion.converted_files_count or 0
380 ) + converted_files_count # type: ignore
381 ingestion.conversion_failures_count = (
382 ingestion.conversion_failures_count or 0
383 ) + conversion_failures_count # type: ignore
384 ingestion.attachments_processed_count = (
385 ingestion.attachments_processed_count or 0
386 ) + attachments_processed_count # type: ignore
387 ingestion.total_conversion_time = (
388 ingestion.total_conversion_time or 0.0
389 ) + total_conversion_time # type: ignore
390 ingestion.updated_at = datetime.now(UTC) # type: ignore
391 else:
392 now = datetime.now(UTC)
393 ingestion = IngestionHistory(
394 source_type=source_type,
395 source=source,
396 last_successful_ingestion=now,
397 status="SUCCESS",
398 document_count=0,
399 converted_files_count=converted_files_count,
400 conversion_failures_count=conversion_failures_count,
401 attachments_processed_count=attachments_processed_count,
402 total_conversion_time=total_conversion_time,
403 created_at=now,
404 updated_at=now,
405 )
406 session.add(ingestion)
407 await session.commit()
410async def get_conversion_metrics(
411 session_factory: AsyncSessionFactory,
412 *,
413 source_type: str,
414 source: str,
415) -> dict[str, int | float]:
416 async with session_factory() as session: # type: ignore
417 result = await session.execute(
418 select(IngestionHistory).filter_by(source_type=source_type, source=source)
419 )
420 ingestion = result.scalar_one_or_none()
421 if ingestion:
422 converted_files: int | None = ingestion.converted_files_count # type: ignore
423 conversion_failures: int | None = ingestion.conversion_failures_count # type: ignore
424 attachments_processed: int | None = ingestion.attachments_processed_count # type: ignore
425 total_time: float | None = ingestion.total_conversion_time # type: ignore
426 return {
427 "converted_files_count": (
428 converted_files if converted_files is not None else 0
429 ),
430 "conversion_failures_count": (
431 conversion_failures if conversion_failures is not None else 0
432 ),
433 "attachments_processed_count": (
434 attachments_processed if attachments_processed is not None else 0
435 ),
436 "total_conversion_time": total_time if total_time is not None else 0.0,
437 }
438 return {
439 "converted_files_count": 0,
440 "conversion_failures_count": 0,
441 "attachments_processed_count": 0,
442 "total_conversion_time": 0.0,
443 }
446async def get_attachment_documents(
447 session_factory: AsyncSessionFactory,
448 *,
449 parent_document_id: str,
450) -> list[DocumentStateRecord]:
451 async with session_factory() as session: # type: ignore
452 result = await session.execute(
453 select(DocumentStateRecord).filter(
454 DocumentStateRecord.parent_document_id == parent_document_id,
455 DocumentStateRecord.is_attachment.is_(True),
456 DocumentStateRecord.is_deleted.is_(False),
457 )
458 )
459 return list(result.scalars().all())
462async def get_converted_documents(
463 session_factory: AsyncSessionFactory,
464 *,
465 source_type: str,
466 source: str,
467 conversion_method: str | None,
468) -> list[DocumentStateRecord]:
469 async with session_factory() as session: # type: ignore
470 query = select(DocumentStateRecord).filter(
471 DocumentStateRecord.source_type == source_type,
472 DocumentStateRecord.source == source,
473 DocumentStateRecord.is_converted.is_(True),
474 DocumentStateRecord.is_deleted.is_(False),
475 )
476 if conversion_method:
477 query = query.filter(
478 DocumentStateRecord.conversion_method == conversion_method
479 )
480 result = await session.execute(query)
481 return list(result.scalars().all())