feat(entities): optional source header on datafabric query/execute (#1808)

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.7"
+version = "0.2.8"
 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/constants/__init__.py

@@ -41,6 +41,7 @@
 HEADER_FOLDER_PATH = "x-uipath-folderpath"
 HEADER_FOLDER_PATH_ENCODED = "x-uipath-folderpath-encoded"
 HEADER_USER_AGENT = "x-uipath-user-agent"
+HEADER_SOURCE = "x-uipath-source"
 HEADER_TENANT_ID = "x-uipath-tenantid"
 HEADER_INTERNAL_TENANT_ID = "x-uipath-internal-tenantid"
 HEADER_INTERNAL_ACCOUNT_ID = "x-uipath-internal-accountid"

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

@@ -1820,7 +1820,9 @@ 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) -> List[Dict[str, Any]]:
+    def query_entity_records(
+        self, sql_query: str, source: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
         """Query entity records using a validated SQL query.
 
         PREVIEW: This method is in preview and may change in future releases.
@@ -1829,6 +1831,9 @@ def query_entity_records(self, sql_query: str) -> List[Dict[str, Any]]:
             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.
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -1841,11 +1846,13 @@ def query_entity_records(self, sql_query: str) -> List[Dict[str, Any]]:
             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)
+        return self._data.query_entity_records(sql_query, source)
 
     @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]]:
+    async def query_entity_records_async(
+        self, sql_query: str, source: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
         """Asynchronously query entity records using a validated SQL query.
 
         PREVIEW: This method is in preview and may change in future releases.
@@ -1854,6 +1861,9 @@ async def query_entity_records_async(self, sql_query: str) -> List[Dict[str, Any
             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.
 
         Notes:
             A routing context is always derived from the configured ``folders_map``
@@ -1866,7 +1876,7 @@ async def query_entity_records_async(self, sql_query: str) -> List[Dict[str, Any
             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)
+        return await self._data.query_entity_records_async(sql_query, source)
 
     @traced(name="entity_upload_attachment", run_type="uipath")
     def upload_attachment(

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 ..constants import HEADER_SOURCE
 from ..errors._enriched_exception import EnrichedException
 from ..orchestrator._folder_service import FolderService
 from ._entity_resolution import RoutingStrategy, create_routing_strategy
@@ -518,16 +519,18 @@ async def retrieve_records_async(
     def query_entity_records(
         self,
         sql_query: str,
+        source: Optional[str] = None,
     ) -> List[Dict[str, Any]]:
         """Internal implementation; see :meth:`EntitiesService.query_entity_records`."""
-        return self._query_entities_for_records(sql_query)
+        return self._query_entities_for_records(sql_query, source)
 
     async def query_entity_records_async(
         self,
         sql_query: str,
+        source: Optional[str] = None,
     ) -> List[Dict[str, Any]]:
         """Async variant of :meth:`query_entity_records`."""
-        return await self._query_entities_for_records_async(sql_query)
+        return await self._query_entities_for_records_async(sql_query, source)
 
     # ------------------------------------------------------------------
     # Attachments
@@ -680,22 +683,28 @@ def validate_entity_batch(
     # Internal helpers — request specs
     # ------------------------------------------------------------------
 
-    def _query_entities_for_records(self, sql_query: str) -> List[Dict[str, Any]]:
+    def _query_entities_for_records(
+        self, sql_query: str, source: Optional[str] = None
+    ) -> 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)
-        response = self.request(spec.method, spec.endpoint, json=spec.json)
+        spec = self._query_entity_records_spec(sql_query, routing_context, source)
+        response = self.request(
+            spec.method, spec.endpoint, json=spec.json, headers=spec.headers
+        )
         return response.json().get("results", [])
 
     async def _query_entities_for_records_async(
-        self, sql_query: str
+        self, sql_query: str, source: Optional[str] = None
     ) -> 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)
-        response = await self.request_async(spec.method, spec.endpoint, json=spec.json)
+        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
+        )
         return response.json().get("results", [])
 
     @staticmethod
@@ -956,17 +965,20 @@ def _retrieve_records_spec(
     def _query_entity_records_spec(
         sql_query: str,
         routing_context: Optional[QueryRoutingOverrideContext] = None,
+        source: Optional[str] = None,
     ) -> 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 {}
         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,6 +459,34 @@ 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(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": []}
+        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"
+        )
+
+        headers = service._data.request.call_args.kwargs.get("headers") or {}
+        assert headers.get("x-uipath-source") == "LOW_CODE_AGENT"
+
+    def test_query_entity_records_omits_source_header_by_default(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        response = MagicMock()
+        response.json.return_value = {"results": []}
+        service._data.request = MagicMock(return_value=response)  # type: ignore[method-assign]
+
+        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
+
     @pytest.mark.anyio
     async def test_query_entity_records_async_rejects_invalid_sql_before_network_call(
         self,