diff options
| author | 2026-09-04 17:09:24 +0200 | |
|---|---|---|
| committer | 2026-09-04 17:09:24 +0200 | |
| commit | 6b6f20a032eb31d747366c8bc270ca81731eaf6a (patch) | |
| tree | ada48080495267bebde4ee1251bbbe77d629745f /proxy/queue.py | |
| parent | 2d22158dd95f76399f5716d556be60c77da04a46 (diff) | |
| download | server-6b6f20a032eb31d747366c8bc270ca81731eaf6a.tar.gz server-6b6f20a032eb31d747366c8bc270ca81731eaf6a.tar.bz2 server-6b6f20a032eb31d747366c8bc270ca81731eaf6a.zip | |
Fixed pre-commit issuesv3
Diffstat (limited to 'proxy/queue.py')
| -rw-r--r-- | proxy/queue.py | 63 |
1 files changed, 27 insertions, 36 deletions
diff --git a/proxy/queue.py b/proxy/queue.py index cfa0f3c..70fe3d0 100644 --- a/proxy/queue.py +++ b/proxy/queue.py | |||
| @@ -2,8 +2,7 @@ import asyncio | |||
| 2 | import json | 2 | import json |
| 3 | import logging | 3 | import logging |
| 4 | import os | 4 | import os |
| 5 | import traceback | 5 | from typing import ClassVar, Self |
| 6 | from typing import Type | ||
| 7 | 6 | ||
| 8 | from ttun_server.redis import RedisConnectionPool | 7 | from ttun_server.redis import RedisConnectionPool |
| 9 | from ttun_server.types import Message | 8 | from ttun_server.types import Message |
| @@ -16,37 +15,37 @@ class BaseProxyQueue: | |||
| 16 | self.identifier = identifier | 15 | self.identifier = identifier |
| 17 | 16 | ||
| 18 | @classmethod | 17 | @classmethod |
| 19 | async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue': | 18 | async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue: |
| 20 | raise NotImplementedError(f'Please implement create_for_identifier') | 19 | raise NotImplementedError("Please implement create_for_identifier") |
| 21 | 20 | ||
| 22 | @classmethod | 21 | @classmethod |
| 23 | async def get_for_identifier(cls, identifier: str) -> Type['self']: | 22 | async def get_for_identifier(cls, identifier: str) -> Self: |
| 24 | assert await cls.has_connection(identifier) | 23 | assert await cls.has_connection(identifier) |
| 25 | return cls(identifier) | 24 | return cls(identifier) |
| 26 | 25 | ||
| 27 | @classmethod | 26 | @classmethod |
| 28 | async def has_connection(cls, identifier) -> bool: | 27 | async def has_connection(cls, identifier) -> bool: |
| 29 | raise NotImplementedError(f'Please implement has_connection') | 28 | raise NotImplementedError("Please implement has_connection") |
| 30 | 29 | ||
| 31 | async def enqueue(self, message: Message): | 30 | async def enqueue(self, message: Message): |
| 32 | raise NotImplementedError(f'Please implement send_request') | 31 | raise NotImplementedError("Please implement send_request") |
| 33 | 32 | ||
| 34 | async def dequeue(self) -> Message: | 33 | async def dequeue(self) -> Message: |
| 35 | raise NotImplementedError(f'Please implement handle_requests') | 34 | raise NotImplementedError("Please implement handle_requests") |
| 36 | 35 | ||
| 37 | async def delete(self): | 36 | async def delete(self): |
| 38 | raise NotImplementedError(f'Please implement delete') | 37 | raise NotImplementedError("Please implement delete") |
| 39 | 38 | ||
| 40 | 39 | ||
| 41 | class MemoryProxyQueue(BaseProxyQueue): | 40 | class MemoryProxyQueue(BaseProxyQueue): |
| 42 | connections: dict[str, asyncio.Queue] = {} | 41 | connections: ClassVar[dict[str, asyncio.Queue]] = {} |
| 43 | 42 | ||
| 44 | @classmethod | 43 | @classmethod |
| 45 | async def has_connection(cls, identifier) -> bool: | 44 | async def has_connection(cls, identifier) -> bool: |
| 46 | return identifier in cls.connections | 45 | return identifier in cls.connections |
| 47 | 46 | ||
| 48 | @classmethod | 47 | @classmethod |
| 49 | async def create_for_identifier(cls, identifier: str) -> 'MemoryProxyQueue': | 48 | async def create_for_identifier(cls, identifier: str) -> MemoryProxyQueue: |
| 50 | instance = cls(identifier) | 49 | instance = cls(identifier) |
| 51 | cls.connections[identifier] = asyncio.Queue() | 50 | cls.connections[identifier] = asyncio.Queue() |
| 52 | 51 | ||
| @@ -66,67 +65,59 @@ class RedisProxyQueue(BaseProxyQueue): | |||
| 66 | def __init__(self, identifier): | 65 | def __init__(self, identifier): |
| 67 | super().__init__(identifier) | 66 | super().__init__(identifier) |
| 68 | 67 | ||
| 69 | self.pubsub = RedisConnectionPool()\ | 68 | self.pubsub = RedisConnectionPool().get_connection().pubsub() |
| 70 | .get_connection()\ | ||
| 71 | .pubsub() | ||
| 72 | 69 | ||
| 73 | self.subscription_queue = asyncio.Queue() | 70 | self.subscription_queue = asyncio.Queue() |
| 74 | 71 | ||
| 75 | @classmethod | 72 | @classmethod |
| 76 | async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue': | 73 | async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue: |
| 77 | instance = cls(identifier) | 74 | instance = cls(identifier) |
| 78 | 75 | ||
| 79 | await instance.pubsub.subscribe(f'request_{identifier}') | 76 | await instance.pubsub.subscribe(f"request_{identifier}") |
| 80 | return instance | 77 | return instance |
| 81 | 78 | ||
| 82 | @classmethod | 79 | @classmethod |
| 83 | async def get_for_identifier(cls, identifier: str) -> 'RedisProxyQueue': | 80 | async def get_for_identifier(cls, identifier: str) -> RedisProxyQueue: |
| 84 | instance: 'RedisProxyQueue' = await super().get_for_identifier(identifier) | 81 | instance: RedisProxyQueue = await super().get_for_identifier(identifier) |
| 85 | 82 | ||
| 86 | await instance.pubsub.subscribe(f'response_{identifier}') | 83 | await instance.pubsub.subscribe(f"response_{identifier}") |
| 87 | 84 | ||
| 88 | return instance | 85 | return instance |
| 89 | 86 | ||
| 90 | @classmethod | 87 | @classmethod |
| 91 | async def has_connection(cls, identifier) -> bool: | 88 | async def has_connection(cls, identifier) -> bool: |
| 92 | logger.debug(await RedisConnectionPool.get_connection().pubsub_channels()) | 89 | logger.debug(await RedisConnectionPool.get_connection().pubsub_channels()) |
| 93 | return f'request_{identifier}' in { | 90 | return f"request_{identifier}" in { |
| 94 | channel.decode() | 91 | channel.decode() |
| 95 | for channel | 92 | for channel in await RedisConnectionPool.get_connection().pubsub_channels() |
| 96 | in await RedisConnectionPool \ | ||
| 97 | .get_connection() \ | ||
| 98 | .pubsub_channels() | ||
| 99 | } | 93 | } |
| 100 | 94 | ||
| 101 | async def wait_for_message(self): | 95 | async def wait_for_message(self): |
| 102 | async for message in self.pubsub.listen(): | 96 | async for message in self.pubsub.listen(): |
| 103 | match message['type']: | 97 | match message["type"]: |
| 104 | case 'subscribe': | 98 | case "subscribe": |
| 105 | continue | 99 | continue |
| 106 | case _: | 100 | case _: |
| 107 | return message['data'] | 101 | return message["data"] |
| 108 | 102 | ||
| 109 | async def enqueue(self, message: Message): | 103 | async def enqueue(self, message: Message): |
| 110 | await RedisConnectionPool \ | 104 | await RedisConnectionPool.get_connection().publish( |
| 111 | .get_connection() \ | 105 | f"request_{self.identifier}", json.dumps(message) |
| 112 | .publish(f'request_{self.identifier}', json.dumps(message)) | 106 | ) |
| 113 | 107 | ||
| 114 | async def dequeue(self) -> Message: | 108 | async def dequeue(self) -> Message: |
| 115 | message = await self.wait_for_message() | 109 | message = await self.wait_for_message() |
| 116 | return json.loads(message) | 110 | return json.loads(message) |
| 117 | 111 | ||
| 118 | async def delete(self): | 112 | async def delete(self): |
| 119 | await self.pubsub.unsubscribe(f'request_{self.identifier}') | 113 | await self.pubsub.unsubscribe(f"request_{self.identifier}") |
| 120 | 114 | ||
| 121 | await RedisConnectionPool.get_connection()\ | 115 | await RedisConnectionPool.get_connection().srem("connections", self.identifier) |
| 122 | .srem('connections', self.identifier) | ||
| 123 | 116 | ||
| 124 | 117 | ||
| 125 | class ProxyQueueMeta(type): | 118 | class ProxyQueueMeta(type): |
| 126 | def __new__(cls, name, superclasses, attributes): | 119 | def __new__(cls, name, superclasses, attributes): |
| 127 | return RedisProxyQueue \ | 120 | return RedisProxyQueue if "REDIS_URL" in os.environ else MemoryProxyQueue |
| 128 | if 'REDIS_URL' in os.environ \ | ||
| 129 | else MemoryProxyQueue | ||
| 130 | 121 | ||
| 131 | 122 | ||
| 132 | class ProxyQueue(BaseProxyQueue, metaclass=ProxyQueueMeta): | 123 | class ProxyQueue(BaseProxyQueue, metaclass=ProxyQueueMeta): |
