import asyncio
import json
import logging
import uuid
from typing import Any, Dict, List, Optional
import aio_pika
from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractRobustConnection
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.config.amq_configuration import AMQConfiguration
[docs]
class UserManagementServiceException(Exception):
"""Exception raised for errors in the User Management Service."""
[docs]
def __init__(self, exception_type: str, message: str, additional_info: Dict[str, Any] = None):
self.exception_type = exception_type
self.message = message
self.additional_info = additional_info or {}
super().__init__(f"{exception_type}: {message}")
[docs]
class UserManagementServiceClient:
"""
Client for the User Management Service that communicates over RabbitMQ.
Supports operations for token, user, group, tenant, and permission services.
"""
# Exchange name for User Management Service
EXCHANGE_NAME = "cleverthis.clevermicro.auth.endpoints"
# Service routing keys
TOKEN_SERVICE = "token-service-v1"
USER_SERVICE = "user-service-v1"
GROUP_SERVICE = "group-service-v1"
TENANT_SERVICE = "tenant-service-v1"
PERMISSION_SERVICE = "permission-service-v1"
# Response queue name
RESPONSE_QUEUE_NAME = "user-management-service-cleverswarm"
[docs]
def __init__(self, configuration: AMQConfiguration, loop: asyncio.AbstractEventLoop = None):
"""
Initialize the User Management Service client.
Args:
configuration: AMQ configuration
loop: Event loop to use (optional)
"""
self.amq_configuration = configuration
self.loop = loop or asyncio.get_event_loop()
self.connection: Optional[AbstractRobustConnection] = None
self.channel: Optional[AbstractRobustChannel] = None
self.exchange: Optional[aio_pika.RobustExchange] = None
self.response_queue: Optional[aio_pika.RobustQueue] = None
self.response_consumer_tag: Optional[str] = None
self.futures: Dict[str, asyncio.Future] = {}
self.initialized = False
[docs]
async def initialize(self) -> None:
"""Initialize the RabbitMQ connection and set up the exchange and response queue."""
if self.initialized:
return
try:
# Connect to RabbitMQ
self.connection = await aio_pika.connect_robust(
host=self.amq_configuration.dispatch.amq_host,
port=self.amq_configuration.dispatch.amq_port,
login=self.amq_configuration.dispatch.rabbit_mq_user,
password=self.amq_configuration.dispatch.rabbit_mq_password,
loop=self.loop,
)
# Create channel
self.channel = await self.connection.channel()
await self.channel.set_qos(prefetch_count=1)
# Declare exchange
self.exchange = await self.channel.declare_exchange(
name=self.EXCHANGE_NAME,
type=aio_pika.ExchangeType.TOPIC,
durable=True,
)
# Declare response queue
self.response_queue = await self.channel.declare_queue(
name=self.RESPONSE_QUEUE_NAME,
durable=True,
exclusive=False,
auto_delete=False,
)
# Start consuming responses
self.response_consumer_tag = await self.response_queue.consume(
self._on_response_callback,
no_ack=False,
)
self.initialized = True
logging_info("User Management Service client initialized")
except Exception as e:
logging_error(f"Failed to initialize User Management Service client: {e}")
if self.connection and not self.connection.is_closed:
await self.connection.close()
raise
async def _on_response_callback(self, message: AbstractIncomingMessage) -> None:
"""
Handle responses from the User Management Service.
Args:
message: The incoming message
"""
try:
correlation_id = message.correlation_id
if not correlation_id:
logging_error("Received response without correlation ID")
await message.ack()
return
future = self.futures.get(correlation_id)
if not future:
logging_error(f"No pending request found for correlation ID: {correlation_id}")
await message.ack()
return
try:
response_data = json.loads(message.body.decode('utf-8'))
logging_debug(f"Received response: {response_data}")
future.set_result(response_data)
except Exception as e:
future.set_exception(e)
await message.ack()
except Exception as e:
logging_error(f"Error processing response: {e}")
await message.ack()
[docs]
async def close(self) -> None:
"""Close the RabbitMQ connection."""
if self.connection and not self.connection.is_closed:
await self.connection.close()
self.initialized = False
logging_info("User Management Service client closed")
[docs]
async def call_service(
self, service: str, method: str, args: Dict[str, Any] = None
) -> Dict[str, Any]:
"""
Call a method on a User Management Service.
Args:
service: Service routing key (e.g., 'token-service-v1')
method: Method name to call
args: Method arguments
Returns:
Response data
Raises:
UserManagementServiceException: If the service returns an error
"""
if not self.initialized:
await self.initialize()
correlation_id = str(uuid.uuid4())
future = self.loop.create_future()
self.futures[correlation_id] = future
request_data = {
"method": method,
"args": args or {},
}
message = Message(
body=json.dumps(request_data).encode('utf-8'),
content_type='application/json',
correlation_id=correlation_id,
reply_to=self.RESPONSE_QUEUE_NAME,
)
try:
await self.exchange.publish(message, routing_key=service)
logging_info(f"Sent request to {service}.{method}: {args}")
# Wait for response
response = await future
# Process response
if response.get("type") == "OK":
return response.get("result", [])
else:
raise UserManagementServiceException(
response.get("exception", "unknown_error"),
response.get("message", "Unknown error"),
response.get("additional_info", {})
)
except asyncio.CancelledError:
raise
except UserManagementServiceException:
raise
except Exception as e:
logging_error(f"Error calling service {service}.{method}: {e}")
raise
finally:
self.futures.pop(correlation_id, None)
[docs]
async def validate_token(self, token: str) -> Dict[str, Any]:
"""
Validate a JWT token.
Args:
token: JWT token to validate
Returns:
Token information
"""
return await self.call_service(
service=self.TOKEN_SERVICE,
method="validateToken",
args={"token": token}
)
[docs]
async def query_user_by_id(self, user_id: str) -> Dict[str, Any]:
"""
Query user information by user ID.
Args:
user_id: User ID
Returns:
User information
"""
return await self.call_service(
service=self.USER_SERVICE,
method="queryUserById",
args={"userId": user_id}
)
[docs]
async def list_users_in_group(self, group_id: str) -> List[Dict[str, Any]]:
"""
List users in a group.
Args:
group_id: Group ID
Returns:
List of users
"""
return await self.call_service(
service=self.GROUP_SERVICE,
method="listUsersInGroup",
args={"groupId": group_id}
)