Added RabbitMQ connection parameters

dbload/cli.py

@@ -301,8 +301,16 @@ def send(actor_name, **kwargs):
     ctx.infuse()
 
     try:
-        from dramatiq import get_broker, actor as register_actor
-        broker = get_broker()
+        from dramatiq import set_broker, actor as register_actor
+        from dramatiq.brokers.rabbitmq import RabbitmqBroker
+
+        broker = RabbitmqBroker(
+            host=config.broker_host,
+            port=config.broker_port,
+            heartbeat=5,
+            blocked_connection_timeout=60
+        )
+        set_broker(broker)
 
         # Decorate all queries and scenarios as actors
         for query_name in ctx.queries:
@@ -339,10 +347,19 @@ def scheduler(actor_name, cron, **kwargs):
     ctx.infuse()
 
     try:
-        from dramatiq import get_broker, actor as register_actor
+        from dramatiq import set_broker, actor as register_actor
+        from dramatiq.brokers.rabbitmq import RabbitmqBroker
         from apscheduler.schedulers.background import BackgroundScheduler
         from apscheduler.triggers.cron import CronTrigger
 
+        broker = RabbitmqBroker(
+            host=config.broker_host,
+            port=config.broker_port,
+            heartbeat=5,
+            blocked_connection_timeout=60
+        )
+        set_broker(broker)
+
         # Decorate all queries and scenarios as actors
         for query_name in ctx.queries:
             register_actor(ctx.queries[query_name].function)
@@ -351,7 +368,6 @@ def scheduler(actor_name, cron, **kwargs):
 
         # Form a schedule if explicitly requested.
         # Otherwise schedule from config will be used
-        broker = get_broker()
         actors = broker.get_declared_actors()
         if actor_name in actors:
             cron = cron or "* * * * *"
@@ -426,12 +442,7 @@ def worker(**kwargs):
     config = get_config(cli_args)
 
     try:
-        from dramatiq import get_broker, actor as register_actor
         from dramatiq.cli import main as dramatiq_main, CPUS
-        # from dramatiq.brokers.rabbitmq import RabbitmqBroker
-
-        # broker = RabbitmqBroker(host="localhost", port="5672", heartbeat=5, blocked_connection_timeout=60)
-        # dramatiq.set_broker(broker)
 
         if dramatiq_args.processes is None:
             dramatiq_args.processes = CPUS

dbload/config.py

@@ -48,7 +48,10 @@ class Config(ilexconf.Config):
         # Predefined simulations
         predefined_simulations=["sap-hana"],
         # Schedule for APScheduler
-        schedule=None
+        schedule=None,
+        # RabbitMQ connection parameters
+        broker_host="localhost",
+        broker_port="5672",
     )
 
     def __init__(self, cli_args):

dbload/resources/scenarios.py

@@ -139,13 +139,24 @@ def create_sale(con):
 
 
 try:
-    import dramatiq
+    from dramatiq import set_broker, actor
+    from dramatiq.brokers.rabbitmq import RabbitmqBroker
+    from dbload import get_config
+
+    config = get_config()
+    broker = RabbitmqBroker(
+        host=config.broker_host,
+        port=config.broker_port,
+        heartbeat=5,
+        blocked_connection_timeout=60
+    )
+    set_broker(broker)
 
     logger.info("Decorating scenarios")
-    dramatiq.actor(update_employee)
-    dramatiq.actor(create_client)
-    dramatiq.actor(update_client)
-    dramatiq.actor(create_sale)
+    actor(update_employee)
+    actor(create_client)
+    actor(update_client)
+    actor(create_sale)
 
 except ImportError:
     pass

demo/sap-hana/README.md

@@ -38,7 +38,7 @@
 1. Launch dbload worker process in a separate terminal window
 
    ```shell
-   dbload worker --predefined sap-hana --process 1 --thread 1
+   dbload worker --predefined sap-hana --processes 1 --threads 1
    ```
 
 1. Launch dbload scheduler process in another terminal window.

demo/sap-hana/dbload.json

@@ -10,7 +10,9 @@
         "password": "HXEHana1"
     },
     "predefined": "sap-hana",
+    "broker_host": "rabbitmq-cluster",
+    "broker_port": "5672",
     "schedule": {
         "create_sale": "* * * * *"
     }
-}
+}