hyperboria/nexus/hub/aioclient/aioclient.py
the-superpirate 681817ceae No description
GitOrigin-RevId: 83514338be1d662518bab9fe5ab6eefef98f229f
2021-04-09 13:29:10 +03:00

73 lines
2.5 KiB
Python

from typing import Optional
from aiokit import AioThing
from grpc.experimental.aio import insecure_channel
from idm.api2.proto.chats_service_pb2 import ChatData as ChatDataPb
from nexus.hub.proto.delivery_service_pb2 import \
StartDeliveryRequest as StartDeliveryRequestPb
from nexus.hub.proto.delivery_service_pb2 import \
StartDeliveryResponse as StartDeliveryResponsePb
from nexus.hub.proto.delivery_service_pb2_grpc import DeliveryStub
from nexus.hub.proto.submitter_service_pb2 import \
SubmitRequest as SubmitRequestPb
from nexus.hub.proto.submitter_service_pb2 import \
SubmitResponse as SubmitResponsePb
from nexus.hub.proto.submitter_service_pb2_grpc import SubmitterStub
from nexus.models.proto.typed_document_pb2 import \
TypedDocument as TypedDocumentPb
class HubGrpcClient(AioThing):
def __init__(
self,
base_url: str,
):
super().__init__()
self.channel = insecure_channel(base_url, [
('grpc.dns_min_time_between_resolutions_ms', 1000),
('grpc.initial_reconnect_backoff_ms', 1000),
('grpc.lb_policy_name', 'round_robin'),
('grpc.min_reconnect_backoff_ms', 1000),
('grpc.max_reconnect_backoff_ms', 2000),
])
self.delivery_stub = DeliveryStub(self.channel)
self.submitter_stub = SubmitterStub(self.channel)
async def start(self):
await self.channel.channel_ready()
async def stop(self):
await self.channel.close()
async def start_delivery(
self,
typed_document_pb: TypedDocumentPb,
chat: ChatDataPb,
request_id: Optional[str],
session_id: Optional[str],
) -> StartDeliveryResponsePb:
return await self.delivery_stub.start_delivery(
StartDeliveryRequestPb(
typed_document=typed_document_pb,
chat=chat,
),
metadata=(('request-id', request_id), ('session-id', session_id))
)
async def submit(
self,
telegram_document: bytes,
telegram_file_id: str,
chat: ChatDataPb,
request_id: Optional[str] = None,
session_id: Optional[str] = None,
) -> SubmitResponsePb:
return await self.submitter_stub.submit(
SubmitRequestPb(
telegram_document=telegram_document,
telegram_file_id=telegram_file_id,
chat=chat,
),
metadata=(('request-id', request_id), ('session-id', session_id))
)