|
5 | 5 | from functools import lru_cache |
6 | 6 | from hashlib import sha256 |
7 | 7 | from logging import getLogger |
8 | | -from typing import TYPE_CHECKING, Any |
| 8 | +from typing import TYPE_CHECKING, Any, cast |
9 | 9 |
|
10 | | -from sqlalchemy import func, or_, select, update |
| 10 | +from sqlalchemy import CursorResult, func, or_, select, update |
11 | 11 | from sqlalchemy.exc import SQLAlchemyError |
12 | 12 | from sqlalchemy.orm import load_only |
13 | 13 | from typing_extensions import NotRequired, Self, override |
@@ -231,6 +231,7 @@ async def add_batch_of_requests( |
231 | 231 |
|
232 | 232 | async with self.get_session() as session: |
233 | 233 | result = await session.execute(stmt) |
| 234 | + result = cast('CursorResult', result) if not isinstance(result, CursorResult) else result |
234 | 235 | existing_requests = {req.request_id: req for req in result.scalars()} |
235 | 236 | state = await self._get_state(session) |
236 | 237 | insert_values: list[dict] = [] |
@@ -498,9 +499,12 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | |
498 | 499 | ) |
499 | 500 | async with self.get_session() as session: |
500 | 501 | result = await session.execute(stmt) |
| 502 | + result = cast('CursorResult', result) if not isinstance(result, CursorResult) else result |
| 503 | + |
501 | 504 | if result.rowcount == 0: |
502 | 505 | logger.warning(f'Request {request.unique_key} not found in database.') |
503 | 506 | return None |
| 507 | + |
504 | 508 | await self._update_metadata( |
505 | 509 | session, |
506 | 510 | **_QueueMetadataUpdateParams( |
@@ -550,6 +554,8 @@ async def reclaim_request( |
550 | 554 | stmt = stmt.values(sequence_number=new_sequence, time_blocked_until=None, client_key=None) |
551 | 555 |
|
552 | 556 | result = await session.execute(stmt) |
| 557 | + result = cast('CursorResult', result) if not isinstance(result, CursorResult) else result |
| 558 | + |
553 | 559 | if result.rowcount == 0: |
554 | 560 | logger.warning(f'Request {request.unique_key} not found in database.') |
555 | 561 | return None |
|
0 commit comments