Coverage for src/qdrant_loader/core/state/checkpoint_manager.py: 71%
98 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Checkpoint manager for resumable ingestion (WS-2 feature)."""
3from dataclasses import dataclass
4from datetime import UTC, datetime
6from sqlalchemy import and_, select
7from sqlalchemy.ext.asyncio import AsyncSession
9from qdrant_loader.core.state.models import IngestionCheckpoint
10from qdrant_loader.utils.logging import LoggingConfig
12logger = LoggingConfig.get_logger(__name__)
15@dataclass
16class Checkpoint:
17 """Represents a checkpoint for resumable ingestion."""
19 project_id: str
20 source_type: str
21 source: str
22 cursor_kind: str # page_token | jql_window | git_commit | since_ts
23 cursor_value: str
24 batch_index: int = 0
25 updated_at: datetime | None = None
27 def to_dict(self) -> dict:
28 """Convert checkpoint to dictionary for serialization."""
29 return {
30 "project_id": self.project_id,
31 "source_type": self.source_type,
32 "source": self.source,
33 "cursor_kind": self.cursor_kind,
34 "cursor_value": self.cursor_value,
35 "batch_index": self.batch_index,
36 "updated_at": self.updated_at.isoformat() if self.updated_at else None,
37 }
39 @classmethod
40 def from_db(cls, db_record: IngestionCheckpoint) -> "Checkpoint":
41 """Create a Checkpoint instance from database record."""
42 return cls(
43 project_id=db_record.project_id,
44 source_type=db_record.source_type,
45 source=db_record.source,
46 cursor_kind=db_record.cursor_kind,
47 cursor_value=db_record.cursor_value,
48 batch_index=db_record.batch_index,
49 updated_at=db_record.updated_at,
50 )
53class CheckpointManager:
54 """Manages ingestion checkpoints for resumable runs."""
56 def __init__(self, session: AsyncSession):
57 """Initialize checkpoint manager.
59 Args:
60 session: SQLAlchemy async session for database operations
61 """
62 self.session = session
64 async def get_checkpoint(
65 self, project_id: str, source_type: str, source: str
66 ) -> Checkpoint | None:
67 """Get the latest checkpoint for a source.
69 Args:
70 project_id: Project ID
71 source_type: Type of source (jira, confluence, etc.)
72 source: Source name/identifier
74 Returns:
75 Checkpoint if exists, None otherwise
76 """
77 try:
78 stmt = select(IngestionCheckpoint).where(
79 and_(
80 IngestionCheckpoint.project_id == project_id,
81 IngestionCheckpoint.source_type == source_type,
82 IngestionCheckpoint.source == source,
83 )
84 )
85 result = await self.session.execute(stmt)
86 db_record = result.scalar_one_or_none()
88 if db_record:
89 checkpoint = Checkpoint.from_db(db_record)
90 logger.info(
91 "Found checkpoint for resumption",
92 project_id=project_id,
93 source_type=source_type,
94 source=source,
95 cursor_kind=checkpoint.cursor_kind,
96 batch_index=checkpoint.batch_index,
97 )
98 return checkpoint
99 else:
100 logger.debug(
101 "No checkpoint found for source",
102 project_id=project_id,
103 source_type=source_type,
104 source=source,
105 )
106 return None
107 except Exception as e:
108 await self.session.rollback()
109 logger.error(
110 "Failed to retrieve checkpoint",
111 project_id=project_id,
112 source_type=source_type,
113 source=source,
114 error=str(e),
115 error_type=type(e).__name__,
116 )
117 raise
119 async def save_checkpoint(self, checkpoint: Checkpoint) -> None:
120 """Save or update a checkpoint after successful batch upsert.
122 Args:
123 checkpoint: Checkpoint to save
125 Raises:
126 Exception: If database operation fails
127 """
128 try:
129 # Set updated_at to now if not provided
130 if checkpoint.updated_at is None:
131 checkpoint.updated_at = datetime.now(UTC)
133 # Try to fetch existing checkpoint
134 stmt = select(IngestionCheckpoint).where(
135 and_(
136 IngestionCheckpoint.project_id == checkpoint.project_id,
137 IngestionCheckpoint.source_type == checkpoint.source_type,
138 IngestionCheckpoint.source == checkpoint.source,
139 )
140 )
141 result = await self.session.execute(stmt)
142 existing = result.scalar_one_or_none()
144 if existing:
145 # Update existing checkpoint
146 existing.cursor_kind = checkpoint.cursor_kind
147 existing.cursor_value = checkpoint.cursor_value
148 existing.batch_index = checkpoint.batch_index
149 existing.updated_at = checkpoint.updated_at
150 await self.session.merge(existing)
151 logger.debug(
152 "Updated checkpoint",
153 project_id=checkpoint.project_id,
154 source_type=checkpoint.source_type,
155 source=checkpoint.source,
156 cursor_value=checkpoint.cursor_value,
157 batch_index=checkpoint.batch_index,
158 )
159 else:
160 # Create new checkpoint
161 db_checkpoint = IngestionCheckpoint(
162 project_id=checkpoint.project_id,
163 source_type=checkpoint.source_type,
164 source=checkpoint.source,
165 cursor_kind=checkpoint.cursor_kind,
166 cursor_value=checkpoint.cursor_value,
167 batch_index=checkpoint.batch_index,
168 updated_at=checkpoint.updated_at,
169 )
170 self.session.add(db_checkpoint)
171 logger.debug(
172 "Created new checkpoint",
173 project_id=checkpoint.project_id,
174 source_type=checkpoint.source_type,
175 source=checkpoint.source,
176 cursor_value=checkpoint.cursor_value,
177 batch_index=checkpoint.batch_index,
178 )
180 await self.session.commit()
181 except Exception as e:
182 await self.session.rollback()
183 logger.error(
184 "Failed to save checkpoint",
185 project_id=checkpoint.project_id,
186 source_type=checkpoint.source_type,
187 source=checkpoint.source,
188 error=str(e),
189 error_type=type(e).__name__,
190 )
191 raise
193 async def clear_checkpoint(
194 self, project_id: str, source_type: str, source: str
195 ) -> None:
196 """Clear checkpoint after successful clean run completion.
198 Args:
199 project_id: Project ID
200 source_type: Type of source (jira, confluence, etc.)
201 source: Source name/identifier
203 Raises:
204 Exception: If database operation fails
205 """
206 try:
207 stmt = select(IngestionCheckpoint).where(
208 and_(
209 IngestionCheckpoint.project_id == project_id,
210 IngestionCheckpoint.source_type == source_type,
211 IngestionCheckpoint.source == source,
212 )
213 )
214 result = await self.session.execute(stmt)
215 db_record = result.scalar_one_or_none()
217 if db_record:
218 await self.session.delete(db_record)
219 await self.session.commit()
220 logger.info(
221 "Cleared checkpoint on successful completion",
222 project_id=project_id,
223 source_type=source_type,
224 source=source,
225 )
226 else:
227 logger.debug(
228 "No checkpoint to clear",
229 project_id=project_id,
230 source_type=source_type,
231 source=source,
232 )
233 except Exception as e:
234 await self.session.rollback()
235 logger.error(
236 "Failed to clear checkpoint",
237 project_id=project_id,
238 source_type=source_type,
239 source=source,
240 error=str(e),
241 error_type=type(e).__name__,
242 )
243 raise
245 async def get_all_checkpoints(self) -> list[Checkpoint]:
246 """Get all active checkpoints (useful for monitoring/debugging).
248 Returns:
249 List of all checkpoints in the database
250 """
251 try:
252 stmt = select(IngestionCheckpoint)
253 result = await self.session.execute(stmt)
254 db_records = result.scalars().all()
255 return [Checkpoint.from_db(record) for record in db_records]
256 except Exception as e:
257 await self.session.rollback()
258 logger.error(
259 "Failed to retrieve all checkpoints",
260 error=str(e),
261 error_type=type(e).__name__,
262 )
263 raise
265 async def clear_all_checkpoints(self) -> None:
266 """Clear all checkpoints (useful for admin/reset operations).
268 Raises:
269 Exception: If database operation fails
270 """
271 try:
272 stmt = select(IngestionCheckpoint)
273 result = await self.session.execute(stmt)
274 records = result.scalars().all()
275 for record in records:
276 await self.session.delete(record)
277 await self.session.commit()
278 logger.info("Cleared all ingestion checkpoints", count=len(records))
279 except Exception as e:
280 await self.session.rollback()
281 logger.error(
282 "Failed to clear all checkpoints",
283 error=str(e),
284 error_type=type(e).__name__,
285 )
286 raise