mirror of
https://github.com/xcat2/confluent.git
synced 2026-08-03 16:07:00 +00:00
Clean up YAML dump/restore code
Rename the format parameter to fmt to stop shadowing the builtin, pass the already parsed key data dict directly to _restore_keys instead of reserializing it, use yaml.safe_dump for symmetry with the safe loader, and consolidate the five repeated per-format dump blocks into one helper.
This commit is contained in:
Markus Hilger
@@ -185,14 +185,14 @@ if args[0] in ('restore', 'merge'):
|
||||
skipped = {'nodes': [], 'nodegroups': []}
|
||||
|
||||
# Use the format parameter based on the --yaml option
|
||||
format = 'yaml' if options.yaml else 'json'
|
||||
fmt = 'yaml' if options.yaml else 'json'
|
||||
|
||||
dp = cfm.restore_db_from_directory(
|
||||
dumpdir, password,
|
||||
merge="skip" if args[0] == 'merge' else False,
|
||||
merge="skip" if args[0] == 'merge' else False,
|
||||
skipped=skipped,
|
||||
format=format)
|
||||
|
||||
fmt=fmt)
|
||||
|
||||
asyncio.run(dp)
|
||||
if skipped['nodes']:
|
||||
skippedn = ','.join(skipped['nodes'])
|
||||
@@ -233,9 +233,9 @@ elif args[0] == 'dump':
|
||||
os.makedirs(dumpdir)
|
||||
|
||||
# Use the format parameter based on the --yaml option
|
||||
format = 'yaml' if options.yaml else 'json'
|
||||
fmt = 'yaml' if options.yaml else 'json'
|
||||
dp = cfm.dump_db_to_directory(dumpdir, password, options.redact,
|
||||
options.skipkeys, format=format)
|
||||
options.skipkeys, fmt=fmt)
|
||||
asyncio.run(dp)
|
||||
|
||||
|
||||
|
||||
@@ -3124,30 +3124,30 @@ _RestrictedYamlLoader.add_implicit_resolver(
|
||||
list('-+0123456789.'))
|
||||
|
||||
|
||||
async def restore_db_from_directory(location, password, merge=False, skipped=None, format='json'):
|
||||
async def restore_db_from_directory(location, password, merge=False, skipped=None, fmt='json'):
|
||||
"""Restore database from a directory
|
||||
|
||||
|
||||
:param location: Directory containing the configuration
|
||||
:param password: Password to decrypt sensitive data
|
||||
:param merge: If True, merge with existing configuration
|
||||
:param skipped: List of elements to skip during restore
|
||||
:param format: Format of the files ('json' [default] or 'yaml')
|
||||
:param fmt: Format of the files ('json' [default] or 'yaml')
|
||||
"""
|
||||
if format not in ('json', 'yaml'):
|
||||
if fmt not in ('json', 'yaml'):
|
||||
raise ValueError("Format must be 'json' or 'yaml'")
|
||||
|
||||
kdd = None
|
||||
try:
|
||||
keys_file = os.path.join(location, f'keys.{format}')
|
||||
keys_file = os.path.join(location, f'keys.{fmt}')
|
||||
with open(keys_file, 'r') as cfgfile:
|
||||
keydata = cfgfile.read()
|
||||
if format == 'json':
|
||||
if fmt == 'json':
|
||||
kdd = json.loads(keydata)
|
||||
else:
|
||||
kdd = yaml.load(keydata, _RestrictedYamlLoader)
|
||||
if kdd is None:
|
||||
raise ValueError(f"Invalid or empty YAML content in {keys_file}")
|
||||
|
||||
|
||||
if merge:
|
||||
if 'cryptkey' in kdd:
|
||||
kdd['cryptkey'] = _parse_key(kdd['cryptkey'], password)
|
||||
@@ -3156,47 +3156,43 @@ async def restore_db_from_directory(location, password, merge=False, skipped=Non
|
||||
else:
|
||||
kdd['integritykey'] = None # GCM
|
||||
else:
|
||||
if format == 'json':
|
||||
_restore_keys(keydata, password)
|
||||
else:
|
||||
# Convert YAML to JSON string for _restore_keys
|
||||
_restore_keys(json.dumps(kdd), password)
|
||||
_restore_keys(kdd, password)
|
||||
kdd = None
|
||||
except IOError as e:
|
||||
if e.errno == 2:
|
||||
otherformat = 'json' if format == 'yaml' else 'yaml'
|
||||
if os.path.exists(os.path.join(location, f'keys.{otherformat}')):
|
||||
otherfmt = 'json' if fmt == 'yaml' else 'yaml'
|
||||
if os.path.exists(os.path.join(location, f'keys.{otherfmt}')):
|
||||
raise Exception(
|
||||
f'Cannot find keys.{format}, but keys.{otherformat} '
|
||||
f'exists; this appears to be a {otherformat} format dump')
|
||||
f'Cannot find keys.{fmt}, but keys.{otherfmt} '
|
||||
f'exists; this appears to be a {otherfmt} format dump')
|
||||
raise Exception("Cannot restore without keys, this may be a "
|
||||
"redacted dump")
|
||||
if not merge:
|
||||
try:
|
||||
globals_file = os.path.join(location, f'globals.{format}')
|
||||
globals_file = os.path.join(location, f'globals.{fmt}')
|
||||
with open(globals_file, 'r') as globin:
|
||||
if format == 'json':
|
||||
if fmt == 'json':
|
||||
moreglobals = json.load(globin)
|
||||
else:
|
||||
moreglobals = yaml.load(globin, _RestrictedYamlLoader)
|
||||
if moreglobals is None:
|
||||
raise ValueError(f"Invalid or empty YAML content in {globals_file}")
|
||||
|
||||
|
||||
for globvar in moreglobals:
|
||||
set_global(globvar, moreglobals[globvar])
|
||||
except IOError as e:
|
||||
if e.errno != 2:
|
||||
raise
|
||||
try:
|
||||
collective_file = os.path.join(location, f'collective.{format}')
|
||||
collective_file = os.path.join(location, f'collective.{fmt}')
|
||||
with open(collective_file, 'r') as collin:
|
||||
if format == 'json':
|
||||
if fmt == 'json':
|
||||
collective = json.load(collin)
|
||||
else:
|
||||
collective = yaml.load(collin, _RestrictedYamlLoader)
|
||||
if collective is None:
|
||||
raise ValueError(f"Invalid or empty YAML content in {collective_file}")
|
||||
|
||||
|
||||
_cfgstore['collective'] = {}
|
||||
for coll in collective:
|
||||
await add_collective_member(coll, collective[coll]['address'],
|
||||
@@ -3204,10 +3200,10 @@ async def restore_db_from_directory(location, password, merge=False, skipped=Non
|
||||
except IOError as e:
|
||||
if e.errno != 2:
|
||||
raise
|
||||
main_file = os.path.join(location, f'main.{format}')
|
||||
main_file = os.path.join(location, f'main.{fmt}')
|
||||
with open(main_file, 'r') as cfgfile:
|
||||
cfgdata = cfgfile.read()
|
||||
if format == 'yaml':
|
||||
if fmt == 'yaml':
|
||||
# Convert YAML to JSON string for _load_from_json
|
||||
yaml_data = yaml.load(cfgdata, _RestrictedYamlLoader)
|
||||
if yaml_data is None:
|
||||
@@ -3216,67 +3212,49 @@ async def restore_db_from_directory(location, password, merge=False, skipped=Non
|
||||
await ConfigManager(tenant=None)._load_from_json(cfgdata, merge=merge, keydata=kdd, skipped=skipped)
|
||||
ConfigManager.wait_for_sync(True)
|
||||
|
||||
async def dump_db_to_directory(location, password, redact=None, skipkeys=False, format='json'):
|
||||
async def dump_db_to_directory(location, password, redact=None, skipkeys=False, fmt='json'):
|
||||
"""Dump database to a directory
|
||||
|
||||
|
||||
:param location: Directory to store the configuration
|
||||
:param password: Password to protect sensitive data
|
||||
:param redact: If True, redact sensitive data
|
||||
:param skipkeys: If True, skip dumping keys
|
||||
:param format: Format to use for dumping ('json' [default] or 'yaml')
|
||||
:param fmt: Format to use for dumping ('json' [default] or 'yaml')
|
||||
"""
|
||||
if format not in ('json', 'yaml'):
|
||||
if fmt not in ('json', 'yaml'):
|
||||
raise ValueError("Format must be 'json' or 'yaml'")
|
||||
|
||||
# Handle keys file
|
||||
if not redact and not skipkeys:
|
||||
with open(os.path.join(location, f'keys.{format}'), 'w') as cfgfile:
|
||||
if format == 'json':
|
||||
cfgfile.write(_dump_keys(password))
|
||||
else:
|
||||
keydata = _dump_keys(password, dojson=False)
|
||||
yaml.dump(keydata, cfgfile, default_flow_style=False)
|
||||
cfgfile.write('\n')
|
||||
|
||||
# Handle main config
|
||||
main_data = await ConfigManager(tenant=None)._dump_to_json(redact=redact)
|
||||
with open(os.path.join(location, f'main.{format}'), 'wb' if format == 'json' else 'w') as cfgfile:
|
||||
if format == 'json':
|
||||
cfgfile.write(main_data)
|
||||
cfgfile.write(b'\n')
|
||||
else:
|
||||
# Convert JSON to Python object, then dump as YAML
|
||||
yaml.dump(json.loads(main_data.decode('utf-8')), cfgfile, default_flow_style=False)
|
||||
|
||||
# Handle collective data
|
||||
if 'collective' in _cfgstore:
|
||||
with open(os.path.join(location, f'collective.{format}'), 'w') as cfgfile:
|
||||
if format == 'json':
|
||||
cfgfile.write(json.dumps(_cfgstore['collective']))
|
||||
|
||||
def writecfg(name, jsondata):
|
||||
# jsondata is serialized JSON (str or bytes), written as-is for
|
||||
# json format or converted for yaml
|
||||
with open(os.path.join(location, f'{name}.{fmt}'), 'w') as cfgfile:
|
||||
if fmt == 'json':
|
||||
cfgfile.write(confluent.util.stringify(jsondata))
|
||||
cfgfile.write('\n')
|
||||
else:
|
||||
yaml.dump(_cfgstore['collective'], cfgfile, default_flow_style=False)
|
||||
|
||||
yaml.safe_dump(json.loads(jsondata), cfgfile,
|
||||
default_flow_style=False)
|
||||
|
||||
# Handle keys file
|
||||
if not redact and not skipkeys:
|
||||
writecfg('keys', _dump_keys(password))
|
||||
# Handle main config
|
||||
writecfg('main',
|
||||
await ConfigManager(tenant=None)._dump_to_json(redact=redact))
|
||||
# Handle collective data
|
||||
if 'collective' in _cfgstore:
|
||||
writecfg('collective', json.dumps(_cfgstore['collective']))
|
||||
# Handle globals
|
||||
bkupglobals = get_globals()
|
||||
if bkupglobals:
|
||||
with open(os.path.join(location, f'globals.{format}'), 'w') as globout:
|
||||
if format == 'json':
|
||||
json.dump(bkupglobals, globout)
|
||||
else:
|
||||
yaml.dump(bkupglobals, globout, default_flow_style=False)
|
||||
|
||||
writecfg('globals', json.dumps(bkupglobals))
|
||||
# Handle tenants
|
||||
try:
|
||||
for tenant in os.listdir(
|
||||
os.path.join(ConfigManager._cfgdir, '/tenants/')):
|
||||
tenant_data = await ConfigManager(tenant=tenant)._dump_to_json(redact=redact)
|
||||
with open(os.path.join(location, 'tenants', tenant, f'main.{format}'), 'wb' if format == 'json' else 'w') as cfgfile:
|
||||
if format == 'json':
|
||||
cfgfile.write(tenant_data)
|
||||
cfgfile.write(b'\n')
|
||||
else:
|
||||
yaml.dump(json.loads(tenant_data.decode('utf-8')), cfgfile, default_flow_style=False)
|
||||
writecfg(os.path.join('tenants', tenant, 'main'),
|
||||
await ConfigManager(tenant=tenant)._dump_to_json(
|
||||
redact=redact))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
Reference in New Issue
Block a user