Coverage for src/qdrant_loader/core/worker/queue.py: 83%
167 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1from __future__ import annotations
3import asyncio
4import json
5from datetime import UTC, datetime, timedelta
6from typing import Any, Protocol
8from sqlalchemy import or_, select, update
9from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
11from qdrant_loader.core.state.models import Job
14class JobQueue(Protocol):
15 """Queue protocol to allow backend swaps without changing worker logic."""
17 async def enqueue(self, job_type: str, payload: dict[str, Any]) -> Job:
18 """Create and persist a new pending job."""
20 async def claim_next(
21 self, lease_seconds: int = 60, job_types: list[str] | None = None
22 ) -> Job | None:
23 """Atomically claim the next visible pending job.
25 Args:
26 lease_seconds: Visibility lease duration.
27 job_types: If provided, only claim jobs whose type is in this list.
28 """
30 def notify(self) -> asyncio.Event:
31 """Return an event that fires when a job becomes available for claiming.
33 The event is raised whenever:
34 - A new job is enqueued (status=PENDING).
35 - A job is released for retry (status goes back to PENDING).
37 Worker loops await this event (with timeout for visibility timeout reclaim)
38 to avoid constant polling. Long-poll backends (SQS) get this for free.
39 """
41 async def mark_done(self, job_id: int, claim_attempt: int) -> bool:
42 """Mark a claimed job as completed if claim ownership still matches."""
44 async def mark_failed(
45 self, job_id: int, error_message: str, claim_attempt: int
46 ) -> bool:
47 """Mark a claimed job as failed if claim ownership still matches."""
49 async def release_for_retry(
50 self,
51 job_id: int,
52 error_message: str,
53 claim_attempt: int,
54 retry_after_seconds: int = 0,
55 ) -> bool:
56 """Release a claimed job back to pending for a later retry."""
58 async def extend_visibility(
59 self, job_id: int, lease_seconds: int, claim_attempt: int
60 ) -> bool:
61 """Extend the visibility deadline of a RUNNING job by lease_seconds.
63 Used to prevent lease expiration during long-running handler execution.
64 Returns True if successfully extended, False if claim ownership changed.
65 """
67 async def list(
68 self, status: str | None = None, limit: int = 100, offset: int = 0
69 ) -> list[Job]:
70 """List jobs with optional status filter and pagination (offset/limit).
72 Args:
73 status: Filter by job status (e.g., 'pending', 'running'). None = all statuses.
74 limit: Max jobs per page (default 100).
75 offset: Pagination offset; skip first N results.
77 Returns:
78 List of Job objects ordered by (enqueued_at, id). May return <limit results.
79 """
81 async def reset_to_pending(self, job_id: int) -> bool:
82 """Reset a failed or done job back to pending so it can be retried."""
84 async def cancel(self, job_id: int) -> bool:
85 """Cancel a pending job (sets status to CANCELLED)."""
88class SQLiteJobQueue:
89 """SQLite-backed job queue implementation using SQLAlchemy async sessions."""
91 PENDING = "pending"
92 RUNNING = "running"
93 DONE = "done"
94 FAILED = "failed"
95 CANCELLED = "cancelled"
97 def __init__(
98 self,
99 session_factory: async_sessionmaker[AsyncSession],
100 db_op_lock: asyncio.Lock,
101 ):
102 self._session_factory = session_factory
103 self._pending_event = asyncio.Event()
104 # Serialize queue DB operations for SQLite-backed state stores.
105 # Lock is injected by the shared owner (e.g. StateManager) so
106 # multiple SQLiteJobQueue instances that share a backend also share
107 # the same DB-operation lock.
108 self._db_op_lock = db_op_lock
110 def notify(self) -> asyncio.Event:
111 """Return the event used to signal when jobs become available."""
112 return self._pending_event
114 async def enqueue(self, job_type: str, payload: dict[str, Any]) -> Job:
115 async with self._db_op_lock:
116 now = datetime.now(UTC)
117 job = Job(
118 type=job_type,
119 payload_json=json.dumps(payload, ensure_ascii=False, sort_keys=True),
120 status=self.PENDING,
121 enqueued_at=now,
122 attempts=0,
123 started_at=None,
124 finished_at=None,
125 last_error=None,
126 visibility_deadline=None,
127 )
129 async with self._session_factory() as session:
130 try:
131 session.add(job)
132 await session.commit()
133 await session.refresh(job)
134 self._pending_event.set()
135 return job
136 except Exception:
137 await session.rollback()
138 raise
140 async def claim_next(
141 self, lease_seconds: int = 60, job_types: list[str] | None = None
142 ) -> Job | None:
143 if lease_seconds < 0:
144 raise ValueError("lease_seconds must be non-negative")
146 type_filter = Job.type.in_(job_types) if job_types else None
148 async with self._db_op_lock:
149 # Compute timestamps inside the lock so they reflect the actual
150 # moment the claim executes, not the moment the caller enqueued
151 # the coroutine (which may have waited for the lock).
152 now = datetime.now(UTC)
153 visibility_deadline = now + timedelta(seconds=lease_seconds)
154 claimable_filter = or_(
155 (Job.status == self.PENDING)
156 & (
157 (Job.visibility_deadline.is_(None))
158 | (Job.visibility_deadline <= now)
159 ),
160 (Job.status == self.RUNNING) & (Job.visibility_deadline <= now),
161 )
163 candidate_query = select(Job.id).where(claimable_filter)
164 if type_filter is not None:
165 candidate_query = candidate_query.where(type_filter)
167 candidate_job_id_subquery = (
168 candidate_query.order_by(Job.enqueued_at.asc(), Job.id.asc())
169 .limit(1)
170 .scalar_subquery()
171 )
173 async with self._session_factory() as session:
174 try:
175 where_clauses = [
176 Job.id == candidate_job_id_subquery,
177 claimable_filter,
178 ]
179 if type_filter is not None:
180 where_clauses.append(type_filter)
182 result = await session.execute(
183 update(Job)
184 .where(*where_clauses)
185 .values(
186 status=self.RUNNING,
187 started_at=now,
188 finished_at=None,
189 visibility_deadline=visibility_deadline,
190 attempts=Job.attempts + 1,
191 last_error=None,
192 )
193 .returning(Job.id)
194 )
196 # SQLite requires RETURNING cursors to be finalized before commit.
197 try:
198 claimed_job_id = result.scalar_one_or_none()
199 finally:
200 result.close()
202 if claimed_job_id is None:
203 await session.commit()
204 return None
206 claimed_job = await session.get(Job, claimed_job_id)
207 await session.commit()
208 return claimed_job
209 except Exception:
210 await session.rollback()
211 raise
213 async def mark_done(self, job_id: int, claim_attempt: int) -> bool:
214 async with self._db_op_lock:
215 now = datetime.now(UTC)
216 async with self._session_factory() as session:
217 try:
218 result = await session.execute(
219 update(Job)
220 .where(
221 Job.id == job_id,
222 Job.status == self.RUNNING,
223 Job.attempts == claim_attempt,
224 )
225 .values(
226 status=self.DONE,
227 finished_at=now,
228 visibility_deadline=None,
229 last_error=None,
230 )
231 )
232 updated = result.rowcount > 0
233 await session.commit()
234 return updated
235 except Exception:
236 await session.rollback()
237 raise
239 async def mark_failed(
240 self, job_id: int, error_message: str, claim_attempt: int
241 ) -> bool:
242 async with self._db_op_lock:
243 now = datetime.now(UTC)
244 async with self._session_factory() as session:
245 try:
246 result = await session.execute(
247 update(Job)
248 .where(
249 Job.id == job_id,
250 Job.status == self.RUNNING,
251 Job.attempts == claim_attempt,
252 )
253 .values(
254 status=self.FAILED,
255 finished_at=now,
256 visibility_deadline=None,
257 last_error=error_message,
258 )
259 )
260 updated = result.rowcount > 0
261 await session.commit()
262 return updated
263 except Exception:
264 await session.rollback()
265 raise
267 async def release_for_retry(
268 self,
269 job_id: int,
270 error_message: str,
271 claim_attempt: int,
272 retry_after_seconds: int = 0,
273 ) -> bool:
274 if retry_after_seconds < 0:
275 raise ValueError("retry_after_seconds must be non-negative")
277 async with self._db_op_lock:
278 now = datetime.now(UTC)
279 retry_deadline = (
280 now + timedelta(seconds=retry_after_seconds)
281 if retry_after_seconds > 0
282 else None
283 )
285 async with self._session_factory() as session:
286 try:
287 result = await session.execute(
288 update(Job)
289 .where(
290 Job.id == job_id,
291 Job.status == self.RUNNING,
292 Job.attempts == claim_attempt,
293 )
294 .values(
295 status=self.PENDING,
296 started_at=None,
297 finished_at=None,
298 visibility_deadline=retry_deadline,
299 last_error=error_message,
300 )
301 )
302 updated = result.rowcount > 0
303 await session.commit()
304 if updated:
305 self._pending_event.set()
306 return updated
307 except Exception:
308 await session.rollback()
309 raise
311 async def extend_visibility(
312 self, job_id: int, lease_seconds: int, claim_attempt: int
313 ) -> bool:
314 """Extend the visibility deadline of a RUNNING job by lease_seconds.
316 Used to prevent lease expiration during long-running handler execution.
317 Returns True if successfully extended, False if claim ownership changed.
318 """
319 async with self._db_op_lock:
320 now = datetime.now(UTC)
321 new_deadline = now + timedelta(seconds=lease_seconds)
323 async with self._session_factory() as session:
324 try:
325 result = await session.execute(
326 update(Job)
327 .where(
328 Job.id == job_id,
329 Job.status == self.RUNNING,
330 Job.attempts == claim_attempt,
331 )
332 .values(
333 visibility_deadline=new_deadline,
334 )
335 )
336 updated = result.rowcount > 0
337 await session.commit()
338 return updated
339 except Exception:
340 await session.rollback()
341 raise
343 async def list(
344 self, status: str | None = None, limit: int = 100, offset: int = 0
345 ) -> list[Job]:
346 """List jobs with optional status filter and pagination support.
348 Args:
349 status: Filter by job status (e.g., 'pending', 'running'). None = all statuses.
350 limit: Max jobs to return per page (default 100).
351 offset: Pagination offset; skip first N results.
353 Returns:
354 List of Job objects ordered by (enqueued_at ASC, id ASC). May return <limit results.
356 Example (paginate through all pending jobs):
357 offset = 0
358 while True:
359 jobs = await queue.list(status='pending', limit=1000, offset=offset)
360 if not jobs:
361 break
362 for job in jobs:
363 process(job)
364 if len(jobs) < 1000:
365 break
366 offset += 1000
367 """
368 async with self._db_op_lock:
369 async with self._session_factory() as session:
370 try:
371 stmt = select(Job)
372 if status:
373 stmt = stmt.where(Job.status == status)
374 stmt = (
375 stmt.order_by(Job.enqueued_at.asc(), Job.id.asc())
376 .offset(offset)
377 .limit(limit)
378 )
380 result = await session.execute(stmt)
381 jobs = list(result.scalars().all())
382 await session.commit()
383 return jobs
384 except Exception:
385 await session.rollback()
386 raise
388 async def reset_to_pending(self, job_id: int) -> bool:
389 """Reset a failed/done job back to pending for retry.
391 Preserve attempts so operator retries do not erase retry history.
392 Does not reset CANCELLED jobs (operator must explicitly delete them).
393 """
394 async with self._db_op_lock:
395 async with self._session_factory() as session:
396 try:
397 result = await session.execute(
398 update(Job)
399 .where(
400 Job.id == job_id,
401 Job.status.in_([self.FAILED, self.DONE]),
402 )
403 .values(
404 status=self.PENDING,
405 started_at=None,
406 finished_at=None,
407 visibility_deadline=None,
408 last_error=None,
409 )
410 )
411 updated = result.rowcount > 0
412 await session.commit()
413 return updated
414 except Exception:
415 await session.rollback()
416 raise
418 async def cancel(self, job_id: int) -> bool:
419 """Cancel a pending job."""
420 async with self._db_op_lock:
421 now = datetime.now(UTC)
422 async with self._session_factory() as session:
423 try:
424 result = await session.execute(
425 update(Job)
426 .where(
427 Job.id == job_id,
428 Job.status == self.PENDING,
429 )
430 .values(
431 status=self.CANCELLED,
432 finished_at=now,
433 visibility_deadline=None,
434 last_error=None,
435 )
436 )
437 updated = result.rowcount > 0
438 await session.commit()
439 return updated
440 except Exception:
441 await session.rollback()
442 raise