From 2966ecdb8c49266d491a1f75791868920d68a510 Mon Sep 17 00:00:00 2001
From: Scott Gasch <scott@gasch.org>
Date: Wed, 17 Aug 2022 11:33:58 -0700
Subject: [PATCH] Persist the zookeeper client and connection.  A step towards
 a config watcher.

---
 config.py        | 57 +++++++++++++++++++++++++++++-------------------
 requirements.txt | 18 +++++----------
 2 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/config.py b/config.py
index f3a2dfe..b344631 100644
--- a/config.py
+++ b/config.py
@@ -103,6 +103,11 @@ SAVED_MESSAGES: List[str] = []
 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
 ORIG_ARGV: List[str] = sys.argv.copy()
 
+# A zookeeper client that is lazily created so as to not incur the
+# latency of connecting to zookeeper for programs that are not reading
+# or writing their config data into zookeeper.
+ZK: Optional[KazooClient] = None
+
 
 class OptionalRawFormatter(argparse.HelpFormatter):
     """This formatter has the same bahavior as the normal argparse text
@@ -325,6 +330,7 @@ def _augment_sys_argv_from_environment_variables():
 def _augment_sys_argv_from_loadfile():
     """Internal.  Augment with arguments persisted in a saved file."""
 
+    global ZK
     loadfile = None
     saw_other_args = False
     grab_next_arg = False
@@ -344,20 +350,21 @@ def _augment_sys_argv_from_loadfile():
         zkpath = None
         if loadfile[:3] == 'zk:':
             try:
-                zk = KazooClient(
-                    hosts=scott_secrets.ZOOKEEPER_NODES,
-                    use_ssl=True,
-                    verify_certs=False,
-                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
-                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                )
-                zk.start()
+                if ZK is None:
+                    ZK = KazooClient(
+                        hosts=scott_secrets.ZOOKEEPER_NODES,
+                        use_ssl=True,
+                        verify_certs=False,
+                        keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                        keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                        certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    )
+                    ZK.start()
                 zkpath = loadfile[3:]
                 if not zkpath.startswith('/config/'):
                     zkpath = '/config/' + zkpath
                     zkpath = re.sub(r'//+', '/', zkpath)
-                if not zk.exists(zkpath):
+                if not ZK.exists(zkpath):
                     raise Exception(
                         f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
                     )
@@ -380,7 +387,8 @@ def _augment_sys_argv_from_loadfile():
         newargs = []
         if zkpath:
             try:
-                contents = zk.get(zkpath)[0]
+                assert ZK
+                contents = ZK.get(zkpath)[0]
                 contents = contents.decode()
                 newargs = [
                     arg.strip('\n') for arg in contents.split('\n') if 'config_savefile' not in arg
@@ -404,6 +412,7 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
     global CONFIG_PARSE_CALLED
     if CONFIG_PARSE_CALLED:
         return config
+    global ZK
 
     # If we're about to do the usage message dump, put the main
     # module's argument group last in the list (if possible) so that
@@ -452,18 +461,20 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
                 zkpath = '/config/' + zkpath
                 zkpath = re.sub(r'//+', '/', zkpath)
             try:
-                zk = KazooClient(
-                    hosts=scott_secrets.ZOOKEEPER_NODES,
-                    use_ssl=True,
-                    verify_certs=False,
-                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
-                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                )
-                zk.start()
-                if zk.exists(zkpath):
-                    zk.delete(zkpath)
-                zk.create(zkpath, data.encode())
+                if not ZK:
+                    ZK = KazooClient(
+                        hosts=scott_secrets.ZOOKEEPER_NODES,
+                        use_ssl=True,
+                        verify_certs=False,
+                        keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                        keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                        certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    )
+                    ZK.start()
+                if not ZK.exists(zkpath):
+                    ZK.create(zkpath, data.encode())
+                else:
+                    ZK.set(zkpath, data.encode())
             except Exception as e:
                 raise Exception(f'Failed to create zookeeper path {zkpath}') from e
             SAVED_MESSAGES.append(f'Saved config to zookeeper in {zkpath}')
diff --git a/requirements.txt b/requirements.txt
index 702b24c..a8673e9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,11 @@
-SpeechRecognition
-aiohttp
 antlr4-python3-runtime
 astral
+beautifulsoup4
 bitstring
 black
 cloudpickle
 coverage
-fastapi[all]
+cv2
 flake8
 holidays
 kazoo
@@ -15,19 +14,14 @@ mypy
 nltk
 numpy
 overrides
-prometheus-fastapi-instrumentator
 psutil
 pychromecast
-pyserial
-python-dateutil
 pylint
+python-dateutil
 pytz
 readchar
 requests
-tinytuya
-wheel
-sqlalchemy
-pymysql
-kazoo
+SpeechRecognition
 sphinx
-threadpoolctl
+sqlalchemy
+tinytuya
-- 
2.47.1