80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
"""
|
|
Energy calculation worker script.
|
|
|
|
This script runs an energy calculation worker that consumes tasks
|
|
from energy_calculation_tasks queue and processes them.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add parent directory to path for imports
|
|
sys_path = str(Path(__file__).parent.parent)
|
|
if sys_path not in sys.path:
|
|
sys.path.insert(0, sys_path)
|
|
|
|
from app.queues.rabbitmq_client import create_rabbitmq_client
|
|
from app.queues.consumers import consume_energy_calculation_tasks
|
|
from app.workers.energy_worker import create_energy_worker
|
|
from app.database import get_db
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def main():
|
|
"""Main entry point for energy calculation worker."""
|
|
# Load configuration from environment variables
|
|
rabbitmq_url = os.getenv(
|
|
'RABBITMQ_URL',
|
|
'amqp://guest:guest@localhost:5672'
|
|
)
|
|
|
|
# Create RabbitMQ client
|
|
try:
|
|
logger.info("🔗 Connecting to RabbitMQ...")
|
|
client = create_rabbitmq_client(rabbitmq_url=rabbitmq_url)
|
|
client.connect()
|
|
logger.info("✅ Connected to RabbitMQ")
|
|
except Exception as e:
|
|
logger.error(f"❌ Failed to connect to RabbitMQ: {e}")
|
|
sys.exit(1)
|
|
|
|
# Create energy calculation worker
|
|
worker = create_energy_worker()
|
|
|
|
# Define task callback
|
|
def task_callback(task_data, db):
|
|
"""Callback to process energy calculation tasks."""
|
|
return worker.execute_energy_calculation_task(task_data, db)
|
|
|
|
# Start consuming tasks
|
|
try:
|
|
logger.info("🚀 Energy calculation worker started")
|
|
logger.info("📥 Waiting for tasks from energy_calculation_tasks queue...")
|
|
|
|
consume_energy_calculation_tasks(
|
|
client=client,
|
|
callback=task_callback,
|
|
db_session_factory=get_db
|
|
)
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("⏹️ Energy calculation worker stopped by user")
|
|
except Exception as e:
|
|
logger.error(f"❌ Energy calculation worker error: {e}")
|
|
finally:
|
|
logger.info("🔌 Closing RabbitMQ connection...")
|
|
client.close()
|
|
logger.info("👋 Energy calculation worker shutdown complete")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|