chainstream-sdk 0.1.0__py3-none-any.whl

chainstream/stream/client.py

@@ -0,0 +1,761 @@
+"""Centrifuge WebSocket client implementation for ChainStream Stream API."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, TypeVar
+
+import aiohttp
+
+from chainstream.stream.fields import replace_filter_fields
+from chainstream.stream.models import (
+    DexPoolBalance,
+    MetricType,
+    NewToken,
+    RankingTokenList,
+    RankingType,
+    Resolution,
+    TokenCandle,
+    TokenHolder,
+    TokenLiquidity,
+    TokenMaxLiquidity,
+    TokenMetadata,
+    TokenStat,
+    TokenSupply,
+    TokenTotalLiquidity,
+    TradeActivity,
+    WalletBalance,
+    WalletTokenPnl,
+)
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class Unsubscribe:
+    """Unsubscribe handle for stream subscriptions."""
+
+    channel: str
+    callback_id: int
+    _api: "StreamApi"
+
+    def unsubscribe(self) -> None:
+        """Unsubscribe from the channel."""
+        self._api._remove_listener(self.channel, self.callback_id)
+
+
+class StreamApi:
+    """Stream API client for real-time data subscriptions using Centrifuge protocol."""
+
+    def __init__(self, url: str, access_token: str) -> None:
+        """Initialize the StreamApi.
+
+        Args:
+            url: The WebSocket URL for the stream server.
+            access_token: The access token for authentication.
+        """
+        # Build URL with token
+        if "?" in url:
+            self._url = f"{url}&token={access_token}"
+        else:
+            self._url = f"{url}?token={access_token}"
+
+        self._access_token = access_token
+        self._connected = False
+        self._command_id = 0
+        self._callback_id = 0
+        self._listeners: Dict[str, List[Dict[str, Any]]] = {}
+        self._subscriptions: Dict[str, int] = {}
+        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._read_task: Optional[asyncio.Task] = None
+        self._lock = asyncio.Lock()
+
+    def _next_command_id(self) -> int:
+        """Get the next command ID."""
+        self._command_id += 1
+        return self._command_id
+
+    def _next_callback_id(self) -> int:
+        """Get the next callback ID."""
+        self._callback_id += 1
+        return self._callback_id
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the client is connected."""
+        return self._connected
+
+    async def connect(self) -> None:
+        """Connect to the WebSocket server."""
+        async with self._lock:
+            if self._connected:
+                return
+
+            self._session = aiohttp.ClientSession()
+            self._ws = await self._session.ws_connect(self._url)
+
+            # Send connect command with token
+            connect_cmd = {
+                "id": self._next_command_id(),
+                "connect": {"token": self._access_token},
+            }
+            await self._ws.send_json(connect_cmd)
+
+            self._connected = True
+
+            # Start read task
+            self._read_task = asyncio.create_task(self._read_loop())
+
+            # Give the read task a chance to start
+            await asyncio.sleep(0.1)
+
+    async def _read_loop(self) -> None:
+        """Read messages from the WebSocket."""
+        if not self._ws:
+            return
+
+        try:
+            async for msg in self._ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    try:
+                        response = json.loads(msg.data)
+                        await self._handle_response(response)
+                    except json.JSONDecodeError:
+                        pass  # Ignore malformed messages
+                elif msg.type == aiohttp.WSMsgType.CLOSED:
+                    logger.info("[streaming] connection closed")
+                    self._connected = False
+                    break
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    logger.error("[streaming] connection error")
+                    self._connected = False
+                    break
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.error(f"[streaming] read error: {e}")
+            self._connected = False
+
+    async def _handle_response(self, response: Dict[str, Any]) -> None:
+        """Handle a response from the server."""
+        # Handle push messages (publications)
+        if "push" in response:
+            push = response["push"]
+            channel = push.get("channel", "")
+            pub = push.get("pub", {})
+            if pub and "data" in pub:
+                self._dispatch_message(channel, pub["data"])
+
+        # Handle errors
+        if "error" in response:
+            error = response["error"]
+            logger.error(
+                f"[streaming] error: code={error.get('code')}, message={error.get('message')}"
+            )
+
+        # Connect and subscribe responses are handled silently
+
+    def _dispatch_message(self, channel: str, data: Any) -> None:
+        """Dispatch a message to listeners."""
+        listeners = self._listeners.get(channel, [])
+        for listener in listeners:
+            try:
+                listener["callback"](data)
+            except Exception as e:
+                logger.error(f"[streaming] callback error: {e}")
+
+    async def _send_subscribe(self, channel: str, filter_expr: Optional[str] = None) -> None:
+        """Send a subscribe command."""
+        if not self._ws:
+            return
+
+        cmd: Dict[str, Any] = {
+            "id": self._next_command_id(),
+            "subscribe": {
+                "channel": channel,
+                "delta": "fossil",
+            },
+        }
+
+        if filter_expr:
+            cmd["subscribe"]["filter"] = filter_expr
+
+        await self._ws.send_json(cmd)
+
+    async def _send_unsubscribe(self, channel: str) -> None:
+        """Send an unsubscribe command."""
+        if not self._ws:
+            return
+
+        cmd = {
+            "id": self._next_command_id(),
+            "unsubscribe": {"channel": channel},
+        }
+        await self._ws.send_json(cmd)
+
+    def _add_listener(self, channel: str, callback: Callable[[Any], None]) -> int:
+        """Add a listener for a channel."""
+        callback_id = self._next_callback_id()
+
+        if channel not in self._listeners:
+            self._listeners[channel] = []
+
+        self._listeners[channel].append({"id": callback_id, "callback": callback})
+        return callback_id
+
+    def _remove_listener(self, channel: str, callback_id: int) -> None:
+        """Remove a listener from a channel."""
+        if channel not in self._listeners:
+            return
+
+        self._listeners[channel] = [
+            l for l in self._listeners[channel] if l["id"] != callback_id
+        ]
+
+        # If no more listeners, unsubscribe
+        if not self._listeners[channel]:
+            del self._listeners[channel]
+            asyncio.create_task(self._send_unsubscribe(channel))
+            if channel in self._subscriptions:
+                del self._subscriptions[channel]
+            logger.info(f"[streaming] unsubscribed from channel: {channel}")
+
+    async def disconnect(self) -> None:
+        """Disconnect from the WebSocket server."""
+        async with self._lock:
+            if self._read_task:
+                self._read_task.cancel()
+                try:
+                    await self._read_task
+                except asyncio.CancelledError:
+                    pass
+                self._read_task = None
+
+            if self._ws:
+                await self._ws.close()
+                self._ws = None
+
+            if self._session:
+                await self._session.close()
+                self._session = None
+
+            self._connected = False
+            logger.info("[streaming] disconnected")
+
+    async def subscribe(
+        self,
+        channel: str,
+        callback: Callable[[Any], None],
+        filter_expr: Optional[str] = None,
+        method_name: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to a channel with a raw callback.
+
+        Args:
+            channel: The channel to subscribe to.
+            callback: The callback function to call when data is received.
+            filter_expr: Optional filter expression.
+            method_name: Optional method name for field mapping.
+
+        Returns:
+            An Unsubscribe handle to cancel the subscription.
+        """
+        # Ensure connected
+        if not self._connected:
+            await self.connect()
+
+        # Process filter if method name is provided
+        processed_filter = None
+        if filter_expr and method_name:
+            processed_filter = replace_filter_fields(filter_expr, method_name)
+        elif filter_expr:
+            processed_filter = filter_expr
+
+        # Check if already subscribed to this channel
+        needs_subscribe = channel not in self._subscriptions
+
+        # Add callback
+        callback_id = self._add_listener(channel, callback)
+
+        # Subscribe to channel if not already subscribed
+        if needs_subscribe:
+            await self._send_subscribe(channel, processed_filter)
+            self._subscriptions[channel] = self._next_command_id()
+            logger.info(f"[streaming] subscribed to channel: {channel}")
+
+        return Unsubscribe(channel=channel, callback_id=callback_id, _api=self)
+
+    # ==================== Subscription Methods ====================
+
+    async def subscribe_token_candles(
+        self,
+        chain: str,
+        token_address: str,
+        resolution: Resolution,
+        callback: Callable[[TokenCandle], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token candle data.
+
+        Args:
+            chain: The blockchain (e.g., "sol").
+            token_address: The token address.
+            resolution: The candle resolution.
+            callback: The callback function to call when data is received.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle to cancel the subscription.
+        """
+        channel = f"dex-candle:{chain}_{token_address}_{resolution.value}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            candle = TokenCandle(
+                open=str(data.get("o", "")),
+                close=str(data.get("c", "")),
+                high=str(data.get("h", "")),
+                low=str(data.get("l", "")),
+                volume=str(data.get("v", "")),
+                resolution=str(data.get("r", "")),
+                time=int(data.get("t", 0)),
+                number=int(data.get("n", 0)),
+            )
+            callback(candle)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_token_candles")
+
+    async def subscribe_token_stats(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TokenStat], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token statistics.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-token-stats:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            stat = TokenStat(
+                address=str(data.get("a", "")),
+                timestamp=int(data.get("t", 0)),
+                buys_1m=data.get("b1m"),
+                sells_1m=data.get("s1m"),
+                buyers_1m=data.get("be1m"),
+                sellers_1m=data.get("se1m"),
+                buy_volume_in_usd_1m=data.get("bviu1m"),
+                sell_volume_in_usd_1m=data.get("sviu1m"),
+                price_1m=data.get("p1m"),
+                open_in_usd_1m=data.get("oiu1m"),
+                close_in_usd_1m=data.get("ciu1m"),
+                price=data.get("p"),
+            )
+            callback(stat)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_token_stats")
+
+    async def subscribe_new_token(
+        self,
+        chain: str,
+        callback: Callable[[NewToken], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to new tokens.
+
+        Args:
+            chain: The blockchain.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-new-token:{chain}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            token = NewToken(
+                token_address=str(data.get("a", "")),
+                name=str(data.get("n", "")),
+                symbol=str(data.get("s", "")),
+                decimals=data.get("d"),
+                created_at_ms=int(data.get("cts", 0)),
+            )
+            callback(token)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_new_token")
+
+    async def subscribe_token_trade(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TradeActivity], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token trades.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-trade:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            trade = TradeActivity(
+                token_address=str(data.get("a", "")),
+                timestamp=int(data.get("t", 0)),
+                kind=str(data.get("k", "")),
+                buy_amount=str(data.get("ba", "")),
+                buy_amount_in_usd=str(data.get("baiu", "")),
+                buy_token_address=str(data.get("btma", "")),
+                buy_token_name=str(data.get("btn", "")),
+                buy_token_symbol=str(data.get("bts", "")),
+                buy_wallet_address=str(data.get("bwa", "")),
+                sell_amount=str(data.get("sa", "")),
+                sell_amount_in_usd=str(data.get("saiu", "")),
+                sell_token_address=str(data.get("stma", "")),
+                sell_token_name=str(data.get("stn", "")),
+                sell_token_symbol=str(data.get("sts", "")),
+                sell_wallet_address=str(data.get("swa", "")),
+                tx_hash=str(data.get("h", "")),
+            )
+            callback(trade)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_token_trades")
+
+    async def subscribe_wallet_balance(
+        self,
+        chain: str,
+        wallet_address: str,
+        callback: Callable[[WalletBalance], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to wallet balance.
+
+        Args:
+            chain: The blockchain.
+            wallet_address: The wallet address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-wallet-balance:{chain}_{wallet_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            balance = WalletBalance(
+                wallet_address=str(data.get("a", "")),
+                token_address=str(data.get("ta", "")),
+                token_price_in_usd=str(data.get("tpiu", "")),
+                balance=str(data.get("b", "")),
+                timestamp=int(data.get("t", 0)),
+            )
+            callback(balance)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_wallet_balance"
+        )
+
+    async def subscribe_token_holders(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TokenHolder], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token holders.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-token-holder:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            holder = TokenHolder(
+                token_address=str(data.get("a", "")),
+                timestamp=int(data.get("ts", 0)),
+                holders=data.get("h"),
+                top100_amount=data.get("t100a"),
+                top10_amount=data.get("t10a"),
+                top100_holders=data.get("t100h"),
+                top10_holders=data.get("t10h"),
+                top100_ratio=data.get("t100r"),
+                top10_ratio=data.get("t10r"),
+            )
+            callback(holder)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_token_holders"
+        )
+
+    async def subscribe_token_supply(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TokenSupply], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token supply.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-token-supply:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            supply = TokenSupply(
+                token_address=str(data.get("a", "")),
+                timestamp=int(data.get("ts", 0)),
+                supply=data.get("s"),
+                market_cap_in_usd=data.get("mciu"),
+            )
+            callback(supply)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_token_supply")
+
+    async def subscribe_dex_pool_balance(
+        self,
+        chain: str,
+        pool_address: str,
+        callback: Callable[[DexPoolBalance], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to DEX pool balance.
+
+        Args:
+            chain: The blockchain.
+            pool_address: The pool address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-pool-balance:{chain}_{pool_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            balance = DexPoolBalance(
+                pool_address=str(data.get("a", "")),
+                token_a_address=str(data.get("taa", "")),
+                token_a_liquidity_in_usd=str(data.get("taliu", "")),
+                token_b_address=str(data.get("tba", "")),
+                token_b_liquidity_in_usd=str(data.get("tbliu", "")),
+            )
+            callback(balance)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_dex_pool_balance"
+        )
+
+    async def subscribe_token_max_liquidity(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TokenMaxLiquidity], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token max liquidity.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-token-liquidity:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            liquidity = TokenMaxLiquidity(
+                token_address=str(data.get("a", "")),
+                pool_address=str(data.get("p", "")),
+                liquidity_in_usd=str(data.get("liu", "")),
+                liquidity_in_native=str(data.get("lin", "")),
+                timestamp=int(data.get("ts", 0)),
+            )
+            callback(liquidity)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_token_max_liquidity"
+        )
+
+    async def subscribe_token_total_liquidity(
+        self,
+        chain: str,
+        token_address: str,
+        callback: Callable[[TokenTotalLiquidity], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to token total liquidity.
+
+        Args:
+            chain: The blockchain.
+            token_address: The token address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-token-total-liquidity:{chain}_{token_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            liquidity = TokenTotalLiquidity(
+                token_address=str(data.get("a", "")),
+                liquidity_in_usd=str(data.get("liu", "")),
+                liquidity_in_native=str(data.get("lin", "")),
+                pool_count=int(data.get("pc", 0)),
+                timestamp=int(data.get("ts", 0)),
+            )
+            callback(liquidity)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_token_total_liquidity"
+        )
+
+    async def subscribe_wallet_pnl(
+        self,
+        chain: str,
+        wallet_address: str,
+        callback: Callable[[WalletTokenPnl], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to wallet PnL.
+
+        Args:
+            chain: The blockchain.
+            wallet_address: The wallet address.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-wallet-pnl:{chain}_{wallet_address}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            pnl = WalletTokenPnl(
+                wallet_address=str(data.get("a", "")),
+                token_address=str(data.get("ta", "")),
+                token_price_in_usd=str(data.get("tpiu", "")),
+                timestamp=int(data.get("t", 0)),
+                open_time=int(data.get("ot", 0)),
+                last_time=int(data.get("lt", 0)),
+                close_time=int(data.get("ct", 0)),
+                buy_amount=str(data.get("ba", "")),
+                buy_amount_in_usd=str(data.get("baiu", "")),
+                buy_count=int(data.get("bs", 0)),
+                buy_count_30d=int(data.get("bs30d", 0)),
+                buy_count_7d=int(data.get("bs7d", 0)),
+                sell_amount=str(data.get("sa", "")),
+                sell_amount_in_usd=str(data.get("saiu", "")),
+                sell_count=int(data.get("ss", 0)),
+                sell_count_30d=int(data.get("ss30d", 0)),
+                sell_count_7d=int(data.get("ss7d", 0)),
+                held_duration_timestamp=int(data.get("hdts", 0)),
+                average_buy_price_in_usd=str(data.get("abpiu", "")),
+                average_sell_price_in_usd=str(data.get("aspiu", "")),
+                unrealized_profit_in_usd=str(data.get("upiu", "")),
+                unrealized_profit_ratio=str(data.get("upr", "")),
+                realized_profit_in_usd=str(data.get("rpiu", "")),
+                realized_profit_ratio=str(data.get("rpr", "")),
+                total_realized_profit_in_usd=str(data.get("trpiu", "")),
+                total_realized_profit_ratio=str(data.get("trr", "")),
+            )
+            callback(pnl)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, "subscribe_wallet_pnl")
+
+    async def subscribe_new_tokens_metadata(
+        self,
+        chain: str,
+        callback: Callable[[TokenMetadata], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to new tokens metadata.
+
+        Args:
+            chain: The blockchain.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-new-token-metadata:{chain}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            metadata = TokenMetadata(
+                token_address=str(data.get("a", "")),
+                name=data.get("n"),
+                symbol=data.get("s"),
+                image_url=data.get("iu"),
+                description=data.get("de"),
+                created_at_ms=data.get("cts"),
+            )
+            callback(metadata)
+
+        return await self.subscribe(
+            channel, parse_callback, filter_expr, "subscribe_new_tokens_metadata"
+        )
+
+    async def subscribe_ranking_tokens_list(
+        self,
+        chain: str,
+        ranking_type: RankingType,
+        callback: Callable[[RankingTokenList], None],
+        filter_expr: Optional[str] = None,
+    ) -> Unsubscribe:
+        """Subscribe to ranking tokens list.
+
+        Args:
+            chain: The blockchain.
+            ranking_type: The ranking type.
+            callback: The callback function.
+            filter_expr: Optional filter expression.
+
+        Returns:
+            An Unsubscribe handle.
+        """
+        channel = f"dex-ranking-token-list:{chain}_{ranking_type.value}"
+
+        def parse_callback(data: Dict[str, Any]) -> None:
+            ranking = RankingTokenList()
+            # TODO: Parse nested objects if needed
+            callback(ranking)
+
+        return await self.subscribe(channel, parse_callback, filter_expr, None)