fix config

Author: david-leifker

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

@@ -6,6 +6,7 @@
 import os
 from collections import defaultdict
 from dataclasses import dataclass
+from enum import auto
 from json.decoder import JSONDecodeError
 from typing import (
     TYPE_CHECKING,
@@ -19,6 +20,7 @@
     Union,
 )
 
+import pydantic
 import requests
 from deprecated import deprecated
 from requests.adapters import HTTPAdapter, Retry
@@ -29,6 +31,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,
@@ -80,6 +83,17 @@
 )
 
 
+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),
+)
+
+
 class RequestsSessionConfig(ConfigModel):
     timeout: Union[float, Tuple[float, float], None] = _DEFAULT_TIMEOUT_SEC
 
@@ -203,7 +217,7 @@ 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."
         )
 
         headers = {

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

@@ -32,7 +32,11 @@
 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,
+    DatahubRestEmitter,
+    RestSinkEndpoint,
+)
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.ingestion.graph.config import (
     DatahubClientConfig as DatahubClientConfig,
@@ -141,6 +145,7 @@ 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,
         )
 
         self.server_id = _MISSING_SERVER_ID

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

@@ -20,7 +20,9 @@
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
+    DEFAULT_REST_SINK_ENDPOINT,
     DataHubRestEmitter,
+    RestSinkEndpoint,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -49,11 +51,6 @@
 )
 
 
-class RestSinkEndpoint(ConfigEnum):
-    RESTLI = auto()
-    OPENAPI = auto()
-
-
 class RestSinkMode(ConfigEnum):
     SYNC = auto()
     ASYNC = auto()
@@ -69,15 +66,9 @@ class RestSinkMode(ConfigEnum):
 )
 
 
-_DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
-    RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
-)
-
-
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
-    endpoint: RestSinkEndpoint = _DEFAULT_REST_SINK_ENDPOINT
+    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -184,9 +175,7 @@ def _make_emitter(cls, config: DatahubRestSinkConfig) -> DataHubRestEmitter:
             ca_certificate_path=config.ca_certificate_path,
             client_certificate_path=config.client_certificate_path,
             disable_ssl_verification=config.disable_ssl_verification,
-            openapi_ingestion=True
-            if config.endpoint == RestSinkEndpoint.OPENAPI
-            else False,
+            openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
         )
 
     @property