feat(entities): choice-set awareness for query engine

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.2.27"
+version = "0.2.28"
 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

@@ -2319,7 +2319,11 @@ async def retrieve_records_v3_async(
     @attach_datafabric_error_mapping("query_entity_records")
     @traced(name="entity_query_records", run_type="uipath")
     def query_entity_records(
-        self, sql_query: str, *, relationships_as_scalar: bool = False
+        self,
+        sql_query: str,
+        *,
+        relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Query entity records using a validated SQL query.
 
@@ -2334,6 +2338,10 @@ def query_entity_records(
                 so a query can join on ``relationshipField = Other.Id``. Sent as
                 ``queryOptions.relationshipsAsScalar`` in the request body. Defaults to
                 ``False`` (unchanged behaviour).
+            resolve_choice_sets (bool, optional): When ``True``, choice-set fields
+                in results are returned as key-value pairs (label + NumberId) instead
+                of bare integers. Sent as ``queryOptions.resolveChoiceSets`` in the
+                request body. Defaults to ``False``.
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -2346,12 +2354,18 @@ def query_entity_records(
             ValueError: If the SQL query fails validation (e.g., non-SELECT, missing
                 WHERE/LIMIT, forbidden keywords, subqueries).
         """
-        return self._data.query_entity_records(sql_query, relationships_as_scalar)
+        return self._data.query_entity_records(
+            sql_query, relationships_as_scalar, resolve_choice_sets
+        )
 
     @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, *, relationships_as_scalar: bool = False
+        self,
+        sql_query: str,
+        *,
+        relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Asynchronously query entity records using a validated SQL query.
 
@@ -2366,6 +2380,10 @@ async def query_entity_records_async(
                 so a query can join on ``relationshipField = Other.Id``. Sent as
                 ``queryOptions.relationshipsAsScalar`` in the request body. Defaults to
                 ``False`` (unchanged behaviour).
+            resolve_choice_sets (bool, optional): When ``True``, choice-set fields
+                in results are returned as key-value pairs (label + NumberId) instead
+                of bare integers. Sent as ``queryOptions.resolveChoiceSets`` in the
+                request body. Defaults to ``False``.
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -2379,7 +2397,7 @@ async def query_entity_records_async(
                 WHERE/LIMIT, forbidden keywords, subqueries).
         """
         return await self._data.query_entity_records_async(
-            sql_query, relationships_as_scalar
+            sql_query, relationships_as_scalar, resolve_choice_sets
         )
 
     @traced(name="entity_upload_attachment", run_type="uipath")

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

@@ -22,6 +22,7 @@
 from ..common._config import UiPathApiConfig
 from ..common._execution_context import UiPathExecutionContext
 from ..common._models import Endpoint, RequestSpec
+from ..errors._datafabric_error import DataFabricSqlValidationError
 from ..errors._enriched_exception import EnrichedException
 from ..orchestrator._folder_service import FolderService
 from ._entity_resolution import RoutingStrategy, create_routing_strategy
@@ -554,18 +555,22 @@ def query_entity_records(
         self,
         sql_query: str,
         relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Internal implementation; see :meth:`EntitiesService.query_entity_records`."""
-        return self._query_entities_for_records(sql_query, relationships_as_scalar)
+        return self._query_entities_for_records(
+            sql_query, relationships_as_scalar, resolve_choice_sets
+        )
 
     async def query_entity_records_async(
         self,
         sql_query: str,
         relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Async variant of :meth:`query_entity_records`."""
         return await self._query_entities_for_records_async(
-            sql_query, relationships_as_scalar
+            sql_query, relationships_as_scalar, resolve_choice_sets
         )
 
     # ------------------------------------------------------------------
@@ -720,25 +725,31 @@ def validate_entity_batch(
     # ------------------------------------------------------------------
 
     def _query_entities_for_records(
-        self, sql_query: str, relationships_as_scalar: bool = False
+        self,
+        sql_query: str,
+        relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Synchronously run a validated SQL query through the federated query engine."""
         self._validate_sql_query(sql_query)
         routing_context = self._routing_strategy.resolve()
         spec = self._query_entity_records_spec(
-            sql_query, routing_context, relationships_as_scalar
+            sql_query, routing_context, relationships_as_scalar, resolve_choice_sets
         )
         response = self.request(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
 
     async def _query_entities_for_records_async(
-        self, sql_query: str, relationships_as_scalar: bool = False
+        self,
+        sql_query: str,
+        relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> List[Dict[str, Any]]:
         """Asynchronously run a validated SQL query through the federated query engine."""
         self._validate_sql_query(sql_query)
         routing_context = await self._routing_strategy.resolve_async()
         spec = self._query_entity_records_spec(
-            sql_query, routing_context, relationships_as_scalar
+            sql_query, routing_context, relationships_as_scalar, resolve_choice_sets
         )
         response = await self.request_async(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
@@ -996,15 +1007,21 @@ def _query_entity_records_spec(
         sql_query: str,
         routing_context: Optional[QueryRoutingOverrideContext] = None,
         relationships_as_scalar: bool = False,
+        resolve_choice_sets: bool = False,
     ) -> RequestSpec:
         """Build the POST spec for the federated SQL query endpoint."""
         body: Dict[str, Any] = {"query": sql_query}
         if routing_context:
             body["routingContext"] = routing_context.model_dump(
                 by_alias=True, exclude_none=True
             )
+        query_options: Dict[str, Any] = {}
         if relationships_as_scalar:
-            body["queryOptions"] = {"relationshipsAsScalar": True}
+            query_options["relationshipsAsScalar"] = True
+        if resolve_choice_sets:
+            query_options["resolveChoiceSets"] = True
+        if query_options:
+            body["queryOptions"] = query_options
         return RequestSpec(
             method="POST",
             endpoint=Endpoint("datafabric_/api/v1/query/execute"),
@@ -1256,20 +1273,37 @@ def _extract_batch_response_from_error(
     # ------------------------------------------------------------------
 
     def _validate_sql_query(self, sql_query: str) -> None:
-        """Validate a SQL string for the federated query endpoint client-side."""
+        """Validate a SQL string for the federated query endpoint client-side.
+
+        Raises:
+            DataFabricSqlValidationError: The statement violates the
+                entity-query subset. Its :class:`DataFabricError` category
+                distinguishes a mechanically fixable statement (``BAD_SQL``)
+                from one whose shape the subset cannot express at all
+                (``UNSUPPORTED_CONSTRUCT``), so a retry loop can stop instead
+                of re-trying variants of an impossible approach.
+        """
         query = sql_query.strip().rstrip(";").strip()
         if not query:
-            raise ValueError("SQL query cannot be empty.")
+            raise DataFabricSqlValidationError(
+                "SQL query cannot be empty.", code="SQL_EMPTY"
+            )
 
         statements = sqlparse.parse(query)
         if len(statements) != 1 or not statements[0].tokens:
-            raise ValueError("Only a single SELECT statement is allowed.")
+            raise DataFabricSqlValidationError(
+                "Only a single SELECT statement is allowed.",
+                code="SQL_MULTIPLE_STATEMENTS",
+            )
 
         stmt = statements[0]
         stmt_type = stmt.get_type()
 
         if stmt_type != "SELECT":
-            raise ValueError("Only SELECT statements are allowed.")
+            raise DataFabricSqlValidationError(
+                "Only SELECT statements are allowed.",
+                code="SQL_STATEMENT_NOT_SELECT",
+            )
 
         keywords = set()
         for token in stmt.flatten():
@@ -1278,46 +1312,65 @@ def _validate_sql_query(self, sql_query: str) -> None:
 
         for kw in _FORBIDDEN_DML:
             if kw in keywords:
-                raise ValueError(f"SQL keyword '{kw}' is not allowed.")
+                raise DataFabricSqlValidationError(
+                    f"SQL keyword '{kw}' is not allowed.",
+                    code="SQL_KEYWORD_NOT_ALLOWED",
+                )
 
         for kw in _FORBIDDEN_DDL:
             if kw in keywords:
-                raise ValueError(f"SQL keyword '{kw}' is not allowed.")
+                raise DataFabricSqlValidationError(
+                    f"SQL keyword '{kw}' is not allowed.",
+                    code="SQL_KEYWORD_NOT_ALLOWED",
+                )
 
         for kw in _DISALLOWED_KEYWORDS:
             if kw in keywords:
-                raise ValueError(
-                    f"SQL construct '{kw}' is not allowed in entity queries."
+                raise DataFabricSqlValidationError(
+                    f"SQL construct '{kw}' is not allowed in entity queries.",
+                    code="SQL_CONSTRUCT_NOT_ALLOWED",
                 )
 
         if self._has_subquery(stmt):
-            raise ValueError("Subqueries are not allowed.")
+            raise DataFabricSqlValidationError(
+                "Subqueries are not allowed.", code="SQL_SUBQUERY_NOT_ALLOWED"
+            )
 
         has_where = any(isinstance(t, Where) for t in stmt.tokens)
         has_limit = "LIMIT" in keywords
         has_from = "FROM" in keywords
 
         if not has_from:
-            raise ValueError("Queries must include a FROM clause.")
+            raise DataFabricSqlValidationError(
+                "Queries must include a FROM clause.", code="SQL_MISSING_FROM"
+            )
 
         projection = self._projection_tokens(stmt)
 
         if self._projection_has_count_star(projection):
-            raise ValueError(
-                "COUNT(*) is not supported. Use COUNT(column_name) instead."
+            raise DataFabricSqlValidationError(
+                "COUNT(*) is not supported. Use COUNT(column_name) instead.",
+                code="SQL_COUNT_STAR_NOT_SUPPORTED",
             )
 
         has_aggregate = self._projection_has_aggregate(projection)
 
         if not has_where and not has_limit and not has_aggregate:
-            raise ValueError("Queries without WHERE must include a LIMIT clause.")
+            raise DataFabricSqlValidationError(
+                "Queries without WHERE must include a LIMIT clause.",
+                code="SQL_LIMIT_REQUIRED",
+            )
 
         has_bare_wildcard = self._projection_has_bare_wildcard(projection)
         if has_bare_wildcard:
-            raise ValueError("SELECT * is not allowed. Specify column names instead.")
+            raise DataFabricSqlValidationError(
+                "SELECT * is not allowed. Specify column names instead.",
+                code="SQL_SELECT_STAR_NOT_ALLOWED",
+            )
         if not has_where and self._projection_column_count(projection) > 4:
-            raise ValueError(
-                "Selecting more than 4 columns without filtering is not allowed."
+            raise DataFabricSqlValidationError(
+                "Selecting more than 4 columns without filtering is not allowed.",
+                code="SQL_TOO_MANY_COLUMNS",
             )
 
     @staticmethod

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

@@ -205,7 +205,11 @@ class FieldMetadata(BaseModel):
     field_display_type: Optional[str] = Field(
         default=None, alias="fieldDisplayType"
     )  # Should be FieldDisplayType enum
-    choiceset_id: Optional[str] = Field(default=None, alias="choicesetId")
+    choiceset_id: Optional[str] = Field(
+        default=None,
+        validation_alias=AliasChoices("choiceSetId", "choicesetId", "ChoiceSetId"),
+        alias="choiceSetId",
+    )
     default_value: Optional[str] = Field(default=None, alias="defaultValue")
     is_attachment: bool = Field(alias="isAttachment")
     is_rbac_enabled: bool = Field(alias="isRbacEnabled")

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

@@ -59,6 +59,48 @@ def is_retryable(self) -> bool:
     def is_bad_sql(self) -> bool:
         return self.category == DataFabricErrorCategory.BAD_SQL
 
+    @property
+    def is_unsupported_construct(self) -> bool:
+        """True when the entity-query subset cannot express this query shape.
+
+        Distinct from :attr:`is_bad_sql`: a bad statement can be fixed by
+        rewriting the SQL, whereas an unsupported construct means retrying a
+        variant of the same approach will fail again.
+        """
+        return self.category == DataFabricErrorCategory.UNSUPPORTED_CONSTRUCT
+
+    @staticmethod
+    def from_exception(exc: BaseException) -> DataFabricError | None:
+        """Extract a DataFabricError from any Data Fabric query failure.
+
+        Covers both origins of a failed query so callers need one branch:
+        client-side validation rejections raised before the request, and
+        server-side errors returned by the query engine.
+
+        Returns None if the exception is not a Data Fabric query failure.
+        """
+        if isinstance(exc, DataFabricSqlValidationError):
+            return exc.error
+        from ._enriched_exception import EnrichedException as _EnrichedException
+
+        if isinstance(exc, _EnrichedException):
+            return DataFabricError.from_enriched_exception(exc)
+        return None
+
+    @staticmethod
+    def from_validation(code: str, message: str) -> DataFabricError:
+        """Build a DataFabricError for a client-side validation rejection.
+
+        These never reach the query engine, so there is no trace id; the code
+        is classified through the same table as server-returned codes.
+        """
+        return DataFabricError(
+            code=code,
+            message=message,
+            trace_id=None,
+            category=classify_error_code(code),
+        )
+
     @staticmethod
     def from_enriched_exception(exc: EnrichedException) -> DataFabricError | None:
         """Extract a DataFabricError from an EnrichedException, if applicable.
@@ -107,3 +149,25 @@ def from_response_body(body: dict[str, Any]) -> DataFabricError:
             trace_id=trace_id,
             category=classify_error_code(code),
         )
+
+
+class DataFabricSqlValidationError(ValueError):
+    """A SQL statement rejected by client-side entity-query validation.
+
+    A thin carrier: the classification callers act on is the
+    :class:`DataFabricError` on :attr:`error`, the same type server-side
+    failures produce. Remains a :class:`ValueError` subclass so existing
+    callers catching ``ValueError`` are unaffected; reach the structured form
+    with :meth:`DataFabricError.from_exception`.
+    """
+
+    def __init__(self, message: str, *, code: str) -> None:
+        """Initialise the error.
+
+        Args:
+            message: Human-readable rejection reason.
+            code: Stable code for this rejection, classified into a
+                :class:`DataFabricErrorCategory` by the shared code table.
+        """
+        super().__init__(message)
+        self.error = DataFabricError.from_validation(code=code, message=message)