amq-adapter-python/amqp/adapter/cleverthis_service_adapter.py

import asyncio
import json
import logging
import os
import re
import threading

from aiohttp import ClientSession, ClientResponse
from fastapi import HTTPException
from typing import Dict, List, Tuple, Callable, Coroutine, AnyStr, Any
from urllib.parse import parse_qs
from xml.etree import ElementTree

from amqp.adapter.amq_response_factory import AMQResponseFactory
from amqp.service.multipart_parser import CleverMultiPartParser
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQMessage, AMQResponse

# Track the number of messages currently being processed
current_parallel_executions = 0
lock = threading.Lock()
parallel_workers = os.getenv("PARALLEL_WORKERS", default=10)


def parse_request_body(message: AMQMessage):
    """
    Parses the HTTP POST request body based on the MIME type.

    Args:
        message (AMQMessage): Contains raw request and MIME headers.

    Returns:
        dict: Parsed data as a dictionary.
    """
    content_type = message.headers.get('Content-Type', '')
    if not content_type:
        raise ValueError("Content-Type header is required")

    # Parse JSON
    if content_type == "application/json":
        try:
            return json.loads(message.body())
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON: {e}")

    # Parse URL-encoded form data
    elif content_type == "application/x-www-form-urlencoded":
        try:
            return {k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()}
        except Exception as e:
            raise ValueError(f"Invalid URL-encoded form data: {e}")

    # Parse multipart/form-data
    elif content_type.startswith("multipart/form-data"):
        try:
            parser: CleverMultiPartParser = CleverMultiPartParser(message)
            return parser.form_data
        except Exception as e:
            raise ValueError(f"Invalid multipart/form-data: {e}")

    # Parse plain text
    elif content_type == "text/plain":
        return {"text": message.body()}

    # Parse XML
    elif content_type == "application/xml":
        try:
            root = ElementTree.fromstring(message.body())
            return {elem.tag: elem.text for elem in root}
        except ElementTree.ParseError as e:
            raise ValueError(f"Invalid XML: {e}")

    else:
        raise ValueError(f"Unsupported Content-Type: {content_type}")


# =============  POST Helper function =============
async def call_cleverthis_post_put_api(
        message: AMQMessage,
        args: Any,
        api_method: Callable[[Any, Any], Coroutine],
        success_code: int,
        authenticated_user: Any
) -> AMQResponse:
    """
    Invokes the provided API method with the given arguments and authenticated user,
    and returns the response as AMQResponse.
    :param message: AMQMessage containing the details of the call and the payload
    :param args: map of the arguments to pass to the API method
    :param api_method: Callable - the method to invoke. Last parameter is the authenticated user
    :param success_code: HTTP code to be returned when operation suceeds
    :param authenticated_user: Authenticated user. The actual class depends on the service
    :return: AMQResponse class
    """
    try:
        coroutine_call: Coroutine = api_method(args, authenticated_user)
        result: str = await coroutine_call
        return AMQResponseFactory.of(
            message.id, success_code, message.headers.get("Content-Type", "application/json"),
            result.encode('utf-8'), message.trace_info
        )
    except HTTPException as http_error:
        return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
                                     str(http_error.detail).encode('utf-8'), message.trace_info)
    except Exception as error:
        return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)


# =============  GET Helper function =============
async def call_cleverthis_get_api(
        message: AMQMessage,
        pattern: str,
        api_method: Callable[[Tuple[AnyStr | Any, ...], Any], Coroutine],
        authenticated_user: Any
) -> AMQResponse:
    """
    invokes provided Callable that is expected to be a GET API method, and returns the response as AMQResponse
    :param message: AMQMessage containing the details of the call
    :param pattern: If there are path variables, use this pattern to extract them
    :param api_method: Callable - the method to invoke. Last parameter is the authenticated user
    :param authenticated_user: Authenticated user. The actual class depends on the service
    :return:
    """
    try:
        match = re.search(pattern, message.path)
        if match:
            crt: Coroutine = api_method(match.groups(), authenticated_user)
            result: str = await crt
            return AMQResponseFactory.of(message.id, 200, "text/plain", result.encode('utf-8'), message.trace_info)
        else:
            _msg = f"Unexpected URL: {message.path}, expected pattern: {pattern}"
            return AMQResponseFactory.of(message.id, 404, "text/plain", _msg.encode('utf-8'), message.trace_info)
    except HTTPException as http_error:
        return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
                                     str(http_error.detail).encode('utf-8'), message.trace_info)
    except Exception as error:
        return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)


# =============  CleverThisServiceAdapter =============
class CleverThisServiceAdapter:
    """
    A base class for CleverThis service adapter that handles the incoming AMQP messages and routes them to the
    appropriate CleverThisService API endpoint.
    """
    def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
        self.amq_configuration = amq_configuration
        self.session = session
        self.service_host = self.amq_configuration.dispatch.service_host
        self.service_port = self.amq_configuration.dispatch.service_port

    async def get_current_user(self, token: str):
        """
        To Be Overriden in subclass for given CleverThis Service, to return the current user based on the token,
        in format appropriate for the service.
        :return: A user object or None if the user is not authenticated.
        """
        return None

    def sync_on_message(self, message: AMQMessage) -> AMQResponse:
        """
        Synchronous version of on_message method. This is a wrapper around the async on_message method.
        :param message: AMQMessage
        :return: AMQResponse
        """
        return asyncio.run(self.on_message(message))

    async def on_message(self, message) -> AMQResponse:
        """
        Called when a message is received from the AMQP.
        :param message:
        :return: AMQP response class with operation status code and optional error details.
        """
        DEBUG_NO_AUTH = 1
        global current_parallel_executions

        # Increment the counter for parallel executions
        with lock:
            current_parallel_executions += 1
            if current_parallel_executions > parallel_workers - 2:
                logging.warn("Warning: CApacity close to depleted!")

        _auth: str = message.headers.get('Authorization', '')
        _auth_user: Any | None = None
        if 'Bearer' in _auth:
            try:
                token = _auth.split(' ')[1]
                _auth_user = await self.get_current_user(token)
            except Exception as error:
                logging.error(f"Error getting user: {error}")
        amq_response: AMQResponse | None = None
        if _auth_user or DEBUG_NO_AUTH == 1:
            method = message.method.lower()
            if method == 'get':
                amq_response = await self.get(message, _auth_user)
            elif method == 'put':
                amq_response = await self.put(message, _auth_user)
            elif method == 'post':
                amq_response = await self.post(message, _auth_user)
            elif method == 'head':
                amq_response = await self.head(message, _auth_user)
            elif method == 'delete':
                amq_response = await self.delete(message, _auth_user)
            else:
                raise ValueError(f"Unexpected HTTP method: {message.method}")
        else:
            amq_response = AMQResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)
        with lock:
            current_parallel_executions -= 1
        return amq_response

    # ==================================================================================================
    # =====================  C l e v e r T h i s    A P I    m e t h o d s  ============================
    # ==================================================================================================

    async def get(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
        """
        Handles the GET requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
        It will attempt to call the REST endpoint directly using HTTP protocol.
        :param message: message from the AMQP that contains the GET request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """

        # unmatched path - try to invoke the service directly on this path
        url = f"http://{self.service_host}:{self.service_port}{message.path}"
        headers = {**message.headers, **message.trace_info}
        logging.info(f" [*] SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}")
        if self.session:
            _http_resp: ClientResponse = await self.session.get(url, headers=headers)
            return AMQResponseFactory.of(message.id, _http_resp.status, _http_resp.content_type,
                                         await _http_resp.read(),
                                         message.trace_info)
        return AMQResponseFactory.of(message.id, 404, "text/plain",
                                     str("Destination location does not exists").encode('utf-8'), message.trace_info)

    async def put(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
        """
        Handles the PUT requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
        It will attempt to call the REST endpoint directly using HTTP protocol.
        :param message: message from the AMQP that contains the PUT request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        return await self.handle_possible_form(
            self.session.put(
                f"http://{self.service_host}:{self.service_port}{message.path}",
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        )

    async def post(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
        """
        Handles the POST requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
        It will attempt to call the REST endpoint directly using HTTP protocol.
        :param message: message from the AMQP that contains the POST request, including body and headers
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        return await self.handle_possible_form(
            message,
            self.session.post(
                f"http://{self.service_host}:{self.service_port}{message.path}",
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        )

    async def handle_possible_form(self, message: AMQMessage, request_coroutine, auth_user: Any) -> AMQResponse:
        args = parse_request_body(message)
        return await request_coroutine

    async def head(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
        """
        Handles the HEAD requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
        It will attempt to call the REST endpoint directly using HTTP protocol.
        :param message: message from the AMQP that contains the POST request, including body and headers
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        raise NotImplementedError

    async def delete(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
        """
        Handles the DELETE requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
        It will attempt to call the REST endpoint directly using HTTP protocol.
        :param message: message from the AMQP that contains the POST request, including body and headers
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        raise NotImplementedError

    def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
        for header, values in message_headers.items():
            if header.lower() == 'content-type':
                return values[0].split(';')[0]
        return 'application/json'