Persist the zookeeper client and connection. A step towards a config watcher.
[python_utils.git] / config.py
index c5813a81145764c05d7af29ce32a07da4ef36ef8..b344631003bb29dc186c0d7db8f2631b8c7ac4bc 100644 (file)
--- a/config.py
+++ b/config.py
@@ -41,7 +41,20 @@ Usage:
 
     If you set this up and remember to invoke config.parse(), all commandline
     arguments will play nicely together.  This is done automatically for you
-    if you're using the bootstrap module's initialize wrapper.::
+    if you're using the :meth:`bootstrap.initialize` decorator on
+    your program's entry point.  See :meth:`python_modules.bootstrap.initialize`
+    for more details.::
+
+        import bootstrap
+
+        @bootstrap.initialize
+        def main():
+            whatever
+
+        if __name__ == '__main__':
+            main()
+
+    Either way, you'll get this behavior from the commandline::
 
         % main.py -h
         usage: main.py [-h]
@@ -72,9 +85,14 @@ import argparse
 import logging
 import os
 import pprint
+import re
 import sys
 from typing import Any, Dict, List, Optional
 
+from kazoo.client import KazooClient
+
+import scott_secrets
+
 # This module is commonly used by others in here and should avoid
 # taking any unnecessary dependencies back on them.
 
@@ -85,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
@@ -241,6 +264,19 @@ def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
     return reordered_action_groups
 
 
+def print_usage() -> None:
+    """Prints the normal help usage message out."""
+    ARGS.print_help()
+
+
+def usage() -> str:
+    """
+    Returns:
+        program usage help text as a string.
+    """
+    return ARGS.format_usage()
+
+
 def _augment_sys_argv_from_environment_variables():
     """Internal.  Look at the system environment for variables that match
     arg names.  This is done via some munging such that:
@@ -256,7 +292,7 @@ def _augment_sys_argv_from_environment_variables():
 
     """
 
-    usage_message = ARGS.format_usage()
+    usage_message = usage()
     optional = False
     var = ''
     for x in usage_message.split():
@@ -294,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
@@ -310,10 +347,36 @@ def _augment_sys_argv_from_loadfile():
             saw_other_args = True
 
     if loadfile is not None:
-        if not os.path.exists(loadfile):
+        zkpath = None
+        if loadfile[:3] == 'zk:':
+            try:
+                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):
+                    raise Exception(
+                        f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
+                    )
+            except Exception as e:
+                raise Exception(
+                    f'ERROR: Error talking with zookeeper while looking for {loadfile}'
+                ) from e
+        elif not os.path.exists(loadfile):
             raise Exception(
                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
             )
+
         if saw_other_args:
             msg = f'Augmenting commandline arguments with those from {loadfile}.'
         else:
@@ -321,9 +384,22 @@ def _augment_sys_argv_from_loadfile():
         print(msg, file=sys.stderr)
         SAVED_MESSAGES.append(msg)
 
-        with open(loadfile, 'r') as rf:
-            newargs = rf.readlines()
-        newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
+        newargs = []
+        if zkpath:
+            try:
+                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
+                ]
+            except Exception as e:
+                raise Exception(f'Error reading {zkpath} from zookeeper.') from e
+            SAVED_MESSAGES.append(f'Loaded config from zookeeper from {zkpath}')
+        else:
+            with open(loadfile, 'r') as rf:
+                newargs = rf.readlines()
+            newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
         sys.argv += newargs
 
 
@@ -336,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
@@ -377,8 +454,33 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
     # Check for savefile and populate it if requested.
     savefile = config['config_savefile']
     if savefile and len(savefile) > 0:
-        with open(savefile, 'w') as wf:
-            wf.write("\n".join(ORIG_ARGV[1:]))
+        data = '\n'.join(ORIG_ARGV[1:])
+        if savefile[:3] == 'zk:':
+            zkpath = savefile[3:]
+            if not zkpath.startswith('/config/'):
+                zkpath = '/config/' + zkpath
+                zkpath = re.sub(r'//+', '/', zkpath)
+            try:
+                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}')
+        else:
+            with open(savefile, 'w') as wf:
+                wf.write(data)
 
     # Also dump the config on stderr if requested.
     if config['config_dump']: