feat: Enhance API responses with trace ID

Author: glondhe-mdsol

dataconnect/client.py

@@ -34,6 +34,11 @@ def __init__(self, service: DataConnectService) -> None:
         """Initialize the client with an injected service implementation."""
         self._service = service
 
+    @property
+    def trace_id(self) -> str | None:
+        """Trace ID from the most recent sequential request, when available."""
+        return getattr(self._service, "trace_id", None)
+
     @classmethod
     def connect(
         cls,
@@ -67,7 +72,7 @@ def fetch_data(
         dataset_uuid: UUID,
         first_n_rows: int | None = None,
     ) -> pd.DataFrame:
-        """Fetch data frames for a given dataset UUID."""
+        """Fetch data for a given dataset UUID."""
         return self._service.fetch_data(dataset_uuid, first_n_rows)
 
     def get_datasets(
@@ -121,8 +126,8 @@ def dry_publish(
 
         Returns:
             A :class:`DryPublishResult` containing the server's validation
-            outcome, including per-field validity flags, error messages, and
-            an optional ``invalid_records`` DataFrame.
+            outcome, including per-field validity flags, error messages, an
+                optional ``invalid_records`` DataFrame.
         """
         return self._service.dry_publish(
             project_token=project_token,

dataconnect/exceptions.py

@@ -48,6 +48,7 @@ class DataConnectError(Exception):
     message: str
     timestamp: str | None = None
     details: list[ErrorDetail] | None = None
+    trace_id: str | None = None
 
     def __str__(self) -> str:
         lines = [
@@ -63,6 +64,10 @@ def __str__(self) -> str:
             for detail in self.details:
                 lines.append(str(detail))
 
+        if self.trace_id is not None:
+            indent = "    " if self.details else ""
+            lines.append(f"{indent}Trace ID: {self.trace_id}")
+
         return "\n".join(lines)
 
 

dataconnect/service/default.py

@@ -74,6 +74,10 @@ class DefaultDataConnectService(DataConnectService):
     def __init__(self, transport: Transport) -> None:
         self._transport = transport
 
+    @property
+    def trace_id(self) -> str | None:
+        return getattr(self._transport, "trace_id", None)
+
     # DataConnectService
 
     def get_studies(self, search_study_name: str | None = None) -> StudiesResult:
@@ -92,7 +96,8 @@ def get_studies(self, search_study_name: str | None = None) -> StudiesResult:
         request = request.append_body({"search_study_name": search_study_name})
 
         try:
-            resources = self._transport.list_resources(request)
+            result = self._transport.list_resources(request)
+            resources = result
             total_records = resources[0].total_records if resources else 0
             studies = [resource_to_study(r) for r in resources]
             return StudiesResult(total_records=total_records, studies=studies)
@@ -106,7 +111,7 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]:
             dataset_uuid: UUID of the dataset whose versions are requested.
 
         Returns:
-            A list of :class:`DatasetVersion` objects for the given dataset.
+            The dataset versions, newest first.
 
         Raises:
             ValidationError: If *dataset_uuid* is not a valid UUID (upstream).
@@ -115,14 +120,15 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]:
         request = ResourceQuery(action=_ACTION_LIST_DATASET_VERSIONS).append_body({"dataset_uuid": str(dataset_uuid)})
 
         try:
-            resources = self._transport.list_resources(request)
+            result = self._transport.list_resources(request)
 
             # Return Sorted dataset versions in descending order (newest first) based on the dataset_version field.
-            return sorted(
-                (resource_to_dataset_version(r) for r in resources),
+            versions = sorted(
+                (resource_to_dataset_version(r) for r in result),
                 key=lambda dv: dv.dataset_version,
                 reverse=True,
             )
+            return versions
         except Exception as ex:
             raise translate_error(ex) from ex
 
@@ -169,7 +175,8 @@ def get_datasets(
         )
 
         try:
-            resources = self._transport.list_resources(request)
+            result = self._transport.list_resources(request)
+            resources = result
             items = []
             for resource in resources:
                 dataset = resource_to_dataset(resource)
@@ -347,7 +354,7 @@ def get_datetime_formats(
         request = DatetimeFormatsRequest(project_token=project_token, format_type=format_type)
 
         try:
-            raw_formats = self._transport.get_datetime_formats(request)
+            formats_response = self._transport.get_datetime_formats(request)
         except TransportError as ex:
             raise translate_error(ex) from ex
 
@@ -356,7 +363,7 @@ def get_datetime_formats(
                 format=fmt,
                 type="datetime" if "HH:mm" in fmt else "date",
             )
-            for fmt in raw_formats
+            for fmt in formats_response
         ]
 
         return DatetimeFormatsResult(formats=formats)

dataconnect/service/error_handler.py

@@ -43,26 +43,54 @@ def translate_error(ex: Exception) -> DataConnectError:
 
     if isinstance(ex, TransportAuthenticationError):
         return AuthenticationError(
-            error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details
+            error_code=ex.error_code,
+            message=ex.message,
+            timestamp=ex.timestamp,
+            details=error_details,
+            trace_id=ex.trace_id,
         )
 
     if isinstance(ex, TransportAuthorizationError):
         return AuthorizationError(
-            error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details
+            error_code=ex.error_code,
+            message=ex.message,
+            timestamp=ex.timestamp,
+            details=error_details,
+            trace_id=ex.trace_id,
         )
 
     if isinstance(ex, TransportValidationError):
         return ValidationError(
-            error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details
+            error_code=ex.error_code,
+            message=ex.message,
+            timestamp=ex.timestamp,
+            details=error_details,
+            trace_id=ex.trace_id,
         )
 
     if isinstance(ex, TransportNotFoundError):
         return NotFoundError(
-            error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details
+            error_code=ex.error_code,
+            message=ex.message,
+            timestamp=ex.timestamp,
+            details=error_details,
+            trace_id=ex.trace_id,
         )
 
     if isinstance(ex, TransportServerError):
-        return ServerError(error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details)
+        return ServerError(
+            error_code=ex.error_code,
+            message=ex.message,
+            timestamp=ex.timestamp,
+            details=error_details,
+            trace_id=ex.trace_id,
+        )
 
     # Non-specific transport error
-    return DataConnectError(error_code=ex.error_code, message=ex.message, timestamp=ex.timestamp, details=error_details)
+    return DataConnectError(
+        error_code=ex.error_code,
+        message=ex.message,
+        timestamp=ex.timestamp,
+        details=error_details,
+        trace_id=ex.trace_id,
+    )

dataconnect/transport/arrow_flight/error_handler.py

@@ -161,10 +161,30 @@ def _normalize_enodia_error(error_message: str) -> str:
         return error_message
 
 
+def extract_error_payload(error_message: str) -> dict | None:
+    """Find a structured error payload in direct or gRPC-wrapped Flight text."""
+    normalized_message = _normalize_enodia_error(error_message)
+    messages = (error_message,) if normalized_message == error_message else (error_message, normalized_message)
+
+    for message in messages:
+        for delimiter in re.finditer(r"::", message):
+            payload_text = message[delimiter.end() :]
+            for brace in re.finditer(r"\{", payload_text):
+                json_part = _extract_json_object(payload_text[brace.start() :])
+                try:
+                    payload = json.loads(json_part)
+                except json.JSONDecodeError:
+                    continue
+                if isinstance(payload, dict):
+                    return payload
+
+    return None
+
+
 _UNKNOWN_ERROR = "Unknown error"
 
 
-def parse_dataconnect_error(ex: Exception) -> TransportError:
+def parse_dataconnect_error(ex: Exception, trace_id: str | None = None) -> TransportError:
     """Parse a raw exception into a typed ``TransportError``.
 
     Extracts a structured ``PREFIX::JSON`` payload from the exception message,
@@ -175,85 +195,88 @@ def parse_dataconnect_error(ex: Exception) -> TransportError:
     the message cannot be parsed or does not match the expected format.
     """
     try:
+        if isinstance(ex, TransportError):
+            if ex.trace_id is None:
+                ex.trace_id = trace_id
+            return ex
+
         error_message = str(ex)
 
-        # Normalize enodia authentication errors into PREFIX::JSON format
-        error_message = _normalize_enodia_error(error_message)
-
-        if "::" in error_message:
-            delimiter_pos = error_message.index("::")
-            json_part_raw = error_message[delimiter_pos + 2 :]
-
-            if json_part_raw:
-                json_part = _extract_json_object(json_part_raw)
-
-                error_data: dict = json.loads(json_part)
-
-                parsed_details: list[ErrorDetail] | None = None
-                raw_details = error_data.get("details")
-                if isinstance(raw_details, list):
-                    standard_keys = {"field", "message", "expected"}
-                    parsed_details = [
-                        ErrorDetail(
-                            field=item.get("field"),
-                            message=item.get("message"),
-                            expected=item.get("expected"),
-                            extra={k: v for k, v in item.items() if k not in standard_keys},
-                        )
-                        for item in raw_details
-                        if isinstance(item, dict)
-                    ]
-
-                error_code = error_data.get("error_code", "SDK_ERROR")
-
-                if error_code.startswith("AUTH_"):
-                    return TransportAuthenticationError(
-                        error_code=error_code,
-                        message=error_data.get("message") or _UNKNOWN_ERROR,
-                        timestamp=error_data.get("timestamp"),
-                        details=parsed_details,
+        error_data = extract_error_payload(error_message)
+        if error_data is not None:
+            parsed_details: list[ErrorDetail] | None = None
+            raw_details = error_data.get("details")
+            if isinstance(raw_details, list):
+                standard_keys = {"field", "message", "expected"}
+                parsed_details = [
+                    ErrorDetail(
+                        field=item.get("field"),
+                        message=item.get("message"),
+                        expected=item.get("expected"),
+                        extra={k: v for k, v in item.items() if k not in standard_keys},
                     )
+                    for item in raw_details
+                    if isinstance(item, dict)
+                ]
 
-                if error_code.startswith("AUTHZ_"):
-                    return TransportAuthorizationError(
-                        error_code=error_code,
-                        message=error_data.get("message") or _UNKNOWN_ERROR,
-                        timestamp=error_data.get("timestamp"),
-                        details=parsed_details,
-                    )
+            error_code = error_data.get("error_code", "SDK_ERROR")
+            payload_trace_id = error_data.get("trace_id")
+            parsed_trace_id = payload_trace_id if isinstance(payload_trace_id, str) and payload_trace_id else trace_id
 
-                if error_code.startswith("VAL_"):
-                    return TransportValidationError(
-                        error_code=error_code,
-                        message=error_data.get("message") or _UNKNOWN_ERROR,
-                        timestamp=error_data.get("timestamp"),
-                        details=parsed_details,
-                    )
+            if error_code.startswith("AUTH_"):
+                return TransportAuthenticationError(
+                    error_code=error_code,
+                    message=error_data.get("message") or _UNKNOWN_ERROR,
+                    timestamp=error_data.get("timestamp"),
+                    details=parsed_details,
+                    trace_id=parsed_trace_id,
+                )
 
-                if error_code.startswith("RES_"):
-                    return TransportNotFoundError(
-                        error_code=error_code,
-                        message=error_data.get("message") or _UNKNOWN_ERROR,
-                        timestamp=error_data.get("timestamp"),
-                        details=parsed_details,
-                    )
+            if error_code.startswith("AUTHZ_"):
+                return TransportAuthorizationError(
+                    error_code=error_code,
+                    message=error_data.get("message") or _UNKNOWN_ERROR,
+                    timestamp=error_data.get("timestamp"),
+                    details=parsed_details,
+                    trace_id=parsed_trace_id,
+                )
 
-                if error_code.startswith("INT_"):
-                    return TransportServerError(
-                        error_code=error_code,
-                        message=error_data.get("message") or _UNKNOWN_ERROR,
-                        timestamp=error_data.get("timestamp"),
-                        details=parsed_details,
-                    )
+            if error_code.startswith("VAL_"):
+                return TransportValidationError(
+                    error_code=error_code,
+                    message=error_data.get("message") or _UNKNOWN_ERROR,
+                    timestamp=error_data.get("timestamp"),
+                    details=parsed_details,
+                    trace_id=parsed_trace_id,
+                )
 
-                return TransportError(
+            if error_code.startswith("RES_"):
+                return TransportNotFoundError(
                     error_code=error_code,
                     message=error_data.get("message") or _UNKNOWN_ERROR,
                     timestamp=error_data.get("timestamp"),
                     details=parsed_details,
+                    trace_id=parsed_trace_id,
                 )
 
-        return TransportError(error_code="SDK_ERROR", message=error_message)
+            if error_code.startswith("INT_"):
+                return TransportServerError(
+                    error_code=error_code,
+                    message=error_data.get("message") or _UNKNOWN_ERROR,
+                    timestamp=error_data.get("timestamp"),
+                    details=parsed_details,
+                    trace_id=parsed_trace_id,
+                )
+
+            return TransportError(
+                error_code=error_code,
+                message=error_data.get("message") or _UNKNOWN_ERROR,
+                timestamp=error_data.get("timestamp"),
+                details=parsed_details,
+                trace_id=parsed_trace_id,
+            )
+
+        return TransportError(error_code="SDK_ERROR", message=error_message, trace_id=trace_id)
 
     except Exception as ex:
-        return TransportError(error_code="SDK_ERROR", message=str(ex))
+        return TransportError(error_code="SDK_ERROR", message=str(ex), trace_id=trace_id)