Coverage for src/qdrant_loader/core/state/state_manager.py: 55%
242 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""
2State management service for tracking document ingestion state.
3"""
5import asyncio
6from datetime import UTC, datetime
7from typing import TYPE_CHECKING
9from sqlalchemy import func, select
11from qdrant_loader.config.source_config import SourceConfig
12from qdrant_loader.config.state import IngestionStatus, StateManagementConfig
13from qdrant_loader.core.document import Document
14from qdrant_loader.core.state import transitions as _transitions
15from qdrant_loader.core.state.models import DocumentStateRecord, IngestionHistory
16from qdrant_loader.core.state.session import create_tables as _create_tables
17from qdrant_loader.core.state.session import dispose_engine as _dispose_engine
18from qdrant_loader.core.state.session import (
19 initialize_engine_and_session as _init_engine_session,
20)
21from qdrant_loader.core.state.utils import generate_sqlite_aiosqlite_url as _gen_url
22from qdrant_loader.utils.logging import LoggingConfig
24if TYPE_CHECKING:
25 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
27logger = LoggingConfig.get_logger(__name__)
30class StateManager:
31 """Manages state for document ingestion."""
33 def __init__(self, config: StateManagementConfig):
34 """Initialize the state manager with configuration."""
35 self.config = config
36 self._initialized = False
37 self._is_sqlite_backend = False
38 self._engine: AsyncEngine | None = None
39 self._session_factory: async_sessionmaker[AsyncSession] | None = None
40 self._queue_db_op_lock = asyncio.Lock()
41 self.logger = LoggingConfig.get_logger(__name__)
43 @property
44 def is_initialized(self) -> bool:
45 """Public accessor for initialization state used by callers/tests."""
46 return self._initialized
48 @property
49 def session_factory(self) -> "async_sessionmaker[AsyncSession]":
50 """Return session factory if initialized, else raise a clear error."""
51 if self._session_factory is None:
52 raise RuntimeError("State manager session factory is not initialized")
53 return self._session_factory
55 @property
56 def queue_db_op_lock(self) -> asyncio.Lock:
57 """Shared lock for queue DB operations on this state backend."""
58 return self._queue_db_op_lock
60 async def get_session(self) -> "AsyncSession":
61 """Return an async session context manager, initializing if needed.
63 This method allows callers to use:
64 async with await state_manager.get_session() as session:
65 ...
66 """
67 if not self._initialized:
68 await self.initialize()
69 if self._session_factory is None:
70 raise RuntimeError("State manager session factory is not available")
71 return self._session_factory()
73 async def create_session(self) -> "AsyncSession":
74 """Alias for get_session for backward compatibility."""
75 return await self.get_session()
77 async def __aenter__(self):
78 """Async context manager entry."""
79 self.logger.debug("=== StateManager.__aenter__() called ===")
80 self.logger.debug(f"Current initialization state: {self._initialized}")
82 # Initialize if not already initialized
83 if not self._initialized:
84 self.logger.debug("StateManager not initialized, calling initialize()")
85 await self.initialize()
86 else:
87 self.logger.debug("StateManager already initialized")
89 return self
91 async def __aexit__(self, exc_type, exc_val, _exc_tb):
92 """Async context manager exit."""
93 await self.dispose()
95 async def initialize(self) -> None:
96 """Initialize the database and create tables if they don't exist."""
97 if self._initialized:
98 self.logger.debug("StateManager already initialized, skipping")
99 return
101 try:
102 self.logger.debug("Starting StateManager initialization")
104 # Process database path with enhanced Windows debugging
105 db_path_str = self.config.database_path
106 self.logger.debug(f"Original database path: {db_path_str}")
108 # Handle special databases and generate URL
109 database_url = _gen_url(db_path_str)
110 self.logger.debug(f"Generated database URL: {database_url}")
111 self._is_sqlite_backend = database_url.startswith("sqlite")
113 # Create database engine and session factory
114 self.logger.debug("Creating database engine and session factory")
115 self._engine, self._session_factory = _init_engine_session(self.config)
116 self.logger.debug("Engine and session factory created successfully")
118 # Create tables
119 self.logger.debug("Creating database tables")
120 await _create_tables(self._engine)
121 self.logger.debug("Database tables created successfully")
123 self._initialized = True
124 self.logger.debug("StateManager initialization completed successfully")
126 except Exception as e:
127 self.logger.error(f"StateManager initialization failed: {e}", exc_info=True)
128 # Ensure we clean up any partial initialization
129 if hasattr(self, "_engine") and self._engine:
130 try:
131 await _dispose_engine(self._engine)
132 except Exception as cleanup_error:
133 self.logger.error(
134 f"Failed to cleanup engine during error handling: {cleanup_error}"
135 )
136 self._initialized = False
137 raise
139 async def dispose(self):
140 """Clean up resources."""
141 if self._engine:
142 self.logger.debug("Disposing database engine")
143 await _dispose_engine(self._engine)
144 self._engine = None
145 self._session_factory = None
146 self._initialized = False
147 self.logger.debug("StateManager resources disposed")
149 async def update_last_ingestion(
150 self,
151 source_type: str,
152 source: str,
153 status: str = IngestionStatus.SUCCESS,
154 error_message: str | None = None,
155 document_count: int = 0,
156 project_id: str | None = None,
157 ) -> None:
158 """Update and get the last successful ingestion time for a source."""
159 self.logger.debug(
160 f"Updating last ingestion for {source_type}:{source} (project: {project_id})"
161 )
162 try:
163 await _transitions.update_last_ingestion(
164 self._session_factory, # type: ignore[arg-type]
165 source_type=source_type,
166 source=source,
167 status=status,
168 error_message=error_message,
169 document_count=document_count,
170 project_id=project_id,
171 )
172 except Exception as e:
173 self.logger.error(
174 f"Error updating last ingestion for {source_type}:{source}: {str(e)}",
175 exc_info=True,
176 )
177 raise
179 async def get_last_ingestion(
180 self, source_type: str, source: str, project_id: str | None = None
181 ) -> IngestionHistory | None:
182 """Get the last ingestion record for a source."""
183 self.logger.debug(
184 f"Getting last ingestion for {source_type}:{source} (project: {project_id})"
185 )
186 try:
187 return await _transitions.get_last_ingestion(
188 self._session_factory, # type: ignore[arg-type]
189 source_type=source_type,
190 source=source,
191 project_id=project_id,
192 )
193 except Exception as e:
194 self.logger.error(
195 f"Error getting last ingestion for {source_type}:{source}: {str(e)}",
196 exc_info=True,
197 )
198 raise
200 async def get_project_document_count(self, project_id: str) -> int:
201 """Get the count of non-deleted documents for a project.
203 Returns 0 on failure to avoid breaking CLI status output.
204 """
205 try:
206 session_factory = getattr(self, "_session_factory", None)
207 if session_factory is None:
208 ctx = await self.get_session()
209 else:
210 ctx = (
211 session_factory() if callable(session_factory) else session_factory
212 )
213 async with ctx as session: # type: ignore
214 result = await session.execute(
215 select(func.count(DocumentStateRecord.id))
216 .filter_by(project_id=project_id)
217 .filter_by(is_deleted=False)
218 )
219 count = result.scalar() or 0
220 return count
221 except Exception as e: # pragma: no cover - fallback path
222 self.logger.error(
223 f"Error getting project document count for {project_id}: {str(e)}",
224 exc_info=True,
225 )
226 return 0
228 async def get_project_latest_ingestion(self, project_id: str) -> str | None:
229 """Get the latest ingestion timestamp (ISO) for a project.
231 Returns None on failure or when no ingestion exists.
232 """
233 try:
234 session_factory = getattr(self, "_session_factory", None)
235 if session_factory is None:
236 ctx = await self.get_session()
237 else:
238 ctx = (
239 session_factory() if callable(session_factory) else session_factory
240 )
241 async with ctx as session: # type: ignore
242 result = await session.execute(
243 select(IngestionHistory.last_successful_ingestion)
244 .filter_by(project_id=project_id)
245 .order_by(IngestionHistory.last_successful_ingestion.desc())
246 .limit(1)
247 )
248 timestamp = result.scalar_one_or_none()
249 return timestamp.isoformat() if timestamp else None
250 except Exception as e: # pragma: no cover - fallback path
251 self.logger.error(
252 f"Error getting project latest ingestion for {project_id}: {str(e)}",
253 exc_info=True,
254 )
255 return None
257 async def mark_document_deleted(
258 self,
259 source_type: str,
260 source: str,
261 document_id: str,
262 project_id: str | None = None,
263 ) -> None:
264 """Mark a document as deleted."""
265 self.logger.debug(
266 f"Marking document as deleted: {source_type}:{source}:{document_id} (project: {project_id})"
267 )
268 try:
269 await _transitions.mark_document_deleted(
270 self._session_factory, # type: ignore[arg-type]
271 source_type=source_type,
272 source=source,
273 document_id=document_id,
274 project_id=project_id,
275 )
276 except Exception as e:
277 self.logger.error(
278 f"Error marking document as deleted {source_type}:{source}:{document_id}: {str(e)}",
279 exc_info=True,
280 )
281 raise
283 async def mark_documents_deleted_atomic(
284 self,
285 deleted_documents: list["Document"],
286 qdrant_manager,
287 project_id: str | None = None,
288 ) -> list[str]:
289 """Atomically mark documents as deleted in state and delete their points in Qdrant.
291 Transaction ordering:
292 1. Mark state records as is_deleted=True and COMMIT to DB immediately
293 2. Delete points from Qdrant (separate operation, outside DB transaction)
294 3. If Qdrant delete fails, vectors remain orphaned but state is correctly marked
295 (orphan vectors are recoverable by re-running delete; orphan state is silent corruption)
297 Returns the list of document IDs that were marked and deleted.
299 WS-3 DESIGN NOTE: Current design prefers orphan vectors (safe, recoverable) over orphan
300 state (dangerous, silent). If Qdrant fails to delete, the operation should be re-enqueued
301 for retry as an idempotent operation.
302 """
303 if not self._initialized:
304 raise RuntimeError("StateManager not initialized. Call initialize() first.")
306 session_factory = getattr(self, "_session_factory", None)
307 if session_factory is None:
308 raise RuntimeError("State manager session factory is not available")
310 document_ids_to_delete: list[str] = []
312 # STEP 1: Commit state changes to DB first
313 async with session_factory() as session: # type: ignore
314 tx = await session.begin()
315 try:
316 now = datetime.now(UTC)
317 for doc in deleted_documents:
318 query = select(DocumentStateRecord).filter(
319 DocumentStateRecord.source_type == doc.source_type,
320 DocumentStateRecord.source == doc.source,
321 DocumentStateRecord.document_id == doc.id,
322 )
323 if project_id is not None:
324 query = query.filter(
325 DocumentStateRecord.project_id == project_id
326 )
327 result = await session.execute(query)
328 state = result.scalar_one_or_none()
329 if state:
330 state.is_deleted = True # type: ignore
331 state.updated_at = now # type: ignore
332 document_ids_to_delete.append(doc.id)
334 # Commit state changes immediately
335 await tx.commit()
336 self.logger.info(
337 f"Marked {len(document_ids_to_delete)} documents as deleted in state DB"
338 )
339 except Exception as e:
340 # Rollback on any DB error
341 try:
342 await tx.rollback()
343 except Exception as rb_err:
344 self.logger.error(
345 f"Failed to rollback transaction after error: {rb_err}",
346 exc_info=True,
347 )
348 self.logger.error(
349 f"Failed to mark documents deleted in DB: {str(e)}",
350 exc_info=True,
351 )
352 raise
354 # STEP 2: Delete from Qdrant (separate operation, after state is committed)
355 # If this fails, vectors remain but state is correctly marked as deleted
356 if document_ids_to_delete:
357 try:
358 await qdrant_manager.delete_points_by_document_id(
359 document_ids_to_delete
360 )
361 self.logger.info(
362 f"Deleted {len(document_ids_to_delete)} documents' points from Qdrant"
363 )
364 except Exception as e:
365 # Qdrant delete failed, but state is already committed
366 # Log the failure; the operation should be re-enqueued for retry (idempotent)
367 self.logger.error(
368 f"Failed to delete points from Qdrant (state still marked as deleted): {str(e)}",
369 exc_info=True,
370 )
371 # Re-raise to signal caller that cleanup should be retried
372 raise
374 return document_ids_to_delete
376 async def get_document_state_record(
377 self,
378 source_type: str,
379 source: str,
380 document_id: str,
381 project_id: str | None = None,
382 ) -> DocumentStateRecord | None:
383 """Get the state of a document."""
384 self.logger.debug(
385 f"Getting document state for {source_type}:{source}:{document_id} (project: {project_id})"
386 )
387 try:
388 return await _transitions.get_document_state_record(
389 self._session_factory, # type: ignore[arg-type]
390 source_type=source_type,
391 source=source,
392 document_id=document_id,
393 project_id=project_id,
394 )
395 except Exception as e:
396 self.logger.error(
397 f"Error getting document state for {source_type}:{source}:{document_id}: {str(e)}",
398 exc_info=True,
399 )
400 raise
402 async def get_document_state_records_by_ids(
403 self,
404 source_type: str,
405 source: str,
406 document_ids: list[str],
407 project_id: str | None = None,
408 ) -> list[DocumentStateRecord]:
409 """Get multiple document state records for a given source in a single query."""
410 self.logger.debug(
411 f"Getting document state records for {source_type}:{source} (batch of {len(document_ids)})"
412 )
413 try:
414 return await _transitions.get_document_state_records_by_ids(
415 self._session_factory, # type: ignore[arg-type]
416 source_type=source_type,
417 source=source,
418 document_ids=document_ids,
419 project_id=project_id,
420 )
421 except Exception as e:
422 self.logger.error(
423 f"Error getting document state records by ids for {source_type}:{source}: {str(e)}",
424 exc_info=True,
425 )
426 raise
428 async def get_document_state_records(
429 self, source_config: SourceConfig, since: datetime | None = None
430 ) -> list[DocumentStateRecord]:
431 """Get all document states for a source, optionally filtered by date."""
432 self.logger.debug(
433 f"Getting document state records for {source_config.source_type}:{source_config.source}"
434 )
435 try:
436 return await _transitions.get_document_state_records(
437 self._session_factory, # type: ignore[arg-type]
438 source_type=source_config.source_type,
439 source=source_config.source,
440 since=since,
441 )
442 except Exception as e:
443 self.logger.error(
444 f"Error getting document state records for {source_config.source_type}:{source_config.source}: {str(e)}",
445 exc_info=True,
446 )
447 raise
449 async def update_document_state(
450 self, document: Document, project_id: str | None = None
451 ) -> DocumentStateRecord:
452 """Update the state of a document."""
453 if not self._initialized:
454 raise RuntimeError("StateManager not initialized. Call initialize() first.")
456 self.logger.debug(
457 f"Updating document state for {document.source_type}:{document.source}:{document.id} (project: {project_id})"
458 )
459 try:
460 if self._is_sqlite_backend:
461 async with self._queue_db_op_lock:
462 return await _transitions.update_document_state(
463 self._session_factory, # type: ignore[arg-type]
464 document=document,
465 project_id=project_id,
466 )
468 return await _transitions.update_document_state(
469 self._session_factory, # type: ignore[arg-type]
470 document=document,
471 project_id=project_id,
472 )
473 except Exception as e:
474 self.logger.error(
475 "Failed to update document state",
476 extra={
477 "project_id": project_id,
478 "document_id": document.id,
479 "error": str(e),
480 "error_type": type(e).__name__,
481 },
482 )
483 raise
485 async def update_document_states_batch(
486 self, documents: list[Document], project_id: str | None = None
487 ) -> list[tuple[Document, DocumentStateRecord | None, Exception | None]]:
488 """Update state for multiple documents in one session/commit.
490 Unlike calling ``update_document_state`` once per document, this
491 commits once for the whole batch (each document's write is isolated
492 by its own SAVEPOINT, so one failure doesn't affect the others).
493 Returns per-document ``(document, record_or_None, exception_or_None)``
494 so callers can report success/failure exactly as before.
495 """
496 if not self._initialized:
497 raise RuntimeError("StateManager not initialized. Call initialize() first.")
499 self.logger.debug(
500 f"Updating document state for {len(documents)} documents in one batch "
501 f"(project: {project_id})"
502 )
503 if self._is_sqlite_backend:
504 async with self._queue_db_op_lock:
505 return await _transitions.update_document_states_batch(
506 self._session_factory, # type: ignore[arg-type]
507 documents=documents,
508 project_id=project_id,
509 )
511 return await _transitions.update_document_states_batch(
512 self._session_factory, # type: ignore[arg-type]
513 documents=documents,
514 project_id=project_id,
515 )
517 async def update_conversion_metrics(
518 self,
519 source_type: str,
520 source: str,
521 converted_files_count: int = 0,
522 conversion_failures_count: int = 0,
523 attachments_processed_count: int = 0,
524 total_conversion_time: float = 0.0,
525 ) -> None:
526 """Update file conversion metrics for a source."""
527 self.logger.debug(f"Updating conversion metrics for {source_type}:{source}")
528 try:
529 await _transitions.update_conversion_metrics(
530 self._session_factory, # type: ignore[arg-type]
531 source_type=source_type,
532 source=source,
533 converted_files_count=converted_files_count,
534 conversion_failures_count=conversion_failures_count,
535 attachments_processed_count=attachments_processed_count,
536 total_conversion_time=total_conversion_time,
537 )
538 except Exception as e:
539 self.logger.error(
540 f"Error updating conversion metrics for {source_type}:{source}: {str(e)}",
541 exc_info=True,
542 )
543 raise
545 async def get_conversion_metrics(
546 self, source_type: str, source: str
547 ) -> dict[str, int | float]:
548 """Get file conversion metrics for a source."""
549 self.logger.debug(f"Getting conversion metrics for {source_type}:{source}")
550 try:
551 return await _transitions.get_conversion_metrics(
552 self._session_factory, # type: ignore[arg-type]
553 source_type=source_type,
554 source=source,
555 )
556 except Exception as e:
557 self.logger.error(
558 f"Error getting conversion metrics for {source_type}:{source}: {str(e)}",
559 exc_info=True,
560 )
561 raise
563 async def get_attachment_documents(
564 self, parent_document_id: str
565 ) -> list[DocumentStateRecord]:
566 """Get all attachment documents for a parent document."""
567 self.logger.debug(
568 f"Getting attachment documents for parent {parent_document_id}"
569 )
570 try:
571 return await _transitions.get_attachment_documents(
572 self._session_factory, # type: ignore[arg-type]
573 parent_document_id=parent_document_id,
574 )
575 except Exception as e:
576 self.logger.error(
577 f"Error getting attachment documents for {parent_document_id}: {str(e)}",
578 exc_info=True,
579 )
580 raise
582 async def get_converted_documents(
583 self, source_type: str, source: str, conversion_method: str | None = None
584 ) -> list[DocumentStateRecord]:
585 """Get all converted documents for a source, optionally filtered by conversion method."""
586 self.logger.debug(f"Getting converted documents for {source_type}:{source}")
587 try:
588 return await _transitions.get_converted_documents(
589 self._session_factory, # type: ignore[arg-type]
590 source_type=source_type,
591 source=source,
592 conversion_method=conversion_method,
593 )
594 except Exception as e:
595 self.logger.error(
596 f"Error getting converted documents for {source_type}:{source}: {str(e)}",
597 exc_info=True,
598 )
599 raise
601 async def close(self):
602 """Close all database connections."""
603 if hasattr(self, "_engine") and self._engine is not None:
604 self.logger.debug("Closing database connections")
605 await _dispose_engine(self._engine)
606 self.logger.debug("Database connections closed")