fix config

Author: david-leifker

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

@@ -8,6 +8,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timedelta
+from enum import auto
 from json.decoder import JSONDecodeError
 from typing import (
     TYPE_CHECKING,
@@ -21,6 +22,7 @@
     Union,
 )
 
+import pydantic
 import requests
 from deprecated import deprecated
 from requests.adapters import HTTPAdapter, Retry
@@ -31,6 +33,7 @@
 from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url, get_or_else
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import (
+    ConfigEnum,
     ConfigModel,
     ConfigurationError,
     OperationalError,
@@ -94,6 +97,28 @@
 )
 
 
+class RestTraceMode(ConfigEnum):
+    ENABLED = auto()
+    DISABLED = auto()
+
+
+class RestSinkEndpoint(ConfigEnum):
+    RESTLI = auto()
+    OPENAPI = auto()
+
+
+DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
+    RestSinkEndpoint,
+    os.getenv("DATAHUB_REST_SINK_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
+)
+
+
+DEFAULT_REST_TRACE_MODE = pydantic.parse_obj_as(
+    RestTraceMode,
+    os.getenv("DATAHUB_REST_TRACE_MODE", RestTraceMode.DISABLED),
+)
+
+
 class RequestsSessionConfig(ConfigModel):
     timeout: Union[float, Tuple[float, float], None] = _DEFAULT_TIMEOUT_SEC
 
@@ -220,10 +245,10 @@ def __init__(
         self._session = requests.Session()
 
         logger.debug(
-            f"Using {'OpenAPI' if openapi_ingestion else 'Restli'} for ingestion."
+            f"Using {'OpenAPI' if self._openapi_ingestion else 'Restli'} for ingestion."
         )
 
-        if default_trace_mode:
+        if self._default_trace_mode:
             logger.debug("Using API Tracing for ingestion.")
 
         headers = {

metadata-ingestion/src/datahub/ingestion/graph/client.py

@@ -32,7 +32,13 @@
 from datahub.emitter.aspect import TIMESERIES_ASPECT_MAP
 from datahub.emitter.mce_builder import DEFAULT_ENV, Aspect
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.rest_emitter import DatahubRestEmitter
+from datahub.emitter.rest_emitter import (
+    DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_TRACE_MODE,
+    DatahubRestEmitter,
+    RestSinkEndpoint,
+    RestTraceMode,
+)
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.ingestion.graph.config import (
     DatahubClientConfig as DatahubClientConfig,
@@ -141,6 +147,8 @@ def __init__(self, config: DatahubClientConfig) -> None:
             ca_certificate_path=self.config.ca_certificate_path,
             client_certificate_path=self.config.client_certificate_path,
             disable_ssl_verification=self.config.disable_ssl_verification,
+            openapi_ingestion=DEFAULT_REST_SINK_ENDPOINT == RestSinkEndpoint.OPENAPI,
+            default_trace_mode=DEFAULT_REST_TRACE_MODE == RestTraceMode.ENABLED,
         )
 
         self.server_id = _MISSING_SERVER_ID

metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py

@@ -20,7 +20,11 @@
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
+    DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_TRACE_MODE,
     DataHubRestEmitter,
+    RestSinkEndpoint,
+    RestTraceMode,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -49,16 +53,6 @@
 )
 
 
-class RestTraceMode(ConfigEnum):
-    ENABLED = auto()
-    DISABLED = auto()
-
-
-class RestSinkEndpoint(ConfigEnum):
-    RESTLI = auto()
-    OPENAPI = auto()
-
-
 class RestSinkMode(ConfigEnum):
     SYNC = auto()
     ASYNC = auto()
@@ -74,22 +68,10 @@ class RestSinkMode(ConfigEnum):
 )
 
 
-_DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
-    RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
-)
-
-
-_DEFAULT_REST_TRACE_MODE = pydantic.parse_obj_as(
-    RestTraceMode,
-    os.getenv("DATAHUB_REST_TRACE_MODE", RestTraceMode.DISABLED),
-)
-
-
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
-    endpoint: RestSinkEndpoint = _DEFAULT_REST_SINK_ENDPOINT
-    default_trace_mode: RestTraceMode = _DEFAULT_REST_TRACE_MODE
+    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
+    default_trace_mode: RestTraceMode = DEFAULT_REST_TRACE_MODE
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS