kuhl-haus-mdp 0.0.1__py3-none-any.whl

kuhl_haus/mdp/integ/massive_data_listener.py

@@ -0,0 +1,140 @@
+import asyncio
+from datetime import datetime, timezone
+from zoneinfo import ZoneInfo
+from logging import Logger
+from typing import Awaitable, Callable, Optional, List, Union
+
+from massive import WebSocketClient
+from massive.websocket import Feed, Market, WebSocketMessage
+
+
+class MassiveDataListener:
+    connection_status: dict
+    ws_connection: Union[WebSocketClient, None]
+    ws_coroutine: Union[asyncio.Task, None]
+    feed: Feed
+    market: Market
+    subscriptions: List[str]
+    raw: bool
+    verbose: bool
+    max_reconnects: Optional[int]
+    secure: bool
+
+    def __init__(
+        self,
+        logger: Logger,
+        message_handler: Union[
+            Callable[[List[WebSocketMessage]], Awaitable],
+            Callable[[Union[str, bytes]], Awaitable],
+        ],
+        api_key: str,
+        feed: Feed,
+        market: Market,
+        subscriptions: List[str],
+        raw: bool = False,
+        verbose: bool = False,
+        max_reconnects: Optional[int] = 5,
+        secure: bool = True,
+        **kwargs,
+    ):
+        self.logger = logger
+        self.message_handler = message_handler
+        self.api_key = api_key
+        self.feed = feed
+        self.market = market
+        self.subscriptions = subscriptions
+        self.raw = raw
+        self.verbose = verbose
+        self.max_reconnects = max_reconnects
+        self.secure = secure
+        self.kwargs = kwargs
+        self.connection_status = {
+            "connected": False,
+            "feed": feed,
+            "market": market,
+            "subscriptions": subscriptions,
+        }
+
+    async def start(self):
+        """Start WebSocket client"""
+        try:
+            self.logger.info("Instantiating WebSocket client...")
+            self.ws_connection = WebSocketClient(
+                api_key=self.api_key,
+                feed=self.feed,
+                market=self.market,
+                raw=self.raw,
+                verbose=self.verbose,
+                subscriptions=self.subscriptions,
+                max_reconnects=self.max_reconnects,
+                secure=self.secure,
+                **self.kwargs,
+            )
+            self.logger.info("Scheduling WebSocket client task...")
+            self.ws_coroutine = asyncio.create_task(self.async_task())
+        except Exception as e:
+            self.logger.error(f"Error starting WebSocket client: {e}")
+            await self.stop()
+
+    async def stop(self):
+        """Stop WebSocket client"""
+        try:
+            self.logger.info("Shutting down WebSocket client...")
+            self.ws_coroutine.cancel()
+            await asyncio.sleep(1)
+            self.logger.info("unsubscribing from all feeds...")
+            self.ws_connection.unsubscribe_all()
+            await asyncio.sleep(1)
+            self.logger.info("closing connection...")
+            await self.ws_connection.close()
+            self.logger.info("done.")
+        except Exception as e:
+            self.logger.error(f"Error stopping WebSocket client: {e}")
+        self.connection_status["connected"] = False
+        self.ws_connection = None
+        self.ws_coroutine = None
+
+    async def restart(self):
+        """Restart WebSocket client"""
+        try:
+            self.logger.info("Stopping WebSocket client...")
+            await self.stop()
+            self.logger.info("done")
+            await asyncio.sleep(1)
+            self.logger.info("Starting WebSocket client...")
+            await self.start()
+            self.logger.info("done")
+        except Exception as e:
+            self.logger.error(f"Error restarting WebSocket client: {e}")
+
+    async def async_task(self):
+        """Main task that runs the WebSocket client"""
+        try:
+            self.logger.info("Connecting to market data provider...")
+            self.connection_status["connected"] = True
+            await asyncio.gather(
+                self.ws_connection.connect(self.message_handler),
+                return_exceptions=True
+            )
+            self.connection_status["connected"] = False
+            self.logger.info("Disconnected from market data provider...")
+            pending_restart = True
+            while pending_restart:
+                # Get current time in UTC, then convert to Eastern Time
+                utc_now = datetime.now(timezone.utc)
+                et_now = utc_now.astimezone(ZoneInfo("America/New_York"))
+
+                # Check if within trading hours: Mon-Fri, 04:00-19:59 ET
+                is_weekday = et_now.weekday() < 5
+                is_trading_hours = 4 <= et_now.hour < 20
+
+                if is_weekday and is_trading_hours:
+                    self.logger.info(f"Reconnecting at {et_now.strftime('%H:%M:%S %Z')}...")
+                    await self.start()
+                    pending_restart = False
+                else:
+                    self.logger.info(f"Outside market hours ({et_now.strftime('%H:%M:%S %Z')}), sleeping 5 min...")
+                    await asyncio.sleep(300)
+        except Exception as e:
+            self.logger.error(f"Fatal error in WebSocket client: {e}")
+            await self.stop()

kuhl_haus/mdp/integ/massive_data_processor.py

@@ -0,0 +1,236 @@
+import asyncio
+import json
+import logging
+from typing import Dict
+
+import aio_pika
+import redis
+import redis.asyncio as aioredis
+from aio_pika.abc import AbstractIncomingMessage
+from aio_pika.exceptions import AMQPConnectionError
+
+from kuhl_haus.mdp.analyzers.massive_data_analyzer import MassiveDataAnalyzer
+from kuhl_haus.mdp.models.market_data_analyzer_result import MarketDataAnalyzerResult
+from kuhl_haus.mdp.integ.web_socket_message_serde import WebSocketMessageSerde
+
+
+class MassiveDataProcessor:
+    queue_name: str
+    mdq_connected: bool
+    mdc_connected: bool
+    processed: int
+    duplicated: int
+    decoding_error: int
+    dropped: int
+    error: int
+
+    def __init__(
+        self,
+        rabbitmq_url: str,
+        queue_name: str,
+        redis_url: str,
+        prefetch_count: int = 100,  # Higher for async throughput
+        max_concurrent_tasks: int = 500,  # Concurrent processing limit
+    ):
+        self.rabbitmq_url = rabbitmq_url
+        self.queue_name = queue_name
+        self.redis_url = redis_url
+        self.prefetch_count = prefetch_count
+        self.max_concurrent_tasks = max_concurrent_tasks
+
+        # Connection objects
+        self.rmq_connection = None
+        self.rmq_channel = None
+        self.redis_client = None
+
+        # Analyzer
+        self.analyzer = MassiveDataAnalyzer()
+
+        # Concurrency control
+        self.semaphore = asyncio.Semaphore(max_concurrent_tasks)
+        self.processing_tasks = set()
+
+        # State
+        self.running = False
+        self.logger = logging.getLogger(__name__)
+
+        # Metrics
+        self.processed = 0
+        self.duplicated = 0
+        self.error = 0
+        self.decoding_error = 0
+        self.dropped = 0
+        self.mdq_connected = False
+        self.mdc_connected = False
+
+    async def connect(self, force: bool = False):
+        """Establish async connections to RabbitMQ and Redis"""
+
+        if not self.mdq_connected or force:
+            # RabbitMQ connection
+            try:
+                self.rmq_connection = await aio_pika.connect_robust(
+                    self.rabbitmq_url,
+                    heartbeat=60,
+                    timeout=30,  # Add connection timeout
+                )
+                self.rmq_channel = await self.rmq_connection.channel()
+                await self.rmq_channel.set_qos(prefetch_count=self.prefetch_count)
+                await self.rmq_channel.get_queue(self.queue_name, ensure=False)
+                self.mdq_connected = True
+                self.logger.info(f"Connected to RabbitMQ queue: {self.queue_name}")
+            except Exception as e:
+                self.logger.error(f"Failed to connect to RabbitMQ: {e}")
+                raise
+
+        if not self.mdc_connected or force:
+            # Redis connection pool
+            try:
+                self.redis_client = aioredis.from_url(
+                    self.redis_url,
+                    encoding="utf-8",
+                    decode_responses=True,
+                    max_connections=1000,
+                    socket_connect_timeout=10,  # Add timeout
+                )
+
+                # Test Redis connection
+                await self.redis_client.ping()
+                self.mdc_connected = True
+                self.logger.debug(f"Connected to Redis: {self.redis_url}")
+            except Exception as e:
+                self.logger.error(f"Failed to connect to Redis: {e}")
+                # Cleanup RabbitMQ connection on Redis failure
+                await self.rmq_channel.close()
+                await self.rmq_connection.close()
+                raise
+
+    async def _process_message(self, message: AbstractIncomingMessage):
+        """Process single message with concurrency control"""
+        async with self.semaphore:
+            try:
+                async with message.process():
+                    # Parse message
+                    web_socket_message = json.loads(message.body.decode())
+                    data = WebSocketMessageSerde.to_dict(web_socket_message)
+
+                    # Delegate to analyzer (async)
+                    analyzer_results = await self.analyzer.analyze_data(data)
+                    if analyzer_results:
+                        self.processed += 1
+                        for analyzer_result in analyzer_results:
+                            # Cache in Redis
+                            await self._cache_result(analyzer_result)
+
+                        self.logger.debug(f"Processed message {message.delivery_tag}")
+                    else:
+                        # Empty result - drop message
+                        self.dropped += 1
+                        self.logger.debug(
+                            f"Analyzer returned empty for {message.delivery_tag}"
+                        )
+            except aio_pika.exceptions.MessageProcessError as e:
+                self.logger.error(f"Message processing error: {e}")
+                self.duplicated += 1
+            except json.JSONDecodeError as e:
+                self.logger.error(f"JSON decode error: {e}")
+                self.decoding_error += 1
+            except Exception as e:
+                self.logger.error(f"Processing error: {e}", exc_info=True)
+                self.error += 1
+
+    async def _callback(self, message: AbstractIncomingMessage):
+        """
+        Message callback - spawns processing task
+
+        Note: Tasks tracked for graceful shutdown
+        """
+        task = asyncio.create_task(self._process_message(message))
+        self.processing_tasks.add(task)
+        task.add_done_callback(self.processing_tasks.discard)
+
+    async def _cache_result(self, analyzer_result: MarketDataAnalyzerResult):
+        """
+        Async cache to Redis with pub/sub notification
+
+        Args:
+            result: Processed data dict from analyzer
+            cache_entry_name: Cache key suffix (e.g., symbol)
+        """
+        result_json = json.dumps(analyzer_result.data)
+
+        # Pipeline - no async context manager, no await on queue methods
+        pipe = self.redis_client.pipeline(transaction=False)
+        if analyzer_result.cache_key:
+            if analyzer_result.cache_ttl > 0:
+                pipe.setex(analyzer_result.cache_key, analyzer_result.cache_ttl, result_json)
+            else:
+                pipe.set(analyzer_result.cache_key, result_json)
+        if analyzer_result.publish_key:
+            pipe.publish(analyzer_result.publish_key, result_json)
+
+        await pipe.execute()
+
+        self.logger.debug(f"Cached result for {analyzer_result.cache_key}")
+
+    async def start(self):
+        """Start async message consumption"""
+        retry_count = 0
+        while not self.mdc_connected or not self.mdq_connected:
+            try:
+                await self.connect()
+            except Exception as e:
+                if retry_count < 5:
+                    retry_count += 1
+                    self.logger.error(f"Connection error: {e}, sleeping for {2*retry_count}s")
+                    await asyncio.sleep(2*retry_count)
+                else:
+                    self.logger.error("Failed to connect to RabbitMQ or Redis")
+                    raise
+        if self.mdc_connected and self.mdq_connected:
+            self.running = True
+        else:
+            self.logger.error("Failed to connect to RabbitMQ or Redis")
+            raise RuntimeError("Failed to connect to RabbitMQ or Redis")
+
+        # Get queue
+        queue = await self.rmq_channel.get_queue(self.queue_name)
+
+        self.logger.info("Starting async message consumption")
+
+        # Start consuming with callback
+        await queue.consume(self._callback, no_ack=False)
+
+        # Run until shutdown signal
+        try:
+            while self.running:
+                await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            self.logger.info("Consumption cancelled")
+        finally:
+            await self.stop()
+
+    async def stop(self):
+        """Graceful async shutdown"""
+        self.logger.info("Stopping processor - waiting for pending tasks")
+        self.running = False
+
+        # Wait for all processing tasks to complete
+        if self.processing_tasks:
+            self.logger.info(f"Waiting for {len(self.processing_tasks)} tasks")
+            await asyncio.gather(*self.processing_tasks, return_exceptions=True)
+
+        # Close connections
+        if self.rmq_channel:
+            await self.rmq_channel.close()
+
+        if self.rmq_connection:
+            await self.rmq_connection.close()
+
+        if self.redis_client:
+            await self.redis_client.close()
+
+        self.logger.info(
+            f"Processor stopped - Processed: {self.processed}, "
+            f"Errors: {self.error}"
+        )

kuhl_haus/mdp/integ/massive_data_queues.py

@@ -0,0 +1,124 @@
+from datetime import datetime
+from typing import List, Union
+
+from aio_pika import Connection, Channel, connect_robust, Message
+from aio_pika import DeliveryMode
+from aio_pika.abc import AbstractConnection, AbstractChannel
+from massive.websocket.models import WebSocketMessage
+
+from kuhl_haus.mdp.models.massive_data_queue import MassiveDataQueue
+from kuhl_haus.mdp.helpers.queue_name_resolver import QueueNameResolver
+from kuhl_haus.mdp.integ.web_socket_message_serde import WebSocketMessageSerde
+
+
+class MassiveDataQueues:
+    rabbitmq_url: str
+    queues: List[str]
+    message_ttl: int
+    connection: Union[Connection, AbstractConnection]
+    channel: Union[Channel, AbstractChannel]
+    connection_status: dict
+
+    def __init__(self, logger, rabbitmq_url, message_ttl: int):
+        self.logger = logger
+        self.rabbitmq_url = rabbitmq_url
+        self.queues = [
+            MassiveDataQueue.TRADES.value,
+            MassiveDataQueue.AGGREGATE.value,
+            MassiveDataQueue.QUOTES.value,
+            MassiveDataQueue.HALTS.value,
+            MassiveDataQueue.NEWS.value,
+            MassiveDataQueue.UNKNOWN.value,
+        ]
+        self.message_ttl = message_ttl
+        self.connection_status = {
+            "connected": False,
+            "last_message_time": None,
+            "messages_received": 0,
+            MassiveDataQueue.TRADES.value: 0,
+            MassiveDataQueue.AGGREGATE.value: 0,
+            MassiveDataQueue.QUOTES.value: 0,
+            MassiveDataQueue.HALTS.value: 0,
+            MassiveDataQueue.NEWS.value: 0,
+            MassiveDataQueue.UNKNOWN.value: 0,
+            "unsupported_messages": 0,
+            "reconnect_attempts": 0,
+        }
+
+    async def connect(self):
+        self.connection = await connect_robust(self.rabbitmq_url)
+        self.channel = await self.connection.channel()
+
+        try:
+            for q in self.queues:
+                _ = await self.channel.declare_queue(q, passive=True)  # Don't create, just check
+
+            self.connection_status["connected"] = self.connection is not None and self.channel is not None
+        except Exception as e:
+            self.logger.error(f"Fatal error while processing request: {e}")
+            raise
+
+    async def handle_messages(self, msgs: List[WebSocketMessage]):
+        if not self.channel:
+            self.logger.error("RabbitMQ channel not initialized")
+            raise Exception("RabbitMQ channel not initialized")
+        if not self.connection:
+            self.logger.error("RabbitMQ connection not initialized")
+            raise Exception("RabbitMQ connection not initialized")
+        try:
+            for message in msgs:
+                await self.fanout_to_queues(message)
+        except Exception as e:
+            self.logger.error(f"Fatal error while processing messages: {e}")
+            raise
+
+    async def shutdown(self):
+        self.connection_status["connected"] = False
+        self.logger.info("Closing RabbitMQ channel")
+        await self.channel.close()
+        self.logger.info("RabbitMQ channel closed")
+        self.logger.info("Closing RabbitMQ connection")
+        await self.connection.close()
+        self.logger.info("RabbitMQ connection closed")
+
+    async def setup_queues(self):
+        self.connection = await connect_robust(self.rabbitmq_url)
+        self.channel = await self.connection.channel()
+
+        # Declare queues with message TTL
+        for queue in self.queues:
+            await self.channel.declare_queue(
+                queue,
+                durable=True,
+                arguments={"x-message-ttl": self.message_ttl}  # Messages are deleted after they expire
+            )
+
+            self.logger.info(f"{queue} queue created with {self.message_ttl}ms TTL")
+        self.connection_status["connected"] = self.connection is not None and self.channel is not None
+
+    async def fanout_to_queues(self, message: WebSocketMessage):
+        try:
+            self.logger.debug(f"Received message: {message}")
+            self.connection_status["messages_received"] += 1
+            self.connection_status["last_message_time"] = datetime.now().isoformat()
+
+            serialized_message = WebSocketMessageSerde.serialize(message)
+            self.logger.debug(f"Serialized message: {serialized_message}")
+
+            encoded_message = serialized_message.encode()
+            rabbit_message = Message(
+                body=encoded_message,
+                delivery_mode=DeliveryMode.PERSISTENT,  # Survive broker restart
+                content_type="application/json",
+                timestamp=datetime.now(),
+            )
+
+            # Publish to event-specific queues
+            queue_name = QueueNameResolver.queue_name_for_web_socket_message(message)
+            self.logger.debug(f"Queue name: {queue_name}")
+
+            await self.channel.default_exchange.publish(rabbit_message, routing_key=queue_name)
+            self.connection_status[queue_name] += 1
+
+        except Exception as e:
+            self.logger.error(f"Error publishing to RabbitMQ: {e}")

kuhl_haus/mdp/integ/utils.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def get_massive_api_key():
+    logger.info("Getting Massive API key...")
+    api_key = os.environ.get("MASSIVE_API_KEY")
+    if not api_key:
+        logger.info("MASSIVE_API_KEY environment variable not set; trying Massive API key...")
+        api_key = os.environ.get("POLYGON_API_KEY")
+    if not api_key:
+        logger.info("POLYGON_API_KEY environment variable not set; trying Massive API key file...")
+        api_key_path = '/app/polygon_api_key.txt'
+        with open(api_key_path, 'r') as f:
+            api_key = f.read().strip()
+    if not api_key:
+        logger.error("No Massive API key found")
+        raise ValueError("MASSIVE_API_KEY environment variable not set")
+    logger.info("Done.")
+    return api_key

kuhl_haus/mdp/integ/web_socket_message_serde.py

@@ -0,0 +1,143 @@
+import json
+from argparse import ArgumentTypeError
+from typing import Union
+
+from massive.websocket.models import (
+    WebSocketMessage,
+    EquityAgg,
+    EquityQuote,
+    EquityTrade,
+    LimitUpLimitDown,
+    EventType
+)
+
+
+class WebSocketMessageSerde:
+    @staticmethod
+    def serialize(message: WebSocketMessage) -> str:
+        if isinstance(message, EquityTrade):
+            return WebSocketMessageSerde.serialize_equity_trade(message)
+        elif isinstance(message, EquityAgg):
+            return WebSocketMessageSerde.serialize_equity_agg(message)
+        elif isinstance(message, EquityQuote):
+            return WebSocketMessageSerde.serialize_equity_quote(message)
+        elif isinstance(message, LimitUpLimitDown):
+            return WebSocketMessageSerde.serialize_limit_up_limit_down(message)
+        else:
+            return json.dumps(message)
+
+    @staticmethod
+    def to_dict(message: WebSocketMessage) -> dict:
+        if isinstance(message, EquityTrade):
+            return WebSocketMessageSerde.decode_equity_trade(message)
+        elif isinstance(message, EquityAgg):
+            return WebSocketMessageSerde.decode_equity_agg(message)
+        elif isinstance(message, EquityQuote):
+            return WebSocketMessageSerde.decode_equity_quote(message)
+        elif isinstance(message, LimitUpLimitDown):
+            return WebSocketMessageSerde.decode_limit_up_limit_down(message)
+        else:
+            return json.loads(json.dumps(message))
+
+    @staticmethod
+    def deserialize(serialized_message: str) -> Union[LimitUpLimitDown, EquityAgg, EquityTrade, EquityQuote]:
+        message: dict = json.loads(serialized_message)
+        event_type = message.get("event_type")
+        if event_type == EventType.LimitUpLimitDown.value:
+            return LimitUpLimitDown(**message)
+        elif event_type == EventType.EquityAgg.value:
+            return EquityAgg(**message)
+        elif event_type == EventType.EquityTrade.value:
+            return EquityTrade(**message)
+        elif event_type == EventType.EquityQuote.value:
+            return EquityQuote(**message)
+        else:
+            raise ArgumentTypeError(f"Unsupported message type: {event_type}")
+
+    @staticmethod
+    def decode_limit_up_limit_down(message: LimitUpLimitDown) -> dict:
+        ret: dict = {
+            "event_type": message.event_type,
+            "symbol": message.symbol,
+            "high": message.high_price,
+            "low": message.low_price,
+            "indicators": message.indicators,
+            "tape": message.tape,
+            "timestamp": message.timestamp,
+            "sequence_number": message.sequence_number,
+        }
+        return ret
+
+    @staticmethod
+    def serialize_limit_up_limit_down(message: LimitUpLimitDown) -> str:
+        return json.dumps(WebSocketMessageSerde.decode_limit_up_limit_down(message))
+
+    @staticmethod
+    def decode_equity_agg(message: EquityAgg) -> dict:
+        ret: dict = {
+            "event_type": message.event_type,
+            "symbol": message.symbol,
+            "volume": message.volume,
+            "accumulated_volume": message.accumulated_volume,
+            "official_open_price": message.official_open_price,
+            "vwap": message.vwap,
+            "open": message.open,
+            "close": message.close,
+            "high": message.high,
+            "low": message.low,
+            "aggregate_vwap": message.aggregate_vwap,
+            "average_size": message.average_size,
+            "start_timestamp": message.start_timestamp,
+            "end_timestamp": message.end_timestamp,
+            "otc": message.otc,
+        }
+        return ret
+
+    @staticmethod
+    def serialize_equity_agg(message: EquityAgg) -> str:
+        return json.dumps(WebSocketMessageSerde.decode_equity_agg(message))
+
+    @staticmethod
+    def decode_equity_trade(message: EquityTrade) -> dict:
+        ret: dict = {
+            "event_type": message.event_type,
+            "symbol": message.symbol,
+            "exchange": message.exchange,
+            "id": message.id,
+            "tape": message.tape,
+            "price": message.price,
+            "size": message.size,
+            "conditions": message.conditions,
+            "timestamp": message.timestamp,
+            "sequence_number": message.sequence_number,
+            "trf_id": message.trf_id,
+            "trf_timestamp": message.trf_timestamp
+        }
+        return ret
+
+    @staticmethod
+    def serialize_equity_trade(message: EquityTrade) -> str:
+        return json.dumps(WebSocketMessageSerde.decode_equity_trade(message))
+
+    @staticmethod
+    def decode_equity_quote(message: EquityQuote) -> dict:
+        ret: dict = {
+            "event_type": message.event_type,
+            "symbol": message.symbol,
+            "bid_exchange_id": message.bid_exchange_id,
+            "bid_price": message.bid_price,
+            "bid_size": message.bid_size,
+            "ask_exchange_id": message.ask_exchange_id,
+            "ask_price": message.ask_price,
+            "ask_size": message.ask_size,
+            "condition": message.condition,
+            "indicators": message.indicators,
+            "timestamp": message.timestamp,
+            "sequence_number": message.sequence_number,
+            "tape": message.tape,
+        }
+        return ret
+
+    @staticmethod
+    def serialize_equity_quote(message: EquityQuote) -> str:
+        return json.dumps(WebSocketMessageSerde.decode_equity_quote(message))

kuhl_haus/mdp/models/market_data_analyzer_result.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass()
+class MarketDataAnalyzerResult:
+    """
+    The object passed as data must be JSON serializable.
+
+    Refer to the following for more details:
+    https://docs.python.org/3.12/library/json.html#py-to-json-table
+    """
+    data: Any
+    cache_key: Optional[str] = None
+    cache_ttl: Optional[int] = 0
+    publish_key: Optional[str] = None