From 6b6f20a032eb31d747366c8bc270ca81731eaf6a Mon Sep 17 00:00:00 2001 From: Tom van der Lee Date: Fri, 4 Sep 2026 17:09:24 +0200 Subject: Fixed pre-commit issues --- proxy/__init__.py | 6 ++--- proxy/endpoints.py | 25 ++++++++++----------- proxy/queue.py | 63 +++++++++++++++++++++++------------------------------ proxy/utils.py | 6 ++++- proxy/websockets.py | 48 +++++++++++++++++++++++++--------------- 5 files changed, 78 insertions(+), 70 deletions(-) (limited to 'proxy') diff --git a/proxy/__init__.py b/proxy/__init__.py index fcdb23b..aba8ad2 100644 --- a/proxy/__init__.py +++ b/proxy/__init__.py @@ -7,7 +7,7 @@ from proxy.websockets import WebsocketProxy app = Starlette( debug=True, routes=[ - Route('/{path:path}', Proxy), - WebSocketRoute('/{path:path}', WebsocketProxy), - ] + Route("/{path:path}", Proxy), + WebSocketRoute("/{path:path}", WebsocketProxy), + ], ) diff --git a/proxy/endpoints.py b/proxy/endpoints.py index 0a23e7b..f514226 100644 --- a/proxy/endpoints.py +++ b/proxy/endpoints.py @@ -1,13 +1,13 @@ import logging -from base64 import b64encode, b64decode +from base64 import b64decode, b64encode from uuid import uuid4 from starlette.endpoints import HTTPEndpoint from starlette.requests import Request from starlette.responses import Response -from proxy.utils import get_path_with_query_string from proxy.queue import ProxyQueue +from proxy.utils import get_path_with_query_string from ttun_server.types import HttpMessage, HttpMessageType, HttpRequestData logger = logging.getLogger(__name__) @@ -18,16 +18,15 @@ class HeaderMapping: self._headers = headers def items(self): - for header in self._headers: - yield header + yield from self._headers class Proxy(HTTPEndpoint): async def dispatch(self) -> None: request = Request(self.scope, self.receive) - subdomain = request.path_params['subdomain'] - response = Response(content='Not Found', status_code=404) + subdomain = request.path_params["subdomain"] + response = Response(content="Not Found", status_code=404) identifier = str(uuid4()) response_queue = await ProxyQueue.create_for_identifier(identifier) @@ -35,7 +34,7 @@ class Proxy(HTTPEndpoint): try: request_queue = await ProxyQueue.get_for_identifier(subdomain) - logger.debug('PROXY %s%s ', subdomain, request.url) + logger.debug("PROXY %s%s ", subdomain, request.url) await request_queue.enqueue( HttpMessage( type=HttpMessageType.request.value, @@ -44,17 +43,17 @@ class Proxy(HTTPEndpoint): method=request.method, path=get_path_with_query_string(request), headers=list(request.headers.items()), - body=b64encode(await request.body()).decode() - ) + body=b64encode(await request.body()).decode(), + ), ) ) _response = await response_queue.dequeue() - payload = _response['payload'] + payload = _response["payload"] response = Response( - status_code=payload['status'], - headers=HeaderMapping(payload['headers']), - content=b64decode(payload['body'].encode()) + status_code=payload["status"], + headers=HeaderMapping(payload["headers"]), + content=b64decode(payload["body"].encode()), ) except AssertionError: pass 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 import json import logging import os -import traceback -from typing import Type +from typing import ClassVar, Self from ttun_server.redis import RedisConnectionPool from ttun_server.types import Message @@ -16,37 +15,37 @@ class BaseProxyQueue: self.identifier = identifier @classmethod - async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue': - raise NotImplementedError(f'Please implement create_for_identifier') + async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue: + raise NotImplementedError("Please implement create_for_identifier") @classmethod - async def get_for_identifier(cls, identifier: str) -> Type['self']: + async def get_for_identifier(cls, identifier: str) -> Self: assert await cls.has_connection(identifier) return cls(identifier) @classmethod async def has_connection(cls, identifier) -> bool: - raise NotImplementedError(f'Please implement has_connection') + raise NotImplementedError("Please implement has_connection") async def enqueue(self, message: Message): - raise NotImplementedError(f'Please implement send_request') + raise NotImplementedError("Please implement send_request") async def dequeue(self) -> Message: - raise NotImplementedError(f'Please implement handle_requests') + raise NotImplementedError("Please implement handle_requests") async def delete(self): - raise NotImplementedError(f'Please implement delete') + raise NotImplementedError("Please implement delete") class MemoryProxyQueue(BaseProxyQueue): - connections: dict[str, asyncio.Queue] = {} + connections: ClassVar[dict[str, asyncio.Queue]] = {} @classmethod async def has_connection(cls, identifier) -> bool: return identifier in cls.connections @classmethod - async def create_for_identifier(cls, identifier: str) -> 'MemoryProxyQueue': + async def create_for_identifier(cls, identifier: str) -> MemoryProxyQueue: instance = cls(identifier) cls.connections[identifier] = asyncio.Queue() @@ -66,67 +65,59 @@ class RedisProxyQueue(BaseProxyQueue): def __init__(self, identifier): super().__init__(identifier) - self.pubsub = RedisConnectionPool()\ - .get_connection()\ - .pubsub() + self.pubsub = RedisConnectionPool().get_connection().pubsub() self.subscription_queue = asyncio.Queue() @classmethod - async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue': + async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue: instance = cls(identifier) - await instance.pubsub.subscribe(f'request_{identifier}') + await instance.pubsub.subscribe(f"request_{identifier}") return instance @classmethod - async def get_for_identifier(cls, identifier: str) -> 'RedisProxyQueue': - instance: 'RedisProxyQueue' = await super().get_for_identifier(identifier) + async def get_for_identifier(cls, identifier: str) -> RedisProxyQueue: + instance: RedisProxyQueue = await super().get_for_identifier(identifier) - await instance.pubsub.subscribe(f'response_{identifier}') + await instance.pubsub.subscribe(f"response_{identifier}") return instance @classmethod async def has_connection(cls, identifier) -> bool: logger.debug(await RedisConnectionPool.get_connection().pubsub_channels()) - return f'request_{identifier}' in { + return f"request_{identifier}" in { channel.decode() - for channel - in await RedisConnectionPool \ - .get_connection() \ - .pubsub_channels() + for channel in await RedisConnectionPool.get_connection().pubsub_channels() } async def wait_for_message(self): async for message in self.pubsub.listen(): - match message['type']: - case 'subscribe': + match message["type"]: + case "subscribe": continue case _: - return message['data'] + return message["data"] async def enqueue(self, message: Message): - await RedisConnectionPool \ - .get_connection() \ - .publish(f'request_{self.identifier}', json.dumps(message)) + await RedisConnectionPool.get_connection().publish( + f"request_{self.identifier}", json.dumps(message) + ) async def dequeue(self) -> Message: message = await self.wait_for_message() return json.loads(message) async def delete(self): - await self.pubsub.unsubscribe(f'request_{self.identifier}') + await self.pubsub.unsubscribe(f"request_{self.identifier}") - await RedisConnectionPool.get_connection()\ - .srem('connections', self.identifier) + await RedisConnectionPool.get_connection().srem("connections", self.identifier) class ProxyQueueMeta(type): def __new__(cls, name, superclasses, attributes): - return RedisProxyQueue \ - if 'REDIS_URL' in os.environ \ - else MemoryProxyQueue + return RedisProxyQueue if "REDIS_URL" in os.environ else MemoryProxyQueue class ProxyQueue(BaseProxyQueue, metaclass=ProxyQueueMeta): diff --git a/proxy/utils.py b/proxy/utils.py index 2b80f43..3fff388 100644 --- a/proxy/utils.py +++ b/proxy/utils.py @@ -3,5 +3,9 @@ from starlette.requests import HTTPConnection def get_path_with_query_string(connection: HTTPConnection) -> str: path = connection.url.path - query_string = '?' + connection.scope['query_string'].decode() if connection.scope['query_string'] else '' + query_string = ( + "?" + connection.scope["query_string"].decode() + if connection.scope["query_string"] + else "" + ) return f"{path}{query_string}" diff --git a/proxy/websockets.py b/proxy/websockets.py index ec65e6d..a09e261 100644 --- a/proxy/websockets.py +++ b/proxy/websockets.py @@ -9,17 +9,22 @@ from uuid import uuid4 from starlette.endpoints import WebSocketEndpoint from starlette.websockets import WebSocket -from proxy.utils import get_path_with_query_string from proxy.queue import ProxyQueue -from ttun_server.types import WebsocketMessage, WebsocketMessageType, WebsocketConnectData, WebsocketMessageData, \ - WebsocketDisconnectData +from proxy.utils import get_path_with_query_string +from ttun_server.types import ( + WebsocketConnectData, + WebsocketDisconnectData, + WebsocketMessage, + WebsocketMessageData, + WebsocketMessageType, +) logger = logging.getLogger(__name__) -logger.setLevel('DEBUG') +logger.setLevel("DEBUG") class WebsocketProxy(WebSocketEndpoint): - encoding = 'json' + encoding = "json" websocket_listen_task = None def __init__(self, *args, **kwargs): @@ -28,16 +33,20 @@ class WebsocketProxy(WebSocketEndpoint): @asynccontextmanager async def proxy(self, websocket: WebSocket, message: WebsocketMessage): - [subdomain, *_] = websocket.url.hostname.split('.') + [subdomain, *_] = websocket.url.hostname.split(".") - expect_ack = WebsocketMessageType(message['type']) == WebsocketMessageType.connect + expect_ack = ( + WebsocketMessageType(message["type"]) == WebsocketMessageType.connect + ) try: request_queue = await ProxyQueue.get_for_identifier(subdomain) await request_queue.enqueue(message) if expect_ack: - response_queue = await ProxyQueue.create_for_identifier(message["identifier"]) + response_queue = await ProxyQueue.create_for_identifier( + message["identifier"] + ) yield await response_queue.dequeue() await response_queue.delete() else: @@ -51,7 +60,9 @@ class WebsocketProxy(WebSocketEndpoint): while True: message: WebsocketMessage = await response_queue.dequeue() logger.debug(message) - await websocket.send_text(b64decode(message['payload']['body'].encode()).decode()) + await websocket.send_text( + b64decode(message["payload"]["body"].encode()).decode() + ) async def on_connect(self, websocket: WebSocket) -> None: message = WebsocketMessage( @@ -60,18 +71,21 @@ class WebsocketProxy(WebSocketEndpoint): payload=WebsocketConnectData( path=get_path_with_query_string(websocket), headers=[ - (k.decode(), v.decode()) - for k, v - in websocket.scope['headers'] + (k.decode(), v.decode()) for k, v in websocket.scope["headers"] ], - ) + ), ) async with self.proxy(websocket, message) as m: - if m is not None and WebsocketMessageType(m['type']) == WebsocketMessageType.ack: + if ( + m is not None + and WebsocketMessageType(m["type"]) == WebsocketMessageType.ack + ): await super().on_connect(websocket) - self.websocket_listen_task = asyncio.create_task(self.listen_for_messages(websocket)) + self.websocket_listen_task = asyncio.create_task( + self.listen_for_messages(websocket) + ) def callback(*args, **kwargs): self.websocket_listen_task = None @@ -92,7 +106,7 @@ class WebsocketProxy(WebSocketEndpoint): identifier=self.id, payload=WebsocketMessageData( body=b64encode(data_bytes).decode(), - ) + ), ) async with self.proxy(websocket, message): @@ -104,7 +118,7 @@ class WebsocketProxy(WebSocketEndpoint): identifier=self.id, payload=WebsocketDisconnectData( close_code=close_code, - ) + ), ) async with self.proxy(websocket, message): -- cgit v1.3