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

@@ -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:

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,35 +161,90 @@ 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
 
         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
 
         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