DeOwl
96548e60ea
Changes: - added rabbitmq hearbeat, consumer and publisher - fixed vqe to run with rabbitmq - added frontend via fasthtml - added proper .env configuration
98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
import os
|
|
from typing import Optional
|
|
|
|
import aio_pika
|
|
from aio_pika.abc import (
|
|
AbstractChannel,
|
|
AbstractRobustConnection,
|
|
)
|
|
from connections.keycloak import (
|
|
get_valid_access_token,
|
|
)
|
|
from connections.local_files import (
|
|
load_client_info,
|
|
)
|
|
|
|
|
|
class RabbitMQManager:
|
|
"""Singleton manager for RabbitMQ connection and channels."""
|
|
|
|
_instance: Optional["RabbitMQManager"] = None
|
|
_connection: Optional[AbstractRobustConnection] = None
|
|
_consumer_channel: Optional[AbstractChannel] = None
|
|
_publisher_channel: Optional[AbstractChannel] = None
|
|
_heartbeat_channel: Optional[AbstractChannel] = None
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
async def connect(self):
|
|
"""Establish the main connection if not already connected."""
|
|
client_info = load_client_info()
|
|
if not client_info:
|
|
return None
|
|
access_token = get_valid_access_token()
|
|
if not access_token:
|
|
return None
|
|
if self._connection is None or self._connection.is_closed:
|
|
print("Creating new RabbitMQ connection...")
|
|
self._connection = await aio_pika.connect_robust(
|
|
host=os.environ["RABBITMQ_HOST"],
|
|
port=int(os.environ["RABBITMQ_PORT"]),
|
|
login=str(client_info["system_id"]),
|
|
password=access_token,
|
|
virtualhost="/",
|
|
)
|
|
print("RabbitMQ connection established")
|
|
return self._connection
|
|
|
|
async def get_consumer_channel(self):
|
|
"""Get channel for consuming messages."""
|
|
if self._connection:
|
|
if self._consumer_channel is None or self._consumer_channel.is_closed:
|
|
self._consumer_channel = await self._connection.channel()
|
|
await self._consumer_channel.set_qos(prefetch_count=1, global_=True)
|
|
print("Consumer channel created")
|
|
return self._consumer_channel
|
|
return None
|
|
|
|
async def get_publisher_channel(self):
|
|
"""Get channel for publishing regular messages."""
|
|
if self._connection:
|
|
if self._publisher_channel is None or self._publisher_channel.is_closed:
|
|
self._publisher_channel = await self._connection.channel()
|
|
print("Publisher channel created")
|
|
return self._publisher_channel
|
|
return None
|
|
|
|
async def get_heartbeat_channel(self):
|
|
"""Get channel for heartbeat publishing."""
|
|
if self._connection:
|
|
if self._heartbeat_channel is None or self._heartbeat_channel.is_closed:
|
|
self._heartbeat_channel = await self._connection.channel()
|
|
print("Heartbeat channel created")
|
|
return self._heartbeat_channel
|
|
return None
|
|
|
|
async def close(self):
|
|
"""Close all channels and the main connection gracefully."""
|
|
print("Closing RabbitMQ channels and connection...")
|
|
|
|
for channel in [
|
|
self._consumer_channel,
|
|
self._publisher_channel,
|
|
self._heartbeat_channel,
|
|
]:
|
|
if channel and not channel.is_closed:
|
|
await channel.close()
|
|
|
|
if self._connection and not self._connection.is_closed:
|
|
await self._connection.close()
|
|
print("RabbitMQ connection closed")
|
|
|
|
|
|
# Create global singleton instance
|
|
rabbitmq_manager = RabbitMQManager()
|