feat(entities): relationships_as_scalar query option for FK joins (#1818)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: milind-jain-uipath

packages/uipath-platform/pyproject.toml

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

@@ -1821,7 +1821,7 @@ async def retrieve_records_async(
     @attach_datafabric_error_mapping("query_entity_records")
     @traced(name="entity_query_records", run_type="uipath")
     def query_entity_records(
-        self, sql_query: str, source: Optional[str] = None
+        self, sql_query: str, *, relationships_as_scalar: bool = False
     ) -> List[Dict[str, Any]]:
         """Query entity records using a validated SQL query.
 
@@ -1831,9 +1831,11 @@ def query_entity_records(
             sql_query (str): A SQL SELECT query to execute against Data Service entities.
                 Only SELECT statements are allowed. Queries without WHERE must include
                 a LIMIT clause. Subqueries and multi-statement queries are not permitted.
-            source (str, optional): Value for the ``x-uipath-source`` header on the
-                query/execute request, identifying the calling surface (e.g.
-                ``"LOW_CODE_AGENT"``). Omitted from the request when not set.
+            relationships_as_scalar (bool, optional): When ``True``, relationship (FK)
+                fields are typed as their underlying scalar id instead of ``VARIANT``,
+                so a query can join on ``relationshipField = Other.Id``. Sent as
+                ``queryOptions.relationshipsAsScalar`` in the request body. Defaults to
+                ``False`` (unchanged behaviour).
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -1846,12 +1848,12 @@ 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, source)
+        return self._data.query_entity_records(sql_query, relationships_as_scalar)
 
     @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, source: Optional[str] = None
+        self, sql_query: str, *, relationships_as_scalar: bool = False
     ) -> List[Dict[str, Any]]:
         """Asynchronously query entity records using a validated SQL query.
 
@@ -1861,9 +1863,11 @@ async def query_entity_records_async(
             sql_query (str): A SQL SELECT query to execute against Data Service entities.
                 Only SELECT statements are allowed. Queries without WHERE must include
                 a LIMIT clause. Subqueries and multi-statement queries are not permitted.
-            source (str, optional): Value for the ``x-uipath-source`` header on the
-                query/execute request, identifying the calling surface (e.g.
-                ``"LOW_CODE_AGENT"``). Omitted from the request when not set.
+            relationships_as_scalar (bool, optional): When ``True``, relationship (FK)
+                fields are typed as their underlying scalar id instead of ``VARIANT``,
+                so a query can join on ``relationshipField = Other.Id``. Sent as
+                ``queryOptions.relationshipsAsScalar`` in the request body. Defaults to
+                ``False`` (unchanged behaviour).
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -1876,7 +1880,9 @@ async def query_entity_records_async(
             ValueError: If the SQL query fails validation (e.g., non-SELECT, missing
                 WHERE/LIMIT, forbidden keywords, subqueries).
         """
-        return await self._data.query_entity_records_async(sql_query, source)
+        return await self._data.query_entity_records_async(
+            sql_query, relationships_as_scalar
+        )
 
     @traced(name="entity_upload_attachment", run_type="uipath")
     def upload_attachment(

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

@@ -22,7 +22,6 @@
 from ..common._config import UiPathApiConfig
 from ..common._execution_context import UiPathExecutionContext
 from ..common._models import Endpoint, RequestSpec
-from ..constants import HEADER_SOURCE
 from ..errors._enriched_exception import EnrichedException
 from ..orchestrator._folder_service import FolderService
 from ._entity_resolution import RoutingStrategy, create_routing_strategy
@@ -519,18 +518,20 @@ async def retrieve_records_async(
     def query_entity_records(
         self,
         sql_query: str,
-        source: Optional[str] = None,
+        relationships_as_scalar: bool = False,
     ) -> List[Dict[str, Any]]:
         """Internal implementation; see :meth:`EntitiesService.query_entity_records`."""
-        return self._query_entities_for_records(sql_query, source)
+        return self._query_entities_for_records(sql_query, relationships_as_scalar)
 
     async def query_entity_records_async(
         self,
         sql_query: str,
-        source: Optional[str] = None,
+        relationships_as_scalar: bool = False,
     ) -> List[Dict[str, Any]]:
         """Async variant of :meth:`query_entity_records`."""
-        return await self._query_entities_for_records_async(sql_query, source)
+        return await self._query_entities_for_records_async(
+            sql_query, relationships_as_scalar
+        )
 
     # ------------------------------------------------------------------
     # Attachments
@@ -684,27 +685,27 @@ def validate_entity_batch(
     # ------------------------------------------------------------------
 
     def _query_entities_for_records(
-        self, sql_query: str, source: Optional[str] = None
+        self, sql_query: str, relationships_as_scalar: 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, source)
-        response = self.request(
-            spec.method, spec.endpoint, json=spec.json, headers=spec.headers
+        spec = self._query_entity_records_spec(
+            sql_query, routing_context, relationships_as_scalar
         )
+        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, source: Optional[str] = None
+        self, sql_query: str, relationships_as_scalar: 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, source)
-        response = await self.request_async(
-            spec.method, spec.endpoint, json=spec.json, headers=spec.headers
+        spec = self._query_entity_records_spec(
+            sql_query, routing_context, relationships_as_scalar
         )
+        response = await self.request_async(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
 
     @staticmethod
@@ -965,20 +966,20 @@ def _retrieve_records_spec(
     def _query_entity_records_spec(
         sql_query: str,
         routing_context: Optional[QueryRoutingOverrideContext] = None,
-        source: Optional[str] = None,
+        relationships_as_scalar: 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
             )
-        headers = {HEADER_SOURCE: source} if source else {}
+        if relationships_as_scalar:
+            body["queryOptions"] = {"relationshipsAsScalar": True}
         return RequestSpec(
             method="POST",
             endpoint=Endpoint("datafabric_/api/v1/query/execute"),
             json=body,
-            headers=headers,
         )
 
     @staticmethod

packages/uipath-platform/tests/services/test_entities_service.py

@@ -459,7 +459,7 @@ def test_query_entity_records_calls_request_for_valid_sql(
         assert result == [{"id": 1}, {"id": 2}]
         service._data.request.assert_called_once()
 
-    def test_query_entity_records_sets_source_header_when_provided(
+    def test_query_entity_records_sets_relationships_as_scalar_option_when_true(
         self,
         service: EntitiesService,
     ) -> None:
@@ -468,13 +468,14 @@ def test_query_entity_records_sets_source_header_when_provided(
         service._data.request = MagicMock(return_value=response)  # type: ignore[method-assign]
 
         service.query_entity_records(
-            "SELECT id FROM Customers WHERE id > 0", source="LOW_CODE_AGENT"
+            "SELECT id FROM Customers WHERE id > 0", relationships_as_scalar=True
         )
 
-        headers = service._data.request.call_args.kwargs.get("headers") or {}
-        assert headers.get("x-uipath-source") == "LOW_CODE_AGENT"
+        call_kwargs = service._data.request.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert body["queryOptions"] == {"relationshipsAsScalar": True}
 
-    def test_query_entity_records_omits_source_header_by_default(
+    def test_query_entity_records_omits_query_options_by_default(
         self,
         service: EntitiesService,
     ) -> None:
@@ -484,8 +485,18 @@ def test_query_entity_records_omits_source_header_by_default(
 
         service.query_entity_records("SELECT id FROM Customers WHERE id > 0")
 
-        headers = service._data.request.call_args.kwargs.get("headers") or {}
-        assert "x-uipath-source" not in headers
+        call_kwargs = service._data.request.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert "queryOptions" not in body
+
+    def test_query_entity_records_flag_is_keyword_only(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        # A second positional arg (as an old ``source`` call would pass) must be
+        # rejected rather than silently coerced into ``relationships_as_scalar``.
+        with pytest.raises(TypeError):
+            service.query_entity_records("SELECT id FROM Customers LIMIT 10", True)
 
     @pytest.mark.anyio
     async def test_query_entity_records_async_rejects_invalid_sql_before_network_call(
@@ -518,6 +529,23 @@ async def test_query_entity_records_async_calls_request_for_valid_sql(
         assert result == [{"id": "c1"}]
         service._data.request_async.assert_called_once()
 
+    @pytest.mark.anyio
+    async def test_query_entity_records_async_sets_relationships_as_scalar_option(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": []}
+        service._data.request_async = AsyncMock(return_value=response)  # type: ignore[method-assign]
+
+        await service.query_entity_records_async(
+            "SELECT id FROM Customers WHERE id > 0", relationships_as_scalar=True
+        )
+
+        call_kwargs = service._data.request_async.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert body["queryOptions"] == {"relationshipsAsScalar": True}
+
     def test_query_entity_records_builds_routing_context_from_folders_map(
         self,
         config: UiPathApiConfig,