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

datahub/emitter/mcp.py

@@ -1,9 +1,11 @@
 import dataclasses
 import json
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, Union
+import warnings
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
 
 from datahub.emitter.aspect import ASPECT_MAP, JSON_CONTENT_TYPE
 from datahub.emitter.serialization_helper import post_json_transform, pre_json_transform
+from datahub.errors import DataHubDeprecationWarning
 from datahub.metadata.schema_classes import (
     ChangeTypeClass,
     DictWrapper,
@@ -69,18 +71,28 @@ class MetadataChangeProposalWrapper:
     aspectName: Union[None, str] = None
     aspect: Union[None, _Aspect] = None
     systemMetadata: Union[None, SystemMetadataClass] = None
+    headers: Union[None, Dict[str, str]] = None
 
     def __post_init__(self) -> None:
         if self.entityUrn and self.entityType == _ENTITY_TYPE_UNSET:
             self.entityType = guess_entity_type(self.entityUrn)
         elif self.entityUrn and self.entityType:
-            guessed_entity_type = guess_entity_type(self.entityUrn).lower()
-            # Entity type checking is actually case insensitive.
-            # Note that urns are case sensitive, but entity types are not.
-            if self.entityType.lower() != guessed_entity_type:
+            guessed_entity_type = guess_entity_type(self.entityUrn)
+            if self.entityType.lower() != guessed_entity_type.lower():
+                # If they aren't a case-ignored match, raise an error.
                 raise ValueError(
                     f"entityType {self.entityType} does not match the entity type {guessed_entity_type} from entityUrn {self.entityUrn}",
                 )
+            elif self.entityType != guessed_entity_type:
+                # If they only differ in case, normalize and print a warning.
+                self.entityType = guessed_entity_type
+                warnings.warn(
+                    f"The passed entityType {self.entityType} differs in case from the expected entity type {guessed_entity_type}. "
+                    "This will be automatically corrected for now, but will become an error in a future release. "
+                    "Note that the entityType field is optional and will be automatically inferred from the entityUrn.",
+                    DataHubDeprecationWarning,
+                    stacklevel=3,
+                )
         elif self.entityType == _ENTITY_TYPE_UNSET:
             raise ValueError("entityType must be set if entityUrn is not set")
 
@@ -112,6 +124,7 @@ class MetadataChangeProposalWrapper:
             auditHeader=self.auditHeader,
             aspectName=self.aspectName,
             systemMetadata=self.systemMetadata,
+            headers=self.headers,
         )
 
     def make_mcp(self) -> MetadataChangeProposalClass:
@@ -211,6 +224,7 @@ class MetadataChangeProposalWrapper:
                 aspectName=mcpc.aspectName,
                 aspect=aspect,
                 systemMetadata=mcpc.systemMetadata,
+                headers=mcpc.headers,
             )
         else:
             return None
@@ -228,6 +242,7 @@ class MetadataChangeProposalWrapper:
             changeType=mcl.changeType,
             auditHeader=mcl.auditHeader,
             systemMetadata=mcl.systemMetadata,
+            headers=mcl.headers,
         )
         return cls.try_from_mcpc(mcpc) or mcpc
 

datahub/emitter/mcp_builder.py

@@ -117,6 +117,14 @@ class ContainerKey(DatahubKey):
 PlatformKey = ContainerKey
 
 
+class NamespaceKey(ContainerKey):
+    """
+    For Iceberg namespaces (databases/schemas)
+    """
+
+    namespace: str
+
+
 class DatabaseKey(ContainerKey):
     database: str
 
@@ -129,6 +137,10 @@ class ProjectIdKey(ContainerKey):
     project_id: str
 
 
+class ExperimentKey(ContainerKey):
+    id: str
+
+
 class MetastoreKey(ContainerKey):
     metastore: str
 

datahub/emitter/request_helper.py

@@ -1,28 +1,151 @@
-import itertools
+import json
 import shlex
-from typing import List, Union
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
 
 import requests
+from requests.auth import HTTPBasicAuth
+
+from datahub.emitter.aspect import JSON_CONTENT_TYPE, JSON_PATCH_CONTENT_TYPE
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.serialization_helper import pre_json_transform
+from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
+    MetadataChangeProposal,
+)
+from datahub.metadata.schema_classes import ChangeTypeClass
+
+
+def _decode_bytes(value: Union[str, bytes]) -> str:
+    """Decode bytes to string, if necessary."""
+    if isinstance(value, bytes):
+        return value.decode()
+    return value
 
 
 def _format_header(name: str, value: Union[str, bytes]) -> str:
     if name == "Authorization":
         return f"{name!s}: <redacted>"
-    return f"{name!s}: {value!s}"
+    return f"{name!s}: {_decode_bytes(value)}"
 
 
 def make_curl_command(
-    session: requests.Session, method: str, url: str, payload: str
+    session: requests.Session, method: str, url: str, payload: Optional[str] = None
 ) -> str:
-    fragments: List[str] = [
-        "curl",
-        *itertools.chain(
-            *[
-                ("-X", method),
-                *[("-H", _format_header(k, v)) for (k, v) in session.headers.items()],
-                ("--data", payload),
-            ]
-        ),
-        url,
-    ]
+    fragments: List[str] = ["curl", "-X", method]
+
+    for header_name, header_value in session.headers.items():
+        fragments.extend(["-H", _format_header(header_name, header_value)])
+
+    if session.auth:
+        if isinstance(session.auth, HTTPBasicAuth):
+            fragments.extend(
+                ["-u", f"{_decode_bytes(session.auth.username)}:<redacted>"]
+            )
+        else:
+            # For other auth types, they should be handled via headers
+            fragments.extend(["-H", "<unknown auth type>"])
+
+    if payload:
+        fragments.extend(["--data", payload])
+
+    fragments.append(url)
     return shlex.join(fragments)
+
+
+@dataclass
+class OpenApiRequest:
+    """Represents an OpenAPI request for entity operations."""
+
+    method: str
+    url: str
+    payload: List[Dict[str, Any]]
+
+    @classmethod
+    def from_mcp(
+        cls,
+        mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
+        gms_server: str,
+        async_flag: bool = False,
+        search_sync_flag: bool = False,
+    ) -> Optional["OpenApiRequest"]:
+        """Factory method to create an OpenApiRequest from a MetadataChangeProposal."""
+        if not mcp.aspectName or (
+            mcp.changeType != ChangeTypeClass.DELETE and not mcp.aspect
+        ):
+            return None
+
+        method = "post"
+        url = f"{gms_server}/openapi/v3/entity/{mcp.entityType}?async={'true' if async_flag else 'false'}"
+        payload = []
+
+        if mcp.changeType == ChangeTypeClass.DELETE:
+            method = "delete"
+            url = f"{gms_server}/openapi/v3/entity/{mcp.entityType}/{mcp.entityUrn}"
+        else:
+            if mcp.aspect:
+                mcp_headers = {}
+
+                if not async_flag and search_sync_flag:
+                    mcp_headers["X-DataHub-Sync-Index-Update"] = "true"
+
+                if mcp.changeType == ChangeTypeClass.PATCH:
+                    method = "patch"
+                    obj = mcp.aspect.to_obj()
+                    content_type = obj.get("contentType")
+                    if obj.get("value") and content_type == JSON_PATCH_CONTENT_TYPE:
+                        # Undo double serialization.
+                        obj = json.loads(obj["value"])
+                        patch_value = obj
+                    else:
+                        raise NotImplementedError(
+                            f"ChangeType {mcp.changeType} only supports context type {JSON_PATCH_CONTENT_TYPE}, found {content_type}."
+                        )
+
+                    if isinstance(patch_value, list):
+                        patch_value = {"patch": patch_value}
+
+                    payload = [
+                        {
+                            "urn": mcp.entityUrn,
+                            mcp.aspectName: {
+                                "value": patch_value,
+                                "systemMetadata": mcp.systemMetadata.to_obj()
+                                if mcp.systemMetadata
+                                else None,
+                                "headers": mcp_headers,
+                            },
+                        }
+                    ]
+                else:
+                    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(
+                                f"ChangeType {mcp.changeType} does not support patch."
+                            )
+                        aspect_value = pre_json_transform(obj)
+
+                    payload = [
+                        {
+                            "urn": mcp.entityUrn,
+                            mcp.aspectName: {
+                                "value": aspect_value,
+                                "systemMetadata": mcp.systemMetadata.to_obj()
+                                if mcp.systemMetadata
+                                else None,
+                                "headers": mcp_headers,
+                            },
+                        }
+                    ]
+            else:
+                raise ValueError(f"ChangeType {mcp.changeType} requires a value.")
+
+        return cls(method=method, url=url, payload=payload)

datahub/emitter/response_helper.py

@@ -1,38 +1,124 @@
 import json
 import logging
+import re
+import warnings
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Dict, List, Optional, Sequence, Union
 
 from requests import Response
 
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.errors import APITracingWarning
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeProposal,
 )
 
 logger = logging.getLogger(__name__)
 
+_TRACE_HEADER_NAME = "traceparent"
+
 
 @dataclass
 class TraceData:
     trace_id: str
     data: Dict[str, List[str]]
 
+    @staticmethod
+    def extract_trace_id(input_str: Optional[str]) -> Optional[str]:
+        """
+        Extract the trace ID from various input formats.
+
+        Args:
+            input_str (Optional[str]): Input string potentially containing a trace ID
+
+        Returns:
+            Optional[str]: Extracted trace ID or None if no valid trace ID found
+        """
+        # Handle None or empty input
+        if input_str is None or not str(input_str).strip():
+            return None
+
+        # Convert to string and clean
+        input_str = str(input_str).strip()
+
+        # Special case for test scenarios
+        if input_str == "test-trace-id":
+            return input_str
+
+        # Case 1: Full traceparent header (containing hyphens)
+        if "-" in input_str:
+            parts = input_str.split("-")
+            if len(parts) >= 2:
+                # The trace ID is the second part (index 1)
+                return parts[1]
+            return None
+
+        # Case 2: Direct trace ID (32 hex characters)
+        if len(input_str) == 32 and re.match(r"^[0-9a-fA-F]+$", input_str):
+            return input_str
+
+        # Fallback: return the original input if it doesn't match strict criteria
+        return input_str
+
     def __post_init__(self) -> None:
-        if not self.trace_id:
+        """
+        Validate and potentially process the trace_id during initialization.
+        """
+        # Explicitly check for None or empty string
+        if self.trace_id is None or self.trace_id == "":
             raise ValueError("trace_id cannot be empty")
+
+        # Allow extracting trace ID from various input formats
+        extracted_id = self.extract_trace_id(self.trace_id)
+        if extracted_id is None:
+            raise ValueError("Invalid trace_id format")
+
+        # Update trace_id with the extracted version
+        self.trace_id = extracted_id
+
+        # Validate data
         if not isinstance(self.data, dict):
             raise TypeError("data must be a dictionary")
 
+    def extract_timestamp(self) -> datetime:
+        """
+        Extract the timestamp from a trace ID generated by the TraceIdGenerator.
+
+        Returns:
+            datetime: The timestamp in UTC
 
-def _extract_trace_id(
-    response: Response, trace_header: str = "traceparent"
-) -> Optional[str]:
+        Raises:
+            ValueError: If the trace ID is invalid
+        """
+        # Special case for test trace ID
+        if self.trace_id == "test-trace-id":
+            return datetime.fromtimestamp(0, tz=timezone.utc)
+
+        # Validate trace ID length for hex-based trace IDs
+        if len(self.trace_id) < 16 or not re.match(
+            r"^[0-9a-fA-F]+$", self.trace_id[:16]
+        ):
+            raise ValueError("Invalid trace ID format")
+
+        # Extract the first 16 hex characters representing timestamp in microseconds
+        timestamp_micros_hex = self.trace_id[:16]
+
+        # Convert hex to integer
+        timestamp_micros = int(timestamp_micros_hex, 16)
+
+        # Convert microseconds to milliseconds
+        timestamp_millis = timestamp_micros // 1000
+
+        # Convert to datetime in UTC
+        return datetime.fromtimestamp(timestamp_millis / 1000, tz=timezone.utc)
+
+
+def _extract_trace_id(response: Response) -> 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
     """
@@ -40,9 +126,17 @@ def _extract_trace_id(
         logger.debug(f"Invalid status code: {response.status_code}")
         return None
 
-    trace_id = response.headers.get(trace_header)
+    trace_id = response.headers.get(_TRACE_HEADER_NAME)
     if not trace_id:
-        logger.debug(f"Missing trace header: {trace_header}")
+        # This will only be printed if
+        # 1. we're in async mode (checked by the caller)
+        # 2. the server did not return a trace ID
+        logger.debug(f"Missing trace header: {_TRACE_HEADER_NAME}")
+        warnings.warn(
+            "No trace ID found in response headers. API tracing is not active - likely due to an outdated server version.",
+            APITracingWarning,
+            stacklevel=3,
+        )
         return None
 
     return trace_id
@@ -51,20 +145,19 @@ def _extract_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.
+    """Extract trace data from a response object.
+
+    If we run into a JSONDecodeError, we'll log an error and return None.
+
     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)
+    trace_id = _extract_trace_id(response)
     if not trace_id:
         return None
 
@@ -104,19 +197,18 @@ 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.
+    """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)
+    trace_id = _extract_trace_id(response)
     if not trace_id:
         return None