feat(platform): add Data Fabric error classification and extractor

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.89"
+version = "0.1.90"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath-platform/src/uipath/platform/entities/_entities_service.py

@@ -24,6 +24,7 @@
 from ..common._bindings import _resource_overwrites
 from ..common._config import UiPathApiConfig
 from ..common._execution_context import UiPathExecutionContext
+from ..errors._datafabric_error import attach_datafabric_error_mapping
 from ..orchestrator._folder_service import FolderService
 from ._entity_data_service import EntityDataService, FileContent
 from ._entity_resolution import (
@@ -1788,6 +1789,7 @@ async def retrieve_records_async(
             limit=limit,
         )
 
+    @attach_datafabric_error_mapping("query_entity_records")
     @traced(name="entity_query_records", run_type="uipath")
     def query_entity_records(self, sql_query: str) -> List[Dict[str, Any]]:
         """Query entity records using a validated SQL query.
@@ -1812,6 +1814,7 @@ def query_entity_records(self, sql_query: str) -> List[Dict[str, Any]]:
         """
         return self._data.query_entity_records(sql_query)
 
+    @attach_datafabric_error_mapping("query_entity_records_async")
     @traced(name="entity_query_records", run_type="uipath")
     async def query_entity_records_async(self, sql_query: str) -> List[Dict[str, Any]]:
         """Asynchronously query entity records using a validated SQL query.

packages/uipath-platform/src/uipath/platform/errors/__init__.py

@@ -22,16 +22,20 @@
 from ._context_grounding_index_not_found_exception import (
     ContextGroundingIndexNotFoundError,
 )
+from ._datafabric_error import DataFabricError
 from ._enriched_exception import EnrichedException, ExtractedErrorInfo
 from ._folder_not_found_exception import FolderNotFoundException
 from ._ingestion_in_progress_exception import IngestionInProgressException
 from ._operation_failed_exception import OperationFailedException
 from ._operation_not_complete_exception import OperationNotCompleteException
 from ._secret_missing_error import SecretMissingError
 from ._unsupported_data_source_exception import UnsupportedDataSourceException
+from .datafabric_error_codes import DataFabricErrorCategory
 
 __all__ = [
     "BaseUrlMissingError",
+    "DataFabricError",
+    "DataFabricErrorCategory",
     "BatchTransformFailedException",
     "BatchTransformNotCompleteException",
     "ContextGroundingIndexNotFoundError",

packages/uipath-platform/src/uipath/platform/errors/_datafabric_error.py

@@ -0,0 +1,109 @@
+"""Data Fabric query engine error classification.
+
+Maps error codes from the DF query engine invoking the "query_execute" endpoint to actionable
+categories so that callers (e.g. the agent SQL sub-graph) can decide
+whether to retry, ask the LLM to fix the SQL, or surface an infra error.
+
+The server error response JSON has the shape:
+    {"error": "<message>", "code": "<ERROR_CODE>", "traceId": "<uuid>"}
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from ._extractors._helpers import extract_service_prefix
+from .datafabric_error_codes import (
+    _QUERY_ENTITY_RECORDS_ERROR_CODES,
+    DataFabricErrorCategory,
+    classify_error_code,
+)
+
+if TYPE_CHECKING:
+    from ._enriched_exception import EnrichedException
+
+
+TCallable = TypeVar("TCallable", bound=Callable[..., Any])
+
+
+def attach_datafabric_error_mapping(
+    method_name: str,
+) -> Callable[[TCallable], TCallable]:
+    """Attach Data Fabric error metadata to a query method."""
+
+    def decorator(func: TCallable) -> TCallable:
+        func.__uipath_datafabric_method__ = method_name  # type: ignore[attr-defined]
+        func.__uipath_datafabric_error_codes__ = (  # type: ignore[attr-defined]
+            _QUERY_ENTITY_RECORDS_ERROR_CODES
+        )
+        return func
+
+    return decorator
+
+
+@dataclass(frozen=True)
+class DataFabricError:
+    """Structured error parsed from a DF query engine response."""
+
+    code: str | None
+    message: str | None
+    trace_id: str | None
+    category: DataFabricErrorCategory
+
+    @property
+    def is_retryable(self) -> bool:
+        return self.category == DataFabricErrorCategory.RETRYABLE
+
+    @property
+    def is_bad_sql(self) -> bool:
+        return self.category == DataFabricErrorCategory.BAD_SQL
+
+    @staticmethod
+    def from_enriched_exception(exc: EnrichedException) -> DataFabricError | None:
+        """Extract a DataFabricError from an EnrichedException, if applicable.
+
+        Returns None if the exception is not from a Data Fabric endpoint.
+        """
+        if extract_service_prefix(exc.url) != "datafabric_":
+            return None
+
+        info = exc.error_info
+        code = info.error_code if info else None
+        message = info.message if info else None
+        trace_id = info.trace_id if info else None
+
+        return DataFabricError(
+            code=code,
+            message=message,
+            trace_id=trace_id,
+            category=classify_error_code(code),
+        )
+
+    @staticmethod
+    def from_response_body(body: dict[str, Any]) -> DataFabricError:
+        """Parse a DataFabricError directly from a response body dict."""
+        raw_code = body.get("code")
+        code = (
+            str(raw_code)
+            if raw_code is not None and not isinstance(raw_code, (dict, list))
+            else None
+        )
+        message = body.get("error")
+        if not isinstance(message, str):
+            message = (
+                body.get("message") if isinstance(body.get("message"), str) else None
+            )
+        trace_id = body.get("traceId")
+        if not isinstance(trace_id, str):
+            trace_id = (
+                body.get("requestId")
+                if isinstance(body.get("requestId"), str)
+                else None
+            )
+        return DataFabricError(
+            code=code,
+            message=message,
+            trace_id=trace_id,
+            category=classify_error_code(code),
+        )

packages/uipath-platform/src/uipath/platform/errors/_extractors/_datafabric.py

@@ -0,0 +1,21 @@
+"""Data Fabric query engine error payload extractor.
+
+DF returns: {"error": "<message>", "code": "<ERROR_CODE>", "traceId": "<uuid>"}
+"""
+
+from typing import Any
+
+from .._enriched_exception import ExtractedErrorInfo
+from ._helpers import get_str_field, get_typed_field
+
+
+def extract_datafabric(body: dict[str, Any]) -> ExtractedErrorInfo:
+    message = get_typed_field(body, str, "error", "message")
+    error_code = get_str_field(body, "code", "errorCode")
+    trace_id = get_typed_field(body, str, "traceId", "requestId")
+
+    return ExtractedErrorInfo(
+        message=message,
+        error_code=error_code,
+        trace_id=trace_id,
+    )

packages/uipath-platform/src/uipath/platform/errors/_extractors/_router.py

@@ -14,6 +14,7 @@
 from .._enriched_exception import ExtractedErrorInfo
 from ._agenthub import extract_agenthub
 from ._apps import extract_apps
+from ._datafabric import extract_datafabric
 from ._elements import extract_elements
 from ._generic import extract_generic
 from ._helpers import extract_service_prefix, is_llm_path
@@ -29,6 +30,7 @@
     "agenthub_": extract_agenthub,
     "agentsruntime_": extract_agenthub,
     "apps_": extract_apps,
+    "datafabric_": extract_datafabric,
     "elements_": extract_elements,
     "llmopstenant_": extract_llmops,
 }

packages/uipath-platform/src/uipath/platform/errors/datafabric_error_codes.py

@@ -0,0 +1,69 @@
+"""Data Fabric query-engine error code constants."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+_RETRYABLE_CODES: frozenset[str] = frozenset(
+    {
+        "EXECUTION_TIMEOUT",
+        "SQLITE_BUSY",
+        "EXECUTION_INTERRUPTED",
+    }
+)
+
+_BAD_SQL_CODES: frozenset[str] = frozenset(
+    {
+        "SQL_PARSING",
+        "SQL_VALIDATION",
+    }
+)
+
+_INFRASTRUCTURE_CODES: frozenset[str] = frozenset(
+    {
+        "SQLITE_MEMORY_FULL",
+        "EPHEMERAL_STORAGE_ERROR",
+        "INTERNAL_ERROR",
+        "FQS_ERROR",
+    }
+)
+
+_DATA_ISSUE_CODES: frozenset[str] = frozenset(
+    {
+        "FRAGMENT_EXECUTION_FAILURE",
+        "CONTEXT_CREATION",
+        "UNKNOWN_ENTITY",
+        "EXECUTION_ERROR",
+        "RESULT_TOO_LARGE",
+    }
+)
+
+_QUERY_ENTITY_RECORDS_ERROR_CODES: frozenset[str] = frozenset(
+    {*_RETRYABLE_CODES, *_BAD_SQL_CODES, *_INFRASTRUCTURE_CODES, *_DATA_ISSUE_CODES}
+)
+
+
+class DataFabricErrorCategory(str, Enum):
+    """Actionable error category for Data Fabric query failures."""
+
+    RETRYABLE = "retryable"
+    BAD_SQL = "bad_sql"
+    INFRASTRUCTURE = "infrastructure"
+    DATA_ISSUE = "data_issue"
+    UNKNOWN = "unknown"
+
+
+def classify_error_code(code: str | None) -> DataFabricErrorCategory:
+    """Classify a DF error code string into an actionable category."""
+    if not code:
+        return DataFabricErrorCategory.UNKNOWN
+    upper = code.upper()
+    if upper in _RETRYABLE_CODES:
+        return DataFabricErrorCategory.RETRYABLE
+    if upper in _BAD_SQL_CODES:
+        return DataFabricErrorCategory.BAD_SQL
+    if upper in _INFRASTRUCTURE_CODES:
+        return DataFabricErrorCategory.INFRASTRUCTURE
+    if upper in _DATA_ISSUE_CODES:
+        return DataFabricErrorCategory.DATA_ISSUE
+    return DataFabricErrorCategory.UNKNOWN