Persist the zookeeper client and connection. A step towards a config watcher.

[python_utils.git] / config.py

diff --git a/config.py b/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 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]
 
         % main.py -h
         usage: main.py [-h]


@@ -72,9 +85,14 @@ import argparse
 import logging
 import os
 import pprint
 import logging
 import os
 import pprint

+import re

 import sys
 from typing import Any, Dict, List, Optional
 
 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.
 
 # 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()
 
 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
 
 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
 
 
     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:
 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():
     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."""
 
 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
     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:
             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.'
             )
             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:
         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)
 
         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
 
 
         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 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
 
     # 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:
     # 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']:
 
     # Also dump the config on stderr if requested.
     if config['config_dump']: