kiba-core 0.5.3.dev50__tar.gz → 0.5.3.dev52__tar.gz

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/CHANGELOG.md

@@ -32,6 +32,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [MINOR] Added `call` and `multicall` to `EthClientInterface`
 - [MINOR] Truncate discord notification messages to fit api
 - [MINOR] Added `blockNumber` to `EthClient.multicall` function
+- [MINOR] Added `SqlMessageQueue` and `CosmosMessageQueue` to support database backed queues
+- [MINOR] Added `shouldAllowFailures` to `EthClient.multicall` to allow individual calls to fail without reverting the whole batch
 
 ### Changed
 

kiba_core-0.5.3.dev52/PKG-INFO

@@ -0,0 +1,41 @@
+Metadata-Version: 2.4
+Name: kiba-core
+Version: 0.5.3.dev52
+Summary: Kiba Labs' python utilities
+Project-URL: repository, https://github.com/kibalabs/core-py
+Requires-Python: <4,>=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: aiofiles>=25.1.0
+Requires-Dist: pydantic>=2.13.4
+Requires-Dist: orjson>=3.11.9
+Provides-Extra: api
+Requires-Dist: fastapi>=0.141.1; extra == "api"
+Requires-Dist: uvicorn[standard]>=0.51.0; extra == "api"
+Provides-Extra: core-api
+Requires-Dist: starlette>=1.6.0; extra == "core-api"
+Requires-Dist: uvicorn[standard]>=0.51.0; extra == "core-api"
+Provides-Extra: storage
+Requires-Dist: aiobotocore>=3.9.0; extra == "storage"
+Provides-Extra: queues
+Requires-Dist: aiobotocore>=3.9.0; extra == "queues"
+Provides-Extra: queue-sqs
+Requires-Dist: aiobotocore>=3.9.0; extra == "queue-sqs"
+Provides-Extra: queue-aqs
+Requires-Dist: azure-storage-queue[aio]>=12.17.0; extra == "queue-aqs"
+Provides-Extra: queue-cosmos
+Requires-Dist: azure-cosmos>=4.16.3; extra == "queue-cosmos"
+Provides-Extra: database-psql
+Requires-Dist: sqlalchemy[asyncio]>=2.0.52; extra == "database-psql"
+Requires-Dist: asyncpg>=0.31.0; extra == "database-psql"
+Provides-Extra: database-sqlite
+Requires-Dist: sqlalchemy[asyncio]>=2.0.52; extra == "database-sqlite"
+Requires-Dist: aiosqlite>=0.22.1; extra == "database-sqlite"
+Provides-Extra: requester
+Requires-Dist: httpx>=0.28.1; extra == "requester"
+Provides-Extra: web3
+Requires-Dist: web3>=7.16.0; extra == "web3"
+Provides-Extra: types
+Requires-Dist: types-aiobotocore[essential]>=3.9.0; extra == "types"
+Requires-Dist: types-aiofiles>=25.1.0.20260518; extra == "types"
+
+# Core-Py

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/api/authorizer.py

@@ -1,9 +1,7 @@
 import functools
 import typing
 from collections.abc import AsyncIterator
-from typing import ParamSpec
 
-from mypy_extensions import Arg
 from pydantic import BaseModel
 
 from core import logging
@@ -13,7 +11,6 @@ from core.exceptions import UnauthorizedException
 from core.http.basic_authentication import BasicAuthentication
 from core.http.jwt import Jwt
 
-_P = ParamSpec('_P')
 _AnyReturn = typing.Awaitable[typing.Any] | AsyncIterator[typing.Any]  # type: ignore[explicit-any]
 
 
@@ -43,19 +40,18 @@ async def _authorize_bearer_jwt[ApiRequest: BaseModel](request: KibaApiRequest[A
 
 def authorize_bearer_jwt[ApiRequest: BaseModel](  # type: ignore[explicit-any]
     authorizer: Authorizer,
-) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
-    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]], typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]]:
+    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]) -> typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]:  # 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)
-            result = func(request=request)
+            result = func(request)
             # NOTE(krishan711): this is here to support streaming responses which return an async generator
             if hasattr(result, '__aiter__'):
                 return result
             return await result
 
-        # TODO(krishan711): figure out correct typing here
-        return async_wrapper  # type: ignore[return-value]
+        return async_wrapper
 
     return decorator
 
@@ -72,26 +68,25 @@ async def get_basic_authentication_from_authorization_signature[ApiRequest: Base
     except UnauthorizedException:
         raise
     except BaseException as exception:  # noqa: BLE001
-        logging.exception(exception)  # type: ignore[arg-type]
+        logging.exception(exception)  # type: ignore[arg-type, ty:invalid-argument-type]
         raise ForbiddenException(message='AUTH_INVALID')
     return BasicAuthentication(username=signerId, password=signatureString)
 
 
 def authorize_signature[ApiRequest: BaseModel](  # type: ignore[explicit-any]
     authorizer: SignatureAuthorizer,
-) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
-    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]], typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]]:
+    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]) -> typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]:  # type: ignore[explicit-any]
         @functools.wraps(func)
         async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
             request.authBasic = await get_basic_authentication_from_authorization_signature(request=request, authorizer=authorizer)
-            result = func(request=request)
+            result = func(request)
             # NOTE(krishan711): this is here to support streaming responses which return an async generator
             if hasattr(result, '__aiter__'):
                 return result
             return await result
 
-        # TODO(krishan711): figure out correct typing here
-        return async_wrapper  # type: ignore[return-value]
+        return async_wrapper
 
     return decorator
 
@@ -112,8 +107,8 @@ class StaticTokenAuthorizer(TokenAuthorizer):
 
 def authorize_token[ApiRequest: BaseModel](  # type: ignore[explicit-any]
     authorizer: TokenAuthorizer,
-) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
-    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]], typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]]:
+    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], _AnyReturn]) -> typing.Callable[[KibaApiRequest[ApiRequest]], typing.Any]:  # type: ignore[explicit-any]
         @functools.wraps(func)
         async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
             authorization = request.headers.get('Authorization')
@@ -122,13 +117,12 @@ def authorize_token[ApiRequest: BaseModel](  # type: ignore[explicit-any]
             if not authorization.startswith('Token '):
                 raise ForbiddenException(message='AUTH_INVALID')
             await authorizer.validate_token(authorization[6:])
-            result = func(request=request)
+            result = func(request)
             # NOTE(krishan711): this is here to support streaming responses which return an async generator
             if hasattr(result, '__aiter__'):
                 return result
             return await result
 
-        # TODO(krishan711): figure out correct typing here
-        return async_wrapper  # type: ignore[return-value]
+        return async_wrapper
 
     return decorator

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/api/json_route.py

@@ -1,8 +1,6 @@
 import functools
 import typing
-from typing import ParamSpec
 
-from mypy_extensions import Arg
 from pydantic import BaseModel
 from pydantic import ValidationError
 
@@ -13,14 +11,12 @@ from core.exceptions import InternalServerErrorException
 from core.util import json_util
 from core.util.typing_util import JsonObject
 
-_P = ParamSpec('_P')
-
 
 def json_route[ApiRequest: BaseModel, ApiResponse: BaseModel](
     requestType: typing.Type[ApiRequest],
     responseType: typing.Type[ApiResponse],
-) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], typing.Awaitable[ApiResponse]]], typing.Callable[_P, KibaJSONResponse]]:
-    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], typing.Awaitable[ApiResponse]]) -> typing.Callable[_P, KibaJSONResponse]:
+) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], typing.Awaitable[ApiResponse]]], typing.Callable[..., typing.Awaitable[KibaJSONResponse]]]:
+    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], typing.Awaitable[ApiResponse]]) -> typing.Callable[..., typing.Awaitable[KibaJSONResponse]]:
         @functools.wraps(func)
         async def async_wrapper(*args: typing.Any) -> KibaJSONResponse:  # type: ignore[explicit-any, misc]
             receivedRequest = args[0]
@@ -42,12 +38,11 @@ def json_route[ApiRequest: BaseModel, ApiResponse: BaseModel](
                 raise BadRequestException(f'Invalid request: {validationErrorMessage}')
             kibaRequest: KibaApiRequest[ApiRequest] = KibaApiRequest(scope=receivedRequest.scope, receive=receivedRequest._receive, send=receivedRequest._send)  # noqa: SLF001
             kibaRequest.data = requestParams
-            receivedResponse = await func(request=kibaRequest)
+            receivedResponse = await func(kibaRequest)
             if not isinstance(receivedResponse, responseType):
                 raise InternalServerErrorException(f'Expected response to be of type {responseType}, got {type(receivedResponse)}')
             return KibaJSONResponse(content=receivedResponse.model_dump())
 
-        # TODO(krishan711): figure out correct typing here
-        return async_wrapper  # type: ignore[return-value]
+        return async_wrapper
 
     return decorator

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/api/streaming_json_route.py

@@ -2,7 +2,6 @@ import functools
 import inspect
 import typing
 from collections.abc import AsyncIterator
-from typing import ParamSpec
 
 from pydantic import BaseModel
 from pydantic import ValidationError
@@ -14,8 +13,6 @@ from core.exceptions import InternalServerErrorException
 from core.util import json_util
 from core.util.typing_util import JsonObject
 
-_P = ParamSpec('_P')
-
 
 async def _convert_to_json_generator[T: BaseModel](response_iterator: AsyncIterator[T], expectedType: typing.Type[T]) -> AsyncIterator[bytes]:
     async for content in response_iterator:
@@ -27,8 +24,8 @@ async def _convert_to_json_generator[T: BaseModel](response_iterator: AsyncItera
 def streaming_json_route[ApiRequest: BaseModel, ApiResponse: BaseModel](
     requestType: typing.Type[ApiRequest],
     responseType: typing.Type[ApiResponse],
-) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], AsyncIterator[ApiResponse]]], typing.Callable[_P, StreamingResponse]]:
-    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], AsyncIterator[ApiResponse]]) -> typing.Callable[_P, StreamingResponse]:
+) -> typing.Callable[[typing.Callable[[KibaApiRequest[ApiRequest]], AsyncIterator[ApiResponse]]], typing.Callable[..., typing.Awaitable[StreamingResponse]]]:
+    def decorator(func: typing.Callable[[KibaApiRequest[ApiRequest]], AsyncIterator[ApiResponse]]) -> typing.Callable[..., typing.Awaitable[StreamingResponse]]:
         @functools.wraps(func)
         async def async_wrapper(*args: typing.Any, **kwargs: typing.Any) -> StreamingResponse:  # type: ignore[explicit-any, misc]
             receivedRequest = kwargs.get('request', args[0] if args else None)
@@ -58,7 +55,6 @@ def streaming_json_route[ApiRequest: BaseModel, ApiResponse: BaseModel](
             # NOTE(krishan711): we set content-encoding to identity to prevent gzip from trying to process it (cos it buffers all the content)
             return StreamingResponse(content=wrappedGenerator, media_type='application/x-ndjson', headers={'Content-Encoding': 'identity'})
 
-        # TODO(krishan711): figure out correct typing here
-        return async_wrapper  # type: ignore[return-value]
+        return async_wrapper
 
     return decorator

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/logging.py

@@ -4,7 +4,6 @@ import logging
 import os
 import re
 import sys
-import typing
 from collections.abc import Callable
 from collections.abc import Collection
 from logging import Formatter
@@ -153,7 +152,7 @@ def init_basic_logging(showDebug: bool = False) -> None:
         init_logger(logger=logger, loggingLevel=loggingLevel, handler=handler)
 
 
-def _serialize_numeric_value(value: Union[None, float]) -> str:
+def _serialize_numeric_value(value: Union[float, None]) -> str:
     if value is None:
         return ''
     roundedNumber = round(value, 6)
@@ -169,7 +168,7 @@ def stat(name: str, key: str, value: float = 1) -> None:
         nameValue = _serialize_string_value(value=str(name))
         keyValue = _serialize_string_value(value=str(key))
         statValue = _serialize_numeric_value(value=value)
-        STAT_LOGGER.log(level=logging.INFO, msg='', extra=typing.cast(dict[str, str], {'statName': nameValue, 'statKey': keyValue, 'statValue': statValue}))
+        STAT_LOGGER.log(level=logging.INFO, msg='', extra={'statName': nameValue, 'statKey': keyValue, 'statValue': statValue})
 
 
 # TODO(krishan711): make pathPattern mandatory in next major release
@@ -184,7 +183,7 @@ def api(action: str, path: str, query: str, pathPattern: str | None = None, resp
         API_LOGGER.log(
             level=logging.INFO,
             msg='',
-            extra=typing.cast(dict[str, str], {'apiAction': actionString, 'apiPath': pathString, 'apiPathPattern': pathPatternString, 'apiQuery': queryString, 'apiResponse': responseString or '', 'apiDuration': durationString or ''}),
+            extra={'apiAction': actionString, 'apiPath': pathString, 'apiPathPattern': pathPatternString, 'apiQuery': queryString, 'apiResponse': responseString or '', 'apiDuration': durationString or ''},
         )
 
 
@@ -204,11 +203,11 @@ def _log(level: int, msg: str, *args: Any, **kwargs: Any) -> None:  # type: igno
 
 
 def critical(msg: str, *args: Any, **kwargs: Any) -> None:  # type: ignore[explicit-any]
-    _log(level=CRITICAL, msg=msg, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(CRITICAL, msg, *args, **kwargs)
 
 
 def error(msg: str, *args: Any, **kwargs: Any) -> None:  # type: ignore[explicit-any]
-    _log(level=ERROR, msg=msg, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(ERROR, msg, *args, **kwargs)
 
 
 def exception(msg: Union[str, Exception], *args: Any, exc_info: bool = True, **kwargs: Any) -> None:  # type: ignore[explicit-any]
@@ -218,19 +217,19 @@ def exception(msg: Union[str, Exception], *args: Any, exc_info: bool = True, **k
         resolvedMessage = str(msg)
     else:
         resolvedMessage = msg
-    _log(level=ERROR, msg=resolvedMessage, exc_info=exc_info, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(ERROR, resolvedMessage, *args, exc_info=exc_info, **kwargs)
 
 
 def warning(msg: str, *args: Any, **kwargs: Any) -> None:  # type: ignore[explicit-any]
-    _log(level=WARNING, msg=msg, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(WARNING, msg, *args, **kwargs)
 
 
 def info(msg: str, *args: Any, **kwargs: Any) -> None:  # type: ignore[explicit-any]
-    _log(level=INFO, msg=msg, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(INFO, msg, *args, **kwargs)
 
 
 def debug(msg: str, *args: Any, **kwargs: Any) -> None:  # type: ignore[explicit-any]
-    _log(level=DEBUG, msg=msg, *args, **kwargs)  # type: ignore[misc]  # noqa: B026
+    _log(DEBUG, msg, *args, **kwargs)
 
 
 def basicConfig(**kwargs: Any) -> None:  # type: ignore[explicit-any]  # noqa: N802  # noqa: B026

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/queues/aqs.py

@@ -12,19 +12,20 @@ from core.queues.model import Message
 from core.util import list_util
 
 
-class AqsMessage(Message):  # type: ignore[explicit-any]
+class AqsMessage(Message):
     aqsId: str
     popReceipt: str | None
 
     @classmethod
     def from_aqs_message(cls, aqsMessage: RawAqsMessage) -> AqsMessage:
-        message = Message.parse_raw(aqsMessage.content)
+        message = Message.model_validate_json(aqsMessage.content)
         return cls(
             command=message.command,
             content=message.content,
             requestId=message.requestId,
             postCount=message.postCount,
             postDate=message.postDate,
+            deduplicationId=message.deduplicationId,
             aqsId=aqsMessage.id,
             popReceipt=aqsMessage.pop_receipt,
         )
@@ -53,7 +54,7 @@ class AqsMessageQueue(MessageQueue[AqsMessage]):
         if not self._aqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to send messages')
         message.prepare_for_send()
-        await self._aqsClient.send_message(visibility_timeout=delaySeconds, content=message.json())
+        await self._aqsClient.send_message(visibility_timeout=delaySeconds, content=message.model_dump_json())
 
     async def send_messages(self, messages: Sequence[Message], delaySeconds: int = 0) -> None:
         if not self._aqsClient:

kiba_core-0.5.3.dev52/core/queues/cosmos.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import time
+import uuid
+from collections.abc import Sequence
+from typing import Any
+
+from azure.core import MatchConditions
+from azure.cosmos import exceptions as cosmos_exceptions
+from azure.cosmos.aio import ContainerProxy
+
+from core.queues.message_queue import MessageQueue
+from core.queues.model import Message
+from core.util.typing_util import JsonObject
+
+
+class CosmosMessage(Message):
+    id: str
+    leaseId: str
+    etag: str
+
+    @classmethod
+    def from_item(cls, item: JsonObject) -> CosmosMessage:
+        message = Message.model_validate(
+            {
+                'command': item['command'],
+                'content': item['content'],
+                'requestId': item.get('requestId'),
+                'postCount': item.get('postCount'),
+                'postDate': item.get('postDate'),
+                'deduplicationId': item.get('deduplicationId'),
+            }
+        )
+        return cls(
+            command=message.command,
+            content=message.content,
+            requestId=message.requestId,
+            postCount=message.postCount,
+            postDate=message.postDate,
+            deduplicationId=message.deduplicationId,
+            id=str(item['id']),
+            leaseId=str(item['leaseId']),
+            etag=str(item['_etag']),
+        )
+
+
+class CosmosMessageQueue(MessageQueue[CosmosMessage]):
+    def __init__(self, container: ContainerProxy, queueName: str, pollIntervalSeconds: float = 1.0) -> None:
+        self.container = container
+        self.queueName = queueName
+        self.pollIntervalSeconds = pollIntervalSeconds
+
+    async def connect(self) -> None:
+        # The caller supplies a shared ContainerProxy and owns CosmosClient lifecycle.
+        return None
+
+    async def disconnect(self) -> None:
+        # The caller supplies a shared ContainerProxy and owns CosmosClient lifecycle.
+        return None
+
+    @staticmethod
+    def _item_body(message: Message, queueName: str, itemId: str, visibleDate: float, createdDate: float) -> dict[str, Any]:  # type: ignore[explicit-any]
+        return {
+            'id': itemId,
+            'itemType': 'message',
+            'queueName': queueName,
+            'command': message.command,
+            'content': message.content,
+            'requestId': message.requestId,
+            'postCount': message.postCount,
+            'postDate': message.postDate.isoformat() if message.postDate else None,
+            'deduplicationId': message.deduplicationId,
+            'leaseId': None,
+            'visibleDate': visibleDate,
+            'createdDate': createdDate,
+        }
+
+    async def send_message(self, message: Message, delaySeconds: int = 0) -> None:
+        await self.send_messages(messages=[message], delaySeconds=delaySeconds)
+
+    async def send_messages(self, messages: Sequence[Message], delaySeconds: int = 0) -> None:
+        if not messages:
+            return
+        now = time.time()
+        visibleDate = now + delaySeconds
+        for message in messages:
+            message.prepare_for_send()
+            itemId = uuid.uuid4().hex
+            item = self._item_body(message=message, queueName=self.queueName, itemId=itemId, visibleDate=visibleDate, createdDate=now)
+            if message.deduplicationId is None:
+                # The SDK derives the create partition key from the body's queueName field.
+                await self.container.create_item(body=item)
+                continue
+
+            deduplicationKey = f'{self.queueName}:{message.deduplicationId}'
+            deduplicationItemId = f'dedup-{hashlib.sha256(deduplicationKey.encode()).hexdigest()}'
+            deduplicationItem = {
+                'id': deduplicationItemId,
+                'itemType': 'deduplication',
+                'queueName': self.queueName,
+                'deduplicationId': message.deduplicationId,
+                'messageId': itemId,
+            }
+            try:
+                await self.container.execute_item_batch(
+                    batch_operations=[
+                        ('create', (deduplicationItem,)),
+                        ('create', (item,)),
+                    ],
+                    partition_key=self.queueName,
+                )
+            except cosmos_exceptions.CosmosBatchOperationError as exception:
+                if exception.error_index != 0 or exception.status_code != 409:  # noqa: PLR2004
+                    raise
+
+    async def get_message(self, expectedProcessingSeconds: int = 300, longPollSeconds: int = 0) -> CosmosMessage | None:
+        messages = await self.get_messages(limit=1, expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds)
+        return messages[0] if messages else None
+
+    async def get_messages(self, limit: int = 1, expectedProcessingSeconds: int = 300, longPollSeconds: int = 0) -> list[CosmosMessage]:
+        if limit <= 0:
+            return []
+        deadline = time.monotonic() + longPollSeconds
+        while True:
+            messages = await self._claim_messages(limit=limit, expectedProcessingSeconds=expectedProcessingSeconds)
+            if messages or time.monotonic() >= deadline:
+                return messages
+            await asyncio.sleep(min(self.pollIntervalSeconds, max(deadline - time.monotonic(), 0)))
+
+    async def _claim_messages(self, limit: int, expectedProcessingSeconds: int) -> list[CosmosMessage]:
+        now = time.time()
+        query = 'SELECT * FROM c WHERE c.itemType = @itemType AND c.queueName = @queueName AND c.visibleDate <= @now ORDER BY c.createdDate'
+        candidates = self.container.query_items(
+            query=query,
+            parameters=[
+                {'name': '@itemType', 'value': 'message'},
+                {'name': '@queueName', 'value': self.queueName},
+                {'name': '@now', 'value': now},
+            ],
+            partition_key=self.queueName,
+            # Read ahead so ETag conflicts caused by concurrent workers do not unnecessarily
+            # under-fill this batch; candidates that another worker claims are skipped.
+            max_item_count=max(limit * 2, 10),
+        )
+        messages: list[CosmosMessage] = []
+        async for candidate in candidates:
+            item = {key: value for key, value in candidate.items() if not key.startswith('_')}
+            leaseId = uuid.uuid4().hex
+            item['leaseId'] = leaseId
+            item['visibleDate'] = now + expectedProcessingSeconds
+            try:
+                claimedItem = await self.container.replace_item(
+                    item=candidate['id'],
+                    body=item,
+                    etag=candidate['_etag'],
+                    match_condition=MatchConditions.IfNotModified,
+                    # The async SDK derives this from the body, but explicitly pass it in
+                    # request options so the replace is targeted to this logical partition.
+                    request_options={'partitionKey': self.queueName},
+                )
+            except cosmos_exceptions.CosmosHttpResponseError as exception:
+                if exception.status_code == 412:  # noqa: PLR2004
+                    continue
+                raise
+            messages.append(CosmosMessage.from_item(item=claimedItem))
+            if len(messages) >= limit:
+                break
+        return messages
+
+    async def delete_message(self, message: CosmosMessage) -> None:
+        if message.deduplicationId is None:
+            await self.container.delete_item(
+                item=message.id,
+                partition_key=self.queueName,
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+            return
+
+        deduplicationKey = f'{self.queueName}:{message.deduplicationId}'
+        deduplicationItemId = f'dedup-{hashlib.sha256(deduplicationKey.encode()).hexdigest()}'
+        await self.container.execute_item_batch(
+            batch_operations=[
+                ('delete', (message.id,), {'if_match_etag': message.etag}),
+                ('delete', (deduplicationItemId,)),
+            ],
+            partition_key=self.queueName,
+        )

{kiba_core-0.5.3.dev50 → kiba_core-0.5.3.dev52}/core/queues/model.py

@@ -1,22 +1,24 @@
 from __future__ import annotations
 
 import datetime
+import hashlib
 import typing
 import uuid
-from typing import Any
 
 from pydantic import BaseModel
 from pydantic.fields import ModelPrivateAttr
 
 from core.util import date_util
+from core.util.typing_util import JsonObject
 
 
-class Message(BaseModel):  # type: ignore[explicit-any]
+class Message(BaseModel):
     command: str
-    content: dict[str, Any]  # type: ignore[explicit-any]
+    content: JsonObject
     requestId: str | None
     postCount: int | None
     postDate: datetime.datetime | None
+    deduplicationId: str | None = None
 
     def prepare_for_send(self) -> None:
         self.requestId = self.requestId or str(uuid.uuid4()).replace('-', '')
@@ -31,6 +33,10 @@ class MessageContent(BaseModel):
     def get_command(cls) -> str:
         return typing.cast(str, typing.cast(ModelPrivateAttr, cls._COMMAND).get_default())
 
+    def get_deduplication_id(self) -> str | None:
+        payload = f'{self.get_command()}:{self.model_dump_json()}'
+        return hashlib.sha256(payload.encode()).hexdigest()
+
     def to_message(self) -> Message:
         return Message(
             command=self.get_command(),
@@ -38,4 +44,10 @@ class MessageContent(BaseModel):
             requestId=None,
             postCount=None,
             postDate=None,
+            deduplicationId=self.get_deduplication_id(),
         )
+
+
+class NonDeduplicatedMessageContent(MessageContent):
+    def get_deduplication_id(self) -> str | None:
+        return None