feat: add market data architecture and complete migration through build 039

app/tools/dzengi_probe/websocket_probe.py

@@ -0,0 +1,134 @@
+# app/tools/dzengi_probe/websocket_probe.py
+
+from __future__ import annotations
+
+import asyncio
+import json
+from dataclasses import dataclass
+from typing import Any
+
+import websockets
+from websockets.typing import Subprotocol
+
+from app.tools.dzengi_probe.config import DzengiProbeConfig
+from app.tools.dzengi_probe.response_store import append_jsonl, save_json
+
+
+@dataclass(frozen=True)
+class WebSocketProbeResult:
+    stream: str
+    output_file: str
+    ok: bool
+    messages_saved: int
+    error: str | None = None
+
+
+class DzengiWebSocketProbe:
+    def __init__(self, config: DzengiProbeConfig) -> None:
+        self.config = config
+
+    async def probe_market_data(self, *, max_messages: int = 10) -> WebSocketProbeResult:
+        output_file = (
+            f"websocket/marketData.subscribe/"
+            f"{self.config.symbol_file_name}.jsonl"
+        )
+
+        request_file = (
+            f"websocket/marketData.subscribe/"
+            f"{self.config.symbol_file_name}.request.json"
+        )
+
+        subscribe_message: dict[str, Any] = {
+            "correlationId": "probe-market-data-subscribe",
+            "destination": "marketData.subscribe",
+            "payload": {
+                "symbols": [self.config.symbol],
+            },
+        }
+
+        save_json(self.config.output_dir / request_file, subscribe_message)
+
+        return await self._subscribe_and_save(
+            stream="marketData.subscribe",
+            subscribe_message=subscribe_message,
+            output_file=output_file,
+            max_messages=max_messages,
+        )
+
+    async def _subscribe_and_save(
+        self,
+        *,
+        stream: str,
+        subscribe_message: dict[str, Any],
+        output_file: str,
+        max_messages: int,
+    ) -> WebSocketProbeResult:
+        output_path = self.config.output_dir / output_file
+        messages_saved = 0
+
+        if output_path.exists():
+            output_path.unlink()
+
+        try:
+            async with websockets.connect(
+                self.config.ws_url,
+                extra_headers={
+                    "Origin": self.config.base_url,
+                    "Content-Type": "application/json",
+                },
+                subprotocols=[Subprotocol("json")],
+                ping_interval=20,
+                ping_timeout=20,
+                close_timeout=5,
+            ) as websocket:
+                await websocket.send(json.dumps(subscribe_message))
+
+                while messages_saved < max_messages:
+                    raw_message = await asyncio.wait_for(
+                        websocket.recv(),
+                        timeout=self.config.timeout_seconds,
+                    )
+
+                    try:
+                        payload = json.loads(raw_message)
+                    except json.JSONDecodeError:
+                        payload = {"raw": raw_message}
+
+                    print(payload)
+                    append_jsonl(output_path, payload)
+                    messages_saved += 1
+
+            return WebSocketProbeResult(
+                stream=stream,
+                output_file=str(output_path),
+                ok=True,
+                messages_saved=messages_saved,
+            )
+
+        except Exception as exc:
+            return WebSocketProbeResult(
+                stream=stream,
+                output_file=str(output_path),
+                ok=False,
+                messages_saved=messages_saved,
+                error=str(exc),
+            )
+
+    async def probe_depth_request(self) -> WebSocketProbeResult:
+        output_file = f"websocket/depth/{self.config.symbol_file_name}.jsonl"
+
+        request_message: dict[str, Any] = {
+            "correlationId": "probe-depth-request",
+            "destination": f"/api/{self.config.api_version}/depth",
+            "payload": {
+                "symbol": self.config.symbol,
+                "limit": int(self.config.depth_limit),
+            },
+        }
+
+        return await self._subscribe_and_save(
+            stream="wss:/api/v1/depth",
+            subscribe_message=request_message,
+            output_file=output_file,
+            max_messages=1,
+        )