feat(ingest): support for env variable in cli (#3215)

docs/how/delete-metadata.md

@@ -15,6 +15,15 @@ will allow you to customize the datahub instance you are communicating with.
 
 _Note: Provide your GMS instance's host when the prompt asks you for the DataHub host._
 
+Alternatively, you can set the following env variables if you don't want to use a config file
+```
+DATAHUB_SKIP_CONFIG=True
+DATAHUB_GMS_HOST=http://localhost:8080
+DATAHUB_GMS_TOKEN=
+```
+
+The env variables take precendence over what is in the config.
+
 ## Delete By Urn
 
 To delete all the data related to a single entity, run

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

@@ -1,18 +1,26 @@
 import json
+import logging
 import os.path
 import sys
 import typing
 from datetime import datetime
-from typing import Optional
+from typing import List, Optional, Tuple
 
 import click
 import requests
 import yaml
 from pydantic import BaseModel, ValidationError
 
+log = logging.getLogger(__name__)
+
+DEFAULT_GMS_HOST = "http://localhost:8080"
 CONDENSED_DATAHUB_CONFIG_PATH = "~/.datahubenv"
 DATAHUB_CONFIG_PATH = os.path.expanduser(CONDENSED_DATAHUB_CONFIG_PATH)
 
+ENV_SKIP_CONFIG = "DATAHUB_SKIP_CONFIG"
+ENV_METADATA_HOST = "DATAHUB_GMS_HOST"
+ENV_METADATA_TOKEN = "DATAHUB_GMS_TOKEN"
+
 
 class GmsConfig(BaseModel):
     server: str
@@ -35,18 +43,23 @@ def write_datahub_config(host: str, token: Optional[str]) -> None:
     return None
 
 
-def get_session_and_host():
-    session = requests.Session()
+def should_skip_config() -> bool:
+    try:
+        return os.environ[ENV_SKIP_CONFIG] == "True"
+    except KeyError:
+        return False
 
-    gms_host = "http://localhost:8080"
-    gms_token = None
+
+def ensure_datahub_config() -> None:
     if not os.path.isfile(DATAHUB_CONFIG_PATH):
         click.secho(
             f"No {CONDENSED_DATAHUB_CONFIG_PATH} file found, generating one for you...",
             bold=True,
         )
-        write_datahub_config(gms_host, gms_token)
+        write_datahub_config(DEFAULT_GMS_HOST, None)
+
 
+def get_details_from_config():
     with open(DATAHUB_CONFIG_PATH, "r") as stream:
         try:
             config_json = yaml.safe_load(stream)
@@ -63,8 +76,38 @@ def get_session_and_host():
 
             gms_host = gms_config.server
             gms_token = gms_config.token
+            return gms_host, gms_token
         except yaml.YAMLError as exc:
             click.secho(f"{DATAHUB_CONFIG_PATH} malformatted, error: {exc}", bold=True)
+    return None, None
+
+
+def get_details_from_env() -> Tuple[Optional[str], Optional[str]]:
+    return os.environ.get(ENV_METADATA_HOST), os.environ.get(ENV_METADATA_TOKEN)
+
+
+def first_non_null(ls: List[Optional[str]]) -> Optional[str]:
+    return next((el for el in ls if el is not None and el.strip() != ""), None)
+
+
+def get_session_and_host():
+    session = requests.Session()
+
+    gms_host_env, gms_token_env = get_details_from_env()
+    if should_skip_config():
+        gms_host = gms_host_env
+        gms_token = gms_token_env
+    else:
+        ensure_datahub_config()
+        gms_host_conf, gms_token_conf = get_details_from_config()
+        gms_host = first_non_null([gms_host_env, gms_host_conf])
+        gms_token = first_non_null([gms_token_env, gms_token_conf])
+
+    if gms_host is None or gms_host.strip() == "":
+        log.error(
+            f"GMS Host is not set. Use datahub init command or set {ENV_METADATA_HOST} env var"
+        )
+        return None, None
 
     session.headers.update(
         {

metadata-ingestion/tests/unit/test_cli_utils.py

@@ -0,0 +1,11 @@
+from datahub.cli import cli_utils
+
+
+def test_first_non_null():
+    assert cli_utils.first_non_null([]) is None
+    assert cli_utils.first_non_null([None]) is None
+    assert cli_utils.first_non_null([None, "1"]) == "1"
+    assert cli_utils.first_non_null([None, "1", "2"]) == "1"
+    assert cli_utils.first_non_null(["3", "1", "2"]) == "3"
+    assert cli_utils.first_non_null(["", "1", "2"]) == "1"
+    assert cli_utils.first_non_null([" ", "1", "2"]) == "1"