Remove redis connection from top level

augur/api/routes/user.py

@@ -16,7 +16,8 @@
 from augur.api.util import ssl_required
 
 from augur.application.db.models import User, UserSessionToken, RefreshToken
-from augur.tasks.init.redis_connection import redis_connection as redis
+from augur.tasks.init.redis_connection import get_redis_connection
+
 from ..server import app
 
 logger = logging.getLogger(__name__)
@@ -64,6 +65,8 @@ def user_authorize():
     code = secrets.token_hex()
     username = current_user.login_name
 
+    redis = get_redis_connection()
+
     redis.set(code, username, ex=300)
 
     return jsonify({"status": "Validated", "code": code})
@@ -81,6 +84,8 @@ def generate_session(application):
 
     if "code" not in grant_type:
         return jsonify({"status": "Invalid grant type"})
+
+    redis = get_redis_connection()
 
     username = redis.get(code)
     redis.delete(code)

augur/application/cli/_cli_util.py

@@ -5,14 +5,17 @@
 import signal
 from urllib.parse import urlparse
 
-from augur.tasks.init.redis_connection import redis_connection 
+from augur.tasks.init.redis_connection import get_redis_connection
 
 def clear_redis_caches(logger):
     """Clears the redis databases that celery and redis use."""
 
     logger.info("Flushing all redis databases this instance was using")
     celery_purge_command = "celery -A augur.tasks.init.celery_app.celery_app purge -f"
     subprocess.call(celery_purge_command.split(" "))
+
+
+    redis_connection = get_redis_connection()
     redis_connection.flushdb()
 
 

augur/application/cli/backend.py

@@ -22,7 +22,7 @@
 from augur.tasks.github.util.github_api_key_handler import GithubApiKeyHandler
 from augur.tasks.gitlab.gitlab_api_key_handler import GitlabApiKeyHandler
 from augur.tasks.data_analysis.contributor_breadth_worker.contributor_breadth_worker import contributor_breadth_model
-from augur.tasks.init.redis_connection import redis_connection 
+from augur.tasks.init.redis_connection import get_redis_connection 
 from augur.application.db.models import UserRepo
 from augur.application.db.session import DatabaseSession
 from augur.application.logs import AugurLogger
@@ -117,7 +117,7 @@
     logger.info(f'Augur is running at: {"http" if development else "https"}://{host}:{port}')
     logger.info(f"The API is available at '{api_response.json()['route']}'")
 
     processes = start_celery_worker_processes(float(worker_vmem_cap), disable_collection)
 
     celery_beat_schedule_db = os.getenv("CELERYBEAT_SCHEDULE_DB", "celerybeat-schedule.db")
     if os.path.exists(celery_beat_schedule_db):
@@ -265,7 +265,7 @@
     """
     Sends SIGTERM to all Augur server & worker processes
     """
     logger = logging.getLogger("augur.cli")
 
     augur_stop(signal.SIGTERM, logger, ctx.obj.engine)
 
@@ -278,7 +278,7 @@
     """
     Stop collection tasks if they are running, block until complete
     """
     processes = get_augur_processes()
 
     stopped = []
 
@@ -288,7 +288,7 @@
             stopped.append(p)
             p.terminate()
 
     if not len(stopped):
         logger.info("No collection processes found")
         return
 
@@ -297,7 +297,7 @@
 
     killed = []
     while True:
         for i in range(len(alive)):
             if alive[i].status() == psutil.STATUS_ZOMBIE:
                 logger.info(f"KILLING ZOMBIE: {alive[i].pid}")
                 alive[i].kill()
@@ -309,7 +309,7 @@
         for i in reversed(killed):
             alive.pop(i)
 
         if not len(alive):
             break
 
         logger.info(f"Waiting on [{', '.join(str(p.pid for p in alive))}]")
@@ -326,11 +326,11 @@
     """
     Sends SIGKILL to all Augur server & worker processes
     """
     logger = logging.getLogger("augur.cli")
     augur_stop(signal.SIGKILL, logger, ctx.obj.engine)
 
 
 def augur_stop(signal, logger, engine):
     """
     Stops augur with the given signal, 
     and cleans up collection if it was running
@@ -346,7 +346,7 @@
         cleanup_collection_status_and_rabbit(logger, engine)
 
 
 def cleanup_collection_status_and_rabbit(logger, engine):
     clear_redis_caches()
 
     connection_string = get_value("RabbitMQ", "connection_string")
@@ -363,6 +363,8 @@
     logger.info("Flushing all redis databases this instance was using")
     celery_purge_command = "celery -A augur.tasks.init.celery_app.celery_app purge -f"
     subprocess.call(celery_purge_command.split(" "))
+
+    redis_connection = get_redis_connection()
     redis_connection.flushdb()
 
 def clear_all_message_queues(connection_string):

augur/application/cli/cache.py

@@ -4,7 +4,7 @@
 """
 import click
 
-from augur.tasks.init.redis_connection import redis_connection as redis_conn
+from augur.tasks.init.redis_connection import get_redis_connection
 from augur.application.logs import AugurLogger
 from augur.application.cli import test_connection, test_db_connection 
 
@@ -30,6 +30,7 @@ def clear_all():
 
         if user_input in ("y", "Y", "Yes", "yes"):
             logger.info("Clearing call redis databases")
+            redis_conn = get_redis_connection()
             redis_conn.flushall()
             return
 
@@ -46,4 +47,6 @@ def clear():
     """Clears the redis cache specified in the config"""
 
     print("Clearing redis cache that is specified in the config")
+
+    redis_conn = get_redis_connection()
     redis_conn.flushdb()

augur/application/cli/tasks.py

@@ -13,7 +13,6 @@
 
 from augur import instance_id
 from augur.application.logs import AugurLogger
-from augur.tasks.init.redis_connection import redis_connection
 from augur.application.cli import test_connection, test_db_connection 
 from augur.application.cli.backend import clear_rabbitmq_messages, raise_open_file_limit
 

augur/tasks/init/redis_connection.py

@@ -2,6 +2,11 @@
 import redis
 from augur.tasks.init import get_redis_conn_values
 
-redis_db_number, redis_conn_string = get_redis_conn_values()
 
-redis_connection= redis.from_url(f'{redis_conn_string}{redis_db_number+2}', decode_responses=True)
+def get_redis_connection():
+
+    redis_db_number, redis_conn_string = get_redis_conn_values()
+
+    redis_connection= redis.from_url(f'{redis_conn_string}{redis_db_number+2}', decode_responses=True)
+
+    return redis_connection

augur/tasks/util/redis_list.py

@@ -4,7 +4,7 @@
 from typing import Iterable, Any, Union
 
 from collections.abc import MutableSequence
-from augur.tasks.init.redis_connection import redis_connection as redis
+from augur.tasks.init.redis_connection import get_redis_connection
 from augur import instance_id
 
 
@@ -25,6 +25,7 @@ def __init__(self, list_name: str):
         """
 
         self.redis_list_key = f"{instance_id}_{list_name}"
+        self.redis = get_redis_connection()
 
 
     def __len__(self) -> int:
@@ -34,7 +35,7 @@ def __len__(self) -> int:
             length of redis list
         """
 
-        return redis.llen(self.redis_list_key) 
+        return self.redis.llen(self.redis_list_key) 
 
 
     def __iter__(self) -> Any:
@@ -63,7 +64,7 @@ def __getitem__(self, index: int) -> Any:
             Item at requested index
         """
 
-        item = redis.lindex(self.redis_list_key, index)
+        item = self.redis.lindex(self.redis_list_key, index)
         if item is None:
             return None
 
@@ -94,10 +95,10 @@ def __setitem__(self, index: int, data: Union[Any, Iterable[Any]]):
         """
 
 
-        if redis.exists(self.redis_list_key):
-            redis.lset(self.redis_list_key, index, data)
+        if self.redis.exists(self.redis_list_key):
+            self.redis.lset(self.redis_list_key, index, data)
         else:
-            redis.rpush(self.redis_list_key, data)
+            self.redis.rpush(self.redis_list_key, data)
 
 
     def __delitem__(self, index: int) -> None:
@@ -110,10 +111,10 @@ def __delitem__(self, index: int) -> None:
             index: index of desired item in list
         """
 
-        items_before = redis.lpop(self.redis_list_key, index+1)
+        items_before = self.redis.lpop(self.redis_list_key, index+1)
         items_before.pop() 
         if items_before:
-            redis.lpush(self.redis_list_key, *items_before)
+            self.redis.lpush(self.redis_list_key, *items_before)
 
     # def contains(self, value: Any):
     #     """Determiens whether the paramater value is in the list
@@ -143,8 +144,8 @@ def insert(self, index: int, value: Any):
             data: value that will be added
         """
 
-        value_before = redis.lindex(self.redis_list_key, index-1)
-        redis.linsert(self.redis_list_key, "after", value_before, value)
+        value_before = self.redis.lindex(self.redis_list_key, index-1)
+        self.redis.linsert(self.redis_list_key, "after", value_before, value)
 
     def append(self, value):
         """Adds the paramater value to the end of the list
@@ -153,7 +154,7 @@ def append(self, value):
             value: value that is added to end of list
         """
 
-        redis.rpush(self.redis_list_key, value)
+        self.redis.rpush(self.redis_list_key, value)
 
 
     def pop(self, index: int = None):
@@ -170,7 +171,7 @@ def pop(self, index: int = None):
             # This will get a random index from the list and remove it, 
             # decreasing the likelihood of everyone using the same key all the time
             #redis.rpop(self.redis_list_key)
-            redis.spop(self.redis_list_key)
+            self.redis.spop(self.redis_list_key)
 
         else:
             # calls __delitem__
@@ -184,7 +185,7 @@ def remove(self, value):
             value: item being removed
         """
 
-        redis.lrem(self.redis_list_key, 0, value)
+        self.redis.lrem(self.redis_list_key, 0, value)
 
 
     def extend(self, values: Iterable[Any]):
@@ -195,27 +196,27 @@ def extend(self, values: Iterable[Any]):
         """
 
         if values:
-            redis.rpush(self.redis_list_key, *values)
+            self.redis.rpush(self.redis_list_key, *values)
 
 
     def clear(self):
         """Removes all items from the redis list"""
 
-        redis.delete(self.redis_list_key)
+        self.redis.delete(self.redis_list_key)
 
 
     def print_values(self):
         """Prints all the values in the redis list"""
 
-        key_list_length = redis.llen(self.redis_list_key) 
+        key_list_length = self.redis.llen(self.redis_list_key) 
 
         for i in range(key_list_length):
-            print(redis.lindex(self.redis_list_key, i))
+            print(self.redis.lindex(self.redis_list_key, i))
 
 if __name__ == "__main__":
 
-    redis.flushdb()
     redis_list = RedisList("list")
+    redis_list.redis.flushdb()
 
     # redis_list.append(5)
     # redis_list.append(15)
@@ -225,5 +226,5 @@ def print_values(self):
     print("List values")
     redis_list.print_values()
 
-    redis.delete("list")
+    redis_list.redis.delete("list")
 

augur/tasks/util/redis_scalar.py

@@ -4,7 +4,7 @@
 from typing import Iterable, Any, Union
 
 from collections.abc import MutableSequence
-from augur.tasks.init.redis_connection import redis_connection as redis
+from augur.tasks.init.redis_connection import get_redis_connection
 from augur import instance_id
 from redis import exceptions
 import numbers
@@ -17,14 +17,15 @@ def __init__(self, scalar_name: str, default_value: int = 0, override_existing:
         self._scalar_name = scalar_name
 
         self.__value = default_value
+        self.redis = get_redis_connection()
 
         #Check redis to see if key exists in cache
-        if 1 != redis.exists(self.redis_scalar_key) or override_existing:
+        if 1 != self.redis.exists(self.redis_scalar_key) or override_existing:
             #Set value
-            redis.set(self.redis_scalar_key,self.__value)
+            self.redis.set(self.redis_scalar_key,self.__value)
         else:
             #else get the value 
-            self.__value = int(float(redis.get(self.redis_scalar_key)))
+            self.__value = int(float(self.redis.get(self.redis_scalar_key)))
 
     @property
     def value(self):
@@ -34,4 +35,4 @@ def value(self):
     def value(self, otherVal):
         if isinstance(otherVal, numbers.Number):
             self.__value = otherVal
-            redis.set(self.redis_scalar_key,self.__value)    
+            self.redis.set(self.redis_scalar_key,self.__value)    

keyman/KeyClient.py

@@ -1,4 +1,4 @@
-from augur.tasks.init.redis_connection import redis_connection as conn
+from augur.tasks.init.redis_connection import get_redis_connection
 from redis.client import PubSub
 from logging import Logger
 from os import getpid
@@ -29,8 +29,10 @@ def __init__(self, platform: str, logger: Logger):
         if not platform:
             raise ValueError("Platform must not be empty")
 
-        self.stdout = conn
-        self.stdin: PubSub = conn.pubsub(ignore_subscribe_messages = True)
+        self.conn = get_redis_connection()
+
+        self.stdout = self.conn
+        self.stdin: PubSub = self.conn.pubsub(ignore_subscribe_messages = True)
         self.stdin.subscribe(f"{self.REQUEST}-{self.id}")
         self.platform = platform
         self.logger = logger
@@ -150,14 +152,15 @@ class KeyPublisher:
         as the process ID is used for async communication between
         the publisher and the orchestrator.
     """
+
     def __init__(self) -> None:
         # Load channel names and IDs from the spec
         for channel in spec["channels"]:
             # IE: self.ANNOUNCE = "augur-oauth-announce"
             setattr(self, channel["name"], channel["id"])
-            
+        self.conn = get_redis_connection()
         self.id = getpid()
-        self.stdin: PubSub = conn.pubsub(ignore_subscribe_messages = True)
+        self.stdin: PubSub = self.conn.pubsub(ignore_subscribe_messages = True)
         self.stdin.subscribe(f"{self.ANNOUNCE}-{self.id}")
 
     def publish(self, key: str, platform: str):
@@ -172,7 +175,7 @@ def publish(self, key: str, platform: str):
             "key_platform": platform
         }
 
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
     def unpublish(self, key: str, platform: str):
         """ Unpublish a key, and remove it from orchestration
@@ -189,7 +192,7 @@ def unpublish(self, key: str, platform: str):
             "key_platform": platform
         }
 
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
     def wait(self, timeout_seconds = 30, republish = False):
         """ Wait for ACK from the orchestrator
@@ -215,7 +218,7 @@ def wait(self, timeout_seconds = 30, republish = False):
         }
 
         listen_delta = 0.1
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
         # Just wait for and consume the next incoming message
         while timeout_seconds >= 0:
@@ -227,7 +230,7 @@ def wait(self, timeout_seconds = 30, republish = False):
             elif timeout_seconds < listen_delta:
                 break
             elif republish:
-                conn.publish(self.ANNOUNCE, json.dumps(message))
+                self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
             time.sleep(listen_delta)
             timeout_seconds -= listen_delta
@@ -245,7 +248,7 @@ def list_platforms(self):
             "requester_id": self.id
         }
 
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
         reply = next(self.stdin.listen())
 
@@ -272,7 +275,7 @@ def list_keys(self, platform):
             "key_platform": platform
         }
 
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
         reply = next(self.stdin.listen())
 
@@ -302,7 +305,7 @@ def list_invalid_keys(self, platform):
             "key_platform": platform
         }
 
-        conn.publish(self.ANNOUNCE, json.dumps(message))
+        self.conn.publish(self.ANNOUNCE, json.dumps(message))
 
         reply = next(self.stdin.listen())
 
@@ -325,4 +328,4 @@ def shutdown(self):
             were sent prior to this message, and will then shut down
             immediately upon processing of the shutdown command
         """
-        conn.publish(self.ANNOUNCE, json.dumps({"type": "SHUTDOWN"}))
+        self.conn.publish(self.ANNOUNCE, json.dumps({"type": "SHUTDOWN"}))