acryl-datahub 1.0.0rc16__py3-none-any.whl → 1.0.0rc18__py3-none-any.whl

{acryl_datahub-1.0.0rc16.dist-info → acryl_datahub-1.0.0rc18.dist-info}/RECORD

@@ -1,6 +1,6 @@
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=E9Ai-iBRxYw3ysQIRJYlsLsIOcjgGBzeqp6OJwZHeAg,322
+datahub/_version.py,sha256=cadwaIVSYjldq1gLBbOrAkoALM9-SUSsE4xECju3hJw,322
 datahub/entrypoints.py,sha256=2TYgHhs3sCxJlojIHjqfxzt3_ImPwPzq4vBtsUuMqu4,8885
 datahub/errors.py,sha256=w6h8b27j9XlmPbTwqpu7-wgiTrXlHzcnUOnJ_iOrwzo,520
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -93,7 +93,7 @@ datahub/cli/specific/structuredproperties_cli.py,sha256=Rb06XJKxqda5RSUo188B90Wn
 datahub/cli/specific/user_cli.py,sha256=jGAokb1NRu8obs6P2g4OL2NQdFgpUBa9De55TBBtun0,1897
 datahub/configuration/__init__.py,sha256=5TN3a7CWNsLRHpdj-sv2bxKWF2IslvJwE6EpNMFrIS4,123
 datahub/configuration/_config_enum.py,sha256=ul2hr5gMmdLvBINicFkMNMi1ApmnmZSwNdUYYted5nk,1447
-datahub/configuration/common.py,sha256=bt_kiy2blqHbxbG-aM_8RNAZoIKMfaMzOZhtknnyLXg,10410
+datahub/configuration/common.py,sha256=GKRNgRCBsFUmZb6huD6DPNYUKNRxUbcHfFazHLte7wk,10626
 datahub/configuration/config_loader.py,sha256=hRzPFxkz-w9IqkpSa5vwCzSra1p49DyfeJNeyqGa8-4,6827
 datahub/configuration/connection_resolver.py,sha256=n4-6MwMiOEDgTouxO0SMjTILKVhJPo6-naE6FuR5qMs,1516
 datahub/configuration/datetimes.py,sha256=nayNc0mmlVKH6oVv9ud6C1dDUiZPGabW-YZxvrkosPg,2870
@@ -124,7 +124,8 @@ datahub/emitter/mcp.py,sha256=hAAYziDdkwjazQU0DtWMbQWY8wS09ACrKJbqxoWXdgc,9637
 datahub/emitter/mcp_builder.py,sha256=Q1bX2BthNvZ7ae71XYF6ICoiN8IOqaAd_h3zOct57Q0,11752
 datahub/emitter/mcp_patch_builder.py,sha256=u7cpW6DkiN7KpLapmMaXgL_FneoN69boxiANbVgMdSI,4564
 datahub/emitter/request_helper.py,sha256=33ORG3S3OVy97_jlWBRn7yUM5XCIkRN6WSdJvN7Ofcg,670
-datahub/emitter/rest_emitter.py,sha256=yJ_QCVe4K-ILXQOhS7CiTHG5Gw2xu4H9mscAnOvfUY4,23633
+datahub/emitter/response_helper.py,sha256=lRMvzF-RPHNkN_ONl-N2uJjKh5XtRFrofrdGibVGn2U,4509
+datahub/emitter/rest_emitter.py,sha256=LyJuTZicSxzyLlwCyVzecjuyFDH0HkvQmSh037OLBc8,29777
 datahub/emitter/serialization_helper.py,sha256=q12Avmf70Vy4ttQGMJoTKlE5EsybMKNg2w3MQeZiHvk,3652
 datahub/emitter/sql_parsing_builder.py,sha256=Cr5imZrm3dYDSCACt5MFscgHCtVbHTD6IjUmsvsKoEs,11991
 datahub/emitter/synchronized_file_emitter.py,sha256=s4ATuxalI4GDAkrZTaGSegxBdvvNPZ9jRSdtElU0kNs,1805
@@ -169,7 +170,7 @@ datahub/ingestion/glossary/classifier.py,sha256=daLxnVv_JlfB_jBOxH5LrU_xQRndrsGo
 datahub/ingestion/glossary/classifier_registry.py,sha256=yFOYLQhDgCLqXYMG3L1BquXafeLcZDcmp8meyw6k9ts,307
 datahub/ingestion/glossary/datahub_classifier.py,sha256=O7wm6gQT1Jf2QSKdWjJQbS5oSzJwplXzfza26Gdq5Mg,7555
 datahub/ingestion/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/graph/client.py,sha256=DpGenZBQ5yziaDRNzKDSlMnE2GWoJe_yk2MdyU3UnLM,65551
+datahub/ingestion/graph/client.py,sha256=4F-IMnz4DwIdLPL5yXMRqqlB3uonN2OHOUIKKHByS7o,65680
 datahub/ingestion/graph/config.py,sha256=_oha8Je7P80ZmrkZUAaRHyYbdMmTkMI5JkYjEP2Ri1Q,751
 datahub/ingestion/graph/connections.py,sha256=9462L0ZWGKURyypAln25eMPhK3pcufBar9tNDoqspXs,741
 datahub/ingestion/graph/entity_versioning.py,sha256=nrcNz0Qm6kpE6oTu_mrYUQDx14KPspBTc6R9SyFUY6c,6901
@@ -188,7 +189,7 @@ datahub/ingestion/sink/blackhole.py,sha256=-jYcWo4i8q7312bCIoHrGr7nT9JdPvA7c4jvS
 datahub/ingestion/sink/console.py,sha256=TZfhA0Ec2eNCrMH7RRy2JOdUE-U-hkoIQrPm1CmKLQs,591
 datahub/ingestion/sink/datahub_kafka.py,sha256=_cjuXu5I6G0zJ2UK7hMbaKjMPZXeIwRMgm7CVeTiNtc,2578
 datahub/ingestion/sink/datahub_lite.py,sha256=7u2aWm7ENLshKHl-PkjJg6Mrw4bWs8sTfKIBz4mm8Ak,1879
-datahub/ingestion/sink/datahub_rest.py,sha256=KLUFteqGPmMvKaMbZG055uBYNyNUDkt_ziuJcjaNl1o,12781
+datahub/ingestion/sink/datahub_rest.py,sha256=4hvMDUxHMJXGgk3Iy7fcYGKixjvVd9DHD03X-F3kOg0,12976
 datahub/ingestion/sink/file.py,sha256=SxXJPJpkIGoaqRjCcSmj2ZE3xE4rLlBABBGwpTj5LWI,3271
 datahub/ingestion/sink/sink_registry.py,sha256=JRBWx8qEYg0ubSTyhqwgSWctgxwyp6fva9GoN2LwBao,490
 datahub/ingestion/source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -212,10 +213,11 @@ datahub/ingestion/source/openapi_parser.py,sha256=K3Z4aLXBQX8eR7tvk9iZakOjghjezx
 datahub/ingestion/source/preset.py,sha256=bbh0ZWiAZMy2zuJDmaRY07_OuGJ9tdtKjwvIxqbY5II,3964
 datahub/ingestion/source/pulsar.py,sha256=u5F8QnCLJsht5-7XCiUTsnfhCPIpKVB_l32CgMCU-As,20187
 datahub/ingestion/source/redash.py,sha256=YxjSad-X_wPmxYH8dJmFz_VCFhiLTCTSlK99WdvcYiA,30653
-datahub/ingestion/source/salesforce.py,sha256=d56tfYqg1rGDvMkLznmBJII55B1Zs8XTaQrrW-wHdLo,32679
+datahub/ingestion/source/salesforce.py,sha256=CQtDFv1OsbC1vyzNbKOc6GxhFQ5GdYj45hgAF0-oIcw,40487
 datahub/ingestion/source/source_registry.py,sha256=a2mLjJPLkSI-gYCTb_7U7Jo4D8jGknNQ_yScPIihXFk,1208
 datahub/ingestion/source/sql_queries.py,sha256=Ip7UZub7fgMh7P5jL_zJPY7lSkc9GGTy8GJ8lqZrcsE,9502
 datahub/ingestion/source/superset.py,sha256=WrpCiZEC17cmFGcfUTTqUdnKASq7ZpT0ih-4xqB9qt4,30976
+datahub/ingestion/source/vertexai.py,sha256=uOtIgHwsH--hkAFqspXGoNN-jHip16s6m5lyvwi-jrg,27735
 datahub/ingestion/source/abs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/abs/config.py,sha256=mBQe0JTaP-Rcv4HnMUUySoYbSr4r3jDEMioxaXHnxXU,6709
 datahub/ingestion/source/abs/datalake_profiler_config.py,sha256=Rkf64evufyVGPiE4VK8QAjzBiJFu85tOGMmJ0lJZ2Og,3600
@@ -245,7 +247,7 @@ datahub/ingestion/source/bigquery_v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeu
 datahub/ingestion/source/bigquery_v2/bigquery.py,sha256=fADrYPiQaiJYvLOrltgv8RJIV9VV2y7vjh3s0zHW6Cw,13950
 datahub/ingestion/source/bigquery_v2/bigquery_audit.py,sha256=kEwWhq3ch6WT4q4hcX8-fvQh28KgrNfspFwIytO3vQA,25103
 datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py,sha256=LuGJ6LgPViLIfDQfylxlQ3CA7fZYM5MDt8M-7sfzm84,5096
-datahub/ingestion/source/bigquery_v2/bigquery_config.py,sha256=UwGCngVm6GEUYvXvdquK-inNIruUoTB_pN6qsrjW5c4,26291
+datahub/ingestion/source/bigquery_v2/bigquery_config.py,sha256=66c7VPPcX1Eamm14VoEZzSPngMiwsNZOeymWBBRsznk,24420
 datahub/ingestion/source/bigquery_v2/bigquery_data_reader.py,sha256=DeT3v_Z82__8En0FcZ0kavBAWQoRvSZ5Rppm9eeDAb8,2393
 datahub/ingestion/source/bigquery_v2/bigquery_helper.py,sha256=QER3gY8e_k1_eNVj7cBso7ZzrWl_vO5PYSa6CpvqNx8,1554
 datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py,sha256=9_sfX8BE2vt9RjBMyq27UxCxBaSlD5o3L4gQxrwlPvA,4961
@@ -268,6 +270,7 @@ datahub/ingestion/source/cassandra/cassandra_profiling.py,sha256=DkSIryZNwLei5Pa
 datahub/ingestion/source/cassandra/cassandra_utils.py,sha256=j-LidYkaCTmGnpUVNLsax_c3z32PsQbsbHeYojygd1s,5105
 datahub/ingestion/source/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/common/data_reader.py,sha256=XbSxiRTYrk6seOz0ZjVjzSpGvP8lEjmqXrNI4cdYYmQ,1819
+datahub/ingestion/source/common/gcp_credentials_config.py,sha256=MM902YksVGv-yCmUTimBXF9YreHUDo3PW3o1Wasb2FE,2212
 datahub/ingestion/source/common/subtypes.py,sha256=LCJefUZ9o8yyhNXOy_HJefBOt93Cmn9r3m4VtCiK4iM,2643
 datahub/ingestion/source/data_lake_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/data_lake_common/config.py,sha256=qUk83B01hjuBKHvVz8SmXnVCy5eFj-2-2QLEOrAdbgk,359
@@ -444,7 +447,7 @@ datahub/ingestion/source/snowflake/snowflake_connection.py,sha256=pEw2O9xoTSIWDi
 datahub/ingestion/source/snowflake/snowflake_data_reader.py,sha256=ffR5E2uhD71FUMXd3XOg2rHwrp1rbbGEFTAbqKcmI2s,2195
 datahub/ingestion/source/snowflake/snowflake_lineage_v2.py,sha256=FBmiONx4EGHWV8RNJT6zHZyntKinPFFyd2oKbTUIbhE,21319
 datahub/ingestion/source/snowflake/snowflake_profiler.py,sha256=0DJiSwII6FY34urlBja2FW66NaVvhbBWmG0p7u8Xyrc,7548
-datahub/ingestion/source/snowflake/snowflake_queries.py,sha256=R3QxWtdR8T_8YV_3aqt3rJdto1gAij_mEHlSYKqdCfA,28326
+datahub/ingestion/source/snowflake/snowflake_queries.py,sha256=gX9E1Z_CemAZsuTDmtvqrxY7vBL2da75j7X8Xwhaf8Y,28441
 datahub/ingestion/source/snowflake/snowflake_query.py,sha256=qz_rhRMNCXxHd23bePbb3YxhFgN7eRpV4s6g58hQ5bU,39678
 datahub/ingestion/source/snowflake/snowflake_report.py,sha256=ahea-bwpW6T0iDehGo0Qq_J7wKxPkV61aYHm8bGwDqo,6651
 datahub/ingestion/source/snowflake/snowflake_schema.py,sha256=qkGgk6WdKSPThFjexXHrxUPYiVtzDk2MbGX3b281A4c,26044
@@ -461,7 +464,7 @@ datahub/ingestion/source/sql/clickhouse.py,sha256=uSRy-HKAiGFTHVLoVtGoh23X0O1lwy
 datahub/ingestion/source/sql/cockroachdb.py,sha256=XaD7eae34plU9ISRC6PzYX9q6RdT2qkzjH6CpTOgkx4,1443
 datahub/ingestion/source/sql/druid.py,sha256=IjGZdntb5hubkIzzT9qDRDpyfbckEg2GwRncvC5mDSs,2722
 datahub/ingestion/source/sql/hana.py,sha256=0PIvcX0Rz59NyR7Ag5Bv1MBV_UbJwxl9UAopo_xe_CA,1342
-datahub/ingestion/source/sql/hive.py,sha256=tfRgzatF4cDb3F7gNXF9zEjFOFrcI318K6yGgykW_EQ,30212
+datahub/ingestion/source/sql/hive.py,sha256=n0XCGkNkVAe-TEyXbxlefvohbmtALbWaC1a0_B9rlG8,30670
 datahub/ingestion/source/sql/hive_metastore.py,sha256=HW0zoHKarBYb8oVCy5fHvPOn-pTo25LctW_AusmH0hQ,36252
 datahub/ingestion/source/sql/mariadb.py,sha256=Hm102kmfs_1rd4lsTYhzVMZq5S3B6cyfvpHSzJjqvMw,737
 datahub/ingestion/source/sql/mysql.py,sha256=nDWK4YbqomcJgnit9b8geUGrp_3eix4bt0_k94o7g-0,3350
@@ -1022,9 +1025,9 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.0.0rc16.dist-info/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
-acryl_datahub-1.0.0rc16.dist-info/METADATA,sha256=GLR8w248XWb0HwHx4mWmrC_lVq4qTRyn2YiFAixM1qU,175337
-acryl_datahub-1.0.0rc16.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-acryl_datahub-1.0.0rc16.dist-info/entry_points.txt,sha256=U1e5ZwqPX1OaIbvGrwvozcdB8SbzFYXQM7plpdLKKeo,9592
-acryl_datahub-1.0.0rc16.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.0.0rc16.dist-info/RECORD,,
+acryl_datahub-1.0.0rc18.dist-info/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.0.0rc18.dist-info/METADATA,sha256=paUotX7MdMGKowqzTD7mt93-f2Hl6uB-og_wV_VbmhU,176898
+acryl_datahub-1.0.0rc18.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+acryl_datahub-1.0.0rc18.dist-info/entry_points.txt,sha256=7-eDilp0OACUtlmmZ-LF6H9MF_SWD_bWHKNG7Dvhhos,9652
+acryl_datahub-1.0.0rc18.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.0.0rc18.dist-info/RECORD,,

{acryl_datahub-1.0.0rc16.dist-info → acryl_datahub-1.0.0rc18.dist-info}/entry_points.txt

@@ -101,6 +101,7 @@ tableau = datahub.ingestion.source.tableau.tableau:TableauSource
 teradata = datahub.ingestion.source.sql.teradata:TeradataSource
 trino = datahub.ingestion.source.sql.trino:TrinoSource
 unity-catalog = datahub.ingestion.source.unity.source:UnityCatalogSource
+vertexai = datahub.ingestion.source.vertexai:VertexAISource
 vertica = datahub.ingestion.source.sql.vertica:VerticaSource
 
 [datahub.ingestion.transformer.plugins]

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.0.0rc16"
+__version__ = "1.0.0rc18"
 
 
 def is_dev_mode() -> bool:

datahub/configuration/common.py

@@ -198,6 +198,14 @@ class IgnorableError(MetaError):
     """An error that can be ignored."""
 
 
+class TraceTimeoutError(OperationalError):
+    """Failure to complete an API Trace within the timeout."""
+
+
+class TraceValidationError(OperationalError):
+    """Failure to complete the expected write operation."""
+
+
 @runtime_checkable
 class ExceptionWithProps(Protocol):
     def get_telemetry_props(self) -> Dict[str, Any]: ...

datahub/emitter/response_helper.py

@@ -0,0 +1,145 @@
+import json
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Union
+
+from requests import Response
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
+    MetadataChangeProposal,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TraceData:
+    trace_id: str
+    data: Dict[str, List[str]]
+
+    def __post_init__(self) -> None:
+        if not self.trace_id:
+            raise ValueError("trace_id cannot be empty")
+        if not isinstance(self.data, dict):
+            raise TypeError("data must be a dictionary")
+
+
+def _extract_trace_id(
+    response: Response, trace_header: str = "traceparent"
+) -> Optional[str]:
+    """
+    Extract trace ID from response headers.
+    Args:
+        response: HTTP response object
+        trace_header: Name of the trace header to use
+    Returns:
+        Trace ID if found and response is valid, None otherwise
+    """
+    if not 200 <= response.status_code < 300:
+        logger.debug(f"Invalid status code: {response.status_code}")
+        return None
+
+    trace_id = response.headers.get(trace_header)
+    if not trace_id:
+        logger.debug(f"Missing trace header: {trace_header}")
+        return None
+
+    return trace_id
+
+
+def extract_trace_data(
+    response: Response,
+    aspects_to_trace: Optional[List[str]] = None,
+    trace_header: str = "traceparent",
+) -> Optional[TraceData]:
+    """
+    Extract trace data from a response object.
+    Args:
+        response: HTTP response object
+        aspects_to_trace: Optional list of aspect names to extract. If None, extracts all aspects.
+        trace_header: Name of the trace header to use (default: "traceparent")
+    Returns:
+        TraceData object if successful, None otherwise
+    Raises:
+        JSONDecodeError: If response body cannot be decoded as JSON
+    """
+    trace_id = _extract_trace_id(response, trace_header)
+    if not trace_id:
+        return None
+
+    try:
+        json_data = response.json()
+        if not isinstance(json_data, list):
+            logger.debug("JSON data is not a list")
+            return None
+
+        data: Dict[str, List[str]] = {}
+
+        for item in json_data:
+            urn = item.get("urn")
+            if not urn:
+                logger.debug(f"Skipping item without URN: {item}")
+                continue
+
+            if aspects_to_trace is None:
+                aspect_names = [
+                    k for k, v in item.items() if k != "urn" and v is not None
+                ]
+            else:
+                aspect_names = [
+                    field for field in aspects_to_trace if item.get(field) is not None
+                ]
+
+            data[urn] = aspect_names
+
+        return TraceData(trace_id=trace_id, data=data)
+
+    except json.JSONDecodeError as e:
+        logger.error(f"Failed to decode JSON response: {e}")
+        return None
+
+
+def extract_trace_data_from_mcps(
+    response: Response,
+    mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
+    aspects_to_trace: Optional[List[str]] = None,
+    trace_header: str = "traceparent",
+) -> Optional[TraceData]:
+    """
+    Extract trace data from a response object and populate data from provided MCPs.
+    Args:
+        response: HTTP response object used only for trace_id extraction
+        mcps: List of MCP URN and aspect data
+        aspects_to_trace: Optional list of aspect names to extract. If None, extracts all aspects.
+        trace_header: Name of the trace header to use (default: "traceparent")
+    Returns:
+        TraceData object if successful, None otherwise
+    """
+    trace_id = _extract_trace_id(response, trace_header)
+    if not trace_id:
+        return None
+
+    data: Dict[str, List[str]] = {}
+    try:
+        for mcp in mcps:
+            entity_urn = getattr(mcp, "entityUrn", None)
+            aspect_name = getattr(mcp, "aspectName", None)
+
+            if not entity_urn or not aspect_name:
+                logger.debug(f"Skipping MCP with missing URN or aspect name: {mcp}")
+                continue
+
+            if aspects_to_trace is not None and aspect_name not in aspects_to_trace:
+                continue
+
+            if entity_urn not in data:
+                data[entity_urn] = []
+
+            data[entity_urn].append(aspect_name)
+
+        return TraceData(trace_id=trace_id, data=data)
+
+    except AttributeError as e:
+        logger.error(f"Error processing MCPs: {e}")
+        return None

datahub/emitter/rest_emitter.py

@@ -4,8 +4,10 @@ import functools
 import json
 import logging
 import os
+import time
 from collections import defaultdict
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 from enum import auto
 from json.decoder import JSONDecodeError
 from typing import (
@@ -35,11 +37,18 @@ from datahub.configuration.common import (
     ConfigModel,
     ConfigurationError,
     OperationalError,
+    TraceTimeoutError,
+    TraceValidationError,
 )
 from datahub.emitter.aspect import JSON_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.response_helper import (
+    TraceData,
+    extract_trace_data,
+    extract_trace_data_from_mcps,
+)
 from datahub.emitter.serialization_helper import pre_json_transform
 from datahub.ingestion.api.closeable import Closeable
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
@@ -69,6 +78,11 @@ _DEFAULT_RETRY_MAX_TIMES = int(
 
 _DATAHUB_EMITTER_TRACE = get_boolean_env_variable("DATAHUB_EMITTER_TRACE", False)
 
+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
 # for overhead like request headers.
 # This applies to pretty much all calls to GMS.
@@ -83,6 +97,11 @@ BATCH_INGEST_MAX_PAYLOAD_LENGTH = int(
 )
 
 
+class RestTraceMode(ConfigEnum):
+    ENABLED = auto()
+    DISABLED = auto()
+
+
 class RestSinkEndpoint(ConfigEnum):
     RESTLI = auto()
     OPENAPI = auto()
@@ -94,6 +113,13 @@ DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
 )
 
 
+# Supported with v1.0
+DEFAULT_REST_TRACE_MODE = pydantic.parse_obj_as(
+    RestTraceMode,
+    os.getenv("DATAHUB_REST_TRACE_MODE", RestTraceMode.DISABLED),
+)
+
+
 class RequestsSessionConfig(ConfigModel):
     timeout: Union[float, Tuple[float, float], None] = _DEFAULT_TIMEOUT_SEC
 
@@ -185,6 +211,7 @@ class DataHubRestEmitter(Closeable, Emitter):
     _token: Optional[str]
     _session: requests.Session
     _openapi_ingestion: bool
+    _default_trace_mode: bool
 
     def __init__(
         self,
@@ -201,6 +228,7 @@ class DataHubRestEmitter(Closeable, Emitter):
         client_certificate_path: Optional[str] = None,
         disable_ssl_verification: bool = False,
         openapi_ingestion: bool = False,
+        default_trace_mode: bool = False,
     ):
         if not gms_server:
             raise ConfigurationError("gms server is required")
@@ -214,12 +242,16 @@ class DataHubRestEmitter(Closeable, Emitter):
         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."
         )
 
+        if self._default_trace_mode:
+            logger.debug("Using API Tracing for ingestion.")
+
         headers = {
             "X-RestLi-Protocol-Version": "2.0.0",
             "X-DataHub-Py-Cli-Version": nice_version_name(),
@@ -396,13 +428,21 @@ class DataHubRestEmitter(Closeable, Emitter):
         self,
         mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
         async_flag: Optional[bool] = None,
+        trace_flag: Optional[bool] = None,
+        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> None:
         ensure_has_system_metadata(mcp)
 
+        trace_data = None
+
         if self._openapi_ingestion:
             request = self._to_openapi_request(mcp, async_flag, async_default=False)
             if request:
-                self._emit_generic(request[0], payload=request[1])
+                response = self._emit_generic(request[0], payload=request[1])
+
+                if self._should_trace(async_flag, trace_flag):
+                    trace_data = extract_trace_data(response) if response else None
+
         else:
             url = f"{self._gms_server}/aspects?action=ingestProposal"
 
@@ -414,12 +454,25 @@ class DataHubRestEmitter(Closeable, Emitter):
 
             payload = json.dumps(payload_dict)
 
-            self._emit_generic(url, payload)
+            response = self._emit_generic(url, payload)
+
+            if self._should_trace(async_flag, trace_flag):
+                trace_data = (
+                    extract_trace_data_from_mcps(response, [mcp]) if response else None
+                )
+
+        if trace_data:
+            self._await_status(
+                [trace_data],
+                trace_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),
     ) -> int:
         if _DATAHUB_EMITTER_TRACE:
             logger.debug(f"Attempting to emit MCP batch of size {len(mcps)}")
@@ -428,7 +481,7 @@ class DataHubRestEmitter(Closeable, Emitter):
             ensure_has_system_metadata(mcp)
 
         if self._openapi_ingestion:
-            return self._emit_openapi_mcps(mcps, async_flag)
+            return self._emit_openapi_mcps(mcps, async_flag, trace_flag, trace_timeout)
         else:
             return self._emit_restli_mcps(mcps, async_flag)
 
@@ -436,6 +489,8 @@ class DataHubRestEmitter(Closeable, Emitter):
         self,
         mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
         async_flag: Optional[bool] = None,
+        trace_flag: Optional[bool] = None,
+        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
     ) -> int:
         """
         1. Grouping MCPs by their entity URL
@@ -483,6 +538,16 @@ class DataHubRestEmitter(Closeable, Emitter):
                 response = self._emit_generic(url, payload=_Chunk.join(chunk))
                 responses.append(response)
 
+        if self._should_trace(async_flag, trace_flag, async_default=True):
+            trace_data = []
+            for response in responses:
+                data = extract_trace_data(response) if response else None
+                if data is not None:
+                    trace_data.append(data)
+
+            if trace_data:
+                self._await_status(trace_data, trace_timeout)
+
         return len(responses)
 
     def _emit_restli_mcps(
@@ -593,6 +658,99 @@ class DataHubRestEmitter(Closeable, Emitter):
                 "Unable to emit metadata to DataHub GMS", {"message": str(e)}
             ) from e
 
+    def _await_status(
+        self,
+        trace_data: List[TraceData],
+        trace_timeout: Optional[timedelta] = timedelta(seconds=3600),
+    ) -> None:
+        """Verify the status of asynchronous write operations.
+        Args:
+            trace_data: List of trace data to verify
+            trace_timeout: Maximum time to wait for verification.
+        Raises:
+            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")
+
+        try:
+            if not trace_data:
+                logger.debug("No trace data to verify")
+                return
+
+            start_time = datetime.now()
+
+            for trace in trace_data:
+                current_backoff = TRACE_INITIAL_BACKOFF
+
+                while trace.data:
+                    if datetime.now() - start_time > trace_timeout:
+                        raise TraceTimeoutError(
+                            f"Timeout waiting for async write completion after {trace_timeout.total_seconds()} seconds"
+                        )
+
+                    base_url = f"{self._gms_server}/openapi/v1/trace/write"
+                    url = f"{base_url}/{trace.trace_id}?onlyIncludeErrors=false&detailed=true"
+
+                    response = self._emit_generic(url, payload=trace.data)
+                    json_data = response.json()
+
+                    for urn, aspects in json_data.items():
+                        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"Persistence failure for URN '{urn}' aspect '{aspect_name}'. "
+                                    f"Status: {aspect_status}"
+                                )
+                                raise TraceValidationError(error_msg, aspect_status)
+
+                            primary_storage = aspect_status["primaryStorage"][
+                                "writeStatus"
+                            ]
+                            search_storage = aspect_status["searchStorage"][
+                                "writeStatus"
+                            ]
+
+                            # Remove resolved statuses
+                            if (
+                                primary_storage != TRACE_PENDING_STATUS
+                                and search_storage != TRACE_PENDING_STATUS
+                            ):
+                                trace.data[urn].remove(aspect_name)
+
+                        # Remove urns with all statuses resolved
+                        if not trace.data[urn]:
+                            trace.data.pop(urn)
+
+                    # Adjust backoff based on response
+                    if trace.data:
+                        # If we still have pending items, increase backoff
+                        current_backoff = min(
+                            current_backoff * TRACE_BACKOFF_FACTOR, TRACE_MAX_BACKOFF
+                        )
+                        logger.debug(
+                            f"Waiting {current_backoff} seconds before next check"
+                        )
+                        time.sleep(current_backoff)
+
+        except Exception as e:
+            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 __repr__(self) -> str:
         token_str = (
             f" with token: {self._token[:4]}**********{self._token[-4:]}"

datahub/ingestion/graph/client.py

@@ -34,8 +34,10 @@ from datahub.emitter.mce_builder import DEFAULT_ENV, Aspect
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.rest_emitter import (
     DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_TRACE_MODE,
     DatahubRestEmitter,
     RestSinkEndpoint,
+    RestTraceMode,
 )
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.ingestion.graph.config import (
@@ -146,6 +148,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             client_certificate_path=self.config.client_certificate_path,
             disable_ssl_verification=self.config.disable_ssl_verification,
             openapi_ingestion=DEFAULT_REST_SINK_ENDPOINT == RestSinkEndpoint.OPENAPI,
+            default_trace_mode=DEFAULT_REST_TRACE_MODE == RestTraceMode.ENABLED,
         )
 
         self.server_id = _MISSING_SERVER_ID

datahub/ingestion/sink/datahub_rest.py

@@ -21,8 +21,10 @@ from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
     DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_TRACE_MODE,
     DataHubRestEmitter,
     RestSinkEndpoint,
+    RestTraceMode,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -69,6 +71,7 @@ _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
     endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
+    default_trace_mode: RestTraceMode = DEFAULT_REST_TRACE_MODE
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -176,6 +179,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             client_certificate_path=config.client_certificate_path,
             disable_ssl_verification=config.disable_ssl_verification,
             openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
+            default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
         )
 
     @property