amq-adapter-python/amqp/rabbitmq/user_management_service_client.py

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


class UserManagementServiceException(Exception):
    """Exception raised for errors in the User Management Service."""

    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}")


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"

    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

    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()

    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")

    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)

    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}
        )

    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}
        )

    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}
        )