yellowstone-fumarole-client 0.1.0__py3-none-any.whl

yellowstone_fumarole_client/__init__.py

@@ -0,0 +1,297 @@
+import asyncio
+import logging
+from yellowstone_fumarole_client.grpc_connectivity import (
+    FumaroleGrpcConnector,
+)
+from typing import Dict, Optional
+from dataclasses import dataclass
+from yellowstone_fumarole_client.config import FumaroleConfig
+from yellowstone_fumarole_client.runtime.aio import (
+    AsyncioFumeDragonsmouthRuntime,
+    FumaroleSM,
+    DEFAULT_GC_INTERVAL,
+    DEFAULT_SLOT_MEMORY_RETENTION,
+    GrpcSlotDownloader,
+)
+from yellowstone_fumarole_proto.geyser_pb2 import SubscribeRequest, SubscribeUpdate
+from yellowstone_fumarole_proto.fumarole_v2_pb2 import (
+    ControlResponse,
+    VersionRequest,
+    VersionResponse,
+    JoinControlPlane,
+    ControlCommand,
+    ListConsumerGroupsRequest,
+    ListConsumerGroupsResponse,
+    GetConsumerGroupInfoRequest,
+    ConsumerGroupInfo,
+    DeleteConsumerGroupRequest,
+    DeleteConsumerGroupResponse,
+    CreateConsumerGroupRequest,
+    CreateConsumerGroupResponse,
+)
+from yellowstone_fumarole_proto.fumarole_v2_pb2_grpc import FumaroleStub
+import grpc
+
+__all__ = [
+    "FumaroleClient",
+    "FumaroleConfig",
+    "FumaroleSubscribeConfig",
+    "DragonsmouthAdapterSession",
+    "DEFAULT_DRAGONSMOUTH_CAPACITY",
+    "DEFAULT_COMMIT_INTERVAL",
+    "DEFAULT_MAX_SLOT_DOWNLOAD_ATTEMPT",
+    "DEFAULT_CONCURRENT_DOWNLOAD_LIMIT_PER_TCP",
+]
+
+# Constants
+DEFAULT_DRAGONSMOUTH_CAPACITY = 10000
+DEFAULT_COMMIT_INTERVAL = 5.0  # seconds
+DEFAULT_MAX_SLOT_DOWNLOAD_ATTEMPT = 3
+DEFAULT_CONCURRENT_DOWNLOAD_LIMIT_PER_TCP = 10
+
+# Error classes
+
+
+# FumaroleSubscribeConfig
+@dataclass
+class FumaroleSubscribeConfig:
+    """Configuration for subscribing to a dragonsmouth stream."""
+
+    # The maximum number of concurrent download tasks per TCP connection.
+    concurrent_download_limit: int = DEFAULT_CONCURRENT_DOWNLOAD_LIMIT_PER_TCP
+
+    # The interval at which to commit the slot memory.
+    commit_interval: float = DEFAULT_COMMIT_INTERVAL
+
+    # The maximum number of failed slot download attempts before giving up.
+    max_failed_slot_download_attempt: int = DEFAULT_MAX_SLOT_DOWNLOAD_ATTEMPT
+
+    # The maximum number of slots to download concurrently.
+    data_channel_capacity: int = DEFAULT_DRAGONSMOUTH_CAPACITY
+
+    # The interval at which to perform garbage collection on the slot memory.
+    gc_interval: int = DEFAULT_GC_INTERVAL
+
+    # The retention period for slot memory in seconds.
+    slot_memory_retention: int = DEFAULT_SLOT_MEMORY_RETENTION
+
+
+# DragonsmouthAdapterSession
+@dataclass
+class DragonsmouthAdapterSession:
+    """Session for interacting with the dragonsmouth-like stream."""
+
+    # The queue for sending SubscribeRequest update to the dragonsmouth stream.
+    sink: asyncio.Queue
+
+    # The queue for receiving SubscribeUpdate from the dragonsmouth stream.
+    source: asyncio.Queue
+
+    # The task handle for the fumarole runtime.
+    fumarole_handle: asyncio.Task
+
+
+# FumaroleClient
+class FumaroleClient:
+    """Fumarole client for interacting with the Fumarole server."""
+
+    logger = logging.getLogger(__name__)
+
+    def __init__(self, connector: FumaroleGrpcConnector, stub: FumaroleStub):
+        self.connector = connector
+        self.stub = stub
+
+    @staticmethod
+    async def connect(config: config.FumaroleConfig) -> "FumaroleClient":
+        """Connect to the Fumarole server using the provided configuration.
+        Args:
+            config (FumaroleConfig): Configuration for the Fumarole client.
+        """
+        endpoint = config.endpoint
+        connector = FumaroleGrpcConnector(config=config, endpoint=endpoint)
+        FumaroleClient.logger.debug(f"Connecting to {endpoint}")
+        client = await connector.connect()
+        FumaroleClient.logger.debug(f"Connected to {endpoint}")
+        return FumaroleClient(connector=connector, stub=client)
+
+    async def version(self) -> VersionResponse:
+        """Get the version of the Fumarole server."""
+        request = VersionRequest()
+        response = await self.stub.Version(request)
+        return response
+
+    async def dragonsmouth_subscribe(
+        self, consumer_group_name: str, request: SubscribeRequest
+    ) -> DragonsmouthAdapterSession:
+        """Subscribe to a dragonsmouth stream with default configuration.
+
+        Args:
+            consumer_group_name (str): The name of the consumer group.
+            request (SubscribeRequest): The request to subscribe to the dragonsmouth stream.
+        """
+        return await self.dragonsmouth_subscribe_with_config(
+            consumer_group_name, request, FumaroleSubscribeConfig()
+        )
+
+    async def dragonsmouth_subscribe_with_config(
+        self,
+        consumer_group_name: str,
+        request: SubscribeRequest,
+        config: FumaroleSubscribeConfig,
+    ) -> DragonsmouthAdapterSession:
+        """Subscribe to a dragonsmouth stream with custom configuration.
+
+        Args:
+            consumer_group_name (str): The name of the consumer group.
+            request (SubscribeRequest): The request to subscribe to the dragonsmouth stream.
+            config (FumaroleSubscribeConfig): The configuration for the dragonsmouth subscription.
+        """
+        dragonsmouth_outlet = asyncio.Queue(maxsize=config.data_channel_capacity)
+        fume_control_plane_q = asyncio.Queue(maxsize=100)
+
+        initial_join = JoinControlPlane(consumer_group_name=consumer_group_name)
+        initial_join_command = ControlCommand(initial_join=initial_join)
+        await fume_control_plane_q.put(initial_join_command)
+
+        FumaroleClient.logger.debug(
+            f"Sent initial join command: {initial_join_command}"
+        )
+
+        async def control_plane_sink():
+            while True:
+                try:
+                    update = await fume_control_plane_q.get()
+                    yield update
+                except asyncio.QueueShutDown:
+                    break
+
+        fume_control_plane_stream_rx: grpc.aio.StreamStreamCall = self.stub.Subscribe(
+            control_plane_sink()
+        )  # it's actually InterceptedStreamStreamCall, but grpc lib doesn't export it
+
+        control_response: ControlResponse = await fume_control_plane_stream_rx.read()
+        init = control_response.init
+        if init is None:
+            raise ValueError(f"Unexpected initial response: {control_response}")
+
+        # Once we have the initial response, we can spin a task to read from the stream
+        # and put the updates into the queue.
+        # This is a bit of a hack, but we need a Queue not a StreamStreamMultiCallable
+        # because Queue are cancel-safe, while Stream are not, or at least didn't find any docs about it.
+        fume_control_plane_rx_q = asyncio.Queue(maxsize=100)
+
+        async def control_plane_source():
+            while True:
+                try:
+                    async for update in fume_control_plane_stream_rx:
+                        await fume_control_plane_rx_q.put(update)
+                except asyncio.QueueShutDown:
+                    break
+
+        _cp_src_task = asyncio.create_task(control_plane_source())
+
+        FumaroleClient.logger.debug(f"Control response: {control_response}")
+
+        last_committed_offset = init.last_committed_offsets.get(0)
+        if last_committed_offset is None:
+            raise ValueError("No last committed offset")
+
+        sm = FumaroleSM(last_committed_offset, config.slot_memory_retention)
+        subscribe_request_queue = asyncio.Queue(maxsize=100)
+
+        data_plane_client = await self.connector.connect()
+
+        grpc_slot_downloader = GrpcSlotDownloader(
+            client=data_plane_client,
+        )
+
+        rt = AsyncioFumeDragonsmouthRuntime(
+            sm=sm,
+            slot_downloader=grpc_slot_downloader,
+            subscribe_request_update_q=subscribe_request_queue,
+            subscribe_request=request,
+            consumer_group_name=consumer_group_name,
+            control_plane_tx_q=fume_control_plane_q,
+            control_plane_rx_q=fume_control_plane_rx_q,
+            dragonsmouth_outlet=dragonsmouth_outlet,
+            commit_interval=config.commit_interval,
+            gc_interval=config.gc_interval,
+            max_concurrent_download=config.concurrent_download_limit,
+        )
+
+        fumarole_handle = asyncio.create_task(rt.run())
+        FumaroleClient.logger.debug(f"Fumarole handle created: {fumarole_handle}")
+        return DragonsmouthAdapterSession(
+            sink=subscribe_request_queue,
+            source=dragonsmouth_outlet,
+            fumarole_handle=fumarole_handle,
+        )
+
+    async def list_consumer_groups(
+        self,
+    ) -> ListConsumerGroupsResponse:
+        """Lists all consumer groups."""
+        return await self.stub.ListConsumerGroups(ListConsumerGroupsRequest())
+
+    async def get_consumer_group_info(
+        self, consumer_group_name: str
+    ) -> Optional[ConsumerGroupInfo]:
+        """Gets information about a consumer group by name.
+        Returns None if the consumer group does not exist.
+
+        Args:
+            consumer_group_name (str): The name of the consumer group to retrieve information for.
+        """
+        try:
+            return await self.stub.GetConsumerGroupInfo(
+                GetConsumerGroupInfoRequest(consumer_group_name=consumer_group_name)
+            )
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                return None
+            else:
+                raise
+
+    async def delete_consumer_group(
+        self, consumer_group_name: str
+    ) -> DeleteConsumerGroupResponse:
+        """Delete a consumer group by name.
+
+        NOTE: this operation is idempotent, meaning that if the consumer group does not exist, it will not raise an error.
+        Args:
+            consumer_group_name (str): The name of the consumer group to delete.
+        """
+        return await self.stub.DeleteConsumerGroup(
+            DeleteConsumerGroupRequest(consumer_group_name=consumer_group_name)
+        )
+
+    async def delete_all_consumer_groups(
+        self,
+    ) -> DeleteConsumerGroupResponse:
+        """Deletes all consumer groups."""
+        consumer_group_list = await self.list_consumer_groups()
+
+        tasks = []
+
+        async with asyncio.TaskGroup() as tg:
+            for group in consumer_group_list.consumer_groups:
+                cg_name = group.consumer_group_name
+                task = tg.create_task(self.delete_consumer_group(cg_name))
+                tasks.append((cg_name, task))
+
+        # Raise an error if any task fails
+        for cg_name, task in tasks:
+            result = task.result()
+            if not result.success:
+                raise RuntimeError(
+                    f"Failed to delete consumer group {cg_name}: {result.error}"
+                )
+
+    async def create_consumer_group(
+        self, request: CreateConsumerGroupRequest
+    ) -> CreateConsumerGroupResponse:
+        """Creates a new consumer group.
+        Args:
+            request (CreateConsumerGroupRequest): The request to create a consumer group.
+        """
+        return await self.stub.CreateConsumerGroup(request)

yellowstone_fumarole_client/config.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Dict, Optional
+import yaml
+
+
+@dataclass
+class FumaroleConfig:
+    endpoint: str
+    x_token: Optional[str] = None
+    max_decoding_message_size_bytes: int = 512_000_000
+    x_metadata: Dict[str, str] = None
+
+    def __post_init__(self):
+        self.x_metadata = self.x_metadata or {}
+
+    @classmethod
+    def from_yaml(cls, fileobj) -> "FumaroleConfig":
+        data = yaml.safe_load(fileobj)
+        return cls(
+            endpoint=data["endpoint"],
+            x_token=data.get("x-token") or data.get("x_token"),
+            max_decoding_message_size_bytes=data.get(
+                "max_decoding_message_size_bytes", cls.max_decoding_message_size_bytes
+            ),
+            x_metadata=data.get("x-metadata", {}),
+        )

yellowstone_fumarole_client/grpc_connectivity.py

@@ -0,0 +1,197 @@
+import logging
+from typing import Optional
+import grpc
+from yellowstone_fumarole_client.config import FumaroleConfig
+from yellowstone_fumarole_proto.fumarole_v2_pb2_grpc import FumaroleStub
+
+X_TOKEN_HEADER = "x-token"
+
+
+def _triton_sign_request(
+    callback: grpc.AuthMetadataPluginCallback,
+    x_token: Optional[str],
+    error: Optional[Exception],
+):
+    # WARNING: metadata is a 1d-tuple (<value>,), the last comma is necessary
+    metadata = ((X_TOKEN_HEADER, x_token),)
+    return callback(metadata, error)
+
+
+class TritonAuthMetadataPlugin(grpc.AuthMetadataPlugin):
+    """Metadata wrapper for raw access token credentials."""
+
+    def __init__(self, x_token: str):
+        self.x_token = x_token
+
+    def __call__(
+        self,
+        context: grpc.AuthMetadataContext,
+        callback: grpc.AuthMetadataPluginCallback,
+    ):
+        return _triton_sign_request(callback, self.x_token, None)
+
+
+def grpc_channel(endpoint: str, x_token=None, compression=None, *grpc_options):
+    options = [("grpc.max_receive_message_length", 111111110), *grpc_options]
+    if x_token is not None:
+        auth = TritonAuthMetadataPlugin(x_token)
+        # ssl_creds allow you to use our https endpoint
+        # grpc.ssl_channel_credentials with no arguments will look through your CA trust store.
+        ssl_creds = grpc.ssl_channel_credentials()
+
+        # call credentials will be sent on each request if setup with composite_channel_credentials.
+        call_creds: grpc.CallCredentials = grpc.metadata_call_credentials(auth)
+
+        # Combined creds will store the channel creds aswell as the call credentials
+        combined_creds = grpc.composite_channel_credentials(ssl_creds, call_creds)
+
+        return grpc.secure_channel(
+            endpoint,
+            credentials=combined_creds,
+            compression=compression,
+            options=options,
+        )
+    else:
+        return grpc.insecure_channel(endpoint, compression=compression, options=options)
+
+
+# Because of a bug in grpcio library, multiple inheritance of ClientInterceptor subclasses does not work.
+# You have to create a new class for each type of interceptor you want to use.
+
+
+class MetadataInterceptor(
+    grpc.aio.UnaryStreamClientInterceptor,
+    grpc.aio.StreamUnaryClientInterceptor,
+    grpc.aio.StreamStreamClientInterceptor,
+    grpc.aio.UnaryUnaryClientInterceptor,
+):
+
+    def __init__(self, metadata):
+        if isinstance(metadata, dict):
+            metadata = metadata.items()
+        self.metadata = list(metadata)
+
+    async def intercept_unary_unary(
+        self, continuation, client_call_details: grpc.aio.ClientCallDetails, request
+    ):
+        # logging.debug("intercept_unary_unary")
+        new_details = client_call_details._replace(
+            metadata=self._merge_metadata(client_call_details.metadata)
+        )
+        return await continuation(new_details, request)
+
+    async def intercept_unary_stream(
+        self, continuation, client_call_details: grpc.aio.ClientCallDetails, request
+    ):
+        # logging.debug("intercept_unary_stream")
+        new_details = client_call_details._replace(
+            metadata=self._merge_metadata(client_call_details.metadata)
+        )
+        return await continuation(new_details, request)
+
+    async def intercept_stream_unary(
+        self, continuation, client_call_details: grpc.aio.ClientCallDetails, request
+    ):
+        # logging.debug("intercept_stream_unary")
+        new_details = client_call_details._replace(
+            metadata=self._merge_metadata(client_call_details.metadata)
+        )
+        return await continuation(new_details, request)
+
+    async def intercept_stream_stream(
+        self, continuation, client_call_details: grpc.aio.ClientCallDetails, request
+    ):
+        # logging.debug("intercept_stream_stream")
+        new_details = client_call_details._replace(
+            metadata=self._merge_metadata(client_call_details.metadata)
+        )
+        return await continuation(new_details, request)
+
+    def unary_stream_interceptor(self) -> grpc.aio.UnaryStreamClientInterceptor:
+        this = self
+
+        class Interceptor(grpc.aio.UnaryStreamClientInterceptor):
+            async def intercept_unary_stream(self, *args):
+                return await this.intercept_unary_stream(*args)
+
+        return Interceptor()
+
+    def stream_unary_interceptor(self) -> grpc.aio.StreamUnaryClientInterceptor:
+        this = self
+
+        class Interceptor(grpc.aio.StreamUnaryClientInterceptor):
+            async def intercept_stream_unary(self, *args):
+                return await this.intercept_stream_unary(*args)
+
+        return Interceptor()
+
+    def stream_stream_interceptor(self) -> grpc.aio.StreamStreamClientInterceptor:
+        this = self
+
+        class Interceptor(grpc.aio.StreamStreamClientInterceptor):
+            async def intercept_stream_stream(self, *args):
+                return await this.intercept_stream_stream(*args)
+
+        return Interceptor()
+
+    def unary_unary_interceptor(self) -> grpc.aio.UnaryUnaryClientInterceptor:
+        this = self
+
+        class Interceptor(grpc.aio.UnaryUnaryClientInterceptor):
+            async def intercept_unary_unary(self, *args):
+                return await this.intercept_unary_unary(*args)
+
+        return Interceptor()
+
+    def interceptors(self) -> list[grpc.aio.ClientInterceptor]:
+        return [
+            self.unary_unary_interceptor(),
+            self.unary_stream_interceptor(),
+            self.stream_unary_interceptor(),
+            self.stream_stream_interceptor(),
+        ]
+
+    def _merge_metadata(self, existing):
+        result = list(existing or []) + self.metadata
+        return result
+
+
+class FumaroleGrpcConnector:
+    logger = logging.getLogger(__name__)
+
+    def __init__(self, config: FumaroleConfig, endpoint: str):
+        self.config = config
+        self.endpoint = endpoint
+
+    async def connect(self, *grpc_options) -> FumaroleStub:
+        options = [("grpc.max_receive_message_length", 111111110), *grpc_options]
+        interceptors = MetadataInterceptor(self.config.x_metadata).interceptors()
+        if self.config.x_token is not None:
+            auth = TritonAuthMetadataPlugin(self.config.x_token)
+            # ssl_creds allow you to use our https endpoint
+            # grpc.ssl_channel_credentials with no arguments will look through your CA trust store.
+            ssl_creds = grpc.ssl_channel_credentials()
+
+            # call credentials will be sent on each request if setup with composite_channel_credentials.
+            call_creds: grpc.CallCredentials = grpc.metadata_call_credentials(auth)
+
+            # Combined creds will store the channel creds aswell as the call credentials
+            combined_creds = grpc.composite_channel_credentials(ssl_creds, call_creds)
+            FumaroleGrpcConnector.logger.debug(
+                "Using secure channel with x-token authentication"
+            )
+            channel = grpc.aio.secure_channel(
+                self.endpoint,
+                credentials=combined_creds,
+                options=options,
+                interceptors=interceptors,
+            )
+        else:
+            FumaroleGrpcConnector.logger.debug(
+                "Using insecure channel without authentication"
+            )
+            channel = grpc.aio.insecure_channel(
+                self.endpoint, options=options, interceptors=interceptors
+            )
+
+        return FumaroleStub(channel)