acryl-datahub 0.15.0.6rc2__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/configuration/common.py

@@ -20,7 +20,7 @@ from pydantic import BaseModel, Extra, ValidationError
 from pydantic.fields import Field
 from typing_extensions import Protocol, Self
 
-from datahub.configuration._config_enum import ConfigEnum as ConfigEnum  # noqa: I250
+from datahub.configuration._config_enum import ConfigEnum as ConfigEnum
 from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
 from datahub.utilities.dedup_list import deduplicate_list
 
@@ -130,7 +130,7 @@ class PermissiveConfigModel(ConfigModel):
     # It is usually used for argument bags that are passed through to third-party libraries.
 
     class Config:
-        if PYDANTIC_VERSION_2:
+        if PYDANTIC_VERSION_2:  # noqa: SIM108
             extra = "allow"
         else:
             extra = Extra.allow
@@ -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/configuration/git.py

@@ -43,9 +43,7 @@ class GitReference(ConfigModel):
 
     @validator("repo", pre=True)
     def simplify_repo_url(cls, repo: str) -> str:
-        if repo.startswith("github.com/"):
-            repo = f"https://{repo}"
-        elif repo.startswith("gitlab.com"):
+        if repo.startswith("github.com/") or repo.startswith("gitlab.com"):
             repo = f"https://{repo}"
         elif repo.count("/") == 1:
             repo = f"https://github.com/{repo}"

datahub/configuration/kafka.py

@@ -44,7 +44,7 @@ class KafkaConsumerConnectionConfig(_KafkaConnectionConfig):
             try:
                 value = CallableConsumerConfig(value).callable_config()
             except Exception as e:
-                raise ConfigurationError(e)
+                raise ConfigurationError(e) from e
         return value
 
 

datahub/emitter/mce_builder.py

@@ -52,7 +52,15 @@ from datahub.metadata.schema_classes import (
     UpstreamLineageClass,
     _Aspect as AspectAbstract,
 )
-from datahub.metadata.urns import DataFlowUrn, DatasetUrn, TagUrn
+from datahub.metadata.urns import (
+    ChartUrn,
+    DashboardUrn,
+    DataFlowUrn,
+    DataJobUrn,
+    DataPlatformUrn,
+    DatasetUrn,
+    TagUrn,
+)
 from datahub.utilities.urn_encoder import UrnEncoder
 
 logger = logging.getLogger(__name__)
@@ -119,7 +127,7 @@ def parse_ts_millis(ts: Optional[float]) -> Optional[datetime]:
 def make_data_platform_urn(platform: str) -> str:
     if platform.startswith("urn:li:dataPlatform:"):
         return platform
-    return f"urn:li:dataPlatform:{platform}"
+    return DataPlatformUrn.create_from_id(platform).urn()
 
 
 def make_dataset_urn(platform: str, name: str, env: str = DEFAULT_ENV) -> str:
@@ -236,7 +244,7 @@ def make_user_urn(username: str) -> str:
     Makes a user urn if the input is not a user or group urn already
     """
     return (
-        f"urn:li:corpuser:{username}"
+        f"urn:li:corpuser:{UrnEncoder.encode_string(username)}"
         if not username.startswith(("urn:li:corpuser:", "urn:li:corpGroup:"))
         else username
     )
@@ -249,7 +257,7 @@ def make_group_urn(groupname: str) -> str:
     if groupname and groupname.startswith(("urn:li:corpGroup:", "urn:li:corpuser:")):
         return groupname
     else:
-        return f"urn:li:corpGroup:{groupname}"
+        return f"urn:li:corpGroup:{UrnEncoder.encode_string(groupname)}"
 
 
 def make_tag_urn(tag: str) -> str:
@@ -301,7 +309,12 @@ def make_data_flow_urn(
 
 
 def make_data_job_urn_with_flow(flow_urn: str, job_id: str) -> str:
-    return f"urn:li:dataJob:({flow_urn},{job_id})"
+    data_flow_urn = DataFlowUrn.from_string(flow_urn)
+    data_job_urn = DataJobUrn.create_from_ids(
+        data_flow_urn=data_flow_urn.urn(),
+        job_id=job_id,
+    )
+    return data_job_urn.urn()
 
 
 def make_data_process_instance_urn(dataProcessInstanceId: str) -> str:
@@ -324,10 +337,11 @@ def make_dashboard_urn(
     platform: str, name: str, platform_instance: Optional[str] = None
 ) -> str:
     # FIXME: dashboards don't currently include data platform urn prefixes.
-    if platform_instance:
-        return f"urn:li:dashboard:({platform},{platform_instance}.{name})"
-    else:
-        return f"urn:li:dashboard:({platform},{name})"
+    return DashboardUrn.create_from_ids(
+        platform=platform,
+        name=name,
+        platform_instance=platform_instance,
+    ).urn()
 
 
 def dashboard_urn_to_key(dashboard_urn: str) -> Optional[DashboardKeyClass]:
@@ -342,10 +356,11 @@ def make_chart_urn(
     platform: str, name: str, platform_instance: Optional[str] = None
 ) -> str:
     # FIXME: charts don't currently include data platform urn prefixes.
-    if platform_instance:
-        return f"urn:li:chart:({platform},{platform_instance}.{name})"
-    else:
-        return f"urn:li:chart:({platform},{name})"
+    return ChartUrn.create_from_ids(
+        platform=platform,
+        name=name,
+        platform_instance=platform_instance,
+    ).urn()
 
 
 def chart_urn_to_key(chart_urn: str) -> Optional[ChartKeyClass]:

datahub/emitter/mcp_builder.py

@@ -36,7 +36,7 @@ from datahub.metadata.schema_classes import (
     SubTypesClass,
     TagAssociationClass,
 )
-from datahub.metadata.urns import StructuredPropertyUrn
+from datahub.metadata.urns import ContainerUrn, StructuredPropertyUrn
 
 # In https://github.com/datahub-project/datahub/pull/11214, we added a
 # new env field to container properties. However, populating this field
@@ -87,6 +87,9 @@ class ContainerKey(DatahubKey):
     def property_dict(self) -> Dict[str, str]:
         return self.dict(by_alias=True, exclude_none=True)
 
+    def as_urn_typed(self) -> ContainerUrn:
+        return ContainerUrn.from_string(self.as_urn())
+
     def as_urn(self) -> str:
         return make_container_urn(guid=self.guid())
 

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