Factories refactor

Author: AlanCoding

dispatcher/__init__.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 
-from dispatcher.main import DispatcherMain
+from dispatcher.factories import from_settings
 
 logger = logging.getLogger(__name__)
 
@@ -12,7 +12,7 @@ def run_service() -> None:
     Before calling this you need to configure by calling dispatcher.config.setup
     """
     loop = asyncio.get_event_loop()
-    dispatcher = DispatcherMain.from_settings()
+    dispatcher = from_settings()
     try:
         loop.run_until_complete(dispatcher.main())
     except KeyboardInterrupt:

dispatcher/brokers/__init__.py

@@ -1,38 +0,0 @@
-import importlib
-from typing import Optional
-from types import ModuleType
-
-from dispatcher.brokers.base import BaseBroker
-from dispatcher.config import settings
-
-
-def get_broker_module(broker_name) -> ModuleType:
-    return importlib.import_module(f'dispatcher.brokers.{broker_name}')
-
-
-def get_async_broker(broker_name, broker_config) -> BaseBroker:
-    broker_module = get_broker_module(broker_name)
-    return broker_module.AsyncBroker(**broker_config)
-
-
-def get_sync_broker(broker_name, broker_config) -> BaseBroker:
-    broker_module = get_broker_module(broker_name)
-    return broker_module.SyncBroker(**broker_config)
-
-
-def get_sync_publisher_from_settings(publish_broker: Optional[str] = None) -> BaseBroker:
-    if publish_broker:
-        pass
-    elif len(settings.brokers) == 1:
-        publish_broker = list(settings.brokers.keys())[0]
-    elif 'default_broker' in settings.publish:
-        publish_broker = settings.publish['default_broker']
-    else:
-        raise RuntimeError(f'Could not determine which broker to publish with between options {list(settings.brokers.keys())}')
-
-    return get_sync_broker(publish_broker, settings.brokers[publish_broker])
-
-
-def get_async_publisher_from_settings() -> BaseBroker:
-    publish_broker = settings.publish['default_broker']
-    return get_async_broker(publish_broker, settings.brokers[publish_broker])

dispatcher/brokers/pg_notify.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterable, Optional, Callable
+from typing import Callable, Iterable, Optional
 
 import psycopg
 
@@ -21,9 +21,22 @@ class PGNotifyBase(BaseBroker):
 
     def __init__(
         self,
+        config: Optional[dict] = None,
         channels: Iterable[str] = ('dispatcher_default',),
         default_publish_channel: Optional[str] = None,
     ) -> None:
+        """
+        channels - listening channels for the service and used for control-and-reply
+        default_publish_channel - if not specified on task level or in the submission
+          by default messages will be sent to this channel.
+          this should be one of the listening channels for messages to be received.
+        """
+        if config:
+            self._config: dict = config.copy()
+            self._config['autocommit'] = True
+        else:
+            self._config = {}
+
         self.channels = channels
         self.default_publish_channel = default_publish_channel
 
@@ -47,6 +60,8 @@ def get_connection_method(self, factory_path: Optional[str] = None) -> Callable:
         else:
             raise RuntimeError('Could not construct connection for lack of config or factory')
 
+    def create_connection(self): ...
+
 
 class AsyncBroker(PGNotifyBase):
     def __init__(
@@ -60,21 +75,17 @@ def __init__(
         if not (config or async_connection_factory or connection):
             raise RuntimeError('Must specify either config or async_connection_factory')
 
-        if config:
-            self._config: dict = config.copy()
-            self._config['autocommit'] = True
-        else:
-            self._config = {}
-
         self._async_connection_factory = async_connection_factory
-        self._connection: Optional[Any] = connection
+        self._connection = connection
 
-        super().__init__(**kwargs)
+        super().__init__(config=config, **kwargs)
 
     async def get_connection(self) -> psycopg.AsyncConnection:
         if not self._connection:
             factory = self.get_connection_method(factory_path=self._async_connection_factory)
-            self._connection = await factory(**self._config)
+            connection = await factory(**self._config)
+            self._connection = connection
+            return connection  # slightly weird due to MyPY
         return self._connection
 
     @staticmethod
@@ -145,20 +156,16 @@ def __init__(
         if not (config or sync_connection_factory or connection):
             raise RuntimeError('Must specify either config or async_connection_factory')
 
-        if config:
-            self._config: Optional[dict] = config.copy()
-            self._config['autocommit'] = True
-        else:
-            self._config = None
-
         self._sync_connection_factory = sync_connection_factory
-        self._connection: Optional[Any] = connection
-        super().__init__(**kwargs)
+        self._connection = connection
+        super().__init__(config=config, **kwargs)
 
     def get_connection(self) -> psycopg.Connection:
         if not self._connection:
             factory = self.get_connection_method(factory_path=self._sync_connection_factory)
-            self._connection = factory(**self._config)
+            connection = factory(**self._config)
+            self._connection = connection
+            return connection
         return self._connection
 
     @staticmethod

dispatcher/config.py

@@ -7,8 +7,8 @@
 
 class DispatcherSettings:
     def __init__(self, config: dict) -> None:
-        self.brokers: list = config.get('brokers', [])
-        self.producers: list = config.get('producers', [])
+        self.brokers: dict = config.get('brokers', [])
+        self.producers: dict = config.get('producers', [])
         self.service: dict = config.get('service', {'max_workers': 3})
         self.publish: dict = config.get('publish', {})
         # TODO: firmly planned sections of config for later

dispatcher/control.py

@@ -5,7 +5,8 @@
 import uuid
 from types import SimpleNamespace
 
-from dispatcher.producers.brokered import BrokeredProducer
+from dispatcher.factories import get_async_publisher_from_settings, get_sync_publisher_from_settings
+from dispatcher.producers import BrokeredProducer
 
 logger = logging.getLogger('awx.main.dispatch.control')
 
@@ -55,10 +56,9 @@ def fatal_error_callback(self, *args):
 
 
 class Control(object):
-    def __init__(self, queue, config=None, async_connection=None):
+    def __init__(self, queue, config=None):
         self.queuename = queue
         self.config = config
-        self.async_connection = async_connection
 
     def running(self, *args, **kwargs):
         return self.control_with_reply('running', *args, **kwargs)
@@ -90,11 +90,8 @@ async def acontrol_with_reply_internal(self, producer, send_data, expected_repli
         return [json.loads(payload) for payload in control_callbacks.received_replies]
 
     def make_producer(self, reply_queue):
-        if self.async_connection:
-            conn_kwargs = {'connection': self.async_connection}
-        else:
-            conn_kwargs = {'config': self.config}
-        return BrokeredProducer(broker='pg_notify', channels=[reply_queue], **conn_kwargs)
+        broker = get_async_publisher_from_settings(channels=[reply_queue])
+        return BrokeredProducer(broker, close_on_exit=True)
 
     async def acontrol_with_reply(self, command, expected_replies=1, timeout=1, data=None):
         reply_queue = Control.generate_reply_queue_name()
@@ -118,9 +115,6 @@ def control_with_reply(self, command, expected_replies=1, timeout=1, data=None):
         start = time.time()
         reply_queue = Control.generate_reply_queue_name()
 
-        if (not self.config) and (not self.async_connection):
-            raise RuntimeError('Must use a new psycopg connection to do control-and-reply')
-
         send_data = {'control': command, 'reply_to': reply_queue}
         if data:
             send_data['control_data'] = data
@@ -139,11 +133,10 @@ def control_with_reply(self, command, expected_replies=1, timeout=1, data=None):
 
     # NOTE: this is the synchronous version, only to be used for no-reply
     def control(self, command, data=None):
-        from dispatcher.brokers.pg_notify import publish_message
-
         send_data = {'control': command}
         if data:
             send_data['control_data'] = data
 
         payload = json.dumps(send_data)
-        publish_message(self.queuename, payload, config=self.config)
+        broker = get_sync_publisher_from_settings()
+        broker.publish_message(channel=self.queuename, message=payload)

dispatcher/factories.py

@@ -0,0 +1,94 @@
+import importlib
+from types import ModuleType
+from typing import Iterable, Optional
+
+from dispatcher import producers
+from dispatcher.brokers.base import BaseBroker
+from dispatcher.config import LazySettings
+from dispatcher.config import settings as global_settings
+from dispatcher.main import DispatcherMain
+
+"""
+Creates objects from settings,
+This is kept separate from the settings and the class definitions themselves,
+which is to avoid import dependencies.
+"""
+
+# ---- Service objects ----
+
+
+def get_broker_module(broker_name) -> ModuleType:
+    "Static method to alias import_module so we use a consistent import path"
+    return importlib.import_module(f'dispatcher.brokers.{broker_name}')
+
+
+def get_async_broker(broker_name: str, broker_config: dict, **overrides) -> BaseBroker:
+    """
+    Given the name of the broker in the settings, and the data under that entry in settings,
+    return the asyncio broker object.
+    """
+    broker_module = get_broker_module(broker_name)
+    kwargs = broker_config.copy()
+    kwargs.update(overrides)
+    return broker_module.AsyncBroker(**kwargs)
+
+
+def producers_from_settings(settings: LazySettings = global_settings) -> Iterable[producers.BaseProducer]:
+    producer_objects = []
+    for broker_name, broker_kwargs in settings.brokers.items():
+        broker = get_async_broker(broker_name, broker_kwargs)
+        producer = producers.BrokeredProducer(broker=broker)
+        producer_objects.append(producer)
+
+    for producer_cls, producer_kwargs in settings.producers.items():
+        producer_objects.append(getattr(producers, producer_cls)(**producer_kwargs))
+
+    return producer_objects
+
+
+def from_settings(settings: LazySettings = global_settings) -> DispatcherMain:
+    """
+    Returns the main dispatcher object, used for running the background task service.
+    You could initialize this yourself, but using the shared settings allows for consistency
+    between the service, publisher, and any other interacting processes.
+    """
+    producers = producers_from_settings(settings=settings)
+    return DispatcherMain(settings.service, producers)
+
+
+# ---- Publisher objects ----
+
+
+def get_sync_broker(broker_name, broker_config) -> BaseBroker:
+    """
+    Given the name of the broker in the settings, and the data under that entry in settings,
+    return the synchronous broker object.
+    """
+    broker_module = get_broker_module(broker_name)
+    return broker_module.SyncBroker(**broker_config)
+
+
+def get_sync_publisher_from_settings(publish_broker: Optional[str] = None, settings: LazySettings = global_settings) -> BaseBroker:
+    if publish_broker:
+        pass
+    elif len(settings.brokers) == 1:
+        publish_broker = list(settings.brokers.keys())[0]
+    elif 'default_broker' in settings.publish:
+        publish_broker = settings.publish['default_broker']
+    else:
+        raise RuntimeError(f'Could not determine which broker to publish with between options {list(settings.brokers.keys())}')
+
+    return get_sync_broker(publish_broker, settings.brokers[publish_broker])
+
+
+def get_async_publisher_from_settings(settings: LazySettings = global_settings, **overrides) -> BaseBroker:
+    """
+    An asynchronous publisher is the ideal choice for submitting control-and-reply actions.
+    This returns an asyncio broker of the default publisher type.
+
+    If channels are specified, these completely replace the channel list from settings.
+    For control-and-reply, this will contain only the reply_to channel, to not receive
+    unrelated traffic.
+    """
+    publish_broker = settings.publish['default_broker']
+    return get_async_broker(publish_broker, settings.brokers[publish_broker], **overrides)