Persist the zookeeper client and connection. A step towards a config watcher.
[python_utils.git] / config.py
index 46caad15c9db04e44a37e4acd5ccad74161b0643..b344631003bb29dc186c0d7db8f2631b8c7ac4bc 100644 (file)
--- a/config.py
+++ b/config.py
@@ -1,68 +1,83 @@
 #!/usr/bin/env python3
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Global configuration driven by commandline arguments, environment variables
 and saved configuration files.  This works across several modules.
 
 Usage:
 
-    module.py:
-    ----------
-    import config
-
-    parser = config.add_commandline_args(
-        "Module",
-        "Args related to module doing the thing.",
-    )
-    parser.add_argument(
-        "--module_do_the_thing",
-        type=bool,
-        default=True,
-        help="Should the module do the thing?"
-    )
-
-    main.py:
-    --------
-    import config
-
-    def main() -> None:
+    In your file.py::
+
+        import config
+
         parser = config.add_commandline_args(
-            "Main",
-            "A program that does the thing.",
+            "Module",
+            "Args related to module doing the thing.",
         )
         parser.add_argument(
-            "--dry_run",
+            "--module_do_the_thing",
             type=bool,
-            default=False,
-            help="Should we really do the thing?"
+            default=True,
+            help="Should the module do the thing?"
         )
-        config.parse()   # Very important, this must be invoked!
+
+    In your main.py::
+
+        import config
+
+        def main() -> None:
+            parser = config.add_commandline_args(
+                "Main",
+                "A program that does the thing.",
+            )
+            parser.add_argument(
+                "--dry_run",
+                type=bool,
+                default=False,
+                help="Should we really do the thing?"
+            )
+            config.parse()   # Very important, this must be invoked!
 
     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
 
-    % main.py -h
-    usage: main.py [-h]
-                   [--module_do_the_thing MODULE_DO_THE_THING]
-                   [--dry_run DRY_RUN]
+        if __name__ == '__main__':
+            main()
 
-    Module:
-      Args related to module doing the thing.
+    Either way, you'll get this behavior from the commandline::
 
-      --module_do_the_thing MODULE_DO_THE_THING
-                   Should the module do the thing?
+        % main.py -h
+        usage: main.py [-h]
+                       [--module_do_the_thing MODULE_DO_THE_THING]
+                       [--dry_run DRY_RUN]
 
-    Main:
-      A program that does the thing
+        Module:
+          Args related to module doing the thing.
 
-      --dry_run
-                   Should we really do the thing?
+          --module_do_the_thing MODULE_DO_THE_THING
+                       Should the module do the thing?
+
+        Main:
+          A program that does the thing
+
+          --dry_run
+                       Should we really do the thing?
 
     Arguments themselves should be accessed via
-    config.config['arg_name'].  e.g.
+    :code:`config.config['arg_name']`.  e.g.::
 
-    if not config.config['dry_run']:
-        module.do_the_thing()
+        if not config.config['dry_run']:
+            module.do_the_thing()
 
 """
 
@@ -70,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.
 
@@ -83,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
@@ -90,6 +115,33 @@ class OptionalRawFormatter(argparse.HelpFormatter):
     "RAW|".  In that case, the line breaks are preserved and the text
     is not wrapped.
 
+    Use this, for example, when you need the helptext of an argument
+    to have its spacing preserved exactly, e.g.::
+
+        args.add_argument(
+            '--mode',
+            type=str,
+            default='PLAY',
+            choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
+            metavar='MODE',
+            help='''RAW|Our mode of operation.  One of:
+
+                PLAY = play wordle with me!  Pick a random solution or
+                       specify a solution with --template.
+
+               CHEAT = given a --template and, optionally, --letters_in_word
+                       and/or --letters_to_avoid, return the best guess word;
+
+            AUTOPLAY = given a complete word in --template, guess it step
+                       by step showing work;
+
+            SELFTEST = autoplay every possible solution keeping track of
+                       wins/losses and average number of guesses;
+
+          PRECOMPUTE = populate hash table with optimal guesses.
+            ''',
+        )
+
     """
 
     def _split_lines(self, text, width):
@@ -104,6 +156,11 @@ ARGS = argparse.ArgumentParser(
     formatter_class=OptionalRawFormatter,
     fromfile_prefix_chars="@",
     epilog=f'{PROGRAM_NAME} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
+    # I don't fully understand why but when loaded by sphinx sometimes
+    # the same module is loaded many times causing any arguments it
+    # registers via module-level code to be redefined.  Work around
+    # this iff the program is 'sphinx-build'
+    conflict_handler='resolve' if PROGRAM_NAME == 'sphinx-build' else 'error',
 )
 
 # Keep track of if we've been called and prevent being called more
@@ -120,8 +177,16 @@ config: Dict[str, Any] = {}
 # It would be really nice if this shit worked from interactive python
 
 
-def add_commandline_args(title: str, description: str = ""):
-    """Create a new context for arguments and return a handle."""
+def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
+    """Create a new context for arguments and return a handle.
+
+    Args:
+        title: A title for your module's commandline arguments group.
+        description: A helpful description of your module.
+
+    Returns:
+        An argparse._ArgumentGroup to be populated by the caller.
+    """
     return ARGS.add_argument_group(title, description)
 
 
@@ -161,18 +226,33 @@ group.add_argument(
 
 
 def overwrite_argparse_epilog(msg: str) -> None:
+    """Allows your code to override the default epilog created by
+    argparse.
+
+    Args:
+        msg: The epilog message to substitute for the default.
+    """
     ARGS.epilog = msg
 
 
-def is_flag_already_in_argv(var: str):
-    """Is a particular flag passed on the commandline?"""
+def is_flag_already_in_argv(var: str) -> bool:
+    """Returns true if a particular flag is passed on the commandline?
+
+    Args:
+        var: The flag to search for.
+    """
     for _ in sys.argv:
         if var in _:
             return True
     return False
 
 
-def reorder_arg_action_groups_before_help(entry_module: Optional[str]):
+def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
+    """Internal.  Used to reorder the arguments before dumping out a
+    generated help string such that the main program's arguments come
+    last.
+
+    """
     reordered_action_groups = []
     for grp in ARGS._action_groups:
         if entry_module is not None and entry_module in grp.title:  # type: ignore
@@ -184,8 +264,35 @@ def reorder_arg_action_groups_before_help(entry_module: Optional[str]):
     return reordered_action_groups
 
 
-def augment_sys_argv_from_environment_variables():
-    usage_message = ARGS.format_usage()
+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:
+
+    :code:`--argument_to_match`
+
+    ...is matched by:
+
+    :code:`ARGUMENT_TO_MATCH`
+
+    This allows programmers to set args via shell environment variables
+    in lieu of passing them on the cmdline.
+
+    """
+
+    usage_message = usage()
     optional = False
     var = ''
     for x in usage_message.split():
@@ -220,7 +327,10 @@ def augment_sys_argv_from_environment_variables():
                 env = ''
 
 
-def augment_sys_argv_from_loadfile():
+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
@@ -237,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:
@@ -248,20 +384,35 @@ 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
 
 
 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
     """Main program should call this early in main().  Note that the
-    bootstrap.initialize wrapper takes care of this automatically.
+    :code:`bootstrap.initialize` wrapper takes care of this automatically.
+    This should only be called once per program invocation.
 
     """
     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
@@ -271,17 +422,17 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
         if arg in ('--help', '-h'):
             if entry_module is not None:
                 entry_module = os.path.basename(entry_module)
-            ARGS._action_groups = reorder_arg_action_groups_before_help(entry_module)
+            ARGS._action_groups = _reorder_arg_action_groups_before_help(entry_module)
 
     # Examine the environment for variables that match known flags.
     # For a flag called --example_flag the corresponding environment
     # variable would be called EXAMPLE_FLAG.  If found, hackily add
     # these into sys.argv to be parsed.
-    augment_sys_argv_from_environment_variables()
+    _augment_sys_argv_from_environment_variables()
 
     # Look for loadfile and read/parse it if present.  This also
     # works by jamming these values onto sys.argv.
-    augment_sys_argv_from_loadfile()
+    _augment_sys_argv_from_loadfile()
 
     # Parse (possibly augmented, possibly completely overwritten)
     # commandline args with argparse normally and populate config.
@@ -303,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']:
@@ -315,7 +491,7 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 
 
 def has_been_parsed() -> bool:
-    """Has the global config been parsed yet?"""
+    """Returns True iff the global config has already been parsed"""
     return CONFIG_PARSE_CALLED
 
 
@@ -329,5 +505,6 @@ def dump_config():
 def late_logging():
     """Log messages saved earlier now that logging has been initialized."""
     logger = logging.getLogger(__name__)
+    logger.debug('Original commandline was: %s', ORIG_ARGV)
     for _ in SAVED_MESSAGES:
         logger.debug(_)