kiba-core 0.5.3.dev65__tar.gz → 0.5.3.dev67__tar.gz

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/.github/workflows/pull-request.yml

@@ -45,16 +45,27 @@ jobs:
           json-file-path: ./type-check-results.json
           check-name: core-type-check
           fail-on-error: false
-      - name: Start postgres
+      - name: Start test services
         run: |
           docker network create tests
           docker run --detach --name postgres --network tests --env POSTGRES_PASSWORD=test --env POSTGRES_DB=core postgres:16-alpine
+          docker run --detach --name cosmos --network tests mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator@sha256:2db1f9e74c506bcf6fc347aa937aea1c00fa756061296a5a9efba530ce86ec02 --protocol http
+          docker run --detach --name elasticmq --network tests --publish 9324:9324 softwaremill/elasticmq-native@sha256:e4580ab9ad1bd5cd37b4ba04911bc5ccc8cd2d9ab4de56ece65acee71c24e05c
+          docker run --detach --name azurite --network tests --publish 10001:10001 mcr.microsoft.com/azure-storage/azurite@sha256:830430c1da1a2d537e08f3e6764dd1f5ae00cf0346bcaf625b968ec3f0971fd5 azurite-queue --queueHost 0.0.0.0 --skipApiVersionCheck
+          echo "Waiting for postgres..."
           # NOTE(krishan711): check over TCP: during initdb the image runs a temporary server on the unix socket only, which reports ready and then shuts down
           for i in $(seq 1 60); do docker exec postgres pg_isready --host 127.0.0.1 --username postgres --quiet && break; sleep 1; done
           docker exec postgres pg_isready --host 127.0.0.1 --username postgres
+          echo "Waiting for sqs and aqs emulators..."
+          # NOTE(krishan711): any http response (even an error) means the emulator is accepting requests
+          for i in $(seq 1 60); do curl --silent --output /dev/null http://localhost:9324/ && curl --silent --output /dev/null http://localhost:10001/ && break; sleep 1; done
+          curl --silent --output /dev/null http://localhost:9324/ && curl --silent --output /dev/null http://localhost:10001/
+          echo "Waiting for cosmos emulator..."
+          for i in $(seq 1 120); do docker exec cosmos curl --silent --fail --output /dev/null http://localhost:8080/ready && break; sleep 1; done
+          docker exec cosmos curl --silent --fail --output /dev/null http://localhost:8080/ready
       - name: Run tests
         run: |
-          docker run --name tests --network tests --env CORE_TEST_PSQL_CONNECTION_STRING=postgresql+asyncpg://postgres:test@postgres:5432/core $DOCKER_IMAGE make test-ci
+          docker run --name tests --network tests --env CORE_TEST_PSQL_CONNECTION_STRING=postgresql+asyncpg://postgres:test@postgres:5432/core --env CORE_TEST_COSMOS_ENDPOINT=http://cosmos:8081 --env CORE_TEST_SQS_ENDPOINT=http://elasticmq:9324 --env CORE_TEST_AQS_ENDPOINT=http://azurite:10001/devstoreaccount1 $DOCKER_IMAGE make test-ci
           docker cp tests:/app/test-check-results.json ./test-check-results.json
       - name: Upload test results
         uses: kibalabs/github-action-create-annotations@main

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/CHANGELOG.md

@@ -8,6 +8,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 
 ### Added
+- [MINOR] Added `authorize_bearer_jwt_request` for validating bearer JWTs in `RouteAuthResolver` implementations; it and `authorize_bearer_jwt` now only strip the leading `Bearer ` prefix, let cancellation propagate, and log why a JWT was rejected
+- [MINOR] Made the `route` `authResolver` optional for routes without `auth`; setting `auth` without an `authResolver` raises `ValueError` at registration
+- [MAJOR] Replaced `MessageQueue.delete_message` with `complete_message`, and added abstract `retry_message` (reschedule with a delay) and `fail_message` (apply the queue's failure policy); `MessageQueueProcessor` now reschedules `MessageNeedsReprocessingException` messages with `retry_message`
+- [MINOR] Updated `MessageQueueProcessor` to retry messages that raise `LockedException` up to 3 times (30s linear backoff) before failing them
+- [MINOR] Updated `MessageQueueProcessor` to log instead of raise when a failure notification cannot be sent
+- [MAJOR] Reworked `SqlMessageQueue` to keep message history in `tbl_queue_messages` with a `status` lifecycle (`pending`, `running`, `succeeded`, `failed`, `deduplicated`), `owner`, `lastError`, `startedDate` and `completedDate`; deduplication now only applies against pending messages and keeps the earliest `visibleDate` of the duplicates, every retry, failure retry and lease-expiry reclaim increments `postCount`, failures back off and move to `failed` once `postCount` reaches `maxAttempts`, and `get_message_count` removes finished messages older than `retentionSeconds` (default 7 days) unless `shouldSkipRemovingNonRetained` is set
+- [MINOR] Added `Database.now` to build database-clock timestamps; `SqlMessageQueue` and `SqlLock` now use the database clock instead of the host clock
+- [MINOR] Capped `SqsMessageQueue` delays at SQS's 900 second maximum
+- [MINOR] Updated `CosmosMessageQueue.retry_message` to reschedule the message in place and added emulator-backed tests (run when `CORE_TEST_COSMOS_ENDPOINT` is set)
+- [MINOR] Added `shouldRaiseOnUncommittedWrites` to `Database` to raise instead of log when an isolated context connection is opened after uncommitted writes (for tests and development)
+- [MAJOR] Added abstract `MessageQueue.extend_message_lease` (implemented for SQL, SQS, AQS and Cosmos queues, each returning `False` once the lease is lost; SQS accepts any receipt handle for a message so it can only detect that the message was deleted or is no longer in flight); `MessageQueueProcessor` now keeps a message's lease alive while it is being processed and, if the lease was lost, raises `MessageLeaseLostException` once the job returns (taking priority over any exception from the job, while cancellation still propagates) so the message is neither completed nor retried; `SqsMessageQueue` and `AqsMessageQueue` `retry_message` check the lease is still held before re-sending so a lost lease cannot duplicate the message
+- [MINOR] Added optional `endpointUrl` to `SqsMessageQueue` and `AqsMessageQueue` to target local emulators, and added emulator-backed tests (run when `CORE_TEST_SQS_ENDPOINT` / `CORE_TEST_AQS_ENDPOINT` are set)
+- [MINOR] Added `Lock.ensure_held` to renew a lease and raise `LockedException` if it has been lost, for use before irreversible work; `LockLease` now records its `ttlSeconds`
 - [MINOR] Added `Lock` with a lease-based `SqlLock` implementation (`tbl_locks`) and a `with_lock` context manager that keeps the lease alive; acquiring raises `LockedException` (naming the current holder) if not acquired within `maxWaitSeconds`, and `get_owner` reports who holds a lock
 - [MINOR] Added `process_util.get_process_name` to identify the current process (`<containerName>:<containerId>` in docker, `<hostname>:<pid>` otherwise)
 - [MINOR] Added `Database.create_isolated_context_connection` to run a block in its own committed transaction inside an existing context connection; logs `ISOLATED_CONNECTION_AFTER_UNCOMMITTED_WRITES` when opened while the outer context connection has uncommitted writes

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev65
+Version: 0.5.3.dev67
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/api/authorizer.py

@@ -10,6 +10,7 @@ from core.api.api_request import KibaApiRequest
 from core.api.route_metadata import OpenApiSecurityScheme
 from core.api.route_metadata import update_route_metadata
 from core.exceptions import ForbiddenException
+from core.exceptions import KibaException
 from core.exceptions import UnauthorizedException
 from core.http.basic_authentication import BasicAuthentication
 from core.http.jwt import Jwt
@@ -27,18 +28,19 @@ class SignatureAuthorizer:
         raise NotImplementedError
 
 
-async def _authorize_bearer_jwt[ApiRequest: BaseModel](request: KibaApiRequest[ApiRequest], authorizer: Authorizer) -> Jwt:
+async def authorize_bearer_jwt_request[ApiRequest: BaseModel](request: KibaApiRequest[ApiRequest], authorizer: Authorizer) -> Jwt:
     authorization = request.headers.get('Authorization')
     if not authorization:
         raise ForbiddenException(message='AUTH_NOT_PROVIDED')
     if not authorization.startswith('Bearer '):
         raise ForbiddenException(message='AUTH_INVALID')
-    jwtString = authorization.replace('Bearer ', '')
     try:
-        jwt = await authorizer.validate_jwt(jwtString=jwtString)
-    except BaseException:  # noqa: BLE001
-        raise ForbiddenException(message='AUTH_INVALID')
-    return jwt
+        return await authorizer.validate_jwt(jwtString=authorization.removeprefix('Bearer '))
+    except Exception as exception:
+        # NOTE(krishan711): only log kiba messages as other exception messages (e.g. json decode errors) can contain parts of the token
+        reason = exception.message if isinstance(exception, KibaException) else type(exception).__name__
+        logging.info(f'Bearer JWT rejected: {reason}')
+        raise ForbiddenException(message='AUTH_INVALID') from exception
 
 
 def authorize_bearer_jwt[ApiRequest: BaseModel](  # type: ignore[explicit-any]
@@ -52,7 +54,7 @@ def authorize_bearer_jwt[ApiRequest: BaseModel](  # type: ignore[explicit-any]
 
         @functools.wraps(func)
         async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
-            request.authJwt = await _authorize_bearer_jwt(request=request, authorizer=authorizer)
+            request.authJwt = await authorize_bearer_jwt_request(request=request, authorizer=authorizer)
             result = func(request)
             # NOTE(krishan711): this is here to support streaming responses which return an async generator
             if hasattr(result, '__aiter__'):

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/api/route.py

@@ -42,7 +42,7 @@ def route[ApiRequest: BaseModel, ApiResponse: BaseModel](  # type: ignore[explic
     requestType: typing.Type[ApiRequest],
     responseType: typing.Type[ApiResponse],
     *,
-    authResolver: RouteAuthResolver,
+    authResolver: RouteAuthResolver | None = None,
     isStreaming: bool = False,
     operationId: str | None = None,
     summary: str | None = None,
@@ -52,13 +52,15 @@ def route[ApiRequest: BaseModel, ApiResponse: BaseModel](  # type: ignore[explic
     rateLimit: RateLimitConfig | None = None,
     extensions: Mapping[str, object] | None = None,
 ) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]], typing.Callable[[Request], typing.Awaitable[KibaJSONResponse | StreamingResponse]]]:
-    securitySchemeNames = authResolver.get_route_security_schemes(auth=auth) if auth is not None else []
-
     def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]) -> typing.Callable[[Request], typing.Awaitable[KibaJSONResponse | StreamingResponse]]:  # type: ignore[explicit-any]
         handler: typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn] = func
         if rateLimit is not None:
             handler = rate_limit(rateLimit)(handler)
+        securitySchemeNames: list[str] = []
         if auth is not None:
+            if authResolver is None:
+                raise ValueError(f'Auth policy {auth} requires an authResolver')
+            securitySchemeNames = authResolver.get_route_security_schemes(auth=auth)
             handler = _authorize_route(authResolver, auth)(handler)
         endpointDecorator = streaming_json_route if isStreaming else json_route
         endpoint = endpointDecorator(requestType=requestType, responseType=responseType)(handler)  # type: ignore[arg-type, ty:invalid-argument-type]

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/locks/lock.py

@@ -5,7 +5,8 @@ from abc import ABC
 from collections.abc import AsyncIterator
 
 from core import logging
-from core.locks.model import Lease
+from core.exceptions import LockedException
+from core.locks.model import LockLease
 
 
 class Lock(ABC):
@@ -18,26 +19,31 @@ class Lock(ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> Lease:
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> LockLease:
         raise NotImplementedError
 
     @abc.abstractmethod
-    async def extend(self, lease: Lease, ttlSeconds: int = 60) -> bool:
+    async def extend(self, lease: LockLease, ttlSeconds: int = 60) -> bool:
         raise NotImplementedError
 
     @abc.abstractmethod
-    async def release(self, lease: Lease) -> None:
+    async def release(self, lease: LockLease) -> None:
         raise NotImplementedError
 
     @abc.abstractmethod
     async def get_owner(self, name: str) -> str | None:
         raise NotImplementedError
 
-    async def _keep_alive(self, lease: Lease, ttlSeconds: int) -> None:
+    async def ensure_held(self, lease: LockLease) -> None:
+        if lease.isLost or not await self.extend(lease=lease, ttlSeconds=lease.ttlSeconds):
+            lease.isLost = True
+            raise LockedException(message=f'LOCK_LOST: {lease.name}')
+
+    async def _keep_alive(self, lease: LockLease) -> None:
         while True:
-            await asyncio.sleep(ttlSeconds / 3)
+            await asyncio.sleep(lease.ttlSeconds / 3)
             try:
-                isExtended = await self.extend(lease=lease, ttlSeconds=ttlSeconds)
+                isExtended = await self.extend(lease=lease, ttlSeconds=lease.ttlSeconds)
             except Exception as exception:  # noqa: BLE001
                 # NOTE(krishan711): a transient failure must not kill the keep-alive; the lease simply lapses at expiryDate if it keeps failing
                 logging.error(f'Failed to extend lease {lease.name}:')
@@ -49,9 +55,9 @@ class Lock(ABC):
                 return
 
     @contextlib.asynccontextmanager
-    async def with_lock(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> AsyncIterator[Lease]:
+    async def with_lock(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> AsyncIterator[LockLease]:
         lease = await self.acquire(name=name, ttlSeconds=ttlSeconds, maxWaitSeconds=maxWaitSeconds)
-        keepAliveTask = asyncio.create_task(self._keep_alive(lease=lease, ttlSeconds=ttlSeconds))
+        keepAliveTask = asyncio.create_task(self._keep_alive(lease=lease))
         try:
             yield lease
         finally:

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/locks/model.py

@@ -5,10 +5,11 @@ from pydantic import BaseModel
 from core.util import date_util
 
 
-class Lease(BaseModel):
+class LockLease(BaseModel):
     name: str
     token: str
     expiryDate: datetime.datetime
+    ttlSeconds: int
     isLost: bool = False
 
     def is_valid(self) -> bool:

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/locks/sql.py

@@ -9,7 +9,7 @@ from sqlalchemy.dialects import sqlite as sqlalchemy_sqlite
 from core.exceptions import InternalServerErrorException
 from core.exceptions import LockedException
 from core.locks.lock import Lock
-from core.locks.model import Lease
+from core.locks.model import LockLease
 from core.store.database import Database
 from core.util import date_util
 from core.util import process_util
@@ -43,7 +43,7 @@ class SqlLock(Lock):
     async def disconnect(self) -> None:
         pass
 
-    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> Lease:
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> LockLease:
         deadline = time.monotonic() + maxWaitSeconds
         while True:
             lease = await self._try_acquire(name=name, ttlSeconds=ttlSeconds)
@@ -56,7 +56,7 @@ class SqlLock(Lock):
 
     async def get_owner(self, name: str) -> str | None:
         async with self.database.create_transaction() as connection:
-            query = sqlalchemy.select(self.table.c.owner).where(self.table.c.name == name).where(self.table.c.expiryDate > date_util.datetime_from_now())
+            query = sqlalchemy.select(self.table.c.owner).where(self.table.c.name == name).where(self.table.c.expiryDate > self.database.now())
             result = await self.database.execute(query=query, connection=connection)
             return result.scalar_one_or_none()
 
@@ -67,30 +67,38 @@ class SqlLock(Lock):
             return sqlalchemy_sqlite.insert(self.table)
         raise InternalServerErrorException(message=f'SqlLock does not support dialect: {dialectName}')
 
-    async def _try_acquire(self, name: str, ttlSeconds: int) -> Lease | None:
-        now = date_util.datetime_from_now()
-        lease = Lease(name=name, token=str(uuid.uuid4()), expiryDate=date_util.datetime_from_now(seconds=ttlSeconds))
-        values = {self.table.c.lockToken: lease.token, self.table.c.acquiredDate: now, self.table.c.expiryDate: lease.expiryDate, self.table.c.owner: self.owner}
+    async def _try_acquire(self, name: str, ttlSeconds: int) -> LockLease | None:
+        lease = LockLease(name=name, token=str(uuid.uuid4()), expiryDate=date_util.datetime_from_now(seconds=ttlSeconds), ttlSeconds=ttlSeconds)
+        values = {self.table.c.lockToken: lease.token, self.table.c.acquiredDate: self.database.now(), self.table.c.expiryDate: self.database.now(seconds=ttlSeconds), self.table.c.owner: self.owner}
         async with self.database.create_transaction() as connection:
             upsertQuery = (
-                self._create_insert(dialectName=connection.dialect.name).values({self.table.c.name: name, **values}).on_conflict_do_update(index_elements=[self.table.c.name], set_=values, where=self.table.c.expiryDate <= now).returning(self.table.c.name)
+                self._create_insert(dialectName=connection.dialect.name)
+                .values({self.table.c.name: name, **values})
+                .on_conflict_do_update(index_elements=[self.table.c.name], set_=values, where=self.table.c.expiryDate <= self.database.now())
+                .returning(self.table.c.name)
             )
             result = await self.database.execute(query=upsertQuery, connection=connection)
             isAcquired = result.first() is not None
         return lease if isAcquired else None
 
-    async def extend(self, lease: Lease, ttlSeconds: int = 60) -> bool:
-        now = date_util.datetime_from_now()
+    async def extend(self, lease: LockLease, ttlSeconds: int = 60) -> bool:
         newExpiryDate = date_util.datetime_from_now(seconds=ttlSeconds)
         async with self.database.create_transaction() as connection:
-            updateQuery = sqlalchemy.update(self.table).where(self.table.c.name == lease.name).where(self.table.c.lockToken == lease.token).where(self.table.c.expiryDate > now).values(expiryDate=newExpiryDate).returning(self.table.c.name)
+            updateQuery = (
+                sqlalchemy.update(self.table)
+                .where(self.table.c.name == lease.name)
+                .where(self.table.c.lockToken == lease.token)
+                .where(self.table.c.expiryDate > self.database.now())
+                .values(expiryDate=self.database.now(seconds=ttlSeconds))
+                .returning(self.table.c.name)
+            )
             result = await self.database.execute(query=updateQuery, connection=connection)
             isExtended = result.first() is not None
         if isExtended:
             lease.expiryDate = newExpiryDate
         return isExtended
 
-    async def release(self, lease: Lease) -> None:
+    async def release(self, lease: LockLease) -> None:
         async with self.database.create_transaction() as connection:
             deleteQuery = sqlalchemy.delete(self.table).where(self.table.c.name == lease.name).where(self.table.c.lockToken == lease.token)
             await self.database.execute(query=deleteQuery, connection=connection)

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/queues/aqs.py

@@ -3,7 +3,9 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Sequence
 
+from azure.core.exceptions import HttpResponseError
 from azure.storage.queue import QueueMessage as RawAqsMessage
+from azure.storage.queue import StorageErrorCode
 from azure.storage.queue.aio import QueueClient
 
 from core.exceptions import InternalServerErrorException
@@ -11,6 +13,8 @@ from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
 from core.util import list_util
 
+AQS_RETRY_HOLD_SECONDS = 60
+
 
 class AqsMessage(Message):
     aqsId: str
@@ -32,14 +36,18 @@ class AqsMessage(Message):
 
 
 class AqsMessageQueue(MessageQueue[AqsMessage]):
-    def __init__(self, storageAccountName: str, storageAccountKey: str, queueName: str) -> None:
+    def __init__(self, storageAccountName: str, storageAccountKey: str, queueName: str, endpointUrl: str | None = None) -> None:
         self._storageAccountName = storageAccountName
         self._storageAccountKey = storageAccountKey
         self.queueName = queueName
+        self.endpointUrl = endpointUrl
         self._aqsClient: QueueClient | None = None
 
     async def connect(self) -> None:
-        self._aqsClient = QueueClient.from_connection_string(conn_str=f'DefaultEndpointsProtocol=https;AccountName={self._storageAccountName};AccountKey={self._storageAccountKey}', queue_name=self.queueName)
+        connectionString = f'DefaultEndpointsProtocol=https;AccountName={self._storageAccountName};AccountKey={self._storageAccountKey}'
+        if self.endpointUrl:
+            connectionString += f';QueueEndpoint={self.endpointUrl}'
+        self._aqsClient = QueueClient.from_connection_string(conn_str=connectionString, queue_name=self.queueName)
         if not self._aqsClient:
             raise InternalServerErrorException('Failed to connect to queue')
         await self._aqsClient.get_queue_properties()
@@ -76,11 +84,35 @@ class AqsMessageQueue(MessageQueue[AqsMessage]):
         aqsMessages = [AqsMessage.from_aqs_message(aqsMessage=message) async for message in messagesIterator]
         return aqsMessages
 
-    async def delete_message(self, message: AqsMessage) -> None:
+    async def complete_message(self, message: AqsMessage) -> None:
         if not self._aqsClient:
-            raise InternalServerErrorException('You need to call .connect() before trying to delete messages')
+            raise InternalServerErrorException('You need to call .connect() before trying to complete messages')
         await self._aqsClient.delete_message(message=message.aqsId, pop_receipt=message.popReceipt)
 
+    async def retry_message(self, message: AqsMessage, delaySeconds: int = 0) -> None:
+        # NOTE(krishan711): check the lease is still held before re-sending, otherwise another worker owns the message and re-sending would duplicate it
+        if not await self.extend_message_lease(message=message, expectedProcessingSeconds=AQS_RETRY_HOLD_SECONDS):
+            return
+        await self.send_message(message=message, delaySeconds=delaySeconds)
+        await self.complete_message(message=message)
+
+    async def fail_message(self, message: AqsMessage, errorMessage: str | None) -> None:
+        # NOTE(krishan711): the message reappears after its visibility timeout
+        pass
+
+    async def extend_message_lease(self, message: AqsMessage, expectedProcessingSeconds: int) -> bool:
+        if not self._aqsClient:
+            raise InternalServerErrorException('You need to call .connect() before trying to extend message leases')
+        try:
+            updatedMessage = await self._aqsClient.update_message(message=message.aqsId, pop_receipt=message.popReceipt, visibility_timeout=expectedProcessingSeconds)
+        except HttpResponseError as exception:
+            if getattr(exception, 'error_code', None) in {StorageErrorCode.POP_RECEIPT_MISMATCH, StorageErrorCode.MESSAGE_NOT_FOUND}:
+                return False
+            raise
+        # NOTE(krishan711): azure invalidates the old pop receipt on every update so later calls must use the new one
+        message.popReceipt = updatedMessage.pop_receipt
+        return True
+
     async def get_message_count(self) -> int:
         if not self._aqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to count messages')

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/queues/cosmos.py

@@ -13,6 +13,7 @@ from azure.cosmos.aio import ContainerProxy
 
 from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
+from core.util import date_util
 from core.util.typing_util import JsonObject
 
 
@@ -169,7 +170,7 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
                 break
         return messages
 
-    async def delete_message(self, message: CosmosMessage) -> None:
+    async def complete_message(self, message: CosmosMessage) -> None:
         if message.deduplicationId is None:
             await self.container.delete_item(
                 item=message.id,
@@ -189,6 +190,46 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
             partition_key=self.queueName,
         )
 
+    async def retry_message(self, message: CosmosMessage, delaySeconds: int = 0) -> None:
+        try:
+            await self.container.patch_item(
+                item=message.id,
+                partition_key=self.queueName,
+                patch_operations=[
+                    {'op': 'set', 'path': '/leaseId', 'value': None},
+                    {'op': 'set', 'path': '/visibleDate', 'value': time.time() + delaySeconds},
+                    {'op': 'set', 'path': '/postCount', 'value': (message.postCount or 0) + 1},
+                    {'op': 'set', 'path': '/postDate', 'value': date_util.datetime_from_now().isoformat()},
+                ],
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except cosmos_exceptions.CosmosHttpResponseError as exception:
+            # NOTE(krishan711): 412 means the lease was lost to another worker, which now owns the message
+            if exception.status_code != 412:  # noqa: PLR2004
+                raise
+
+    async def fail_message(self, message: CosmosMessage, errorMessage: str | None) -> None:
+        # NOTE(krishan711): the message reappears once its lease expires
+        pass
+
+    async def extend_message_lease(self, message: CosmosMessage, expectedProcessingSeconds: int) -> bool:
+        try:
+            updatedItem = await self.container.patch_item(
+                item=message.id,
+                partition_key=self.queueName,
+                patch_operations=[{'op': 'set', 'path': '/visibleDate', 'value': time.time() + expectedProcessingSeconds}],
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except cosmos_exceptions.CosmosHttpResponseError as exception:
+            if exception.status_code == 412:  # noqa: PLR2004
+                return False
+            raise
+        # NOTE(krishan711): the patch changes the etag so later calls must use the new one to prove they still hold the lease
+        message.etag = str(updatedItem['_etag'])
+        return True
+
     async def get_message_count(self) -> int:
         counts = self.container.query_items(
             query='SELECT COUNT(1) AS messageCount FROM c WHERE c.itemType = @itemType AND c.queueName = @queueName AND c.visibleDate <= @now',

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/queues/message_queue.py

@@ -31,7 +31,19 @@ class MessageQueue[MessageType: Message](ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    async def delete_message(self, message: MessageType) -> None:
+    async def complete_message(self, message: MessageType) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def retry_message(self, message: MessageType, delaySeconds: int = 0) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def fail_message(self, message: MessageType, errorMessage: str | None) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def extend_message_lease(self, message: MessageType, expectedProcessingSeconds: int) -> bool:
         raise NotImplementedError
 
     @abc.abstractmethod

{kiba_core-0.5.3.dev65 → kiba_core-0.5.3.dev67}/core/queues/message_queue_processor.py

@@ -1,5 +1,7 @@
 import abc
 import asyncio
+import contextlib
+import sys
 import time
 import urllib.parse as urlparse
 import uuid
@@ -8,11 +10,15 @@ from abc import ABC
 from core import logging
 from core.exceptions import InternalServerErrorException
 from core.exceptions import KibaException
+from core.exceptions import LockedException
 from core.notifications.notification_client import NotificationClient
 from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
 from core.util.value_holder import RequestIdHolder
 
+LOCKED_MAX_RETRY_COUNT = 3
+LOCKED_RETRY_DELAY_SECONDS = 30
+
 
 class MessageProcessor(ABC):
     @abc.abstractmethod
@@ -28,6 +34,11 @@ class MessageNeedsReprocessingException(InternalServerErrorException):
         self.originalException = originalException
 
 
+class MessageLeaseLostException(InternalServerErrorException):
+    def __init__(self) -> None:
+        super().__init__(message='MESSAGE_LEASE_LOST')
+
+
 class MessageQueueProcessor[MessageType: Message]:
     def __init__(self, queue: MessageQueue[MessageType], messageProcessor: MessageProcessor, notificationClients: list[NotificationClient], requestIdHolder: RequestIdHolder | None = None) -> None:
         self.queue = queue
@@ -35,7 +46,36 @@ class MessageQueueProcessor[MessageType: Message]:
         self.notificationClients = notificationClients
         self.requestIdHolder = requestIdHolder
 
-    async def _process_message(self, message: MessageType) -> None:
+    async def _handle_failure(self, message: MessageType, exception: Exception, requestId: str) -> int:
+        logging.error('Caught exception whilst processing message:')
+        logging.exception(exception)
+        kibaException = KibaException.from_exception(exception=exception)
+        await self.queue.fail_message(message=message, errorMessage=kibaException.message)
+        for client in self.notificationClients:
+            try:
+                await client.post(messageText=f'Error processing message: {message.command}\n```\n{requestId}\n{message.content}\n{kibaException.message}```')
+            except Exception as notificationException:  # noqa: BLE001
+                logging.error('Failed to send message failure notification:')
+                logging.exception(notificationException)
+        return exception.statusCode if isinstance(exception, KibaException) else 500
+
+    async def _keep_message_alive(self, message: MessageType, expectedProcessingSeconds: int, stopEvent: asyncio.Event) -> bool:
+        while True:
+            with contextlib.suppress(TimeoutError):
+                await asyncio.wait_for(stopEvent.wait(), timeout=expectedProcessingSeconds / 3)
+            if stopEvent.is_set():
+                return False
+            try:
+                isExtended = await self.queue.extend_message_lease(message=message, expectedProcessingSeconds=expectedProcessingSeconds)
+            except Exception as exception:  # noqa: BLE001
+                logging.error(f'Failed to extend message lease {message.command}:')
+                logging.exception(exception)
+                continue
+            if not isExtended:
+                logging.error(f'Lost message lease for {message.command}')
+                return True
+
+    async def _process_message(self, message: MessageType, expectedProcessingSeconds: int) -> None:
         requestId = message.requestId or str(uuid.uuid4()).replace('-', '')
         if self.requestIdHolder:
             self.requestIdHolder.set_value(value=requestId)
@@ -43,20 +83,34 @@ class MessageQueueProcessor[MessageType: Message]:
         logging.api(action='MESSAGE', path=message.command, pathPattern=message.command, query=query)
         startTime = time.time()
         statusCode = 200
+        # NOTE(krishan711): the keep-alive is stopped and awaited (not cancelled) so an in-flight extend finishes before the message is completed, retried or failed.
+        # A lost lease is raised once the job returns rather than interrupting it, as cancelling mid-job could abandon work that was already sent.
+        # It is raised from the finally so it replaces any exception from the job: the message belongs to another worker, so it must not be retried or completed here.
+        # Cancellation (and other non-Exception errors) must still propagate, so those are never replaced.
+        stopKeepAliveEvent = asyncio.Event()
+        keepAliveTask = asyncio.create_task(self._keep_message_alive(message=message, expectedProcessingSeconds=expectedProcessingSeconds, stopEvent=stopKeepAliveEvent))
         try:
-            await self.messageProcessor.process_message(message=message)
-            await self.queue.delete_message(message=message)
+            try:
+                await self.messageProcessor.process_message(message=message)
+            finally:
+                stopKeepAliveEvent.set()
+                isLeaseLost = await keepAliveTask
+                inFlightException = sys.exception()
+                if isLeaseLost and (inFlightException is None or isinstance(inFlightException, Exception)):
+                    raise MessageLeaseLostException
+            await self.queue.complete_message(message=message)
         except MessageNeedsReprocessingException as exception:
             logging.info(msg=f'Scheduling reprocessing for message:{message.command} due to: {exception.originalException!s}')
-            await self.queue.send_message(message=message, delaySeconds=((message.postCount or 0) * exception.delaySeconds))
+            await self.queue.retry_message(message=message, delaySeconds=((message.postCount or 0) * exception.delaySeconds))
+        except LockedException as exception:
+            postCount = message.postCount or 0
+            if postCount <= LOCKED_MAX_RETRY_COUNT:
+                logging.info(msg=f'Scheduling reprocessing for message:{message.command} due to: {exception!s}')
+                await self.queue.retry_message(message=message, delaySeconds=(postCount * LOCKED_RETRY_DELAY_SECONDS))
+            else:
+                statusCode = await self._handle_failure(message=message, exception=exception, requestId=requestId)
         except Exception as exception:  # noqa: BLE001
-            statusCode = exception.statusCode if isinstance(exception, KibaException) else 500
-            logging.error('Caught exception whilst processing message:')
-            logging.exception(exception)
-            kibaException = KibaException.from_exception(exception=exception)
-            for client in self.notificationClients:
-                await client.post(messageText=f'Error processing message: {message.command}\n```\n{requestId}\n{message.content}\n{kibaException.message}```')
-            # TODO(krishan711): should possibly reset the visibility timeout
+            statusCode = await self._handle_failure(message=message, exception=exception, requestId=requestId)
         duration = time.time() - startTime
         logging.api(action='MESSAGE', path=message.command, pathPattern=message.command, query=query, response=statusCode, duration=duration)
         if self.requestIdHolder:
@@ -66,10 +120,10 @@ class MessageQueueProcessor[MessageType: Message]:
         logging.info('Retrieving messages...')
         messages = await self.queue.get_messages(expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds, limit=batchSize)
         if shouldProcessInParallel:
-            await asyncio.gather(*[self._process_message(message=message) for message in messages])
+            await asyncio.gather(*[self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds) for message in messages])
         else:
             for message in messages:
-                await self._process_message(message=message)
+                await self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds)
         return len(messages)
 
     async def execute(self, expectedProcessingSeconds: int = 300, longPollSeconds: int = 20) -> bool: