diff --git a/packages/uipath-platform/pyproject.toml b/packages/uipath-platform/pyproject.toml index 40955931c..01353417d 100644 --- a/packages/uipath-platform/pyproject.toml +++ b/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" diff --git a/packages/uipath-platform/src/uipath/platform/constants/__init__.py b/packages/uipath-platform/src/uipath/platform/constants/__init__.py index a42cf7c84..fc28a10c8 100644 --- a/packages/uipath-platform/src/uipath/platform/constants/__init__.py +++ b/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" diff --git a/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py b/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py index 8ab08f2f1..8b346cf0c 100644 --- a/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py +++ b/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( diff --git a/packages/uipath-platform/src/uipath/platform/entities/_entity_data_service.py b/packages/uipath-platform/src/uipath/platform/entities/_entity_data_service.py index 60cfb99b4..5a98fd7dd 100644 --- a/packages/uipath-platform/src/uipath/platform/entities/_entity_data_service.py +++ b/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,6 +965,7 @@ 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} @@ -963,10 +973,12 @@ def _query_entity_records_spec( 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 diff --git a/packages/uipath-platform/tests/services/test_entities_service.py b/packages/uipath-platform/tests/services/test_entities_service.py index 8ea33b618..547ec5335 100644 --- a/packages/uipath-platform/tests/services/test_entities_service.py +++ b/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, diff --git a/packages/uipath-platform/uv.lock b/packages/uipath-platform/uv.lock index 29a1abba1..779a6d1e1 100644 --- a/packages/uipath-platform/uv.lock +++ b/packages/uipath-platform/uv.lock @@ -1095,7 +1095,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.2.7" +version = "0.2.8" source = { editable = "." } dependencies = [ { name = "httpx" }, diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 730276088..fc5139669 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2741,7 +2741,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.2.7" +version = "0.2.8" source = { editable = "../uipath-platform" } dependencies = [ { name = "httpx" },