From c73a6bb2b4657b8a00f6c456dcbfdcf5911775b4 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Tue, 8 Oct 2024 15:05:20 -0700 Subject: [PATCH] feat(ingest): support `__from_env__` special server value --- metadata-ingestion/src/datahub/cli/config_utils.py | 7 +++++++ metadata-ingestion/src/datahub/emitter/rest_emitter.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/metadata-ingestion/src/datahub/cli/config_utils.py b/metadata-ingestion/src/datahub/cli/config_utils.py index bb85809174ea9..5d9604de7836f 100644 --- a/metadata-ingestion/src/datahub/cli/config_utils.py +++ b/metadata-ingestion/src/datahub/cli/config_utils.py @@ -84,6 +84,13 @@ def _get_config_from_env() -> Tuple[Optional[str], Optional[str]]: return url or host, token +def require_config_from_env() -> Tuple[str, Optional[str]]: + host, token = _get_config_from_env() + if host is None: + raise MissingConfigError("No GMS host was provided in env variables.") + return host, token + + def load_client_config() -> DatahubClientConfig: gms_host_env, gms_token_env = _get_config_from_env() if gms_host_env: diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py index e370ad3562a06..948060c3c4f44 100644 --- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py +++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py @@ -76,6 +76,12 @@ def __init__( ): if not gms_server: raise ConfigurationError("gms server is required") + if gms_server == "__from_env__" and token is None: + # HACK: similar to what we do with system auth, we transparently + # inject the config in here. Ideally this should be done in the + # config loader or by the caller, but it gets the job done for now. + gms_server, token = config_utils.require_config_from_env() + self._gms_server = fixup_gms_url(gms_server) self._token = token self.server_config: Dict[str, Any] = {}