acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/emitter/rest_emitter.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import functools
 import json
 import logging
-import os
+import re
 import time
 from collections import defaultdict
 from dataclasses import dataclass
@@ -20,18 +20,18 @@ from typing import (
     Sequence,
     Tuple,
     Union,
+    overload,
 )
 
 import pydantic
 import requests
-from deprecated import deprecated
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import HTTPError, RequestException
+from typing_extensions import deprecated
 
 from datahub._version import nice_version_name
 from datahub.cli import config_utils
 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,
@@ -40,10 +40,17 @@ from datahub.configuration.common import (
     TraceTimeoutError,
     TraceValidationError,
 )
-from datahub.emitter.aspect import JSON_CONTENT_TYPE
+from datahub.configuration.env_vars import (
+    get_emit_mode,
+    get_emitter_trace,
+    get_rest_emitter_batch_max_payload_bytes,
+    get_rest_emitter_batch_max_payload_length,
+    get_rest_emitter_default_endpoint,
+    get_rest_emitter_default_retry_max_times,
+)
 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,
@@ -51,11 +58,20 @@ from datahub.emitter.response_helper import (
 )
 from datahub.emitter.serialization_helper import pre_json_transform
 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.metadata.schema_classes import (
+    KEY_ASPECT_NAMES,
+    ChangeTypeClass,
+)
+from datahub.utilities.server_config_util import RestServiceConfig, ServiceFeature
 
 if TYPE_CHECKING:
     from datahub.ingestion.graph.client import DataHubGraph
@@ -72,51 +88,77 @@ _DEFAULT_RETRY_STATUS_CODES = [  # Additional status codes to retry on
     504,
 ]
 _DEFAULT_RETRY_METHODS = ["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"]
-_DEFAULT_RETRY_MAX_TIMES = int(
-    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
-)
+_DEFAULT_RETRY_MAX_TIMES = int(get_rest_emitter_default_retry_max_times())
 
-_DATAHUB_EMITTER_TRACE = get_boolean_env_variable("DATAHUB_EMITTER_TRACE", False)
+_DATAHUB_EMITTER_TRACE = get_emitter_trace()
+
+_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
 TRACE_BACKOFF_FACTOR = 2.0  # Double the wait time each attempt
 
-# The limit is 16mb. We will use a max of 15mb to have some space
+# The limit is 16,000,000 bytes. We will use a max of 15mb to have some space
 # for overhead like request headers.
 # This applies to pretty much all calls to GMS.
-INGEST_MAX_PAYLOAD_BYTES = 15 * 1024 * 1024
+INGEST_MAX_PAYLOAD_BYTES = get_rest_emitter_batch_max_payload_bytes()
 
 # This limit is somewhat arbitrary. All GMS endpoints will timeout
 # and return a 500 if processing takes too long. To avoid sending
 # too much to the backend and hitting a timeout, we try to limit
 # the number of MCPs we send in a batch.
-BATCH_INGEST_MAX_PAYLOAD_LENGTH = int(
-    os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_LENGTH", 200)
+BATCH_INGEST_MAX_PAYLOAD_LENGTH = get_rest_emitter_batch_max_payload_length()
+
+
+def preserve_unicode_escapes(obj: Any) -> Any:
+    """Recursively convert unicode characters back to escape sequences"""
+    if isinstance(obj, dict):
+        return {k: preserve_unicode_escapes(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [preserve_unicode_escapes(item) for item in obj]
+    elif isinstance(obj, str):
+        # Convert non-ASCII characters back to \u escapes
+        def escape_unicode(match: Any) -> Any:
+            return f"\\u{ord(match.group(0)):04x}"
+
+        return re.sub(r"[^\x00-\x7F]", escape_unicode, obj)
+    else:
+        return obj
+
+
+class EmitMode(ConfigEnum):
+    # Fully synchronous processing that updates both primary storage (SQL) and search storage (Elasticsearch) before returning.
+    # Provides the strongest consistency guarantee but with the highest cost. Best for critical operations where immediate
+    # searchability and consistent reads are required.
+    SYNC_WAIT = auto()
+    # Synchronously updates the primary storage (SQL) but asynchronously updates search storage (Elasticsearch). Provides
+    # a balance between consistency and performance. Suitable for updates that need to be immediately reflected in direct
+    # entity retrievals but where search index consistency can be slightly delayed.
+    SYNC_PRIMARY = auto()
+    # Queues the metadata change for asynchronous processing and returns immediately. The client continues execution without
+    # waiting for the change to be fully processed. Best for high-throughput scenarios where eventual consistency is acceptable.
+    ASYNC = auto()
+    # Queues the metadata change asynchronously but blocks until confirmation that the write has been fully persisted.
+    # More efficient than fully synchronous operations due to backend parallelization and batching while still providing
+    # strong consistency guarantees. Useful when you need confirmation of successful persistence without sacrificing performance.
+    ASYNC_WAIT = auto()
+
+
+_DEFAULT_EMIT_MODE = pydantic.parse_obj_as(
+    EmitMode,
+    get_emit_mode() or EmitMode.SYNC_PRIMARY,
 )
 
 
-class RestTraceMode(ConfigEnum):
-    ENABLED = auto()
-    DISABLED = auto()
-
-
 class RestSinkEndpoint(ConfigEnum):
     RESTLI = auto()
     OPENAPI = auto()
 
 
-DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
+DEFAULT_REST_EMITTER_ENDPOINT = pydantic.parse_obj_as(
     RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
-)
-
-
-# Supported with v1.0
-DEFAULT_REST_TRACE_MODE = pydantic.parse_obj_as(
-    RestTraceMode,
-    os.getenv("DATAHUB_REST_TRACE_MODE", RestTraceMode.DISABLED),
+    get_rest_emitter_default_endpoint() or RestSinkEndpoint.RESTLI,
 )
 
 
@@ -132,12 +174,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
@@ -185,6 +239,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:
@@ -210,8 +317,8 @@ class DataHubRestEmitter(Closeable, Emitter):
     _gms_server: str
     _token: Optional[str]
     _session: requests.Session
-    _openapi_ingestion: bool
-    _default_trace_mode: bool
+    _openapi_ingestion: Optional[bool]
+    _server_config: RestServiceConfig
 
     def __init__(
         self,
@@ -227,8 +334,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 = False,
-        default_trace_mode: bool = False,
+        openapi_ingestion: Optional[bool] = None,
+        client_mode: Optional[ClientMode] = None,
+        datahub_component: Optional[str] = None,
+        server_config_refresh_interval: Optional[int] = None,
     ):
         if not gms_server:
             raise ConfigurationError("gms server is required")
@@ -240,21 +349,15 @@ 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:
-            logger.debug("Using API Tracing for ingestion.")
+        self._server_config_refresh_interval = server_config_refresh_interval
+        self._config_fetch_time: Optional[float] = None
 
         headers = {
             "X-RestLi-Protocol-Version": "2.0.0",
-            "X-DataHub-Py-Cli-Version": nice_version_name(),
             "Content-Type": "application/json",
         }
         if token:
@@ -300,39 +403,116 @@ 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
+    @property
+    def server_config(self) -> RestServiceConfig:
+        return self.fetch_server_config()
 
-            else:
+    # TODO: This should move to DataHubGraph once it no longer inherits from DataHubRestEmitter
+    def fetch_server_config(self) -> RestServiceConfig:
+        """
+        Fetch configuration from the server if not already loaded.
+
+        Returns:
+            The configuration dictionary
+
+        Raises:
+            ConfigurationError: If there's an error fetching or validating the configuration
+        """
+
+        if (
+            not hasattr(self, "_server_config")
+            or self._server_config is None
+            or (
+                self._server_config_refresh_interval is not None
+                and self._config_fetch_time is not None
+                and (time.time() - self._config_fetch_time)
+                > self._server_config_refresh_interval
+            )
+        ):
+            if self._session is None or self._gms_server is None:
                 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"
+                    "Session and URL are required to load configuration"
                 )
-        else:
-            logger.debug(
-                f"Unable to connect to {url} with status_code: {response.status_code}. Response: {response.text}"
-            )
-            if response.status_code == 401:
-                message = f"Unable to connect to {url} - got an authentication error: {response.text}."
+
+            url = f"{self._gms_server}/config"
+            response = self._session.get(url)
+
+            if response.status_code == 200:
+                raw_config = response.json()
+
+                # Validate that we're connected to the correct service
+                if not raw_config.get("noCode") == "true":
+                    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"
+                    )
+
+                self._server_config = RestServiceConfig(raw_config=raw_config)
+                self._config_fetch_time = time.time()
+                self._post_fetch_server_config()
+
+            else:
+                logger.debug(
+                    f"Unable to connect to {url} with status_code: {response.status_code}. Response: {response.text}"
+                )
+
+                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)
+
+        return self._server_config
+
+    def _post_fetch_server_config(self) -> None:
+        # Determine OpenAPI mode
+        if self._openapi_ingestion is None:
+            # No constructor parameter
+            if (
+                not get_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:
-                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)
+                # The system env is specifying the value
+                self._openapi_ingestion = (
+                    DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI
+                )
+
+    def test_connection(self) -> None:
+        self.fetch_server_config()
+        logger.debug(
+            f"Using {'OpenAPI' if self._openapi_ingestion else 'Restli'} for ingestion."
+        )
+        logger.debug(
+            f"{EmitMode.ASYNC_WAIT} {'IS' if self._should_trace(emit_mode=EmitMode.ASYNC_WAIT, warn=False) else 'IS NOT'} supported."
+        )
 
     def get_server_config(self) -> dict:
-        self.test_connection()
-        return self.server_config
+        return self.server_config.raw_config
+
+    def invalidate_config_cache(self) -> None:
+        """Manually invalidate the configuration cache."""
+        if (
+            hasattr(self, "_server_config")
+            and self._server_config is not None
+            and self._server_config_refresh_interval is not None
+        ):
+            # Set fetch time to beyond TTL in the past to force refresh on next access
+            self._config_fetch_time = (
+                time.time() - self._server_config_refresh_interval - 1
+            )
 
     def to_graph(self) -> "DataHubGraph":
         from datahub.ingestion.graph.client import DataHubGraph
@@ -342,39 +522,24 @@ class DataHubRestEmitter(Closeable, Emitter):
     def _to_openapi_request(
         self,
         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'}"
+        emit_mode: EmitMode,
+    ) -> 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()
-                if obj.get("value") and obj.get("contentType") == JSON_CONTENT_TYPE:
-                    obj = json.loads(obj["value"])
-                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
+            emit_mode: Client emit mode
+
+        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=emit_mode in (EmitMode.ASYNC, EmitMode.ASYNC_WAIT),
+            search_sync_flag=emit_mode == EmitMode.SYNC_WAIT,
+        )
 
     def emit(
         self,
@@ -385,7 +550,7 @@ class DataHubRestEmitter(Closeable, Emitter):
             UsageAggregation,
         ],
         callback: Optional[Callable[[Exception, str], None]] = None,
-        async_flag: Optional[bool] = None,
+        emit_mode: EmitMode = _DEFAULT_EMIT_MODE,
     ) -> None:
         try:
             if isinstance(item, UsageAggregation):
@@ -393,7 +558,7 @@ class DataHubRestEmitter(Closeable, Emitter):
             elif isinstance(
                 item, (MetadataChangeProposal, MetadataChangeProposalWrapper)
             ):
-                self.emit_mcp(item, async_flag=async_flag)
+                self.emit_mcp(item, emit_mode=emit_mode)
             else:
                 self.emit_mce(item)
         except Exception as e:
@@ -421,42 +586,84 @@ class DataHubRestEmitter(Closeable, Emitter):
             "systemMetadata": system_metadata_obj,
         }
         payload = json.dumps(snapshot)
+        if len(payload) > INGEST_MAX_PAYLOAD_BYTES:
+            logger.warning(
+                f"MCE object has size {len(payload)} that exceeds the max payload size of {INGEST_MAX_PAYLOAD_BYTES}, "
+                "so this metadata will likely fail to be emitted."
+            )
 
         self._emit_generic(url, payload)
 
+    @overload
+    @deprecated("Use emit_mode instead of async_flag")
+    def emit_mcp(
+        self,
+        mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
+        *,
+        async_flag: Optional[bool] = None,
+    ) -> None: ...
+
+    @overload
+    def emit_mcp(
+        self,
+        mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
+        *,
+        emit_mode: EmitMode = _DEFAULT_EMIT_MODE,
+        wait_timeout: Optional[timedelta] = timedelta(seconds=3600),
+    ) -> None: ...
+
     def emit_mcp(
         self,
         mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
         async_flag: Optional[bool] = None,
-        trace_flag: Optional[bool] = None,
-        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
+        emit_mode: EmitMode = _DEFAULT_EMIT_MODE,
+        wait_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> None:
+        if async_flag is True:
+            emit_mode = EmitMode.ASYNC
+
         ensure_has_system_metadata(mcp)
 
         trace_data = None
 
         if self._openapi_ingestion:
-            request = self._to_openapi_request(mcp, async_flag, async_default=False)
+            request = self._to_openapi_request(mcp, emit_mode)
             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):
+                if self._should_trace(emit_mode):
                     trace_data = extract_trace_data(response) if response else None
 
         else:
-            url = f"{self._gms_server}/aspects?action=ingestProposal"
+            if mcp.changeType == ChangeTypeClass.DELETE:
+                if mcp.aspectName not in KEY_ASPECT_NAMES:
+                    raise OperationalError(
+                        f"Delete not supported for non key aspect: {mcp.aspectName} for urn: "
+                        f"{mcp.entityUrn}"
+                    )
 
-            mcp_obj = pre_json_transform(mcp.to_obj())
-            payload_dict = {"proposal": mcp_obj}
+                url = f"{self._gms_server}/entities?action=delete"
+                payload_dict = {
+                    "urn": mcp.entityUrn,
+                }
+            else:
+                url = f"{self._gms_server}/aspects?action=ingestProposal"
 
-            if async_flag is not None:
-                payload_dict["async"] = "true" if async_flag else "false"
+                mcp_obj = preserve_unicode_escapes(pre_json_transform(mcp.to_obj()))
+                payload_dict = {
+                    "proposal": mcp_obj,
+                    "async": "true"
+                    if emit_mode in (EmitMode.ASYNC, EmitMode.ASYNC_WAIT)
+                    else "false",
+                }
 
             payload = json.dumps(payload_dict)
 
             response = self._emit_generic(url, payload)
 
-            if self._should_trace(async_flag, trace_flag):
+            if self._should_trace(emit_mode):
                 trace_data = (
                     extract_trace_data_from_mcps(response, [mcp]) if response else None
                 )
@@ -464,15 +671,14 @@ class DataHubRestEmitter(Closeable, Emitter):
         if trace_data:
             self._await_status(
                 [trace_data],
-                trace_timeout,
+                wait_timeout,
             )
 
     def emit_mcps(
         self,
         mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
-        async_flag: Optional[bool] = None,
-        trace_flag: Optional[bool] = None,
-        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
+        emit_mode: EmitMode = _DEFAULT_EMIT_MODE,
+        wait_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> int:
         if _DATAHUB_EMITTER_TRACE:
             logger.debug(f"Attempting to emit MCP batch of size {len(mcps)}")
@@ -481,43 +687,46 @@ class DataHubRestEmitter(Closeable, Emitter):
             ensure_has_system_metadata(mcp)
 
         if self._openapi_ingestion:
-            return self._emit_openapi_mcps(mcps, async_flag, trace_flag, trace_timeout)
+            return self._emit_openapi_mcps(mcps, emit_mode, wait_timeout)
         else:
-            return self._emit_restli_mcps(mcps, async_flag)
+            return self._emit_restli_mcps(mcps, emit_mode)
 
     def _emit_openapi_mcps(
         self,
         mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
-        async_flag: Optional[bool] = None,
-        trace_flag: Optional[bool] = None,
-        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
+        emit_mode: EmitMode,
+        wait_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 emit_mode: the mode to emit the MCPs
+        :param wait_timeout: timeout for blocking queue
         :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)
+            request = self._to_openapi_request(mcp, emit_mode)
             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
@@ -527,18 +736,20 @@ 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):
+        if self._should_trace(emit_mode):
             trace_data = []
             for response in responses:
                 data = extract_trace_data(response) if response else None
@@ -546,28 +757,36 @@ class DataHubRestEmitter(Closeable, Emitter):
                     trace_data.append(data)
 
             if trace_data:
-                self._await_status(trace_data, trace_timeout)
+                self._await_status(trace_data, wait_timeout)
 
         return len(responses)
 
     def _emit_restli_mcps(
         self,
         mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
-        async_flag: Optional[bool] = None,
+        emit_mode: EmitMode,
     ) -> int:
         url = f"{self._gms_server}/aspects?action=ingestProposalBatch"
 
         mcp_objs = [pre_json_transform(mcp.to_obj()) for mcp in mcps]
+        if len(mcp_objs) == 0:
+            return 0
 
         # As a safety mechanism, we need to make sure we don't exceed the max payload size for GMS.
         # If we will exceed the limit, we need to break it up into chunks.
-        mcp_obj_chunks: List[List[str]] = []
-        current_chunk_size = INGEST_MAX_PAYLOAD_BYTES
+        mcp_obj_chunks: List[List[str]] = [[]]
+        current_chunk_size = 0
         for mcp_obj in mcp_objs:
+            mcp_identifier = f"{mcp_obj.get('entityUrn')}-{mcp_obj.get('aspectName')}"
             mcp_obj_size = len(json.dumps(mcp_obj))
             if _DATAHUB_EMITTER_TRACE:
                 logger.debug(
-                    f"Iterating through object with size {mcp_obj_size} (type: {mcp_obj.get('aspectName')}"
+                    f"Iterating through object ({mcp_identifier}) with size {mcp_obj_size}"
+                )
+            if mcp_obj_size > INGEST_MAX_PAYLOAD_BYTES:
+                logger.warning(
+                    f"MCP object {mcp_identifier} has size {mcp_obj_size} that exceeds the max payload size of {INGEST_MAX_PAYLOAD_BYTES}, "
+                    "so this metadata will likely fail to be emitted."
                 )
 
             if (
@@ -580,7 +799,7 @@ class DataHubRestEmitter(Closeable, Emitter):
                 current_chunk_size = 0
             mcp_obj_chunks[-1].append(mcp_obj)
             current_chunk_size += mcp_obj_size
-        if len(mcp_obj_chunks) > 0:
+        if len(mcp_obj_chunks) > 1 or _DATAHUB_EMITTER_TRACE:
             logger.debug(
                 f"Decided to send {len(mcps)} MCP batch in {len(mcp_obj_chunks)} chunks"
             )
@@ -588,16 +807,19 @@ class DataHubRestEmitter(Closeable, Emitter):
         for mcp_obj_chunk in mcp_obj_chunks:
             # TODO: We're calling json.dumps on each MCP object twice, once to estimate
             # the size when chunking, and again for the actual request.
-            payload_dict: dict = {"proposals": mcp_obj_chunk}
-            if async_flag is not None:
-                payload_dict["async"] = "true" if async_flag else "false"
+            payload_dict: dict = {
+                "proposals": mcp_obj_chunk,
+                "async": "true"
+                if emit_mode in (EmitMode.ASYNC, EmitMode.ASYNC_WAIT)
+                else "false",
+            }
 
             payload = json.dumps(payload_dict)
             self._emit_generic(url, payload)
 
         return len(mcp_obj_chunks)
 
-    @deprecated
+    @deprecated("Use emit with a datasetUsageStatistics aspect instead")
     def emit_usage(self, usageStats: UsageAggregation) -> None:
         url = f"{self._gms_server}/usageStats?action=batchIngest"
 
@@ -608,11 +830,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
@@ -625,7 +849,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:
@@ -661,7 +886,7 @@ class DataHubRestEmitter(Closeable, Emitter):
     def _await_status(
         self,
         trace_data: List[TraceData],
-        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
+        wait_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> None:
         """Verify the status of asynchronous write operations.
         Args:
@@ -671,8 +896,8 @@ class DataHubRestEmitter(Closeable, Emitter):
             TraceTimeoutError: If verification fails or times out
             TraceValidationError: Expected write was not completed successfully
         """
-        if trace_timeout is None:
-            raise ValueError("trace_timeout cannot be None")
+        if wait_timeout is None:
+            raise ValueError("wait_timeout cannot be None")
 
         try:
             if not trace_data:
@@ -685,9 +910,9 @@ class DataHubRestEmitter(Closeable, Emitter):
                 current_backoff = TRACE_INITIAL_BACKOFF
 
                 while trace.data:
-                    if datetime.now() - start_time > trace_timeout:
+                    if datetime.now() - start_time > wait_timeout:
                         raise TraceTimeoutError(
-                            f"Timeout waiting for async write completion after {trace_timeout.total_seconds()} seconds"
+                            f"Timeout waiting for async write completion after {wait_timeout.total_seconds()} seconds"
                         )
 
                     base_url = f"{self._gms_server}/openapi/v1/trace/write"
@@ -700,7 +925,7 @@ class DataHubRestEmitter(Closeable, Emitter):
                         for aspect_name, aspect_status in aspects.items():
                             if not aspect_status["success"]:
                                 error_msg = (
-                                    f"Unable to validate async write to DataHub GMS: "
+                                    f"Unable to validate async write {trace.trace_id} ({trace.extract_timestamp()}) to DataHub GMS: "
                                     f"Persistence failure for URN '{urn}' aspect '{aspect_name}'. "
                                     f"Status: {aspect_status}"
                                 )
@@ -739,17 +964,28 @@ class DataHubRestEmitter(Closeable, Emitter):
             logger.error(f"Error during status verification: {str(e)}")
             raise
 
-    def _should_trace(
-        self,
-        async_flag: Optional[bool] = None,
-        trace_flag: Optional[bool] = None,
-        async_default: bool = False,
-    ) -> bool:
-        resolved_trace_flag = (
-            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
-        return resolved_trace_flag and resolved_async_flag
+    def _should_trace(self, emit_mode: EmitMode, warn: bool = True) -> bool:
+        if emit_mode == EmitMode.ASYNC_WAIT:
+            if not bool(self._openapi_ingestion):
+                if warn:
+                    logger.warning(
+                        f"{emit_mode} requested but is only available when using OpenAPI."
+                    )
+                return False
+            elif getattr(
+                self, "server_config", None
+            ) is None or not self.server_config.supports_feature(
+                ServiceFeature.API_TRACING
+            ):
+                if warn:
+                    logger.warning(
+                        f"{emit_mode} requested but is only available with a newer GMS version."
+                    )
+                return False
+            else:
+                return True
+        else:
+            return False
 
     def __repr__(self) -> str:
         token_str = (