|
| 1 | +import asyncio |
| 2 | +import logging |
| 3 | +import os |
| 4 | + |
| 5 | +import websockets |
| 6 | +from channels.exceptions import DenyConnection |
| 7 | +from channels.generic.websocket import AsyncWebsocketConsumer |
| 8 | +from django.contrib.auth.models import AnonymousUser |
| 9 | + |
| 10 | +logger = logging.getLogger(__name__) |
| 11 | + |
| 12 | + |
| 13 | +class DagsterWebSocketProxyConsumer(AsyncWebsocketConsumer): |
| 14 | + |
| 15 | + async def connect(self): |
| 16 | + logger.info(f"WebSocket connection attempt: {self.scope['path']}") |
| 17 | + |
| 18 | + # Authentication check |
| 19 | + if isinstance(self.scope["user"], AnonymousUser): |
| 20 | + logger.error("Authentication required") |
| 21 | + raise DenyConnection("Authentication required") |
| 22 | + |
| 23 | + perm = "common.access_dagster_ui" |
| 24 | + if not self.scope["user"].has_perm(perm): |
| 25 | + logger.error( |
| 26 | + f"User {self.scope['user'].username} lacks permission {perm} for accessing {self.scope['path']}" |
| 27 | + ) |
| 28 | + raise DenyConnection("Permission denied") |
| 29 | + |
| 30 | + logger.info(f"User {self.scope['user'].username} authenticated") |
| 31 | + |
| 32 | + # Build upstream URL |
| 33 | + dagster_url = os.environ.get("DAGSTER_WEBSERVER_URL", "http://localhost:3000") |
| 34 | + dagster_prefix = os.environ.get("DAGSTER_WEBSERVER_PREFIX", "pipelines") |
| 35 | + |
| 36 | + path = self.scope["path"] |
| 37 | + if path.startswith(f"/{dagster_prefix}/"): |
| 38 | + path = path[len(f"/{dagster_prefix}/") :] |
| 39 | + |
| 40 | + # Convert http to ws |
| 41 | + if dagster_url.startswith("https://"): |
| 42 | + ws_url = dagster_url.replace("https://", "wss://", 1) |
| 43 | + else: |
| 44 | + ws_url = dagster_url.replace("http://", "ws://", 1) |
| 45 | + |
| 46 | + # Build target URL |
| 47 | + if dagster_prefix: |
| 48 | + target_url = f"{ws_url}/{dagster_prefix}/{path}" |
| 49 | + else: |
| 50 | + target_url = f"{ws_url}/{path}" |
| 51 | + # Add query string |
| 52 | + if self.scope.get("query_string"): |
| 53 | + target_url += f"?{self.scope['query_string'].decode()}" |
| 54 | + |
| 55 | + logger.info(f"Connecting to upstream: {target_url}") |
| 56 | + |
| 57 | + # Get subprotocols from client |
| 58 | + subprotocols = self.scope.get("subprotocols", []) |
| 59 | + |
| 60 | + try: |
| 61 | + self.websocket = await websockets.connect( |
| 62 | + target_url, |
| 63 | + max_size=52428800, # Increase to 50 MB |
| 64 | + ping_interval=20, |
| 65 | + open_timeout=30, # Default is 10 sec |
| 66 | + subprotocols=subprotocols if subprotocols else None, |
| 67 | + ) |
| 68 | + logger.info("Connected to upstream") |
| 69 | + except Exception as e: |
| 70 | + logger.error(f"Failed to connect: {e}") |
| 71 | + raise DenyConnection(f"Connection to upstream failed: {e}") |
| 72 | + |
| 73 | + await self.accept(self.websocket.subprotocol) |
| 74 | + logger.info(f"Client accepted with subprotocol: {self.websocket.subprotocol}") |
| 75 | + |
| 76 | + self.consumer_task = asyncio.create_task(self.consume_from_upstream()) |
| 77 | + |
| 78 | + async def disconnect(self, close_code): |
| 79 | + logger.info(f"Disconnecting with code {close_code}") |
| 80 | + if hasattr(self, "consumer_task"): |
| 81 | + self.consumer_task.cancel() |
| 82 | + try: |
| 83 | + await self.consumer_task |
| 84 | + except asyncio.CancelledError: |
| 85 | + pass |
| 86 | + if hasattr(self, "websocket"): |
| 87 | + await self.websocket.close() |
| 88 | + |
| 89 | + async def receive(self, text_data=None, bytes_data=None): |
| 90 | + try: |
| 91 | + await self.websocket.send(bytes_data or text_data) |
| 92 | + except Exception as e: |
| 93 | + logger.error(f"Error forwarding to upstream: {e}") |
| 94 | + await self.close() |
| 95 | + |
| 96 | + async def consume_from_upstream(self): |
| 97 | + try: |
| 98 | + async for message in self.websocket: |
| 99 | + if isinstance(message, bytes): |
| 100 | + await self.send(bytes_data=message) |
| 101 | + else: |
| 102 | + await self.send(text_data=message) |
| 103 | + except asyncio.CancelledError: |
| 104 | + pass |
| 105 | + except Exception as e: |
| 106 | + logger.error(f"Error consuming from upstream: {e}") |
| 107 | + await self.close() |
0 commit comments