import asyncio import time from collections.abc import MutableMapping from signal import Signals from typing import Any, Dict, Iterable, Optional, Callable from amqpworker.conf import logger from amqpworker.connections import Connection, ConnectionsMapping from amqpworker.exceptions import InvalidConnection, InvalidRoute from amqpworker.options import Options, DefaultValues, RouteTypes from amqpworker.rabbitmq.entrypoints import AMQPRouteEntryPointImpl from amqpworker.routes import RoutesRegistry, Route from amqpworker.scheduled import ScheduledThreadPoolExecutor from amqpworker.signals.base import Freezable, Signal from concurrent.futures import ThreadPoolExecutor, Executor from amqpworker.signals.handlers.rabbitmq import RabbitMQ from amqpworker.utils import entrypoint class App(MutableMapping, Freezable): shutdown_os_signals = (Signals.SIGINT, Signals.SIGTERM) handlers = (RabbitMQ(),) def __init__(self, connections: Optional[Iterable[Connection]] = None, loop: Optional[Executor] = None): Freezable.__init__(self) self.routes_registry = RoutesRegistry() self.default_route_options: dict = {} self.loop: Optional[Executor] = loop self.scheduler = None self._state: Dict[Any, Any] = self._get_initial_state() self.connections = ConnectionsMapping() if connections: self.connections.add(connections) self._on_startup: Signal = Signal(self) self._on_shutdown: Signal = Signal(self) for handler in self.handlers: self._on_startup.append(handler.startup) self._on_shutdown.append(handler.shutdown) # for signal_ in self.shutdown_os_signals: # signal.signal(signal_, self.shutdown) # self.loop.add_signal_handler(signal, self.shutdown) def _check_frozen(self): if self.frozen: raise RuntimeError( "You shouldn't change the state of a started application" ) def freeze(self): self.connections.freeze() super(App, self).freeze() def get_connection(self, name: str) -> Connection: try: return self.connections[name] except KeyError as e: raise InvalidConnection( f"There's no Connection with name `{name}` registered " f"in `App.connections`" ) from e def _get_initial_state(self) -> Dict[str, Dict]: # fixme: typeignore reason - https://github.com/python/mypy/issues/4537 return {route_type: {} for route_type in RouteTypes} # type: ignore def __getitem__(self, key): return self._state[key] def __setitem__(self, key, value): self._check_frozen() self._state[key] = value def __delitem__(self, key): self._check_frozen() del self._state[key] def __len__(self): return len(self._state) def __iter__(self): return iter(self._state) @entrypoint def run(self): logger.debug({"event": "Booting App..."}) self.startup() try: while True: time.sleep(10) except KeyboardInterrupt: self.shutdown() def startup(self): """ Causes on_startup signal Should be called in the event loop along with the request handler. """ if self.loop is None: self.loop = ThreadPoolExecutor(thread_name_prefix='AMQPWORKER-',max_workers=min(16, sum(1 for _ in self.routes_registry.amqp_routes))) self._on_startup.send(self) def shutdown(self): """ Schedules an on_startup signal Is called automatically when the application receives a SIGINT or SIGTERM """ return self._on_shutdown.send(self) @property def amqp(self) -> AMQPRouteEntryPointImpl: return AMQPRouteEntryPointImpl(self) def route( self, routes: Iterable[str], type: RouteTypes, options: Optional[dict] = None, **kwargs, ): if options is None: options = {} if not isinstance(type, RouteTypes): raise TypeError( f"type parameter is not a valid RouteTypes." f" Found: '{type}'" ) def wrapper(f): handler = f if not callable(handler): raise TypeError( f"Object passed as handler is not callable: {f}" ) if asyncio.iscoroutinefunction(handler) or ( hasattr(handler, '__call__') and asyncio.iscoroutinefunction(handler.__call__)): raise TypeError( f"handler must be not a coroutine: {handler}" ) self.routes_registry[handler] = { "type": type, "routes": routes, "handler": handler, "options": options, "default_options": self.default_route_options, **kwargs, } return f return wrapper def run_on_startup(self, func: Callable[["App"], None]) -> None: """ Registers a coroutine to be awaited for during app startup """ self._on_startup.append(func) def run_on_shutdown(self, func: Callable[["App"], None]) -> None: """ Registers a coroutine to be awaited for during app shutdown """ self._on_shutdown.append(func) def run_every(self, seconds: int, options: Optional[Dict] = None): """ Registers a coroutine to be called with a given interval """ if options is None: options = {} max_concurrency = options.get( Options.MAX_CONCURRENCY, DefaultValues.RUN_EVERY_MAX_CONCURRENCY ) def wrapper(task: Callable[..., None]): runner = ScheduledThreadPoolExecutor(max_workers=max_concurrency) runner.schedule_at_fixed_delay(task, initial_delay=seconds, period=seconds) # runner = ScheduledTaskRunner( # seconds=seconds, # task=task, # app=self, # max_concurrency=max_concurrency, # ) self._on_startup.append(runner.start) self._on_shutdown.append(runner.stop) if "task_runners" not in self: self["task_runners"] = [] self["task_runners"].append(runner) return task return wrapper def get_connection_for_route(self, route_info: Route): route_connection = route_info.connection if route_connection is None: route_connection = route_info.options.get("connection") connections = self.connections.with_type(route_info.type) if route_connection is not None: if isinstance(route_connection, Connection): return route_connection elif isinstance(route_connection, str): return self.get_connection(name=route_connection) else: # pragma: nocover raise InvalidRoute( f"Expected `Route.connection` to be either `str` or " f"`Connection`. Got `{type(route_connection)}` instead." ) elif len(connections) > 1: raise InvalidRoute( f"Invalid route definition for App. You are trying to " f"define a {route_info.type} into an amqpworker.App " f"with multiple connections without specifying which " f"one to use." ) else: try: return connections[0] except IndexError as e: raise InvalidRoute( f"Invalid route definition for App. You are trying to " f"define a {route_info.type} without an " f"Connection registered on App" ) from e