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

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
@@ -41,23 +40,26 @@ from datahub.configuration.common import (
     TraceTimeoutError,
     TraceValidationError,
 )
-from datahub.emitter.aspect import JSON_CONTENT_TYPE, JSON_PATCH_CONTENT_TYPE
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.request_helper import make_curl_command
+from datahub.emitter.request_helper import OpenApiRequest, make_curl_command
 from datahub.emitter.response_helper import (
     TraceData,
     extract_trace_data,
     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
@@ -80,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
@@ -134,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
@@ -187,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:
@@ -212,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,
@@ -229,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")
@@ -244,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:
@@ -258,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:
@@ -304,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
 
@@ -348,43 +431,24 @@ class DataHubRestEmitter(Closeable, Emitter):
         mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
         async_flag: Optional[bool] = None,
         async_default: bool = False,
-    ) -> Optional[Tuple[str, List[Dict[str, Any]]]]:
-        if mcp.aspect and mcp.aspectName:
-            resolved_async_flag = (
-                async_flag if async_flag is not None else async_default
-            )
-            url = f"{self._gms_server}/openapi/v3/entity/{mcp.entityType}?async={'true' if resolved_async_flag else 'false'}"
+    ) -> Optional[OpenApiRequest]:
+        """
+        Convert a MetadataChangeProposal to an OpenAPI request format.
 
-            if isinstance(mcp, MetadataChangeProposalWrapper):
-                aspect_value = pre_json_transform(
-                    mcp.to_obj(simplified_structure=True)
-                )["aspect"]["json"]
-            else:
-                obj = mcp.aspect.to_obj()
-                content_type = obj.get("contentType")
-                if obj.get("value") and content_type == JSON_CONTENT_TYPE:
-                    # Undo double serialization.
-                    obj = json.loads(obj["value"])
-                elif content_type == JSON_PATCH_CONTENT_TYPE:
-                    raise NotImplementedError(
-                        "Patches are not supported for OpenAPI ingestion. Set the endpoint to RESTLI."
-                    )
-                aspect_value = pre_json_transform(obj)
-            return (
-                url,
-                [
-                    {
-                        "urn": mcp.entityUrn,
-                        mcp.aspectName: {
-                            "value": aspect_value,
-                            "systemMetadata": mcp.systemMetadata.to_obj()
-                            if mcp.systemMetadata
-                            else None,
-                        },
-                    }
-                ],
-            )
-        return None
+        Args:
+            mcp: The metadata change proposal
+            async_flag: Optional flag to override async behavior
+            async_default: Default async behavior if not specified
+
+        Returns:
+            An OpenApiRequest object or None if the MCP doesn't have required fields
+        """
+        return OpenApiRequest.from_mcp(
+            mcp=mcp,
+            gms_server=self._gms_server,
+            async_flag=async_flag,
+            async_default=async_default,
+        )
 
     def emit(
         self,
@@ -448,7 +512,9 @@ class DataHubRestEmitter(Closeable, Emitter):
         if self._openapi_ingestion:
             request = self._to_openapi_request(mcp, async_flag, async_default=False)
             if request:
-                response = self._emit_generic(request[0], payload=request[1])
+                response = self._emit_generic(
+                    request.url, payload=request.payload, method=request.method
+                )
 
                 if self._should_trace(async_flag, trace_flag):
                     trace_data = extract_trace_data(response) if response else None
@@ -503,31 +569,36 @@ class DataHubRestEmitter(Closeable, Emitter):
         trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> int:
         """
-        1. Grouping MCPs by their 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)
 
         The Chunk class encapsulates both the items and their byte size tracking
-        Serializing the items only once with json.dumps(request[1]) and reusing that
+        Serializing the items only once with json.dumps(request.payload) and reusing that
         The chunking logic handles edge cases (always accepting at least one item per chunk)
         The joining logic is efficient with a simple string concatenation
 
         :param mcps: metadata change proposals to transmit
         :param async_flag: the mode
+        :param trace_flag: whether to trace the requests
+        :param trace_timeout: timeout for tracing
         :return: number of requests
         """
-        # group by entity url
-        batches: Dict[str, List[_Chunk]] = defaultdict(
+        # Group by entity URL and HTTP method
+        batches: Dict[Tuple[str, str], List[_Chunk]] = defaultdict(
             lambda: [_Chunk(items=[])]
         )  # Initialize with one empty Chunk
 
         for mcp in mcps:
             request = self._to_openapi_request(mcp, async_flag, async_default=True)
             if request:
-                current_chunk = batches[request[0]][-1]  # Get the last chunk
-                # Only serialize once
-                serialized_item = json.dumps(request[1][0])
+                # Create a composite key with both method and URL
+                key = (request.method, request.url)
+                current_chunk = batches[key][-1]  # Get the last chunk
+
+                # Only serialize once - we're serializing a single payload item
+                serialized_item = json.dumps(request.payload[0])
                 item_bytes = len(serialized_item.encode())
 
                 # If adding this item would exceed max_bytes, create a new chunk
@@ -537,15 +608,17 @@ class DataHubRestEmitter(Closeable, Emitter):
                     or len(current_chunk.items) >= BATCH_INGEST_MAX_PAYLOAD_LENGTH
                 ):
                     new_chunk = _Chunk(items=[])
-                    batches[request[0]].append(new_chunk)
+                    batches[key].append(new_chunk)
                     current_chunk = new_chunk
 
                 current_chunk.add_item(serialized_item)
 
         responses = []
-        for url, chunks in batches.items():
+        for (method, url), chunks in batches.items():
             for chunk in chunks:
-                response = self._emit_generic(url, payload=_Chunk.join(chunk))
+                response = self._emit_generic(
+                    url, payload=_Chunk.join(chunk), method=method
+                )
                 responses.append(response)
 
         if self._should_trace(async_flag, trace_flag, async_default=True):
@@ -618,11 +691,13 @@ class DataHubRestEmitter(Closeable, Emitter):
         payload = json.dumps(snapshot)
         self._emit_generic(url, payload)
 
-    def _emit_generic(self, url: str, payload: Union[str, Any]) -> requests.Response:
+    def _emit_generic(
+        self, url: str, payload: Union[str, Any], method: str = "POST"
+    ) -> requests.Response:
         if not isinstance(payload, str):
             payload = json.dumps(payload)
 
-        curl_command = make_curl_command(self._session, "POST", url, payload)
+        curl_command = make_curl_command(self._session, method, url, payload)
         payload_size = len(payload)
         if payload_size > INGEST_MAX_PAYLOAD_BYTES:
             # since we know total payload size here, we could simply avoid sending such payload at all and report a warning, with current approach we are going to cause whole ingestion to fail
@@ -635,7 +710,8 @@ class DataHubRestEmitter(Closeable, Emitter):
             curl_command,
         )
         try:
-            response = self._session.post(url, data=payload)
+            method_func = getattr(self._session, method.lower())
+            response = method_func(url, data=payload) if payload else method_func(url)
             response.raise_for_status()
             return response
         except HTTPError as e:
@@ -759,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/api/source.py

@@ -420,12 +420,9 @@ class Source(Closeable, metaclass=ABCMeta):
         Run in order, first in list is applied first. Be careful with order when overriding.
         """
         browse_path_processor: Optional[MetadataWorkUnitProcessor] = None
-        if (
-            self.ctx.pipeline_config
-            and self.ctx.pipeline_config.flags.generate_browse_path_v2
-        ):
+        if self.ctx.flags.generate_browse_path_v2:
             browse_path_processor = self._get_browse_path_processor(
-                self.ctx.pipeline_config.flags.generate_browse_path_v2_dry_run
+                self.ctx.flags.generate_browse_path_v2_dry_run
             )
 
         auto_lowercase_dataset_urns: Optional[MetadataWorkUnitProcessor] = None

datahub/ingestion/glossary/classification_mixin.py

@@ -319,8 +319,10 @@ def classification_workunit_processor(
                         partial(
                             data_reader.get_sample_data_for_table,
                             table_id,
-                            classification_handler.config.classification.sample_size
-                            * SAMPLE_SIZE_MULTIPLIER,
+                            int(
+                                classification_handler.config.classification.sample_size
+                                * SAMPLE_SIZE_MULTIPLIER
+                            ),
                             **(data_reader_kwargs or {}),
                         )
                         if data_reader

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,11 +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=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()
@@ -193,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
@@ -201,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.
@@ -218,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,
             )
         )
 
@@ -1952,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"""
 
@@ -17,3 +28,6 @@ class DatahubClientConfig(ConfigModel):
     ca_certificate_path: Optional[str] = None
     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/graph/filters.py

@@ -18,7 +18,7 @@ RawSearchFilterRule: TypeAlias = Dict[str, Union[str, bool, List[str]]]
 # This can be put directly into the orFilters parameter in GraphQL.
 RawSearchFilter: TypeAlias = List[Dict[Literal["and"], List[RawSearchFilterRule]]]
 
-# Mirrors our GraphQL enum: https://datahubproject.io/docs/graphql/enums#filteroperator
+# Mirrors our GraphQL enum: https://docs.datahub.com/docs/graphql/enums#filteroperator
 FilterOperator: TypeAlias = Literal[
     "CONTAIN",
     "EQUAL",

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__)

datahub/ingestion/sink/datahub_rest.py

@@ -34,7 +34,7 @@ from datahub.ingestion.api.sink import (
     WriteCallback,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.graph.client import DatahubClientConfig
+from datahub.ingestion.graph.config import ClientMode, DatahubClientConfig
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
@@ -140,11 +140,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
                 f"💥 Failed to connect to DataHub with {repr(self.emitter)}"
             ) from exc
 
-        self.report.gms_version = (
-            gms_config.get("versions", {})
-            .get("acryldata/datahub", {})
-            .get("version", None)
-        )
+        self.report.gms_version = gms_config.service_version
         self.report.mode = self.config.mode
         self.report.max_threads = self.config.max_threads
         logger.debug("Setting env variables to override config")
@@ -180,6 +176,8 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             disable_ssl_verification=config.disable_ssl_verification,
             openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
             default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
+            client_mode=config.client_mode,
+            datahub_component=config.datahub_component,
         )
 
     @property
@@ -190,6 +188,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         # https://github.com/psf/requests/issues/1871#issuecomment-32751346
         thread_local = self._emitter_thread_local
         if not hasattr(thread_local, "emitter"):
+            self.config.client_mode = ClientMode.INGESTION
             thread_local.emitter = DatahubRestSink._make_emitter(self.config)
         return thread_local.emitter