refactor + fix env-based config loading

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

@@ -1,7 +1,5 @@
 import json
 import logging
-import os
-import os.path
 import typing
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
@@ -25,9 +23,6 @@
 
 log = logging.getLogger(__name__)
 
-ENV_DATAHUB_SYSTEM_CLIENT_ID = "DATAHUB_SYSTEM_CLIENT_ID"
-ENV_DATAHUB_SYSTEM_CLIENT_SECRET = "DATAHUB_SYSTEM_CLIENT_SECRET"
-
 # TODO: Many of the methods in this file duplicate logic that already lives
 # in the DataHubGraph client. We should refactor this to use the client instead.
 # For the methods that aren't duplicates, that logic should be moved to the client.
@@ -37,14 +32,6 @@ 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_system_auth() -> Optional[str]:
-    system_client_id = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_ID)
-    system_client_secret = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_SECRET)
-    if system_client_id is not None and system_client_secret is not None:
-        return f"Basic {system_client_id}:{system_client_secret}"
-    return None
-
-
 def parse_run_restli_response(response: requests.Response) -> dict:
     response_json = response.json()
     if response.status_code != 200:
@@ -310,20 +297,16 @@ def command(ctx: click.Context) -> None:
     return command
 
 
-def get_session_login_as(
+def get_frontend_session_login_as(
     username: str, password: str, frontend_url: str
 ) -> requests.Session:
     session = requests.Session()
     headers = {
         "Content-Type": "application/json",
     }
-    system_auth = get_system_auth()
-    if system_auth is not None:
-        session.headers.update({"Authorization": system_auth})
-    else:
-        data = '{"username":"' + username + '", "password":"' + password + '"}'
-        response = session.post(f"{frontend_url}/logIn", headers=headers, data=data)
-        response.raise_for_status()
+    data = '{"username":"' + username + '", "password":"' + password + '"}'
+    response = session.post(f"{frontend_url}/logIn", headers=headers, data=data)
+    response.raise_for_status()
     return session
 
 
@@ -367,7 +350,7 @@ def generate_access_token(
     validity: str = "ONE_HOUR",
 ) -> Tuple[str, str]:
     frontend_url = guess_frontend_url_from_gms_url(gms_url)
-    session = get_session_login_as(
+    session = get_frontend_session_login_as(
         username=username,
         password=password,
         frontend_url=frontend_url,

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

@@ -4,36 +4,133 @@
 
 import logging
 import os
-from typing import Optional
+import sys
+from typing import Optional, Tuple
 
 import click
 import yaml
+from pydantic import BaseModel, ValidationError
 
 from datahub.cli.env_utils import get_boolean_env_variable
+from datahub.ingestion.graph.config import DatahubClientConfig
 
-log = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
-DEFAULT_GMS_HOST = "http://localhost:8080"
 CONDENSED_DATAHUB_CONFIG_PATH = "~/.datahubenv"
 DATAHUB_CONFIG_PATH = os.path.expanduser(CONDENSED_DATAHUB_CONFIG_PATH)
 DATAHUB_ROOT_FOLDER = os.path.expanduser("~/.datahub")
 ENV_SKIP_CONFIG = "DATAHUB_SKIP_CONFIG"
 
+ENV_DATAHUB_SYSTEM_CLIENT_ID = "DATAHUB_SYSTEM_CLIENT_ID"
+ENV_DATAHUB_SYSTEM_CLIENT_SECRET = "DATAHUB_SYSTEM_CLIENT_SECRET"
 
-def persist_datahub_config(config: dict) -> None:
-    with open(DATAHUB_CONFIG_PATH, "w+") as outfile:
-        yaml.dump(config, outfile, default_flow_style=False)
+ENV_METADATA_HOST_URL = "DATAHUB_GMS_URL"
+ENV_METADATA_TOKEN = "DATAHUB_GMS_TOKEN"
+ENV_METADATA_HOST = "DATAHUB_GMS_HOST"
+ENV_METADATA_PORT = "DATAHUB_GMS_PORT"
+ENV_METADATA_PROTOCOL = "DATAHUB_GMS_PROTOCOL"
+
+
+class MissingConfigError(Exception):
+    SHOW_STACK_TRACE = False
+
+
+def get_system_auth() -> Optional[str]:
+    system_client_id = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_ID)
+    system_client_secret = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_SECRET)
+    if system_client_id is not None and system_client_secret is not None:
+        return f"Basic {system_client_id}:{system_client_secret}"
     return None
 
 
-def should_skip_config() -> bool:
+def _should_skip_config() -> bool:
     return get_boolean_env_variable(ENV_SKIP_CONFIG, False)
 
 
-def get_client_config() -> Optional[dict]:
+def persist_raw_datahub_config(config: dict) -> None:
+    with open(DATAHUB_CONFIG_PATH, "w+") as outfile:
+        yaml.dump(config, outfile, default_flow_style=False)
+    return None
+
+
+def get_raw_client_config() -> Optional[dict]:
     with open(DATAHUB_CONFIG_PATH) as stream:
         try:
             return yaml.safe_load(stream)
         except yaml.YAMLError as exc:
             click.secho(f"{DATAHUB_CONFIG_PATH} malformed, error: {exc}", bold=True)
             return None
+
+
+class DatahubConfig(BaseModel):
+    gms: DatahubClientConfig
+
+
+def _get_config_from_env() -> Tuple[Optional[str], Optional[str]]:
+    host = os.environ.get(ENV_METADATA_HOST)
+    port = os.environ.get(ENV_METADATA_PORT)
+    token = os.environ.get(ENV_METADATA_TOKEN)
+    protocol = os.environ.get(ENV_METADATA_PROTOCOL, "http")
+    url = os.environ.get(ENV_METADATA_HOST_URL)
+    if port is not None:
+        url = f"{protocol}://{host}:{port}"
+        return url, token
+    # The reason for using host as URL is backward compatibility
+    # If port is not being used we assume someone is using host env var as URL
+    if url is None and host is not None:
+        logger.warning(
+            f"Do not use {ENV_METADATA_HOST} as URL. Use {ENV_METADATA_HOST_URL} instead"
+        )
+    return url or host, token
+
+
+def load_client_config() -> DatahubClientConfig:
+    gms_host_env, gms_token_env = _get_config_from_env()
+    if gms_host_env:
+        # TODO We should also load system auth credentials here.
+        return DatahubClientConfig(server=gms_host_env, token=gms_token_env)
+
+    if _should_skip_config():
+        raise MissingConfigError(
+            "You have set the skip config flag, but no GMS host or token was provided in env variables."
+        )
+
+    try:
+        _ensure_datahub_config()
+        client_config_dict = get_raw_client_config()
+        datahub_config: DatahubClientConfig = DatahubConfig.parse_obj(
+            client_config_dict
+        ).gms
+
+        return datahub_config
+    except ValidationError as e:
+        click.echo(f"Error loading your {CONDENSED_DATAHUB_CONFIG_PATH}")
+        click.echo(e, err=True)
+        sys.exit(1)
+
+
+def _ensure_datahub_config() -> None:
+    if not os.path.isfile(DATAHUB_CONFIG_PATH):
+        raise MissingConfigError(
+            f"No {CONDENSED_DATAHUB_CONFIG_PATH} file found. Run `datahub init` to create one."
+        )
+
+
+def write_gms_config(
+    host: str, token: Optional[str], merge_with_previous: bool = True
+) -> None:
+    config = DatahubConfig(gms=DatahubClientConfig(server=host, token=token))
+    if merge_with_previous:
+        try:
+            previous_config = get_raw_client_config()
+            assert isinstance(previous_config, dict)
+        except Exception as e:
+            # ok to fail on this
+            previous_config = {}
+            logger.debug(
+                f"Failed to retrieve config from file {DATAHUB_CONFIG_PATH}: {e}. This isn't fatal."
+            )
+        config_dict = {**previous_config, **config.dict()}
+    else:
+        config_dict = config.dict()
+    persist_raw_datahub_config(config_dict)

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

@@ -11,12 +11,12 @@
 
 from datahub.cli.config_utils import (
     DATAHUB_ROOT_FOLDER,
-    get_client_config,
-    persist_datahub_config,
+    DatahubConfig,
+    get_raw_client_config,
+    persist_raw_datahub_config,
 )
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.sink import NoopWriteCallback
-from datahub.ingestion.graph.client import DatahubConfig
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.ingestion.sink.file import FileSink, FileSinkConfig
 from datahub.lite.duckdb_lite_config import DuckDBLiteConfig
@@ -45,7 +45,7 @@ class LiteCliConfig(DatahubConfig):
 
 
 def get_lite_config() -> LiteLocalConfig:
-    client_config_dict = get_client_config()
+    client_config_dict = get_raw_client_config()
     lite_config = LiteCliConfig.parse_obj(client_config_dict)
     return lite_config.lite
 
@@ -309,10 +309,10 @@ def search(
 
 
 def write_lite_config(lite_config: LiteLocalConfig) -> None:
-    cli_config = get_client_config()
+    cli_config = get_raw_client_config()
     assert isinstance(cli_config, dict)
     cli_config["lite"] = lite_config.dict()
-    persist_datahub_config(cli_config)
+    persist_raw_datahub_config(cli_config)
 
 
 @lite.command(context_settings=dict(allow_extra_args=True))

metadata-ingestion/src/datahub/emitter/rest_emitter.py

@@ -10,11 +10,8 @@
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import HTTPError, RequestException
 
-from datahub.cli.cli_utils import (
-    ensure_has_system_metadata,
-    fixup_gms_url,
-    get_system_auth,
-)
+from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url
+from datahub.cli.config_utils import get_system_auth
 from datahub.configuration.common import ConfigurationError, OperationalError
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper

metadata-ingestion/src/datahub/entrypoints.py

@@ -13,7 +13,11 @@
     generate_access_token,
     make_shim_command,
 )
-from datahub.cli.config_utils import DATAHUB_CONFIG_PATH, get_boolean_env_variable
+from datahub.cli.config_utils import (
+    DATAHUB_CONFIG_PATH,
+    get_boolean_env_variable,
+    write_gms_config,
+)
 from datahub.cli.delete_cli import delete
 from datahub.cli.docker_cli import docker
 from datahub.cli.exists_cli import exists
@@ -33,7 +37,7 @@
 from datahub.cli.telemetry import telemetry as telemetry_cli
 from datahub.cli.timeline_cli import timeline
 from datahub.configuration.common import should_show_stack_trace
-from datahub.ingestion.graph.client import get_default_graph, write_gms_config
+from datahub.ingestion.graph.client import get_default_graph
 from datahub.telemetry import telemetry
 from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.logging_manager import configure_logging
@@ -149,7 +153,7 @@ def init(use_password: bool = False) -> None:
             type=str,
             default="",
         )
-    write_gms_config(host, token)
+    write_gms_config(host, token, merge_with_previous=False)
 
     click.echo(f"Written to {DATAHUB_CONFIG_PATH}")