From 3878a12529231c38d23c6c12c912b8ca82128349 Mon Sep 17 00:00:00 2001 From: Argenis Leon Date: Wed, 6 Nov 2019 08:59:02 -0600 Subject: [PATCH] Added support for custom comm object --- optimus/bumblebee.py | 8 +++++--- optimus/optimus.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/optimus/bumblebee.py b/optimus/bumblebee.py index bdd785ac..dfc27330 100644 --- a/optimus/bumblebee.py +++ b/optimus/bumblebee.py @@ -28,18 +28,20 @@ class Comm: Send encrypted message to the Bumblebee """ - def __init__(self, queue_name=None): + def __init__(self, queue_name=None, key=None): # If queue_name was not given try lo load from file if not generate one if queue_name is None: self.queue_name = save_config_key("bumblebee.ini", "DEFAULT", "QueueName", str(uuid.uuid4())) + else: + self.queue_name = queue_name + if key is None: # key is generated as byte convert to base64 so we can saved it in the config file key = Fernet.generate_key() self.key = save_config_key("bumblebee.ini", "DEFAULT", "Key", key.decode()) - else: - self.queue_name = queue_name + self.key = key keys_link = " here".format(FULL_DOMAIN=FULL_DOMAIN, SESSION=self.queue_name, KEY=self.key) diff --git a/optimus/optimus.py b/optimus/optimus.py index f1ca24c0..132f7362 100644 --- a/optimus/optimus.py +++ b/optimus/optimus.py @@ -81,8 +81,8 @@ def __init__(self, session=None, master="local[*]", app_name="optimus", checkpoi if comm is True: Comm.instance = Comm() - # else: - # Comm.instance = comm + else: + Comm.instance = comm if session is None: # Creating Spark Session