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
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
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)'
)
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
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
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}')