From 734e529e23747dc9ad4a8da309d80815e066da48 Mon Sep 17 00:00:00 2001 From: Pallab Pain Date: Tue, 6 Aug 2024 01:03:42 +0530 Subject: [PATCH] feat(configtree): supports overrides when importing config from files This commit adds the support for config overrides while importing configurations from files. Users can specify one or more override files and they will be applied to the base configs in order. It is expected that the overrides have the full key hierarchy such that the command is able to match the right override with the base config. Wrike Ticket: https://www.wrike.com/open.htm?id=1459374448 --- riocli/configtree/etcd.py | 5 + riocli/configtree/import_keys.py | 157 ++++++++++++++++++++++++------- 2 files changed, 126 insertions(+), 36 deletions(-) diff --git a/riocli/configtree/etcd.py b/riocli/configtree/etcd.py index 173f0351..f4a503dd 100644 --- a/riocli/configtree/etcd.py +++ b/riocli/configtree/etcd.py @@ -27,6 +27,11 @@ def import_in_etcd( ) -> None: cli = Etcd3Client(host=endpoint, port=port) + try: + cli.status() + except Exception as e: + raise ConnectionError(f'cannot connect to etcd server at {endpoint}:{port}') + if prefix: cli.delete_prefix(prefix) else: diff --git a/riocli/configtree/import_keys.py b/riocli/configtree/import_keys.py index 61962a49..97c12dd3 100644 --- a/riocli/configtree/import_keys.py +++ b/riocli/configtree/import_keys.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from pathlib import Path -from typing import Optional, Iterable +from typing import Iterable, Optional import click from benedict import benedict @@ -23,7 +23,7 @@ from riocli.configtree.etcd import import_in_etcd from riocli.configtree.revision import Revision from riocli.configtree.util import Metadata, export_to_files -from riocli.constants import Symbols, Colors +from riocli.constants import Colors, Symbols from riocli.utils.spinner import with_spinner @@ -33,22 +33,28 @@ help_headers_color=Colors.YELLOW, help_options_color=Colors.GREEN, ) -@click.option('--commit/--no-commit', 'commit', is_flag=True, type=bool, ) -@click.option('--update-head/--no-update-head', 'update_head', is_flag=True, type=bool) +@click.option('--commit/--no-commit', 'commit', is_flag=True, type=bool, + help='Commit the imported keys to the Config Tree.') +@click.option('--update-head/--no-update-head', 'update_head', is_flag=True, type=bool, + help='Update the HEAD of the Config Tree after importing the keys.') @click.option('--milestone', 'milestone', type=str, - help='Minestone name for the imported revision.') + help='Milestone name for the imported revision.') @click.option('--etcd-endpoint', 'etcd_endpoint', type=str, help='Import keys to local etcd instead of rapyuta.io cloud') @click.option('--export-directory', 'export_directory', type=str, help='Path to the directory for exporting files.') -@click.option('--etcd-port', 'etcd_port', type=int, +@click.option('--export-format', 'export_format', type=click.Choice(['json', 'yaml']), + default='json', help='Format of the exported files.') +@click.option('--etcd-port', 'etcd_port', type=int, default=2379, help='Port for the etcd endpoint') @click.option('--etcd-prefix', 'etcd_prefix', type=str, help='Prefix to use for the key-space') @click.option('--organization', 'with_org', is_flag=True, type=bool, default=False, help='Operate on organization-scoped Config Trees only.') +@click.option('--override', 'overrides', type=click.Path(exists=True), default=None, + multiple=True, help='Override values for keys in the imported files.') @click.argument('tree-name', type=str) -@click.argument('files', type=str, nargs=-1) +@click.argument('files', type=click.Path(exists=True, file_okay=True, dir_okay=True, resolve_path=True), nargs=-1) @click.pass_context @with_spinner(text="Importing keys...") def import_keys( @@ -59,43 +65,67 @@ def import_keys( update_head: bool, milestone: Optional[str], export_directory: Optional[str], + export_format: Optional[str], etcd_endpoint: Optional[str], etcd_port: Optional[int], etcd_prefix: Optional[str], + overrides: Optional[Iterable[str]], with_org: bool, spinner: Yaspin, ) -> None: - """ - Imports keys in a Config tree from YAML files. - """ + """Imports keys from JSON or YAML files. - data = {} + The import command can import keys into an existing Config Tree or + to an ETCD cluster. It can also output the keys to files in JSON or + YAML format. The command can also apply overrides to the keys before + importing them. - for f in files: - file_prefix = Path(f).stem - file_format = 'yaml' - if f.endswith('json'): - file_format = 'json' + The keys are imported from the one or more files provided as arguments. + The supported file formats are JSON and YAML. - data[file_prefix] = benedict(f, format=file_format) - spinner.write( - click.style( - '{} File {} processed.'.format(Symbols.SUCCESS, f), - fg=Colors.CYAN, - ) - ) + Usage Examples: - data, metadata = split_metadata(data) + Import keys from master.json to the sootballs Config Tree along with overrides. Also, commit and update the head of the Config Tree. + + $ rio configtree import sootballs master.json --override overrides.json --commit --update-head + + Import keys from master.json to etcd along with overrides. + + $ rio configtree import sootballs master.json --override overrides.json --etcd-endpoint localhost + + + You can specify more than one override file by providing the files with --override flag multiple times. + + When importing to ETCD, the name of the base JSON or YAML file will be prefixed to the keys. + + Note: If --etcd-endpoint is provided, the keys are imported to the local etcd cluster instead of the rapyuta.io cloud. + """ + data, metadata = _process_files_with_overrides(files, overrides, spinner) if export_directory is not None: - export_to_files(base_dir=export_directory, data=data) + try: + export_to_files(base_dir=export_directory, data=data, file_format=export_format) + spinner.write(click.style( + f'{Symbols.SUCCESS} Keys exported to {export_format} files in {Path(export_directory).absolute()}.', + fg=Colors.GREEN)) + except Exception as e: + spinner.text = click.style(f'Error exporting keys to files: {e}', fg=Colors.RED) + spinner.red.fail(Symbols.ERROR) + raise SystemExit(1) from e data = benedict(data).flatten(separator='/') metadata = benedict(metadata).flatten(separator='/') if etcd_endpoint: - import_in_etcd(data=data, endpoint=etcd_endpoint, port=etcd_port, prefix=etcd_prefix) - return + try: + import_in_etcd(data=data, endpoint=etcd_endpoint, port=etcd_port, prefix=etcd_prefix) + spinner.text = click.style('Keys imported to etcd successfully.', fg=Colors.GREEN) + spinner.green.ok(Symbols.SUCCESS) + return + except Exception as e: + spinner.text = click.style(f'Error importing keys to etcd: {e}', fg=Colors.RED) + spinner.red.fail(Symbols.ERROR) + raise SystemExit(1) from e try: client = new_v2_client(with_project=(not with_org)) @@ -131,22 +161,22 @@ def import_keys( } client.set_revision_config_tree(tree_name, payload) - spinner.text = click.style('Config tree Head updated successfully.', fg=Colors.CYAN) + spinner.text = click.style('Config tree HEAD updated successfully.', fg=Colors.CYAN) spinner.green.ok(Symbols.SUCCESS) - except Exception as e: - spinner.red.text = str(e) + spinner.text = click.style(f'Error importing keys: {e}', fg=Colors.RED) spinner.red.fail(Symbols.ERROR) raise SystemExit(1) from e -def split_metadata(input: Iterable) -> (Iterable, Iterable): - if not isinstance(input, dict): - return input, None +def split_metadata(data: Iterable) -> (Iterable, Iterable): + """Helper function to split data and metadata from the input data.""" + if not isinstance(data, dict): + return data, None content, metadata = {}, {} - for key, value in input.items(): + for key, value in data.items(): if not isinstance(value, dict): content[key] = value continue @@ -154,8 +184,8 @@ def split_metadata(input: Iterable) -> (Iterable, Iterable): potential_content = value.get('value') potential_meta = value.get('metadata') - if len(value) == 2 and potential_content is not None and \ - potential_meta is not None and isinstance(potential_meta, dict): + if (len(value) == 2 and potential_content is not None and + potential_meta is not None and isinstance(potential_meta, dict)): content[key] = potential_content metadata[key] = Metadata(potential_meta) continue @@ -163,3 +193,58 @@ def split_metadata(input: Iterable) -> (Iterable, Iterable): content[key], metadata[key] = split_metadata(value) return content, metadata + + +def _process_files_with_overrides( + files: Iterable[str], + overrides: Iterable[str], + spinner: Yaspin, +) -> (benedict, benedict): + """Helper function to process the files and overrides. + + Reads the base files and splits data and metadata. Then + applies overrides to the data and metadata. + """ + data = {} + + for f in files: + file_prefix = Path(f).stem + file_format = 'yaml' + if f.endswith('json'): + file_format = 'json' + + data[file_prefix] = benedict(f, format=file_format) + spinner.write( + click.style( + '{} File {} processed.'.format(Symbols.SUCCESS, f), + fg=Colors.CYAN, + ) + ) + + data, metadata = split_metadata(data) + + # Process the override files. + override = benedict({}) + + for f in overrides: + file_format = 'yaml' + if f.endswith('json'): + file_format = 'json' + + override.merge(benedict(f, format=file_format).unflatten(separator='/')) + + spinner.write( + click.style( + '{} Override file {} processed.'.format(Symbols.SUCCESS, f), + fg=Colors.CYAN, + ) + ) + + override_data, override_metadata = split_metadata(override) + + # Merge the override data and metadata with + # the original data and metadata. + benedict(data).merge(override_data) + benedict(metadata).merge(override_metadata) + + return data, metadata