Add documentation + session update sub-command

Author: pedro93

docs/cli.md

@@ -48,13 +48,10 @@ datahub --help
 Usage: datahub [OPTIONS] COMMAND [ARGS]...
 
 Options:
-  --debug / --no-debug            Enable debug logging.
-  --log-file FILE                 Enable debug logging.
-  --debug-vars / --no-debug-vars  Show variable values in stack traces. Implies --debug. While we try to avoid
-                                  printing sensitive information like passwords, this may still happen.
-  --version                       Show the version and exit.
-  -dl, --detect-memory-leaks      Run memory leak detection.
-  --help                          Show this message and exit.
+  --debug / --no-debug  Enable debug logging.
+  --log-file FILE       Write debug-level logs to a file.
+  --version             Show the version and exit.
+  --help                Show this message and exit.
 
 Commands:
   actions       <disabled due to missing dependencies>
@@ -76,6 +73,7 @@ Commands:
   migrate       Helper commands for migrating metadata within DataHub.
   properties    A group of commands to interact with structured properties in DataHub.
   put           A group of commands to put metadata in DataHub.
+  session       Manage DataHub session profiles
   state         Managed state stored in DataHub by stateful ingestion.
   telemetry     Toggle telemetry.
   timeline      Get timeline for an entity based on certain categories
@@ -280,16 +278,83 @@ DATAHUB_TELEMETRY_TIMEOUT=10
 DATAHUB_DEBUG=false
 ```
 
-### container
+### session
 
-A group of commands to interact with containers in DataHub.
+**🤝 Version Compatibility:** `acryl-datahub>=0.15.0.5`
+
+The `session` group of commands allows you to manage DataHub sessions to multiple instances using the same cli. The session information will be store in a local file under `~/.datahub/sessions.json`.
+which can be inspected at any time.
 
-e.g. You can use this to apply a tag to all datasets recursively in this container.
 ```shell
-datahub container tag --container-urn "urn:li:container:0e9e46bd6d5cf645f33d5a8f0254bc2d" --tag-urn "urn:li:tag:tag1"
-datahub container domain --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --domain-urn  "urn:li:domain:ajsajo-b832-4ab3-8881-7ed5e991a44c"
-datahub container owner --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --owner-urn  "urn:li:corpGroup:[email protected]"
-datahub container term --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --term-urn  "urn:li:term:PII"
+datahub session --help
+Usage: datahub session [OPTIONS] COMMAND [ARGS]...
+
+  Manage DataHub session profiles
+
+Options:
+  --help  Show this message and exit.
+
+Commands:
+  create  Create profile with which to connect to a DataHub instance
+  delete  Delete a session profile
+  list    List all session profiles
+  save    Save the current active datahubenv config as a session
+  use     Set the active session
+```
+
+Here we detail the sub-commands available under the dataproduct group of commands:
+
+#### create
+
+Use this command to create a new datahub session. This is similar to [datahub init](#init) with an extra step to fill out the name of the *profile* under which to save the session. The profile name is important as it uniquely identifies the name of the session.
+
+```shell
+datahub session create
+```
+
+:::note
+This command has a flag `--use-password`, that can be used to authenticate to the instance using a username/password combo and dynamically generate an access token with an 1 hour validaty for that username.
+:::
+
+#### update
+
+Use this command to update an existing session identified by the profile name.
+A form will then appear for the user to fill out, this will be prefilled with the existing value if it exists or a default exists.
+
+```shell
+datahub session update -p <profile name>
+```
+
+#### delete
+
+Use this command to delete an existing session identified by the profile name.
+
+```shell
+datahub session delete -p <profile name>
+```
+
+#### list
+
+Use this command to list all existing sessions available in `~/.datahub/sessions.json`.
+
+```shell
+datahub session list
+```
+
+#### save
+
+Use this command to save the existing configuration stored in `~/.datahubenv` as a session with the specified name.
+
+```shell
+datahub session save -p <profile name>
+```
+
+#### use
+
+Command used to specify which profile to use. This overwrites whatever is in `~/.datahubenv`.
+
+```shell
+datahub session use -p <profile name>
 ```
 
 ### check
@@ -626,6 +691,18 @@ Use this to delete a Data Product from DataHub. Default to `--soft` which preser
 # > datahub dataproduct delete --urn "urn:li:dataProduct:pet_of_the_week" --hard
 ```
 
+### container
+
+A group of commands to interact with containers in DataHub.
+
+e.g. You can use this to apply a tag to all datasets recursively in this container.
+```shell
+datahub container tag --container-urn "urn:li:container:0e9e46bd6d5cf645f33d5a8f0254bc2d" --tag-urn "urn:li:tag:tag1"
+datahub container domain --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --domain-urn  "urn:li:domain:ajsajo-b832-4ab3-8881-7ed5e991a44c"
+datahub container owner --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --owner-urn  "urn:li:corpGroup:[email protected]"
+datahub container term --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --term-urn  "urn:li:term:PII"
+```
+
 ## Miscellaneous Admin Commands
 
 ### lite (experimental)

metadata-ingestion/src/datahub/cli/session_cli.py

@@ -1,18 +1,76 @@
 import json
+import copy
+from math import log
 import os
-from typing import Dict
-
 import click
+import logging
+from sqlalchemy import null
 from tabulate import tabulate
 
 from datahub.cli.cli_utils import fixup_gms_url, generate_access_token
 from datahub.cli.config_utils import load_client_config, persist_raw_datahub_config
 from datahub.ingestion.graph.config import DatahubClientConfig
 
+logger = logging.getLogger(__name__)
 DATAHUB_SESSIONS_PATH = os.path.expanduser("~/.datahub/sessions.json")
 
+from typing import Optional, Type, Any
+import click
 
-def load_sessions() -> Dict[str, DatahubClientConfig]:
+def parse_value(value: str, expected_type: Optional[Type] = None) -> tuple[Any, Optional[str]]:
+    
+    if not value:
+        return None, None
+    if expected_type is None:
+        return value, None
+     
+    try:
+        if expected_type == str:
+            return value, None
+        elif expected_type == float:
+            return float(value), None
+        elif expected_type == bool:
+            return bool(value), None
+        elif expected_type == list:
+            return [item.strip() for item in value.split(',') if item.strip()], None
+        elif expected_type == dict:
+            pairs = [pair.strip().split(':') for pair in value.split(',') if pair.strip()]
+            return {k.strip(): v.strip() for k, v in pairs}, None
+        else:
+            return None, f"Unsupported type: {expected_type}"
+    except Exception as e:
+        return None, str(e)
+
+
+def dynamic_prompt(prompt_text: str, expected_type: Optional[Type] = None, default: Optional[object] = None) -> Any:
+    type_hint = {
+        str: "text",
+        float: "number", 
+        bool: "boolean",
+        list: "comma-separated values",
+        dict: "key1:value1,key2:value2",
+        None: "any value"
+    }.get(expected_type, "value")
+    
+    if default is not None:    
+       prompt_text += f" [{str(default)}]" 
+    
+    while True:
+        if default is not None:
+            value = click.prompt(prompt_text, default=default, show_default=False)
+        else:
+            value = click.prompt(prompt_text, default='', show_default=False) 
+        result, error = parse_value(value, expected_type)
+        
+        if error and value:  # Only show error and retry if there was input
+            click.echo(f"Error: {error}")
+            retry = click.confirm("Would you like to try again?", default=True)
+            if not retry:
+                return None
+        else:
+            return result
+
+def load_sessions() -> dict[str, DatahubClientConfig]:
     if not os.path.exists(DATAHUB_SESSIONS_PATH):
         return {}
     with open(DATAHUB_SESSIONS_PATH, "r") as f:
@@ -23,7 +81,7 @@ def load_sessions() -> Dict[str, DatahubClientConfig]:
         }
 
 
-def save_sessions(sessions: Dict[str, DatahubClientConfig]) -> None:
+def save_sessions(sessions: dict[str, DatahubClientConfig]) -> None:
     os.makedirs(os.path.dirname(DATAHUB_SESSIONS_PATH), exist_ok=True)
     with open(DATAHUB_SESSIONS_PATH, "w") as f:
         json.dump(
@@ -37,6 +95,19 @@ def session() -> None:
     pass
 
 
+@session.command()
+def list() -> None:
+    """List all session profiles"""
+    sessions = load_sessions()
+    if not sessions:
+        click.echo("No profiles found")
+        return
+
+    headers = ["Profile", "URL"]
+    table_data = [[name, config.server] for name, config in sessions.items()]
+    click.echo(tabulate(table_data, headers=headers))
+
+
 @session.command()
 @click.option(
     "--use-password",
@@ -75,27 +146,50 @@ def create(use_password: bool) -> None:
 
     profile_name = click.prompt("Enter name for profile", type=str)
 
+    # TODO if profile_name already exists, cancel the create, ask the customer to delete the old one or update instead
     config = DatahubClientConfig(server=host, token=token)
     sessions[profile_name] = config
     save_sessions(sessions)
     click.echo(f"Created profile: {profile_name}")
 
 
 @session.command()
-def list() -> None:
-    """List all session profiles"""
+@click.option(
+    "-p",
+    "--profile",
+    type=str,
+    required=True,
+    help="Name of profile to delete",
+)
+def update(profile: str) -> None:
+    """Update a session profile"""
     sessions = load_sessions()
-    if not sessions:
-        click.echo("No profiles found")
+    if profile not in sessions:
+        click.echo(f"Profile {profile} not found")
         return
 
-    headers = ["Profile", "URL"]
-    table_data = [[name, config.server] for name, config in sessions.items()]
-    click.echo(tabulate(table_data, headers=headers))
+    config: DatahubClientConfig = sessions[profile]
+    config_copy: DatahubClientConfig = copy.deepcopy(config)
+    for key, value in config_copy.dict().items():
+        # Get type information for each field in DatahubClientConfig so we can generate
+        # a dynamic, type-safe click prompt for it.
+        field_type = config_copy.__fields__.get(key).type_
+        new_value = dynamic_prompt(f"Enter new value for {key}", field_type, value)
+        config_copy.__setattr__(key, new_value)
 
+    sessions[profile] = config_copy
+    save_sessions(sessions)
+    click.echo(f"Updated profile: {profile}")
+ 
 
 @session.command()
-@click.argument("profile", type=str)
+@click.option(
+    "-p",
+    "--profile",
+    type=str,
+    required=True,
+    help="Name of profile to delete",
+)
 def delete(profile: str) -> None:
     """Delete a session profile"""
     sessions = load_sessions()
@@ -109,7 +203,13 @@ def delete(profile: str) -> None:
 
 
 @session.command()
-@click.argument("profile", type=str)
+@click.option(
+    "-p",
+    "--profile",
+    type=str,
+    required=True,
+    help="Name of profile to use",
+)
 def use(profile: str) -> None:
     """Set the active session"""
     sessions = load_sessions()
@@ -124,6 +224,7 @@ def use(profile: str) -> None:
 
 @session.command()
 @click.option(
+    "-p",
     "--profile",
     type=str,
     required=True,

metadata-ingestion/src/datahub/configuration/common.py

@@ -15,6 +15,7 @@
 )
 
 import pydantic
+from pydantic.fields import ModelField
 from cached_property import cached_property
 from pydantic import BaseModel, Extra, ValidationError
 from pydantic.fields import Field
@@ -122,6 +123,13 @@ def parse_obj_allow_extras(cls, obj: Any) -> Self:
             with unittest.mock.patch.object(cls.Config, "extra", pydantic.Extra.allow):
                 return cls.parse_obj(obj)
 
+    @classmethod
+    def is_field_required(cls, field_name: str) -> bool:
+        field: ModelField | None = cls.__fields__.get(field_name, None)
+        if field: 
+            return field.required
+        return False
+
 
 class PermissiveConfigModel(ConfigModel):
     # A permissive config model that allows extra fields.