From b35d8c1755173db03eb74301ebae38c4766e1f3d Mon Sep 17 00:00:00 2001 From: Martin Kersner Date: Sun, 5 Jul 2026 10:44:39 +0900 Subject: [PATCH 1/3] feat: async WebSocket ticker client (pilot) on generated WS surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validates the codegen WS pipeline end-to-end for Python. Consumes the generated datamaxi/_ws_endpoints.py (WS_CHANNELS/WS_BASE_PATH/WS_AUTH_HEADER) and datamaxi/_ws_models.py (TickerMessage) — both emitted by datamaxi-codegen from the backend WS surface (route table + protobuf + Go view structs). datamaxi/aio/ws.py: AsyncDatamaxiWS + AsyncWSConnection + TickerChannel on the websockets lib (optional [ws] extra). One connection per channel path; SUBSCRIBE /UNSUBSCRIBE/PING protocol; app-level PING keepalive (~30s vs the ~90s proxy idle); reconnect-with-resubscribe; ack filtering; channel-stream routing (caller filters by msg['s']/['e'] — the wire has no channel tag). async with AsyncDatamaxiWS(api_key) as ws: async for msg in ws.ticker.subscribe('BTC-USDT@binance', market='spot'): ... Tests (tests/test_ws.py) run a real in-process websockets server: auth header, path from the generated registry, subscribe protocol, streamed data, ack filtering, multi-symbol multiplexing, keepalive PING. importorskip-guarded; websockets added to requirements-test. Generated _ws_*.py excluded from black/flake8 like _endpoints.py. import datamaxi does not load websockets. Suite: 205 passed. --- datamaxi/_ws_endpoints.py | 112 +++++++++++ datamaxi/_ws_models.py | 184 ++++++++++++++++++ datamaxi/aio/ws.py | 288 +++++++++++++++++++++++++++++ pyproject.toml | 5 +- requirements/requirements-test.txt | 1 + setup.cfg | 4 +- tests/test_ws.py | 140 ++++++++++++++ 7 files changed, 731 insertions(+), 3 deletions(-) create mode 100644 datamaxi/_ws_endpoints.py create mode 100644 datamaxi/_ws_models.py create mode 100644 datamaxi/aio/ws.py create mode 100644 tests/test_ws.py diff --git a/datamaxi/_ws_endpoints.py b/datamaxi/_ws_endpoints.py new file mode 100644 index 0000000..c43ec82 --- /dev/null +++ b/datamaxi/_ws_endpoints.py @@ -0,0 +1,112 @@ +""" +Auto-generated WebSocket channel registry from the datamaxi-backend WS surface. +DO NOT EDIT — regenerate with: make ws-python + +Source: pkg/apiws/app.go route table + protobuf/*.proto + Go view structs. +The subscribe `param` formats come from overrides/ws_channels.json (the one +piece not derivable from those sources — see parse_ws.py). +""" + +WS_BASE_PATH = "/ws/v1" +WS_AUTH_HEADER = "X-DTMX-APIKEY" + +WS_CHANNELS = { + "/announcement/listing": { + "plan": "pro_plus", + "market": None, + "message": "ListingMessage", + "param": None, + "subscribe": True, + "unsubscribe": True, + }, + "/announcement/listing/internal": { + "plan": "pro_plus", + "market": None, + "message": "InternalListingMessage", + "param": None, + "subscribe": True, + "unsubscribe": True, + }, + "/forex": { + "plan": "basic", + "market": None, + "message": "ForexMessage", + "param": "SYMBOL", + "subscribe": True, + "unsubscribe": True, + "param_note": "fx pair e.g. USD-KRW, no exchange", + }, + "/funding-rate": { + "plan": "basic", + "market": None, + "message": "FundingRateMessage", + "param": "SYMBOL@exchange", + "subscribe": True, + "unsubscribe": True, + "param_note": "UPPER symbol + lower exchange; code/doc disagree — verify live", + }, + "/liquidation": { + "plan": "basic", + "market": None, + "message": "LiquidationMessage", + "param": "SYMBOL@exchange", + "subscribe": True, + "unsubscribe": False, + }, + "/liquidation/feed": { + "plan": "basic", + "market": None, + "message": "LiquidationMessage", + "param": None, + "subscribe": False, + "unsubscribe": False, + }, + "/open-interest": { + "plan": "basic", + "market": None, + "message": "OpenInterestMessage", + "param": "SYMBOL@exchange", + "subscribe": True, + "unsubscribe": False, + }, + "/orderbook/futures": { + "plan": "basic", + "market": "futures", + "message": "OrderbookMessage", + "param": "SYMBOL@exchange[@currency@conversionBase]", + "subscribe": True, + "unsubscribe": True, + }, + "/orderbook/spot": { + "plan": "basic", + "market": "spot", + "message": "OrderbookMessage", + "param": "SYMBOL@exchange[@currency@conversionBase]", + "subscribe": True, + "unsubscribe": True, + }, + "/premium": { + "plan": "basic", + "market": None, + "message": "PremiumMessage", + "param": "src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt", + "subscribe": True, + "unsubscribe": True, + }, + "/ticker/futures": { + "plan": "basic", + "market": "futures", + "message": "TickerMessage", + "param": "SYMBOL@exchange[@currency@conversionBase]", + "subscribe": True, + "unsubscribe": True, + }, + "/ticker/spot": { + "plan": "basic", + "market": "spot", + "message": "TickerMessage", + "param": "SYMBOL@exchange[@currency@conversionBase]", + "subscribe": True, + "unsubscribe": True, + }, +} diff --git a/datamaxi/_ws_models.py b/datamaxi/_ws_models.py new file mode 100644 index 0000000..695e415 --- /dev/null +++ b/datamaxi/_ws_models.py @@ -0,0 +1,184 @@ +""" +Auto-generated WebSocket message models from the datamaxi-backend WS surface. +DO NOT EDIT — regenerate with: make ws-python + +Source: pkg/apiws/app.go route table + protobuf/*.proto + Go view structs. +The subscribe `param` formats come from overrides/ws_channels.json (the one +piece not derivable from those sources — see parse_ws.py). +""" + +from __future__ import annotations + +from typing import Any, Dict, List, TypedDict # noqa: F401 + + +# source: gostruct:pkg/apiforex/app.go::Forex +ForexMessage = TypedDict( + "ForexMessage", + { + "s": str, + "d": int, + "r": float, + }, + total=False, +) + + +# source: gostruct:pkg/apifundingrate/types.go::Snapshot +FundingRateMessage = TypedDict( + "FundingRateMessage", + { + "f": float, + "i": int, + "e": str, + "id": str, + "s": str, + "b": str, + "q": str, + "d": int, + "p": int, + }, + total=False, +) + + +# source: gostruct:pkg/apiannouncement/app.go::InternalListing +InternalListingMessage = TypedDict( + "InternalListingMessage", + { + "s": str, + "e": str, + "b": str, + "t": str, + "u": str, + "d": int, + }, + total=False, +) + + +# source: proto:protobuf/liquidation.proto::liquidation +LiquidationMessage = TypedDict( + "LiquidationMessage", + { + "id": str, + "e": str, + "d": str, + "s": str, + "b": str, + "q": str, + "sd": str, + "p": float, + "pusd": float, + "pfiat": float, + "v": float, + "vusd": float, + "vfiat": float, + "pt": Dict[str, Any], + "pa": str, + "src": Any, + }, + total=False, +) + + +# source: gostruct:pkg/apiannouncement/app.go::Listing +ListingMessage = TypedDict( + "ListingMessage", + { + "e": str, + "b": str, + "q": str, + "u": str, + "d": int, + }, + total=False, +) + + +# source: proto:protobuf/open-interest.proto::open_interest +OpenInterestMessage = TypedDict( + "OpenInterestMessage", + { + "id": str, + "e": str, + "d": str, + "s": str, + "b": str, + "q": str, + "oi": float, + "oiusd": float, + "oifiat": float, + "pt": Dict[str, Any], + "pa": str, + }, + total=False, +) + + +# source: proto:protobuf/orderbook.proto::Orderbook +OrderbookMessage = TypedDict( + "OrderbookMessage", + { + "id": str, + "m": str, + "e": str, + "d": str, + "s": str, + "b": str, + "q": str, + "as": List[Any], + "bs": List[Any], + "pt": Dict[str, Any], + "pa": str, + "ms": bool, + "src": Any, + "snap": bool, + }, + total=False, +) + + +# source: gostruct:pkg/apipremium/handlers/dataapipremiumws/handler.go::PremiumOut +PremiumMessage = TypedDict( + "PremiumMessage", + { + "key": str, + "source_exchange": str, + "target_exchange": str, + "token_id": str, + "source_base": str, + "source_quote": str, + "target_quote": str, + "source_market": str, + "target_market": str, + "premium": float, + "source_price": float, + "target_price": float, + "timestamp": int, + }, + total=False, +) + + +# source: gostruct:pkg/apiticker/app.go::View +TickerMessage = TypedDict( + "TickerMessage", + { + "p": float, + "v": float, + "p24h": float, + "pc": float, + "hb": float, + "la": float, + "ud": float, + "ld": float, + "e": str, + "s": str, + "b": str, + "q": str, + "d": int, + "m": str, + }, + total=False, +) diff --git a/datamaxi/aio/ws.py b/datamaxi/aio/ws.py new file mode 100644 index 0000000..4517ba1 --- /dev/null +++ b/datamaxi/aio/ws.py @@ -0,0 +1,288 @@ +"""Async WebSocket client (Phase-1 pilot — ticker). + +Streams real-time data over the DataMaxi+ WebSocket API. Requires the ``ws`` +extra:: + + pip install "datamaxi[ws]" + +Usage:: + + from datamaxi.aio.ws import AsyncDatamaxiWS + + async with AsyncDatamaxiWS(api_key="...") as ws: + async for msg in ws.ticker.subscribe("BTC-USDT@binance", market="spot"): + print(msg["s"], msg["p"]) + +This is a deliberately small pilot (ticker) that consumes the **generated** +WS surface — ``WS_CHANNELS`` / ``WS_BASE_PATH`` / ``WS_AUTH_HEADER`` from +``datamaxi._ws_endpoints`` and the ``TickerMessage`` model from +``datamaxi._ws_models`` (both emitted by datamaxi-codegen from the backend). +The goal is to validate that the generated artifacts are sufficient to drive a +real client before expanding to every channel. + +Routing model: one connection per channel path, multiplexing all subscribed +``params``. ``subscribe()`` returns an async iterator over **every** message on +that channel — the WS protocol tags messages by payload fields (``s``/``e``), +not by a channel id, so callers filter by ``msg["s"]`` when subscribing to +multiple symbols on one connection. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from typing import Any, AsyncIterator, Dict, List, Optional, TYPE_CHECKING + +from datamaxi.__version__ import __version__ +from datamaxi._ws_endpoints import WS_CHANNELS, WS_BASE_PATH, WS_AUTH_HEADER + +if TYPE_CHECKING: + from datamaxi._ws_models import TickerMessage + +_DEFAULT_WS_URL = "wss://api.datamaxiplus.com" +# Send an app-level PING within the ~90s openresty proxy idle timeout. +_KEEPALIVE_INTERVAL = 30.0 +_RECONNECT_BACKOFF = 1.0 +_CLOSED = object() # sentinel pushed to subscriber queues on shutdown + + +def _import_websockets(): + try: + import websockets + except ImportError as exc: # pragma: no cover - exercised via extra + raise ImportError( + "The WebSocket client requires websockets. Install it with: " + "pip install 'datamaxi[ws]'" + ) from exc + return websockets + + +def _derive_ws_url(base_url: Optional[str]) -> str: + if not base_url: + return _DEFAULT_WS_URL + if base_url.startswith("https://"): + return "wss://" + base_url.removeprefix("https://") + if base_url.startswith("http://"): + return "ws://" + base_url.removeprefix("http://") + return base_url + + +class AsyncWSConnection: + """One WebSocket connection to a single channel path. + + Owns the SUBSCRIBE / UNSUBSCRIBE / PING protocol, a reader that fans each + incoming data message out to every subscriber stream, an app-level PING + keepalive, and reconnect-with-resubscribe on a dropped connection. + """ + + def __init__( + self, + url: str, + api_key: Optional[str], + keepalive: float = _KEEPALIVE_INTERVAL, + reconnect: bool = True, + connect_kwargs: Optional[dict] = None, + ): + self._url = url + self._api_key = api_key + self._keepalive = keepalive + self._reconnect = reconnect + self._connect_kwargs = connect_kwargs or {} + self._ws = None + self._websockets = None + self._id = 0 + self._active: set = set() # params to replay on reconnect + self._subscribers: List[asyncio.Queue] = [] + self._reader_task: Optional[asyncio.Task] = None + self._keepalive_task: Optional[asyncio.Task] = None + self._closed = False + + async def start(self) -> None: + await self._open() + self._reader_task = asyncio.create_task(self._reader()) + if self._keepalive: + self._keepalive_task = asyncio.create_task(self._keepalive_loop()) + + async def _open(self) -> None: + self._websockets = _import_websockets() + headers = { + WS_AUTH_HEADER: str(self._api_key), + "User-Agent": "datamaxi/" + __version__, + } + # No Origin header: coder/websocket enforces same-origin otherwise. + self._ws = await self._websockets.connect( + self._url, additional_headers=headers, **self._connect_kwargs + ) + if self._active: # resubscribe after a reconnect + await self._send( + { + "method": "SUBSCRIBE", + "params": sorted(self._active), + "id": self._next_id(), + } + ) + + def _next_id(self) -> int: + self._id += 1 + return self._id + + async def _send(self, obj: dict) -> None: + await self._ws.send(json.dumps(obj)) + + async def subscribe(self, params: List[str]) -> None: + self._active.update(params) + await self._send( + {"method": "SUBSCRIBE", "params": list(params), "id": self._next_id()} + ) + + async def unsubscribe(self, params: List[str]) -> None: + for p in params: + self._active.discard(p) + await self._send( + {"method": "UNSUBSCRIBE", "params": list(params), "id": self._next_id()} + ) + + def stream(self) -> AsyncIterator[Dict[str, Any]]: + """Register a subscriber queue *now* and return an iterator over it.""" + q: asyncio.Queue = asyncio.Queue() + self._subscribers.append(q) + return self._drain(q) + + async def _drain(self, q: asyncio.Queue) -> AsyncIterator[Dict[str, Any]]: + try: + while True: + item = await q.get() + if item is _CLOSED: + return + yield item + finally: + if q in self._subscribers: + self._subscribers.remove(q) + + async def _reader(self) -> None: + while not self._closed: + try: + raw = await self._ws.recv() + except self._websockets.ConnectionClosed: + if self._closed or not self._reconnect: + break + await asyncio.sleep(_RECONNECT_BACKOFF) + try: + await self._open() + except Exception: + continue + else: + continue + msg = json.loads(raw) + # Subscription acks are {"result": [...], "id": N} — not data. + if isinstance(msg, dict) and "result" in msg and "id" in msg: + continue + for q in list(self._subscribers): + q.put_nowait(msg) + for q in list(self._subscribers): + q.put_nowait(_CLOSED) + + async def _keepalive_loop(self) -> None: + while not self._closed: + await asyncio.sleep(self._keepalive) + try: + await self._send({"method": "PING"}) + except Exception: + pass + + async def close(self) -> None: + self._closed = True + for task in (self._reader_task, self._keepalive_task): + if task is not None: + task.cancel() + if self._ws is not None: + await self._ws.close() + for q in list(self._subscribers): + q.put_nowait(_CLOSED) + + +class TickerChannel: + """`ws.ticker` — subscribe to real-time ticker updates. + + Resolves the channel path from the generated ``WS_CHANNELS`` registry by + market (``/ticker/spot`` or ``/ticker/futures``). + """ + + def __init__(self, client: "AsyncDatamaxiWS"): + self._client = client + + def _path(self, market: str) -> str: + path = f"/ticker/{market}" + if path not in WS_CHANNELS: + valid = [p for p in WS_CHANNELS if p.startswith("/ticker/")] + raise ValueError(f"unknown ticker market {market!r}; valid paths: {valid}") + return path + + async def subscribe( + self, *params: str, market: str = "spot" + ) -> AsyncIterator["TickerMessage"]: + """Subscribe to ``params`` (``SYMBOL@exchange[@currency@conversionBase]``). + + Returns an async iterator over ticker messages on the channel. + """ + conn = await self._client._conn(self._path(market)) + stream = conn.stream() # register the queue before SUBSCRIBE (no missed msgs) + await conn.subscribe(list(params)) + return stream + + +class AsyncDatamaxiWS: + """Async WebSocket entrypoint (pilot). Exposes ``ticker``. + + Use as an async context manager so open connections are closed, or call + :meth:`aclose` explicitly. + """ + + def __init__( + self, + api_key: Optional[str] = None, + ws_url: Optional[str] = None, + base_url: Optional[str] = None, + keepalive: float = _KEEPALIVE_INTERVAL, + reconnect: bool = True, + connect_kwargs: Optional[dict] = None, + ): + self.api_key = api_key or os.environ.get("DATAMAXI_API_KEY") + self.ws_url = ws_url or _derive_ws_url(base_url) + self._keepalive = keepalive + self._reconnect = reconnect + self._connect_kwargs = connect_kwargs + self._conns: Dict[str, AsyncWSConnection] = {} + self.ticker = TickerChannel(self) + + async def _conn(self, path: str) -> AsyncWSConnection: + conn = self._conns.get(path) + if conn is None: + url = self.ws_url + WS_BASE_PATH + path + conn = AsyncWSConnection( + url, + self.api_key, + keepalive=self._keepalive, + reconnect=self._reconnect, + connect_kwargs=self._connect_kwargs, + ) + await conn.start() + self._conns[path] = conn + return conn + + async def aclose(self) -> None: + for conn in list(self._conns.values()): + await conn.close() + self._conns.clear() + + async def __aenter__(self) -> "AsyncDatamaxiWS": + return self + + async def __aexit__(self, *exc) -> None: + await self.aclose() + + def __repr__(self) -> str: + return "AsyncDatamaxiWS(ws_url={!r}, has_key={})".format( + self.ws_url, bool(self.api_key) + ) diff --git a/pyproject.toml b/pyproject.toml index 971bd9f..a10fb01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dynamic = ["dependencies", "version"] [project.optional-dependencies] async = ["httpx>=0.27,<1"] +ws = ["websockets>=13,<17"] [tool.setuptools.dynamic] dependencies = {file = ["requirements/common.txt"]} @@ -33,5 +34,5 @@ Homepage = "https://github.com/bisonai/datamaxi-python" Issues = "https://github.com/bisonai/datamaxi-python/issues" [tool.black] -# datamaxi/_endpoints.py is generated by datamaxi-codegen; don't reformat it. -force-exclude = 'datamaxi/_endpoints\.py' +# Generated by datamaxi-codegen; don't reformat. +force-exclude = 'datamaxi/(_endpoints|_ws_endpoints|_ws_models)\.py' diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt index 0a43e44..8f620bd 100644 --- a/requirements/requirements-test.txt +++ b/requirements/requirements-test.txt @@ -3,6 +3,7 @@ pytest-cov>=5 pytest>=8 responses>=0.25 httpx>=0.27,<1 +websockets>=13,<17 black==26.3.1 flake8==7.3.0 wheel==0.46.3 diff --git a/setup.cfg b/setup.cfg index 16aa1b2..2f4a75f 100644 --- a/setup.cfg +++ b/setup.cfg @@ -19,6 +19,8 @@ exclude = .git, build, dist, - datamaxi/_endpoints.py + datamaxi/_endpoints.py, + datamaxi/_ws_endpoints.py, + datamaxi/_ws_models.py max-complexity = 10 ignore = E501, W503, W504 diff --git a/tests/test_ws.py b/tests/test_ws.py new file mode 100644 index 0000000..67b1099 --- /dev/null +++ b/tests/test_ws.py @@ -0,0 +1,140 @@ +"""Local tests for the async WebSocket ticker client (pilot). + +Runs a real in-process `websockets` server implementing the subscribe protocol, +so the client is exercised end-to-end (auth header, path, SUBSCRIBE, streamed +data, ack filtering, keepalive PING). Skipped when the optional `websockets` +dependency is absent. +""" + +import asyncio +import json + +import pytest + +websockets = pytest.importorskip("websockets") + +from datamaxi.aio.ws import AsyncDatamaxiWS # noqa: E402 +from datamaxi._ws_endpoints import WS_CHANNELS, WS_BASE_PATH # noqa: E402 +from datamaxi._ws_models import TickerMessage # noqa: E402 + + +def _serve(handler): + return websockets.serve(handler, "localhost", 0) + + +def _port(server): + return server.sockets[0].getsockname()[1] + + +async def _first(stream, timeout=2.0): + return await asyncio.wait_for(stream.__anext__(), timeout) + + +def _run(coro): + return asyncio.run(coro) + + +def test_ws_consumes_generated_registry_and_model(): + # The client resolves the path from the generated WS_CHANNELS, and the + # generated TickerMessage carries the wire keys. + assert "/ticker/spot" in WS_CHANNELS + assert WS_CHANNELS["/ticker/spot"]["message"] == "TickerMessage" + assert WS_BASE_PATH == "/ws/v1" + assert {"p", "e", "s", "d"} <= set(TickerMessage.__annotations__) + + +def test_ws_ticker_subscribe_streams_and_filters_ack(): + seen = {} + + async def handler(conn): + seen["apikey"] = conn.request.headers.get("X-DTMX-APIKEY") + seen["path"] = conn.request.path + async for raw in conn: + m = json.loads(raw) + if m.get("method") == "SUBSCRIBE": + seen["params"] = m["params"] + await conn.send(json.dumps({"result": m["params"], "id": m["id"]})) + for p in m["params"]: + sym, exch = p.split("@")[0], p.split("@")[1] + await conn.send( + json.dumps({"s": sym, "e": exch, "p": 105.5, "d": 1}) + ) + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.ticker.subscribe("BTC-USDT@binance", market="spot") + return await _first(stream) + + msg = _run(run()) + assert seen["apikey"] == "k" # X-DTMX-APIKEY on the handshake + assert seen["path"] == "/ws/v1/ticker/spot" # base path + channel path + assert seen["params"] == ["BTC-USDT@binance"] # SUBSCRIBE protocol + # first yielded item is data, not the {"result":...,"id":...} ack + assert msg == {"s": "BTC-USDT", "e": "binance", "p": 105.5, "d": 1} + + +def test_ws_multi_symbol_multiplexed_on_one_connection(): + async def handler(conn): + async for raw in conn: + m = json.loads(raw) + if m.get("method") == "SUBSCRIBE": + await conn.send(json.dumps({"result": m["params"], "id": m["id"]})) + for p in m["params"]: + sym = p.split("@")[0] + await conn.send( + json.dumps({"s": sym, "e": "binance", "p": 1.0, "d": 1}) + ) + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.ticker.subscribe( + "BTC-USDT@binance", "ETH-USDT@binance", market="spot" + ) + got = {(await _first(stream))["s"], (await _first(stream))["s"]} + return got + + assert _run(run()) == {"BTC-USDT", "ETH-USDT"} + + +def test_ws_keepalive_sends_ping(): + got_ping = asyncio.Event() + + async def handler(conn): + async for raw in conn: + if json.loads(raw).get("method") == "PING": + got_ping.set() + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", + ws_url=f"ws://localhost:{_port(server)}", + keepalive=0.1, + ) as ws: + await ws.ticker.subscribe("BTC-USDT@binance", market="spot") + await asyncio.wait_for(got_ping.wait(), timeout=2.0) + return True + + assert _run(run()) is True + + +def test_ws_invalid_market_raises(): + async def run(): + async with AsyncDatamaxiWS(api_key="k", ws_url="ws://localhost:1") as ws: + await ws.ticker.subscribe("BTC-USDT@binance", market="bogus") + + with pytest.raises(ValueError): + _run(run()) + + +def test_ws_repr_no_key_leak(): + ws = AsyncDatamaxiWS(api_key="secret", base_url="https://api.datamaxiplus.com") + assert ws.ws_url == "wss://api.datamaxiplus.com" + assert "secret" not in repr(ws) + assert "has_key=True" in repr(ws) From 9acff43ddac3df7e15fffec793ee077495844395 Mon Sep 17 00:00:00 2001 From: Martin Kersner Date: Sun, 5 Jul 2026 11:03:14 +0900 Subject: [PATCH 2/3] fix: filter empty-result subscribe acks in WS client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live testing surfaced that when the accepted param list is empty (e.g. the announcement channel), the server omits `result` and sends just {"id": N}. The ack filter only matched {"result":..., "id":...}, so {"id": N} leaked as a data message. Detect an ack as any dict whose keys are a subset of {"result", "id"} (data payloads always carry other fields — and may themselves include an "id" token, so presence of "id" alone isn't enough). Verified live against the announcement channel. --- datamaxi/aio/ws.py | 8 ++++++-- tests/test_ws.py | 25 +++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/datamaxi/aio/ws.py b/datamaxi/aio/ws.py index 4517ba1..82d7c02 100644 --- a/datamaxi/aio/ws.py +++ b/datamaxi/aio/ws.py @@ -175,8 +175,12 @@ async def _reader(self) -> None: else: continue msg = json.loads(raw) - # Subscription acks are {"result": [...], "id": N} — not data. - if isinstance(msg, dict) and "result" in msg and "id" in msg: + # Subscription acks are {"result": [...], "id": N}; when the accepted + # param list is empty the server omits `result`, leaving just + # {"id": N}. Data payloads always carry other fields (s/e/d/...) — + # note they may also include an "id" (token id), so detect an ack as + # any dict whose keys are a subset of {"result", "id"}. + if isinstance(msg, dict) and set(msg) <= {"result", "id"}: continue for q in list(self._subscribers): q.put_nowait(msg) diff --git a/tests/test_ws.py b/tests/test_ws.py index 67b1099..92db282 100644 --- a/tests/test_ws.py +++ b/tests/test_ws.py @@ -76,6 +76,31 @@ async def run(): assert msg == {"s": "BTC-USDT", "e": "binance", "p": 105.5, "d": 1} +def test_ws_empty_result_ack_is_filtered(): + # When the accepted param list is empty the server omits `result`, sending + # just {"id": N}. That must be filtered, not yielded as data. (A data + # payload can itself carry an "id" token, e.g. {"id": "bitcoin", ...}.) + async def handler(conn): + async for raw in conn: + m = json.loads(raw) + if m.get("method") == "SUBSCRIBE": + await conn.send(json.dumps({"id": m["id"]})) # empty-result ack + await conn.send( + json.dumps({"id": "bitcoin", "s": "BTC-USDT", "p": 1.0}) + ) + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.ticker.subscribe("BTC-USDT@binance", market="spot") + return await _first(stream) + + msg = _run(run()) + assert msg == {"id": "bitcoin", "s": "BTC-USDT", "p": 1.0} # data, not the ack + + def test_ws_multi_symbol_multiplexed_on_one_connection(): async def handler(conn): async for raw in conn: From e45afd52cdfb450ce6f05f4ee23bf25f47559dc9 Mon Sep 17 00:00:00 2001 From: Martin Kersner Date: Sun, 5 Jul 2026 11:16:11 +0900 Subject: [PATCH 3/3] feat: generalize WS client to all data types; exclude orderbook Expand the async WS client from ticker-only to every DataMaxi+ WS data type, driven by the generated WS_CHANNELS registry. Accessors: ticker (market-keyed), forex, premium, funding_rate, open_interest, liquidation (subscribe), liquidation_feed (firehose), announcement / announcement_internal (Pro+). Generic channel types (Subscription / MarketSubscription / Feed) over the same AsyncWSConnection transport. Regenerate _ws_endpoints.py + _ws_models.py from codegen with orderbook excluded (unsupported product, backend removal tracked in Bisonai/datamaxi-backend#7927) -> 10 channels / 8 models. Tests: orderbook-excluded, every channel maps to a generated model + a client accessor, param_format from registry, plus forex-subscribe and feed-stream over the mock server. Live-verified through the client against prod (ticker, open_interest, liquidation_feed, premium). Suite: 211 passed. --- datamaxi/_ws_endpoints.py | 18 +----- datamaxi/_ws_models.py | 23 -------- datamaxi/aio/ws.py | 117 +++++++++++++++++++++++++++++--------- tests/test_ws.py | 71 +++++++++++++++++++++++ 4 files changed, 162 insertions(+), 67 deletions(-) diff --git a/datamaxi/_ws_endpoints.py b/datamaxi/_ws_endpoints.py index c43ec82..267ce65 100644 --- a/datamaxi/_ws_endpoints.py +++ b/datamaxi/_ws_endpoints.py @@ -43,7 +43,7 @@ "param": "SYMBOL@exchange", "subscribe": True, "unsubscribe": True, - "param_note": "UPPER symbol + lower exchange; code/doc disagree — verify live", + "param_note": "UPPER symbol + lower exchange; verified live (handler.go doc comment is wrong — Bisonai/datamaxi-backend#7926)", }, "/liquidation": { "plan": "basic", @@ -69,22 +69,6 @@ "subscribe": True, "unsubscribe": False, }, - "/orderbook/futures": { - "plan": "basic", - "market": "futures", - "message": "OrderbookMessage", - "param": "SYMBOL@exchange[@currency@conversionBase]", - "subscribe": True, - "unsubscribe": True, - }, - "/orderbook/spot": { - "plan": "basic", - "market": "spot", - "message": "OrderbookMessage", - "param": "SYMBOL@exchange[@currency@conversionBase]", - "subscribe": True, - "unsubscribe": True, - }, "/premium": { "plan": "basic", "market": None, diff --git a/datamaxi/_ws_models.py b/datamaxi/_ws_models.py index 695e415..3aa9735 100644 --- a/datamaxi/_ws_models.py +++ b/datamaxi/_ws_models.py @@ -116,29 +116,6 @@ ) -# source: proto:protobuf/orderbook.proto::Orderbook -OrderbookMessage = TypedDict( - "OrderbookMessage", - { - "id": str, - "m": str, - "e": str, - "d": str, - "s": str, - "b": str, - "q": str, - "as": List[Any], - "bs": List[Any], - "pt": Dict[str, Any], - "pa": str, - "ms": bool, - "src": Any, - "snap": bool, - }, - total=False, -) - - # source: gostruct:pkg/apipremium/handlers/dataapipremiumws/handler.go::PremiumOut PremiumMessage = TypedDict( "PremiumMessage", diff --git a/datamaxi/aio/ws.py b/datamaxi/aio/ws.py index 82d7c02..c4c0199 100644 --- a/datamaxi/aio/ws.py +++ b/datamaxi/aio/ws.py @@ -1,4 +1,4 @@ -"""Async WebSocket client (Phase-1 pilot — ticker). +"""Async WebSocket client for every DataMaxi+ WS data type. Streams real-time data over the DataMaxi+ WebSocket API. Requires the ``ws`` extra:: @@ -13,12 +13,18 @@ async for msg in ws.ticker.subscribe("BTC-USDT@binance", market="spot"): print(msg["s"], msg["p"]) -This is a deliberately small pilot (ticker) that consumes the **generated** -WS surface — ``WS_CHANNELS`` / ``WS_BASE_PATH`` / ``WS_AUTH_HEADER`` from -``datamaxi._ws_endpoints`` and the ``TickerMessage`` model from -``datamaxi._ws_models`` (both emitted by datamaxi-codegen from the backend). -The goal is to validate that the generated artifacts are sufficient to drive a -real client before expanding to every channel. + async for oi in ws.open_interest.subscribe("BTC-USDT@binance"): + ... + +Channels (accessors driven by the generated ``WS_CHANNELS`` registry): +``ticker`` (market-keyed), ``forex``, ``premium``, ``funding_rate``, +``open_interest``, ``liquidation`` (subscribe), ``liquidation_feed`` +(firehose, no params), ``announcement`` / ``announcement_internal`` (Pro+). + +Consumes the **generated** WS surface — ``WS_CHANNELS`` / ``WS_BASE_PATH`` / +``WS_AUTH_HEADER`` from ``datamaxi._ws_endpoints`` and the per-channel message +models from ``datamaxi._ws_models`` (both emitted by datamaxi-codegen from the +backend). Orderbook is intentionally excluded (unsupported product). Routing model: one connection per channel path, multiplexing all subscribed ``params``. ``subscribe()`` returns an async iterator over **every** message on @@ -32,14 +38,11 @@ import asyncio import json import os -from typing import Any, AsyncIterator, Dict, List, Optional, TYPE_CHECKING +from typing import Any, AsyncIterator, Dict, List, Optional from datamaxi.__version__ import __version__ from datamaxi._ws_endpoints import WS_CHANNELS, WS_BASE_PATH, WS_AUTH_HEADER -if TYPE_CHECKING: - from datamaxi._ws_models import TickerMessage - _DEFAULT_WS_URL = "wss://api.datamaxiplus.com" # Send an app-level PING within the ~90s openresty proxy idle timeout. _KEEPALIVE_INTERVAL = 30.0 @@ -206,38 +209,87 @@ async def close(self) -> None: q.put_nowait(_CLOSED) -class TickerChannel: - """`ws.ticker` — subscribe to real-time ticker updates. +def _require_channel(path: str) -> str: + if path not in WS_CHANNELS: + raise ValueError( + f"unknown WS channel {path!r}; generated: {sorted(WS_CHANNELS)}" + ) + return path - Resolves the channel path from the generated ``WS_CHANNELS`` registry by - market (``/ticker/spot`` or ``/ticker/futures``). + +class Subscription: + """A single-path subscribable channel (``ws.forex``, ``ws.premium``, ...). + + ``param_format`` echoes the generated registry's subscribe param format for + reference; callers pass the raw param strings (e.g. ``"BTC-USDT@binance"``). """ - def __init__(self, client: "AsyncDatamaxiWS"): + def __init__(self, client: "AsyncDatamaxiWS", path: str): + self._client = client + self._path = _require_channel(path) + + @property + def param_format(self) -> Optional[str]: + return WS_CHANNELS[self._path].get("param") + + async def subscribe(self, *params: str) -> AsyncIterator[Dict[str, Any]]: + """SUBSCRIBE to ``params``; return an async iterator over the channel.""" + conn = await self._client._conn(self._path) + stream = conn.stream() # register the queue before SUBSCRIBE (no missed msgs) + await conn.subscribe(list(params)) + return stream + + async def unsubscribe(self, *params: str) -> None: + conn = await self._client._conn(self._path) + await conn.unsubscribe(list(params)) + + +class MarketSubscription: + """A market-keyed subscribable channel (``ws.ticker``) → ``/``.""" + + def __init__(self, client: "AsyncDatamaxiWS", base: str): self._client = client + self._base = base def _path(self, market: str) -> str: - path = f"/ticker/{market}" + path = f"{self._base}/{market}" if path not in WS_CHANNELS: - valid = [p for p in WS_CHANNELS if p.startswith("/ticker/")] - raise ValueError(f"unknown ticker market {market!r}; valid paths: {valid}") + valid = [p for p in WS_CHANNELS if p.startswith(self._base + "/")] + raise ValueError(f"unknown market {market!r}; valid paths: {valid}") return path async def subscribe( self, *params: str, market: str = "spot" - ) -> AsyncIterator["TickerMessage"]: - """Subscribe to ``params`` (``SYMBOL@exchange[@currency@conversionBase]``). - - Returns an async iterator over ticker messages on the channel. - """ + ) -> AsyncIterator[Dict[str, Any]]: conn = await self._client._conn(self._path(market)) - stream = conn.stream() # register the queue before SUBSCRIBE (no missed msgs) + stream = conn.stream() await conn.subscribe(list(params)) return stream + async def unsubscribe(self, *params: str, market: str = "spot") -> None: + conn = await self._client._conn(self._path(market)) + await conn.unsubscribe(list(params)) + + +class Feed: + """A firehose channel (``ws.liquidation_feed``) — auto-subscribed, no params.""" + + def __init__(self, client: "AsyncDatamaxiWS", path: str): + self._client = client + self._path = _require_channel(path) + + async def stream(self) -> AsyncIterator[Dict[str, Any]]: + conn = await self._client._conn(self._path) + return conn.stream() + class AsyncDatamaxiWS: - """Async WebSocket entrypoint (pilot). Exposes ``ticker``. + """Async WebSocket entrypoint — every DataMaxi+ WS data type. + + Accessors are driven by the generated ``WS_CHANNELS`` registry: + ``ticker`` (market-keyed), ``forex``, ``premium``, ``funding_rate``, + ``open_interest``, ``liquidation`` (subscribe), ``liquidation_feed`` + (firehose), ``announcement`` / ``announcement_internal`` (Pro+). Use as an async context manager so open connections are closed, or call :meth:`aclose` explicitly. @@ -258,7 +310,18 @@ def __init__( self._reconnect = reconnect self._connect_kwargs = connect_kwargs self._conns: Dict[str, AsyncWSConnection] = {} - self.ticker = TickerChannel(self) + + self.ticker = MarketSubscription(self, "/ticker") + self.forex = Subscription(self, "/forex") + self.premium = Subscription(self, "/premium") + self.funding_rate = Subscription(self, "/funding-rate") + self.open_interest = Subscription(self, "/open-interest") + self.liquidation = Subscription(self, "/liquidation") + self.liquidation_feed = Feed(self, "/liquidation/feed") + self.announcement = Subscription(self, "/announcement/listing") + self.announcement_internal = Subscription( + self, "/announcement/listing/internal" + ) async def _conn(self, path: str) -> AsyncWSConnection: conn = self._conns.get(path) diff --git a/tests/test_ws.py b/tests/test_ws.py index 92db282..61de1e3 100644 --- a/tests/test_ws.py +++ b/tests/test_ws.py @@ -13,6 +13,7 @@ websockets = pytest.importorskip("websockets") +import datamaxi._ws_models as _ws_models # noqa: E402 from datamaxi.aio.ws import AsyncDatamaxiWS # noqa: E402 from datamaxi._ws_endpoints import WS_CHANNELS, WS_BASE_PATH # noqa: E402 from datamaxi._ws_models import TickerMessage # noqa: E402 @@ -43,6 +44,76 @@ def test_ws_consumes_generated_registry_and_model(): assert {"p", "e", "s", "d"} <= set(TickerMessage.__annotations__) +def test_ws_orderbook_excluded(): + assert not any("orderbook" in p for p in WS_CHANNELS) + + +def test_ws_every_channel_has_a_generated_model_and_accessor(): + # Every generated channel maps to a real model in _ws_models... + for path, ch in WS_CHANNELS.items(): + assert hasattr(_ws_models, ch["message"]), f"{path}: missing {ch['message']}" + # ...and the client exposes an accessor for every data type. + ws = AsyncDatamaxiWS(api_key="k") + for name in ( + "ticker", + "forex", + "premium", + "funding_rate", + "open_interest", + "liquidation", + "liquidation_feed", + "announcement", + "announcement_internal", + ): + assert hasattr(ws, name), f"client missing ws.{name}" + + +def test_ws_param_format_from_registry(): + ws = AsyncDatamaxiWS(api_key="k") + assert ws.premium.param_format == "src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt" + assert ws.forex.param_format == "SYMBOL" + + +def test_ws_forex_subscribe_streams(): + async def handler(conn): + seen = None + async for raw in conn: + m = json.loads(raw) + if m.get("method") == "SUBSCRIBE": + seen = (conn.request.path, m["params"]) + await conn.send(json.dumps({"result": m["params"], "id": m["id"]})) + await conn.send(json.dumps({"s": "USD-KRW", "d": 1, "r": 1530.0})) + assert seen == ("/ws/v1/forex", ["USD-KRW"]) + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.forex.subscribe("USD-KRW") + return await _first(stream) + + assert _run(run()) == {"s": "USD-KRW", "d": 1, "r": 1530.0} + + +def test_ws_liquidation_feed_streams_without_subscribe(): + async def handler(conn): + # firehose: push immediately, no SUBSCRIBE expected + await conn.send(json.dumps({"s": "RPL-USDT", "sd": "sell", "p": 2.2})) + async for _ in conn: + pass + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.liquidation_feed.stream() + return await _first(stream) + + assert _run(run())["s"] == "RPL-USDT" + + def test_ws_ticker_subscribe_streams_and_filters_ack(): seen = {}