acryl-datahub 1.0.0.3rc10__py3-none-any.whl → 1.0.0.3rc11__py3-none-any.whl

datahub/cli/specific/dataproduct_cli.py

@@ -20,6 +20,7 @@ from datahub.emitter.mce_builder import (
     validate_ownership_type,
 )
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import OwnerClass, OwnershipTypeClass
 from datahub.specific.dataproduct import DataProductPatchBuilder
 from datahub.telemetry import telemetry
@@ -81,7 +82,7 @@ def mutate(file: Path, validate_assets: bool, external_url: str, upsert: bool) -
 
     config_dict = load_file(pathlib.Path(file))
     id = config_dict.get("id") if isinstance(config_dict, dict) else None
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product: DataProduct = DataProduct.from_yaml(file, graph)
         external_url_override = (
             external_url
@@ -162,7 +163,7 @@ def upsert(file: Path, validate_assets: bool, external_url: str) -> None:
 def diff(file: Path, update: bool) -> None:
     """Diff a Data Product file with its twin in DataHub"""
 
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         id: Optional[str] = None
         try:
             data_product_local: DataProduct = DataProduct.from_yaml(file, emitter)
@@ -216,7 +217,7 @@ def delete(urn: str, file: Path, hard: bool) -> None:
         raise click.Abort()
 
     graph: DataHubGraph
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product_urn = (
             urn if urn.startswith("urn:li:dataProduct") else f"urn:li:dataProduct:{urn}"
         )
@@ -248,7 +249,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataProduct:"):
         urn = f"urn:li:dataProduct:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataproduct: DataProduct = DataProduct.from_datahub(graph=graph, id=urn)
             click.secho(
@@ -306,7 +307,7 @@ def set_description(urn: str, description: str, md_file: Path) -> None:
 
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.set_description(description)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "set description")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -342,7 +343,7 @@ def add_owner(urn: str, owner: str, owner_type: str) -> None:
             owner=_get_owner_urn(owner), type=owner_type, typeUrn=owner_type_urn
         )
     )
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add owners")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -360,7 +361,7 @@ def remove_owner(urn: str, owner_urn: str) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_owner(owner=_get_owner_urn(owner_urn))
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove owners")
         for mcp in dataproduct_patcher.build():
             click.echo(json.dumps(mcp.to_obj()))
@@ -382,7 +383,7 @@ def add_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.add_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add assets")
         if validate_assets:
             _abort_if_non_existent_urn(
@@ -409,7 +410,7 @@ def remove_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove assets")
         if validate_assets:
             _abort_if_non_existent_urn(

datahub/cli/specific/dataset_cli.py

@@ -12,6 +12,7 @@ from click_default_group import DefaultGroup
 from datahub.api.entities.dataset.dataset import Dataset, DatasetRetrievalConfig
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -54,7 +55,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataset:"):
         urn = f"urn:li:dataset:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataset: Dataset = Dataset.from_datahub(graph=graph, urn=urn)
             click.secho(
@@ -82,7 +83,7 @@ def add_sibling(urn: str, sibling_urns: Tuple[str]) -> None:
     all_urns.add(urn)
     for sibling_urn in sibling_urns:
         all_urns.add(sibling_urn)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         for _urn in all_urns:
             _emit_sibling(graph, urn, _urn, all_urns)
 
@@ -181,7 +182,7 @@ def sync(file: str, to_datahub: bool, dry_run: bool) -> None:
     dry_run_prefix = "[dry-run]: " if dry_run else ""  # prefix to use in messages
 
     failures: List[str] = []
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         datasets = Dataset.from_yaml(file)
         for dataset in datasets:
             assert (

datahub/cli/specific/forms_cli.py

@@ -7,6 +7,7 @@ from click_default_group import DefaultGroup
 
 from datahub.api.entities.forms.forms import Forms
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -40,7 +41,7 @@ def upsert(file: Path) -> None:
 @telemetry.with_telemetry()
 def get(urn: str, to_file: str) -> None:
     """Get form from DataHub"""
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             form: Forms = Forms.from_datahub(graph=graph, urn=urn)
             click.secho(

datahub/cli/specific/group_cli.py

@@ -10,6 +10,7 @@ from datahub.api.entities.corpgroup.corpgroup import (
 )
 from datahub.cli.specific.file_loader import load_file
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -40,7 +41,7 @@ def upsert(file: Path, override_editable: bool) -> None:
 
     config_dict = load_file(file)
     group_configs = config_dict if isinstance(config_dict, list) else [config_dict]
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         for group_config in group_configs:
             try:
                 datahub_group = CorpGroup.parse_obj(group_config)

datahub/cli/specific/structuredproperties_cli.py

@@ -11,6 +11,7 @@ from datahub.api.entities.structuredproperties.structuredproperties import (
     StructuredProperties,
 )
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.urns.urn import Urn
@@ -33,7 +34,7 @@ def properties() -> None:
 def upsert(file: Path) -> None:
     """Upsert structured properties in DataHub."""
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         StructuredProperties.create(str(file), graph)
 
 
@@ -48,7 +49,7 @@ def get(urn: str, to_file: str) -> None:
     """Get structured properties from DataHub"""
     urn = Urn.make_structured_property_urn(urn)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             structuredproperties: StructuredProperties = (
                 StructuredProperties.from_datahub(graph=graph, urn=urn)
@@ -117,7 +118,7 @@ def list(details: bool, to_file: str) -> None:
         with open(file, "w") as fp:
             yaml.dump(serialized_objects, fp)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if details:
             logger.info(
                 "Listing structured properties with details. Use --no-details for urns only"

datahub/cli/specific/user_cli.py

@@ -8,6 +8,7 @@ from click_default_group import DefaultGroup
 from datahub.api.entities.corpuser.corpuser import CorpUser, CorpUserGenerationConfig
 from datahub.cli.specific.file_loader import load_file
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -38,7 +39,7 @@ def upsert(file: Path, override_editable: bool) -> None:
 
     config_dict = load_file(pathlib.Path(file))
     user_configs = config_dict if isinstance(config_dict, list) else [config_dict]
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         for user_config in user_configs:
             try:
                 datahub_user: CorpUser = CorpUser.parse_obj(user_config)

datahub/cli/state_cli.py

@@ -5,6 +5,7 @@ import click
 from click_default_group import DefaultGroup
 
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -28,7 +29,7 @@ def inspect(pipeline_name: str, platform: str) -> None:
     Only works for state entity removal for now.
     """
 
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.CLI)
     checkpoint = datahub_graph.get_latest_pipeline_checkpoint(pipeline_name, platform)
     if not checkpoint:
         click.secho("No ingestion state found.", fg="red")

datahub/cli/timeline_cli.py

@@ -9,6 +9,7 @@ from requests import Response
 
 from datahub.emitter.mce_builder import dataset_urn_to_key, schema_field_urn_to_key
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.urns.urn import Urn
@@ -63,7 +64,7 @@ def get_timeline(
     diff: bool,
     graph: Optional[DataHubGraph] = None,
 ) -> Any:
-    client = graph if graph else get_default_graph()
+    client = graph if graph else get_default_graph(ClientMode.CLI)
     session = client._session
     host = client.config.server
     if urn.startswith("urn%3A"):

datahub/emitter/rest_emitter.py

@@ -5,7 +5,6 @@ import json
 import logging
 import os
 import time
-import warnings
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timedelta
@@ -50,13 +49,17 @@ from datahub.emitter.response_helper import (
     extract_trace_data_from_mcps,
 )
 from datahub.emitter.serialization_helper import pre_json_transform
-from datahub.errors import APITracingWarning
 from datahub.ingestion.api.closeable import Closeable
+from datahub.ingestion.graph.config import (
+    DATAHUB_COMPONENT_ENV,
+    ClientMode,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.usage import UsageAggregation
+from datahub.utilities.server_config_util import RestServiceConfig, ServiceFeature
 
 if TYPE_CHECKING:
     from datahub.ingestion.graph.client import DataHubGraph
@@ -79,6 +82,8 @@ _DEFAULT_RETRY_MAX_TIMES = int(
 
 _DATAHUB_EMITTER_TRACE = get_boolean_env_variable("DATAHUB_EMITTER_TRACE", False)
 
+_DEFAULT_CLIENT_MODE: ClientMode = ClientMode.SDK
+
 TRACE_PENDING_STATUS = "PENDING"
 TRACE_INITIAL_BACKOFF = 1.0  # Start with 1 second
 TRACE_MAX_BACKOFF = 300.0  # Cap at 5 minutes
@@ -133,12 +138,24 @@ class RequestsSessionConfig(ConfigModel):
     ca_certificate_path: Optional[str] = None
     client_certificate_path: Optional[str] = None
     disable_ssl_verification: bool = False
+    client_mode: Optional[ClientMode] = _DEFAULT_CLIENT_MODE
+    datahub_component: Optional[str] = None
 
     def build_session(self) -> requests.Session:
         session = requests.Session()
 
-        if self.extra_headers:
-            session.headers.update(self.extra_headers)
+        user_agent = self._get_user_agent_string(session)
+
+        base_headers = {
+            "User-Agent": user_agent,
+            "X-DataHub-Client-Mode": self.client_mode.name
+            if self.client_mode
+            else _DEFAULT_CLIENT_MODE.name,
+            "X-DataHub-Py-Cli-Version": nice_version_name(),
+        }
+
+        headers = {**base_headers, **self.extra_headers}
+        session.headers.update(headers)
 
         if self.client_certificate_path:
             session.cert = self.client_certificate_path
@@ -186,6 +203,59 @@ class RequestsSessionConfig(ConfigModel):
 
         return session
 
+    @classmethod
+    def get_client_mode_from_session(
+        cls, session: requests.Session
+    ) -> Optional[ClientMode]:
+        """
+        Extract the ClientMode enum from a requests Session by checking the headers.
+
+        Args:
+            session: The requests.Session object to check
+
+        Returns:
+            The corresponding ClientMode enum value if found, None otherwise
+        """
+        # Check if the session has the X-DataHub-Client-Mode header
+        mode_str = session.headers.get("X-DataHub-Client-Mode")
+
+        if not mode_str:
+            return None
+
+        # Try to convert the string value to enum
+        try:
+            # First ensure we're working with a str value
+            if isinstance(mode_str, bytes):
+                mode_str = mode_str.decode("utf-8")
+
+            # Then find the matching enum value
+            for mode in ClientMode:
+                if mode.name == mode_str:
+                    return mode
+
+            # If we got here, no matching enum was found
+            return None
+        except Exception:
+            # Handle any other errors
+            return None
+
+    def _get_user_agent_string(self, session: requests.Session) -> str:
+        """Generate appropriate user agent string based on client mode"""
+        version = nice_version_name()
+        client_mode = self.client_mode if self.client_mode else _DEFAULT_CLIENT_MODE
+
+        if "User-Agent" in session.headers:
+            user_agent = session.headers["User-Agent"]
+            if isinstance(user_agent, bytes):
+                requests_user_agent = " " + user_agent.decode("utf-8")
+            else:
+                requests_user_agent = " " + user_agent
+        else:
+            requests_user_agent = ""
+
+        # 1.0 refers to the user agent string version
+        return f"DataHub-Client/1.0 ({client_mode.name.lower()}; {self.datahub_component if self.datahub_component else DATAHUB_COMPONENT_ENV}; {version}){requests_user_agent}"
+
 
 @dataclass
 class _Chunk:
@@ -211,8 +281,9 @@ class DataHubRestEmitter(Closeable, Emitter):
     _gms_server: str
     _token: Optional[str]
     _session: requests.Session
-    _openapi_ingestion: bool
+    _openapi_ingestion: Optional[bool]
     _default_trace_mode: bool
+    server_config: RestServiceConfig
 
     def __init__(
         self,
@@ -228,10 +299,10 @@ class DataHubRestEmitter(Closeable, Emitter):
         ca_certificate_path: Optional[str] = None,
         client_certificate_path: Optional[str] = None,
         disable_ssl_verification: bool = False,
-        openapi_ingestion: bool = (
-            DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI
-        ),
+        openapi_ingestion: Optional[bool] = None,
         default_trace_mode: bool = False,
+        client_mode: Optional[ClientMode] = None,
+        datahub_component: Optional[str] = None,
     ):
         if not gms_server:
             raise ConfigurationError("gms server is required")
@@ -243,13 +314,10 @@ class DataHubRestEmitter(Closeable, Emitter):
 
         self._gms_server = fixup_gms_url(gms_server)
         self._token = token
-        self.server_config: Dict[str, Any] = {}
-        self._openapi_ingestion = openapi_ingestion
         self._default_trace_mode = default_trace_mode
         self._session = requests.Session()
-
-        logger.debug(
-            f"Using {'OpenAPI' if self._openapi_ingestion else 'Restli'} for ingestion."
+        self._openapi_ingestion = (
+            openapi_ingestion  # Re-evaluated after test connection
         )
 
         if self._default_trace_mode:
@@ -257,7 +325,6 @@ class DataHubRestEmitter(Closeable, Emitter):
 
         headers = {
             "X-RestLi-Protocol-Version": "2.0.0",
-            "X-DataHub-Py-Cli-Version": nice_version_name(),
             "Content-Type": "application/json",
         }
         if token:
@@ -303,37 +370,54 @@ class DataHubRestEmitter(Closeable, Emitter):
             ca_certificate_path=ca_certificate_path,
             client_certificate_path=client_certificate_path,
             disable_ssl_verification=disable_ssl_verification,
+            client_mode=client_mode,
+            datahub_component=datahub_component,
         )
 
         self._session = self._session_config.build_session()
 
     def test_connection(self) -> None:
         url = f"{self._gms_server}/config"
-        response = self._session.get(url)
-        if response.status_code == 200:
-            config: dict = response.json()
-            if config.get("noCode") == "true":
-                self.server_config = config
-                return
+        try:
+            # Create a config instance with session and URL
+            config = RestServiceConfig(session=self._session, url=url)
+            # Attempt to load config, which will throw ConfigurationError if there's an issue
+            config.fetch_config()
+            self.server_config = config
+
+            # Determine OpenAPI mode
+            if self._openapi_ingestion is None:
+                # No constructor parameter
+                if (
+                    not os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT")
+                    and self._session_config.client_mode == ClientMode.SDK
+                    and self.server_config.supports_feature(ServiceFeature.OPEN_API_SDK)
+                ):
+                    # Enable if SDK client and no environment variable specified
+                    self._openapi_ingestion = True
+                else:
+                    # The system env is specifying the value
+                    self._openapi_ingestion = (
+                        DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI
+                    )
 
-            else:
-                raise ConfigurationError(
-                    "You seem to have connected to the frontend service instead of the GMS endpoint. "
-                    "The rest emitter should connect to DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms). "
-                    "For Acryl users, the endpoint should be https://<name>.acryl.io/gms"
-                )
-        else:
             logger.debug(
-                f"Unable to connect to {url} with status_code: {response.status_code}. Response: {response.text}"
+                f"Using {'OpenAPI' if self._openapi_ingestion else 'Restli'} for ingestion."
             )
-            if response.status_code == 401:
-                message = f"Unable to connect to {url} - got an authentication error: {response.text}."
-            else:
-                message = f"Unable to connect to {url} with status_code: {response.status_code}."
-            message += "\nPlease check your configuration and make sure you are talking to the DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms)."
-            raise ConfigurationError(message)
 
-    def get_server_config(self) -> dict:
+            # Set default tracing for SDK
+            if (
+                self._session_config.client_mode == ClientMode.SDK
+                and self.server_config.supports_feature(ServiceFeature.API_TRACING)
+            ):
+                # Enable tracing if using SDK & server supported
+                self._default_trace_mode = True
+
+        except ConfigurationError as e:
+            # Just re-raise the exception
+            raise e
+
+    def get_server_config(self) -> RestServiceConfig:
         self.test_connection()
         return self.server_config
 
@@ -485,7 +569,7 @@ class DataHubRestEmitter(Closeable, Emitter):
         trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> int:
         """
-        1. Grouping MCPs by their HTTP method and entity URL
+        1. Grouping MCPs by their HTTP method and entity URL and HTTP method
         2. Breaking down large batches into smaller chunks based on both:
          * Total byte size (INGEST_MAX_PAYLOAD_BYTES)
          * Maximum number of items (BATCH_INGEST_MAX_PAYLOAD_LENGTH)
@@ -751,12 +835,6 @@ class DataHubRestEmitter(Closeable, Emitter):
             trace_flag if trace_flag is not None else self._default_trace_mode
         )
         resolved_async_flag = async_flag if async_flag is not None else async_default
-        if resolved_trace_flag and not resolved_async_flag:
-            warnings.warn(
-                "API tracing is only available with async ingestion. For sync mode, API errors will be surfaced as exceptions.",
-                APITracingWarning,
-                stacklevel=3,
-            )
         return resolved_trace_flag and resolved_async_flag
 
     def __repr__(self) -> str:

datahub/entrypoints.py

@@ -37,6 +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
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.logging_manager import configure_logging
@@ -117,7 +118,7 @@ def version(include_server: bool = False) -> None:
     click.echo(f"Models: {model_version_name()}")
     click.echo(f"Python version: {sys.version}")
     if include_server:
-        server_config = get_default_graph().get_config()
+        server_config = get_default_graph(ClientMode.CLI).get_config()
         click.echo(f"Server config: {server_config}")
 
 

datahub/ingestion/graph/client.py

@@ -34,14 +34,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 (
-    DEFAULT_REST_EMITTER_ENDPOINT,
     DEFAULT_REST_TRACE_MODE,
     DatahubRestEmitter,
-    RestSinkEndpoint,
     RestTraceMode,
 )
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.ingestion.graph.config import (
+    ClientMode,
     DatahubClientConfig as DatahubClientConfig,
 )
 from datahub.ingestion.graph.connections import (
@@ -158,13 +157,12 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             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=self.config.openapi_ingestion
-            if self.config.openapi_ingestion is not None
-            else (DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI),
+            openapi_ingestion=self.config.openapi_ingestion,
             default_trace_mode=DEFAULT_REST_TRACE_MODE == RestTraceMode.ENABLED,
+            client_mode=config.client_mode,
+            datahub_component=config.datahub_component,
         )
-
-        self.server_id = _MISSING_SERVER_ID
+        self.server_id: str = _MISSING_SERVER_ID
 
     def test_connection(self) -> None:
         super().test_connection()
@@ -195,7 +193,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         if not self.server_config:
             self.test_connection()
 
-        base_url = self.server_config.get("baseUrl")
+        base_url = self.server_config.raw_config.get("baseUrl")
         if not base_url:
             raise ValueError("baseUrl not found in server config")
         return base_url
@@ -203,6 +201,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     @classmethod
     def from_emitter(cls, emitter: DatahubRestEmitter) -> "DataHubGraph":
         session_config = emitter._session_config
+
         if isinstance(session_config.timeout, tuple):
             # TODO: This is slightly lossy. Eventually, we want to modify the emitter
             # to accept a tuple for timeout_sec, and then we'll be able to remove this.
@@ -220,6 +219,8 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
                 disable_ssl_verification=session_config.disable_ssl_verification,
                 ca_certificate_path=session_config.ca_certificate_path,
                 client_certificate_path=session_config.client_certificate_path,
+                client_mode=session_config.client_mode,
+                datahub_component=session_config.datahub_component,
             )
         )
 
@@ -1954,8 +1955,14 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         super().close()
 
 
-def get_default_graph() -> DataHubGraph:
+@functools.lru_cache(maxsize=None)
+def get_default_graph(
+    client_mode: Optional[ClientMode] = None,
+    datahub_component: Optional[str] = None,
+) -> DataHubGraph:
     graph_config = config_utils.load_client_config()
+    graph_config.client_mode = client_mode
+    graph_config.datahub_component = datahub_component
     graph = DataHubGraph(graph_config)
     graph.test_connection()
     telemetry_instance.set_context(server=graph)

datahub/ingestion/graph/config.py

@@ -1,8 +1,19 @@
+import os
+from enum import Enum, auto
 from typing import Dict, List, Optional
 
 from datahub.configuration.common import ConfigModel
 
 
+class ClientMode(Enum):
+    INGESTION = auto()
+    CLI = auto()
+    SDK = auto()
+
+
+DATAHUB_COMPONENT_ENV: str = os.getenv("DATAHUB_COMPONENT", "datahub").lower()
+
+
 class DatahubClientConfig(ConfigModel):
     """Configuration class for holding connectivity to datahub gms"""
 
@@ -18,3 +29,5 @@ class DatahubClientConfig(ConfigModel):
     client_certificate_path: Optional[str] = None
     disable_ssl_verification: bool = False
     openapi_ingestion: Optional[bool] = None
+    client_mode: Optional[ClientMode] = None
+    datahub_component: Optional[str] = None

datahub/ingestion/run/pipeline.py

@@ -31,6 +31,7 @@ from datahub.ingestion.api.source import Extractor, Source
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.extractor.extractor_registry import extractor_registry
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.reporting.reporting_provider_registry import (
     reporting_provider_registry,
 )
@@ -136,9 +137,8 @@ class CliReport(Report):
 
 
 def _make_default_rest_sink(ctx: PipelineContext) -> DatahubRestSink:
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.INGESTION)
     sink_config = graph._make_rest_sink_config()
-
     return DatahubRestSink(ctx, sink_config)
 
 
@@ -175,6 +175,7 @@ class Pipeline:
             self.graph: Optional[DataHubGraph] = None
             with _add_init_error_context("connect to DataHub"):
                 if self.config.datahub_api:
+                    self.config.datahub_api.client_mode = ClientMode.INGESTION
                     self.graph = exit_stack.enter_context(
                         DataHubGraph(self.config.datahub_api)
                     )

datahub/ingestion/run/pipeline_config.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional
 from pydantic import Field, validator
 
 from datahub.configuration.common import ConfigModel, DynamicTypedConfig
-from datahub.ingestion.graph.client import DatahubClientConfig
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.ingestion.sink.file import FileSinkConfig
 
 logger = logging.getLogger(__name__)