Fix forecasting service
This commit is contained in:
Urtzi Alfaro
@@ -1,5 +1,5 @@
|
||||
# ================================================================
|
||||
# services/forecasting/app/services/messaging.py
|
||||
# services/forecasting/app/services/messaging.py
|
||||
# ================================================================
|
||||
"""
|
||||
Messaging service for event publishing and consuming
|
||||
@@ -10,89 +10,126 @@ import json
|
||||
from typing import Dict, Any
|
||||
import asyncio
|
||||
|
||||
from shared.messaging.rabbitmq import RabbitMQPublisher, RabbitMQConsumer
|
||||
from shared.messaging.rabbitmq import RabbitMQClient
|
||||
from shared.messaging.events import (
|
||||
TrainingCompletedEvent,
|
||||
DataImportedEvent,
|
||||
ForecastGeneratedEvent,
|
||||
)
|
||||
from app.core.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# Global messaging instances
|
||||
publisher = None
|
||||
consumer = None
|
||||
# Global messaging instance
|
||||
rabbitmq_client = None
|
||||
|
||||
async def setup_messaging():
|
||||
"""Initialize messaging services"""
|
||||
global publisher, consumer
|
||||
|
||||
global rabbitmq_client
|
||||
|
||||
try:
|
||||
# Initialize publisher
|
||||
publisher = RabbitMQPublisher(settings.RABBITMQ_URL)
|
||||
await publisher.connect()
|
||||
|
||||
# Initialize consumer
|
||||
consumer = RabbitMQConsumer(settings.RABBITMQ_URL)
|
||||
await consumer.connect()
|
||||
|
||||
rabbitmq_client = RabbitMQClient(settings.RABBITMQ_URL, service_name="forecasting_service")
|
||||
await rabbitmq_client.connect()
|
||||
|
||||
# Set up event handlers
|
||||
await consumer.subscribe("training.model.updated", handle_model_updated)
|
||||
await consumer.subscribe("data.weather.updated", handle_weather_updated)
|
||||
|
||||
# We need to adapt the callback to accept aio_pika.IncomingMessage
|
||||
await rabbitmq_client.consume_events(
|
||||
exchange_name="training.events",
|
||||
queue_name="forecasting_model_update_queue",
|
||||
routing_key="training.completed", # Assuming model updates are part of training.completed events
|
||||
callback=handle_model_updated_message
|
||||
)
|
||||
await rabbitmq_client.consume_events(
|
||||
exchange_name="data.events",
|
||||
queue_name="forecasting_weather_update_queue",
|
||||
routing_key="data.weather.updated", # This needs to match the actual event type if different
|
||||
callback=handle_weather_updated_message
|
||||
)
|
||||
|
||||
logger.info("Messaging setup completed")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to setup messaging", error=str(e))
|
||||
raise
|
||||
|
||||
async def cleanup_messaging():
|
||||
"""Cleanup messaging connections"""
|
||||
global publisher, consumer
|
||||
|
||||
global rabbitmq_client
|
||||
|
||||
try:
|
||||
if consumer:
|
||||
await consumer.close()
|
||||
if publisher:
|
||||
await publisher.close()
|
||||
|
||||
if rabbitmq_client:
|
||||
await rabbitmq_client.disconnect()
|
||||
|
||||
logger.info("Messaging cleanup completed")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error during messaging cleanup", error=str(e))
|
||||
|
||||
async def publish_forecast_completed(data: Dict[str, Any]):
|
||||
"""Publish forecast completed event"""
|
||||
if publisher:
|
||||
await publisher.publish("forecasting.forecast.completed", data)
|
||||
if rabbitmq_client:
|
||||
event = ForecastGeneratedEvent(service_name="forecasting_service", data=data, event_type="forecast.completed")
|
||||
await rabbitmq_client.publish_forecast_event(event_type="completed", forecast_data=event.to_dict())
|
||||
|
||||
async def publish_alert_created(data: Dict[str, Any]):
|
||||
"""Publish alert created event"""
|
||||
if publisher:
|
||||
await publisher.publish("forecasting.alert.created", data)
|
||||
# Assuming 'alert.created' is a type of forecast event, or define a new exchange/publisher method
|
||||
if rabbitmq_client:
|
||||
event = ForecastGeneratedEvent(service_name="forecasting_service", data=data, event_type="alert.created")
|
||||
await rabbitmq_client.publish_forecast_event(event_type="alert.created", forecast_data=event.to_dict())
|
||||
|
||||
async def publish_batch_completed(data: Dict[str, Any]):
|
||||
"""Publish batch forecast completed event"""
|
||||
if publisher:
|
||||
await publisher.publish("forecasting.batch.completed", data)
|
||||
if rabbitmq_client:
|
||||
event = ForecastGeneratedEvent(service_name="forecasting_service", data=data, event_type="forecast.batch.completed")
|
||||
await rabbitmq_client.publish_forecast_event(event_type="batch.completed", forecast_data=event.to_dict())
|
||||
|
||||
# Event handlers
|
||||
|
||||
# Event handler wrappers for aio_pika messages
|
||||
async def handle_model_updated_message(message: Any):
|
||||
async with message.process():
|
||||
try:
|
||||
event_data = json.loads(message.body.decode())
|
||||
# Assuming the actual event data is nested under a 'data' key within the event dictionary
|
||||
await handle_model_updated(event_data.get("data", {}))
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error("Failed to decode model updated message JSON", error=str(e), body=message.body)
|
||||
except Exception as e:
|
||||
logger.error("Error processing model updated message", error=str(e), body=message.body)
|
||||
|
||||
async def handle_weather_updated_message(message: Any):
|
||||
async with message.process():
|
||||
try:
|
||||
event_data = json.loads(message.body.decode())
|
||||
# Assuming the actual event data is nested under a 'data' key within the event dictionary
|
||||
await handle_weather_updated(event_data.get("data", {}))
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error("Failed to decode weather updated message JSON", error=str(e), body=message.body)
|
||||
except Exception as e:
|
||||
logger.error("Error processing weather updated message", error=str(e), body=message.body)
|
||||
|
||||
|
||||
# Original Event handlers (now called from the message wrappers)
|
||||
async def handle_model_updated(data: Dict[str, Any]):
|
||||
"""Handle model updated event from training service"""
|
||||
try:
|
||||
logger.info("Received model updated event",
|
||||
logger.info("Received model updated event",
|
||||
model_id=data.get("model_id"),
|
||||
tenant_id=data.get("tenant_id"))
|
||||
|
||||
|
||||
# Clear model cache for this model
|
||||
# This will be handled by PredictionService
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error handling model updated event", error=str(e))
|
||||
|
||||
async def handle_weather_updated(data: Dict[str, Any]):
|
||||
"""Handle weather data updated event"""
|
||||
try:
|
||||
logger.info("Received weather updated event",
|
||||
logger.info("Received weather updated event",
|
||||
date=data.get("date"))
|
||||
|
||||
|
||||
# Could trigger re-forecasting if needed
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error handling weather updated event", error=str(e))
|
||||
logger.error("Error handling weather updated event", error=str(e))
|
||||
Reference in New Issue
Block a user