Fix token transfers handling for realtime index status

Author: igorsereda

src/dipdup/datasources/datasource.py

@@ -17,6 +17,7 @@
 from dipdup.models import BigMapData
 from dipdup.models import HeadBlockData
 from dipdup.models import OperationData
+from dipdup.models import TokenTransferData
 from dipdup.utils import FormattedLogger
 
 _logger = logging.getLogger('dipdup.datasource')
@@ -25,6 +26,7 @@
 EmptyCallbackT = Callable[[], Awaitable[None]]
 HeadCallbackT = Callable[['IndexDatasource', HeadBlockData], Awaitable[None]]
 OperationsCallbackT = Callable[['IndexDatasource', Tuple[OperationData, ...]], Awaitable[None]]
+TokenTransfersCallbackT = Callable[['IndexDatasource', Tuple[TokenTransferData, ...]], Awaitable[None]]
 BigMapsCallbackT = Callable[['IndexDatasource', Tuple[BigMapData, ...]], Awaitable[None]]
 RollbackCallbackT = Callable[['IndexDatasource', MessageType, int, int], Awaitable[None]]
 
@@ -74,6 +76,7 @@ def __init__(self, url: str, http_config: HTTPConfig, merge_subscriptions: bool
         self._on_disconnected_callbacks: Set[EmptyCallbackT] = set()
         self._on_head_callbacks: Set[HeadCallbackT] = set()
         self._on_operations_callbacks: Set[OperationsCallbackT] = set()
+        self._on_token_transfers_callbacks: Set[TokenTransfersCallbackT] = set()
         self._on_big_maps_callbacks: Set[BigMapsCallbackT] = set()
         self._on_rollback_callbacks: Set[RollbackCallbackT] = set()
         self._subscriptions: SubscriptionManager = SubscriptionManager(merge_subscriptions)
@@ -100,6 +103,9 @@ def call_on_head(self, fn: HeadCallbackT) -> None:
     def call_on_operations(self, fn: OperationsCallbackT) -> None:
         self._on_operations_callbacks.add(fn)
 
+    def call_on_token_transfers(self, fn: TokenTransfersCallbackT) -> None:
+        self._on_token_transfers_callbacks.add(fn)
+
     def call_on_big_maps(self, fn: BigMapsCallbackT) -> None:
         self._on_big_maps_callbacks.add(fn)
 
@@ -120,6 +126,10 @@ async def emit_operations(self, operations: Tuple[OperationData, ...]) -> None:
         for fn in self._on_operations_callbacks:
             await fn(self, operations)
 
+    async def emit_token_transfers(self, token_transfers: Tuple[TokenTransferData, ...]) -> None:
+        for fn in self._on_token_transfers_callbacks:
+            await fn(self, token_transfers)
+
     async def emit_big_maps(self, big_maps: Tuple[BigMapData, ...]) -> None:
         for fn in self._on_big_maps_callbacks:
             await fn(self, big_maps)

src/dipdup/datasources/tzkt/datasource.py

@@ -899,6 +899,7 @@ def _get_ws_client(self) -> SignalRClient:
         self._ws_client.on_error(self._on_error)
 
         self._ws_client.on('operations', partial(self._on_message, MessageType.operation))
+        self._ws_client.on('transfers', partial(self._on_message, MessageType.token_transfer))
         self._ws_client.on('bigmaps', partial(self._on_message, MessageType.big_map))
         self._ws_client.on('head', partial(self._on_message, MessageType.head))
 
@@ -977,6 +978,8 @@ async def _on_message(self, type_: MessageType, message: List[Dict[str, Any]]) -
         for buffered_message in self._buffer.yield_from():
             if buffered_message.type == MessageType.operation:
                 await self._process_operations_data(cast(list, buffered_message.data))
+            elif buffered_message.type == MessageType.token_transfer:
+                await self._process_token_transfers_data(cast(list, buffered_message.data))
             elif buffered_message.type == MessageType.big_map:
                 await self._process_big_maps_data(cast(list, buffered_message.data))
             elif buffered_message.type == MessageType.head:
@@ -997,6 +1000,17 @@ async def _process_operations_data(self, data: List[Dict[str, Any]]) -> None:
         for _level, operations in level_operations.items():
             await self.emit_operations(tuple(operations))
 
+    async def _process_token_transfers_data(self, data: List[Dict[str, Any]]) -> None:
+        """Parse and emit raw token transfers from WS"""
+        level_token_transfers: DefaultDict[int, Deque[TokenTransferData]] = defaultdict(deque)
+
+        for token_transfer_json in data:
+            token_transfer = self.convert_token_transfer(token_transfer_json)
+            level_token_transfers[token_transfer.level].append(token_transfer)
+
+        for _level, token_transfers in level_token_transfers.items():
+            await self.emit_token_transfers(tuple(token_transfers))
+
     async def _process_big_maps_data(self, data: List[Dict[str, Any]]) -> None:
         """Parse and emit raw big map diffs from WS"""
         level_big_maps: DefaultDict[int, Deque[BigMapData]] = defaultdict(deque)

src/dipdup/dipdup.py

@@ -46,6 +46,7 @@
 from dipdup.index import HeadIndex
 from dipdup.index import Index
 from dipdup.index import OperationIndex
+from dipdup.index import TokenTransferIndex
 from dipdup.index import extract_operation_subgroups
 from dipdup.models import BigMapData
 from dipdup.models import Contract
@@ -55,6 +56,7 @@
 from dipdup.models import IndexStatus
 from dipdup.models import OperationData
 from dipdup.models import Schema
+from dipdup.models import TokenTransferData
 from dipdup.prometheus import Metrics
 from dipdup.scheduler import SchedulerManager
 from dipdup.utils import is_importable
@@ -220,6 +222,7 @@ async def _subscribe_to_datasource_events(self) -> None:
                 continue
             datasource.call_on_head(self._on_head)
             datasource.call_on_operations(self._on_operations)
+            datasource.call_on_token_transfers(self._on_token_transfers)
             datasource.call_on_big_maps(self._on_big_maps)
             datasource.call_on_rollback(self._on_rollback)
 
@@ -259,6 +262,11 @@ async def _on_operations(self, datasource: IndexDatasource, operations: Tuple[Op
         for index in operation_indexes:
             index.push_operations(operation_subgroups)
 
+    async def _on_token_transfers(self, datasource: IndexDatasource, token_transfers: Tuple[TokenTransferData, ...]) -> None:
+        token_transfer_indexes = (i for i in self._indexes.values() if isinstance(i, TokenTransferIndex) and i.datasource == datasource)
+        for index in token_transfer_indexes:
+            index.push_token_transfers(token_transfers)
+
     async def _on_big_maps(self, datasource: IndexDatasource, big_maps: Tuple[BigMapData, ...]) -> None:
         big_map_indexes = (i for i in self._indexes.values() if isinstance(i, BigMapIndex) and i.datasource == datasource)
         for index in big_map_indexes:

src/dipdup/index.py

@@ -937,8 +937,14 @@ def __init__(self, ctx: DipDupContext, config: TokenTransferIndexConfig, datasou
         super().__init__(ctx, config, datasource)
         self._queue: Deque[Tuple[TokenTransferData, ...]] = deque()
 
+    def push_token_transfers(self, token_transfers: Tuple[TokenTransferData, ...]) -> None:
+        self._queue.append(token_transfers)
+
+        if Metrics.enabled:
+            Metrics.set_levels_to_realtime(self._config.name, len(self._queue))
+
     async def _synchronize(self, last_level: int, cache: bool = False) -> None:
-        """Fetch operations via Fetcher and pass to message callback"""
+        """Fetch token transfers via Fetcher and pass to message callback"""
         first_level = await self._enter_sync_state(last_level)
         if first_level is None:
             return