amqp-worker/amqpworker/connections.py

import abc
import collections
from collections.abc import KeysView, ValuesView
from typing import (
    Any,
    Counter,
    Dict,
    ItemsView,
    Iterable,
    List,
    Mapping,
    Optional,
    Type,
    Union,
)

from pydantic import BaseModel, validator

from amqpworker.conf import settings
from amqpworker.easyqueue.queue import JsonQueue
from amqpworker.exceptions import InvalidConnection
from amqpworker.options import RouteTypes
from amqpworker.signals.base import Freezable


class Connection(BaseModel, abc.ABC):
    """
    Common ancestral for all Connection classes that auto generates a
    connection name and is responsible for keeping track of new connections on
    the ConnectionsMapping
    """

    route_type: RouteTypes
    name: Optional[str] = None


class ConnectionsMapping(Mapping[str, Connection], Freezable):
    """
    A mapping (Connection.name->Connection) of all available connections that
    also keeps a counter for each connection type
    """

    def __getitem__(self, k: str) -> Connection:
        return self._data[k]

    def __len__(self) -> int:
        return len(self._data)

    def __iter__(self):
        return iter(self._data)

    def __init__(self) -> None:
        Freezable.__init__(self)
        self._data: Dict[str, Connection] = {}
        self.counter: Counter[Type[Connection]] = collections.Counter()

    def __contains__(self, item):
        if isinstance(item, Connection):
            return item in self.values()
        return super(ConnectionsMapping, self).__contains__(item)

    def __setitem__(self, key: str, value: Connection) -> None:
        if self.frozen:
            raise RuntimeError(
                "You shouldn't change the state of ConnectionsMapping "
                "after App startup"
            )

        if key is None:
            key = id(value)

        if key in self:
            raise InvalidConnection(
                f"Invalid connection: `{value}`. "
                f"The name `{key}` already exists in {self.__class__.__name__}"
            )
        self._data[key] = value
        self.counter[value.__class__] += 1

    def __delitem__(self, key: str) -> None:
        if self.frozen:
            raise RuntimeError(
                "You shouldn't change the state of ConnectionsMapping "
                "after App startup"
            )
        del self._data[key]

    def add(self, connections: Iterable[Connection]) -> None:
        for conn in connections:
            self[conn.name] = conn  # type: ignore

    def with_type(self, route_type: RouteTypes) -> List["Connection"]:
        # todo: manter uma segunda estrutura de dados ou aceitar O(n) sempre que chamado?
        return [conn for conn in self.values() if conn.route_type == route_type]


_TYPE_COUNTER: Counter[Type[Connection]] = collections.Counter()

Message = Union[List, Dict]


class AMQPConnection(Connection):
    hostname: str
    port: int
    username: str
    password: str
    route_type = RouteTypes.AMQP_RABBITMQ
    prefetch: int = settings.AMQP_DEFAULT_PREFETCH_COUNT
    heartbeat: int = settings.AMQP_DEFAULT_HEARTBEAT
    name: Optional[str] = None
    connections: Dict[str, JsonQueue] = {}

    class Config:
        arbitrary_types_allowed = True

    @validator("connections", pre=True, always=True, check_fields=False)
    def set_connections(cls, v):
        return v or {}

    def __len__(self) -> int:
        return len(self.connections)

    def __iter__(self):
        return iter(self.connections)

    def __getitem__(self, key: str) -> JsonQueue:
        """
        Gets a JsonQueue instance for a given virtual host

        :param key: The virtual host of the connection
        :return: An instance of the connection
        """
        try:
            return self.connections[key]
        except KeyError:
            conn: JsonQueue = JsonQueue(
                host=self.hostname,
                port=self.port,
                username=self.username,
                password=self.password,
                virtual_host=key,
            )
            self.connections[key] = conn
            return conn

    def keys(self):
        return KeysView(self)

    def items(self):
        return ItemsView(self)

    def values(self):
        return ValuesView(self)

    def register(self, queue: JsonQueue) -> None:
        self.connections[queue.virtual_host] = queue

    def put(
            self,
            routing_key: str,
            data: Any = None,
            serialized_data: Union[str, bytes, None] = None,
            exchange: str = "",
            vhost: str = settings.AMQP_DEFAULT_VHOST,
            properties: Optional[dict] = None,
            mandatory: bool = False,
            immediate: bool = False,
    ):
        conn = self[vhost]
        return conn.put(
            routing_key=routing_key,
            data=data,
            serialized_data=serialized_data,
            exchange=exchange,
            properties=properties,
            mandatory=mandatory,
            immediate=immediate,
        )