X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=config.py;h=b344631003bb29dc186c0d7db8f2631b8c7ac4bc;hb=2966ecdb8c49266d491a1f75791868920d68a510;hp=c5813a81145764c05d7af29ce32a07da4ef36ef8;hpb=1e858172519e9339e4720b8bf9b39b6d9801e305;p=python_utils.git diff --git a/config.py b/config.py index c5813a8..b344631 100644 --- 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']: