fix: make entity query service folder-aware with routing context (#1490)

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

Author: UIPath-Harshit

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.11"
+version = "0.1.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/__init__.py

@@ -10,11 +10,13 @@
     EntityFieldMetadata,
     EntityRecord,
     EntityRecordsBatchResponse,
+    EntityRouting,
     ExternalField,
     ExternalObject,
     ExternalSourceFields,
     FieldDataType,
     FieldMetadata,
+    QueryRoutingOverrideContext,
     ReferenceType,
     SourceJoinCriteria,
 )
@@ -25,12 +27,14 @@
     "EntityField",
     "EntityRecord",
     "EntityFieldMetadata",
+    "EntityRouting",
     "FieldDataType",
     "FieldMetadata",
     "EntityRecordsBatchResponse",
     "ExternalField",
     "ExternalObject",
     "ExternalSourceFields",
+    "QueryRoutingOverrideContext",
     "ReferenceType",
     "SourceJoinCriteria",
 ]

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

@@ -14,6 +14,7 @@
     Entity,
     EntityRecord,
     EntityRecordsBatchResponse,
+    QueryRoutingOverrideContext,
 )
 
 _FORBIDDEN_DML = {"INSERT", "UPDATE", "DELETE", "MERGE", "REPLACE"}
@@ -416,6 +417,7 @@ class CustomerRecord:
     def query_entity_records(
         self,
         sql_query: str,
+        routing_context: Optional[QueryRoutingOverrideContext] = None,
     ) -> List[Dict[str, Any]]:
         """Query entity records using a validated SQL query.
 
@@ -425,6 +427,9 @@ 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.
+            routing_context (Optional[QueryRoutingOverrideContext]): Per-entity routing context
+                for multi-folder queries. When present, included in the request body
+                and takes precedence over the folder header on the backend.
 
         Returns:
             List[Dict[str, Any]]: A list of result records as dictionaries.
@@ -433,12 +438,15 @@ def query_entity_records(
             ValueError: If the SQL query fails validation (e.g., non-SELECT, missing
                 WHERE/LIMIT, forbidden keywords, subqueries).
         """
-        return self._query_entities_for_records(sql_query)
+        return self._query_entities_for_records(
+            sql_query, routing_context=routing_context
+        )
 
     @traced(name="entity_query_records", run_type="uipath")
     async def query_entity_records_async(
         self,
         sql_query: str,
+        routing_context: Optional[QueryRoutingOverrideContext] = None,
     ) -> List[Dict[str, Any]]:
         """Asynchronously query entity records using a validated SQL query.
 
@@ -448,6 +456,9 @@ 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.
+            routing_context (Optional[QueryRoutingOverrideContext]): Per-entity routing context
+                for multi-folder queries. When present, included in the request body
+                and takes precedence over the folder header on the backend.
 
         Returns:
             List[Dict[str, Any]]: A list of result records as dictionaries.
@@ -456,19 +467,29 @@ 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._query_entities_for_records_async(sql_query)
+        return await self._query_entities_for_records_async(
+            sql_query, routing_context=routing_context
+        )
 
-    def _query_entities_for_records(self, sql_query: str) -> List[Dict[str, Any]]:
+    def _query_entities_for_records(
+        self,
+        sql_query: str,
+        *,
+        routing_context: Optional[QueryRoutingOverrideContext] = None,
+    ) -> List[Dict[str, Any]]:
         self._validate_sql_query(sql_query)
-        spec = self._query_entity_records_spec(sql_query)
+        spec = self._query_entity_records_spec(sql_query, routing_context)
         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
+        self,
+        sql_query: str,
+        *,
+        routing_context: Optional[QueryRoutingOverrideContext] = None,
     ) -> List[Dict[str, Any]]:
         self._validate_sql_query(sql_query)
-        spec = self._query_entity_records_spec(sql_query)
+        spec = self._query_entity_records_spec(sql_query, routing_context)
         response = await self.request_async(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
 
@@ -958,11 +979,17 @@ def _list_records_spec(
     def _query_entity_records_spec(
         self,
         sql_query: str,
+        routing_context: Optional[QueryRoutingOverrideContext] = None,
     ) -> RequestSpec:
+        body: Dict[str, Any] = {"query": sql_query}
+        if routing_context:
+            body["routingContext"] = routing_context.model_dump(
+                by_alias=True, exclude_none=True
+            )
         return RequestSpec(
             method="POST",
             endpoint=Endpoint("datafabric_/api/v1/query/execute"),
-            json={"query": sql_query},
+            json=body,
         )
 
     def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec:

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

@@ -322,4 +322,24 @@ class EntityRecordsBatchResponse(BaseModel):
     failure_records: List[EntityRecord] = Field(alias="failureRecords")
 
 
+class EntityRouting(BaseModel):
+    """A single entity-to-folder routing entry for query execution."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    entity_name: str = Field(alias="entityName")
+    folder_id: str = Field(alias="folderId")
+    override_entity_name: Optional[str] = Field(
+        default=None, alias="overrideEntityName"
+    )
+
+
+class QueryRoutingOverrideContext(BaseModel):
+    """Routing context that maps entities to their folders for multi-entity queries."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    entity_routings: List[EntityRouting] = Field(alias="entityRoutings")
+
+
 Entity.model_rebuild()

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

@@ -8,7 +8,7 @@
 from pytest_httpx import HTTPXMock
 
 from uipath.platform import UiPathApiConfig, UiPathExecutionContext
-from uipath.platform.entities import Entity
+from uipath.platform.entities import Entity, EntityRouting, QueryRoutingOverrideContext
 from uipath.platform.entities._entities_service import EntitiesService
 
 
@@ -389,3 +389,80 @@ async def test_query_entity_records_async_calls_request_for_valid_sql(
 
         assert result == [{"id": "c1"}]
         service.request_async.assert_called_once()
+
+    def test_query_entity_records_with_routing_context(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": [{"id": 1}]}
+        service.request = MagicMock(return_value=response)  # type: ignore[method-assign]
+
+        routing = QueryRoutingOverrideContext(
+            entity_routings=[
+                EntityRouting(entity_name="Customers", folder_id="folder-1"),
+                EntityRouting(
+                    entity_name="Orders",
+                    folder_id="folder-2",
+                    override_entity_name="OrdersV2",
+                ),
+            ]
+        )
+
+        result = service.query_entity_records(
+            "SELECT id FROM Customers LIMIT 10", routing_context=routing
+        )
+
+        assert result == [{"id": 1}]
+        call_kwargs = service.request.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert body["query"] == "SELECT id FROM Customers LIMIT 10"
+        assert body["routingContext"] == {
+            "entityRoutings": [
+                {"entityName": "Customers", "folderId": "folder-1"},
+                {
+                    "entityName": "Orders",
+                    "folderId": "folder-2",
+                    "overrideEntityName": "OrdersV2",
+                },
+            ]
+        }
+
+    @pytest.mark.anyio
+    async def test_query_entity_records_async_with_routing_context(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": [{"id": "c1"}]}
+        service.request_async = AsyncMock(return_value=response)  # type: ignore[method-assign]
+
+        routing = QueryRoutingOverrideContext(
+            entity_routings=[
+                EntityRouting(entity_name="Customers", folder_id="folder-1"),
+            ]
+        )
+
+        result = await service.query_entity_records_async(
+            "SELECT id FROM Customers WHERE id = 'c1'",
+            routing_context=routing,
+        )
+
+        assert result == [{"id": "c1"}]
+        call_kwargs = service.request_async.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert "routingContext" in body
+
+    def test_query_entity_records_without_routing_context_omits_key(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": []}
+        service.request = MagicMock(return_value=response)  # type: ignore[method-assign]
+
+        service.query_entity_records("SELECT id FROM Customers WHERE id > 0")
+
+        call_kwargs = service.request.call_args
+        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
+        assert "routingContext" not in body