From 2766413e54a1d0549bc9f7fed3f49e5247430f72 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Fri, 13 Feb 2026 01:12:51 +0530 Subject: [PATCH 1/9] Integration-e2e --- pyproject.toml | 4 + src/uipath/agent/models/agent.py | 7 ++ .../platform/entities/_entities_service.py | 103 +++++++++++++++++- tests/agent/models/test_agent.py | 50 +++++++++ uv.lock | 2 + 5 files changed, 165 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 886c9b2a6..3ad3ce874 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ dev = [ "inflection>=0.5.1", "types-toml>=0.10.8", "pytest-timeout>=2.4.0", + "uipath", ] [tool.hatch.build.targets.wheel] @@ -148,3 +149,6 @@ name = "testpypi" url = "https://test.pypi.org/simple/" publish-url = "https://test.pypi.org/legacy/" explicit = true + +[tool.uv.sources] +uipath = { workspace = true } diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 2941577d3..b5336569f 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -85,6 +85,7 @@ class AgentContextRetrievalMode(str, Enum): STRUCTURED = "Structured" DEEP_RAG = "DeepRAG" BATCH_TRANSFORM = "BatchTransform" + DATA_FABRIC = "DataFabric" UNKNOWN = "Unknown" # fallback branch discriminator @@ -317,6 +318,7 @@ class AgentContextSettings(BaseCfg): AgentContextRetrievalMode.STRUCTURED, AgentContextRetrievalMode.DEEP_RAG, AgentContextRetrievalMode.BATCH_TRANSFORM, + AgentContextRetrievalMode.DATA_FABRIC, AgentContextRetrievalMode.UNKNOWN, ] = Field(alias="retrievalMode") threshold: float = Field(default=0) @@ -336,6 +338,10 @@ class AgentContextSettings(BaseCfg): output_columns: Optional[List[AgentContextOutputColumn]] = Field( None, alias="outputColumns" ) + # Data Fabric specific settings + entity_identifiers: Optional[List[str]] = Field( + None, alias="entityIdentifiers" + ) class AgentContextResourceConfig(BaseAgentResourceConfig): @@ -1162,6 +1168,7 @@ def _normalize_resources(v: Dict[str, Any]) -> None: "structured": "Structured", "deeprag": "DeepRAG", "batchtransform": "BatchTransform", + "datafabric": "DataFabric", "unknown": "Unknown", } diff --git a/src/uipath/platform/entities/_entities_service.py b/src/uipath/platform/entities/_entities_service.py index 2b31b830c..d5b19a2f4 100644 --- a/src/uipath/platform/entities/_entities_service.py +++ b/src/uipath/platform/entities/_entities_service.py @@ -1,4 +1,4 @@ -from typing import Any, List, Optional, Type +from typing import Any, Dict, List, Optional, Type from httpx import Response @@ -389,6 +389,95 @@ class CustomerRecord: EntityRecord.from_data(data=record, model=schema) for record in records_data ] + @traced(name="entity_query_records", run_type="uipath") + def query_entity_records( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = entities_service.query_entity_records( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = entities_service.query_entity_records( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + # Use absolute URL to bypass scoping since org/tenant are embedded in the path + full_url = f"{self._url.base_url}{spec.endpoint}" + response = self.request(spec.method, full_url, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + # Return raw dicts for query results - they may not have Id field + # if SELECT doesn't include it + return records_data + + @traced(name="entity_query_records", run_type="uipath") + async def query_entity_records_async( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Asynchronously query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = await entities_service.query_entity_records_async( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = await entities_service.query_entity_records_async( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + full_url = f"{self._url.base_url}{spec.endpoint}" + response = await self.request_async(spec.method, full_url, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + return records_data + @traced(name="entity_record_insert_batch", run_type="uipath") def insert_records( self, @@ -872,6 +961,18 @@ def _list_records_spec( params=({"start": start, "limit": limit}), ) + def _query_entity_records_spec( + self, + sql_query: str, + ) -> RequestSpec: + # Endpoint includes org/tenant in the path: dataservice_/{org}/{tenant}/datafabric_/api/v1/query/execute + endpoint = f"/dataservice_/{self._url.org_name}/{self._url.tenant_name}/datafabric_/api/v1/query/execute" + return RequestSpec( + method="POST", + endpoint=Endpoint(endpoint), + json={"query": sql_query}, + ) + def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec: return RequestSpec( method="POST", diff --git a/tests/agent/models/test_agent.py b/tests/agent/models/test_agent.py index 63a8a917f..6af2b9663 100644 --- a/tests/agent/models/test_agent.py +++ b/tests/agent/models/test_agent.py @@ -2755,3 +2755,53 @@ def test_is_conversational_false_by_default(self): ) assert config.is_conversational is False + + +class TestDataFabricContextConfig: + """Tests for Data Fabric context resource configuration.""" + + def test_datafabric_retrieval_mode_exists(self): + """Test that DATA_FABRIC retrieval mode is defined.""" + assert AgentContextRetrievalMode.DATA_FABRIC == "DataFabric" + + def test_datafabric_context_config_parses(self): + """Test that Data Fabric context config parses correctly.""" + config = { + "$resourceType": "context", + "name": "Customer Data", + "description": "Query customer and order data", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", + "settings": { + "retrievalMode": "DataFabric", + "resultCount": 100, + "entityIdentifiers": ["customers-key", "orders-key"], + }, + } + + parsed = AgentContextResourceConfig.model_validate(config) + + assert parsed.name == "Customer Data" + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers == ["customers-key", "orders-key"] + + def test_datafabric_context_config_without_entity_identifiers(self): + """Test that entity_identifiers is optional.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "Test", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", + "settings": { + "retrievalMode": "DataFabric", + "resultCount": 10, + }, + } + + parsed = AgentContextResourceConfig.model_validate(config) + + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers is None diff --git a/uv.lock b/uv.lock index 80cd0b9f6..9426a50ef 100644 --- a/uv.lock +++ b/uv.lock @@ -2580,6 +2580,7 @@ dev = [ { name = "termynal" }, { name = "tomli-w" }, { name = "types-toml" }, + { name = "uipath" }, { name = "virtualenv" }, ] @@ -2631,6 +2632,7 @@ dev = [ { name = "termynal", specifier = ">=0.13.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, { name = "types-toml", specifier = ">=0.10.8" }, + { name = "uipath", editable = "." }, { name = "virtualenv", specifier = ">=20.36.1" }, ] From 1ee79552ae2e34b443d2aa63bb25f1ed7a7ebf63 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Fri, 13 Feb 2026 17:49:24 +0530 Subject: [PATCH 2/9] SDK Commit --- src/uipath/platform/entities/_entities_service.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/uipath/platform/entities/_entities_service.py b/src/uipath/platform/entities/_entities_service.py index d5b19a2f4..d7a4ba132 100644 --- a/src/uipath/platform/entities/_entities_service.py +++ b/src/uipath/platform/entities/_entities_service.py @@ -431,11 +431,10 @@ def query_entity_records( response = self.request(spec.method, full_url, json=spec.json, headers=headers) records_data = response.json().get("results", []) - # Return raw dicts for query results - they may not have Id field - # if SELECT doesn't include it return records_data - @traced(name="entity_query_records", run_type="uipath") + + @traced(name="query_entities_async", run_type="uipath") async def query_entity_records_async( self, sql_query: str, @@ -965,7 +964,6 @@ def _query_entity_records_spec( self, sql_query: str, ) -> RequestSpec: - # Endpoint includes org/tenant in the path: dataservice_/{org}/{tenant}/datafabric_/api/v1/query/execute endpoint = f"/dataservice_/{self._url.org_name}/{self._url.tenant_name}/datafabric_/api/v1/query/execute" return RequestSpec( method="POST", From 628c29d98be48c7b6786544d4982a7983e4b6340 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Tue, 17 Feb 2026 08:42:44 +0530 Subject: [PATCH 3/9] Introduced query_multiple_entities method as part of DataFabric integration with agents. Switch data-fabric to private preview under the label as there is major rearchitecture going on our side --- pyproject.toml | 1 + .../platform/entities/_entities_service.py | 176 +++++++++++- tests/datafabric/test_entities_service.py | 117 ++++++++ tests/sdk/services/test_entities_service.py | 262 ------------------ uv.lock | 11 + 5 files changed, 290 insertions(+), 277 deletions(-) create mode 100644 tests/datafabric/test_entities_service.py delete mode 100644 tests/sdk/services/test_entities_service.py diff --git a/pyproject.toml b/pyproject.toml index 3ad3ce874..d016de14d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "mermaid-builder==0.0.3", "graphtty==0.1.6", "applicationinsights>=0.11.10", + "sqlparse>=0.4.4", ] classifiers = [ "Intended Audience :: Developers", diff --git a/src/uipath/platform/entities/_entities_service.py b/src/uipath/platform/entities/_entities_service.py index d7a4ba132..150818e44 100644 --- a/src/uipath/platform/entities/_entities_service.py +++ b/src/uipath/platform/entities/_entities_service.py @@ -1,6 +1,16 @@ from typing import Any, Dict, List, Optional, Type +import sqlparse from httpx import Response +from sqlparse.sql import ( + IdentifierList, + Parenthesis, + Statement, + Token, + TokenList, + Where, +) +from sqlparse.tokens import DML, Comment, Keyword, Punctuation, Wildcard from ..._utils import Endpoint, RequestSpec from ...tracing import traced @@ -11,15 +21,39 @@ EntityRecordsBatchResponse, ) +_FORBIDDEN_SQL_KEYWORDS = { + "INSERT", + "UPDATE", + "DELETE", + "MERGE", + "DROP", + "ALTER", + "CREATE", + "TRUNCATE", + "REPLACE", +} +_DISALLOWED_SQL_OPERATORS = { + "WITH", + "UNION", + "INTERSECT", + "EXCEPT", + "OVER", + "ROLLUP", + "CUBE", + "GROUPING SETS", + "PARTITION BY", +} + class EntitiesService(BaseService): """Service for managing UiPath Data Service entities. - Entities are database tables in UiPath Data Service that can store - structured data for automation processes. + Entities represent business objects that provide structured data storage and access via the Data Service. + This service allows you to retrieve entity metadata, list entities, and query records using SQL. - See Also: - https://docs.uipath.com/data-service/automation-cloud/latest/user-guide/introduction + !!! warning "Preview Feature" + This function is currently experimental. + Behavior and parameters, request and response formats are subject to change in future versions. """ def __init__( @@ -389,8 +423,8 @@ class CustomerRecord: EntityRecord.from_data(data=record, model=schema) for record in records_data ] - @traced(name="entity_query_records", run_type="uipath") - def query_entity_records( + @traced(name="query_multiple_entities", run_type="uipath") + def query_multiple_entities( self, sql_query: str, schema: Optional[Type[Any]] = None, @@ -411,17 +445,18 @@ def query_entity_records( Examples: Basic query:: - records = entities_service.query_entity_records( + records = entities_service.query_multiple_entities( "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" ) Query with specific fields:: - records = entities_service.query_entity_records( + records = entities_service.query_multiple_entities( "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" ) """ - spec = self._query_entity_records_spec(sql_query) + self._validate_sql_query(sql_query) + spec = self._query_multiple_records_spec(sql_query) headers = { "X-UiPath-Internal-TenantName": self._url.tenant_name, "X-UiPath-Internal-AccountName": self._url.org_name, @@ -434,8 +469,8 @@ def query_entity_records( return records_data - @traced(name="query_entities_async", run_type="uipath") - async def query_entity_records_async( + @traced(name="query_multiple_entities_async", run_type="uipath") + async def query_multiple_entities_async( self, sql_query: str, schema: Optional[Type[Any]] = None, @@ -456,17 +491,18 @@ async def query_entity_records_async( Examples: Basic query:: - records = await entities_service.query_entity_records_async( + records = await entities_service.query_multiple_entities_async( "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" ) Query with specific fields:: - records = await entities_service.query_entity_records_async( + records = await entities_service.query_multiple_entities_async( "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" ) """ - spec = self._query_entity_records_spec(sql_query) + self._validate_sql_query(sql_query) + spec = self._query_multiple_entities_spec(sql_query) headers = { "X-UiPath-Internal-TenantName": self._url.tenant_name, "X-UiPath-Internal-AccountName": self._url.org_name, @@ -960,7 +996,7 @@ def _list_records_spec( params=({"start": start, "limit": limit}), ) - def _query_entity_records_spec( + def _query_multiple_entities_spec( self, sql_query: str, ) -> RequestSpec: @@ -999,3 +1035,113 @@ def _delete_batch_spec(self, entity_key: str, record_ids: List[str]) -> RequestS ), json=record_ids, ) + + def _validate_sql_query(self, sql_query: str) -> None: + query = sql_query.strip() + if not query: + raise ValueError("SQL query cannot be empty.") + + statements = [stmt for stmt in sqlparse.parse(query) if stmt.tokens] + if len(statements) != 1: + raise ValueError("Only a single SELECT statement is allowed.") + + statement = statements[0] + if statement.get_type() != "SELECT": + raise ValueError("Only SELECT statements are allowed.") + + normalized_keywords = { + token.normalized + for token in statement.flatten() + if token.ttype in Keyword or token.ttype is DML + } + + for keyword in _FORBIDDEN_SQL_KEYWORDS: + if keyword in normalized_keywords: + raise ValueError(f"SQL keyword '{keyword}' is not allowed.") + + for operator in _DISALLOWED_SQL_OPERATORS: + if operator in normalized_keywords: + raise ValueError( + f"SQL construct '{operator}' is not allowed in entity queries." + ) + + if self._contains_subquery(statement): + raise ValueError("Subqueries are not allowed.") + + has_where = any(isinstance(token, Where) for token in statement.tokens) + has_limit = any( + token.ttype in Keyword and token.normalized == "LIMIT" + for token in statement.flatten() + ) + if not has_where and not has_limit: + raise ValueError("Queries without WHERE must include a LIMIT clause.") + + projection_tokens = self._projection_tokens(statement) + has_wildcard_projection = any( + token.ttype is Wildcard + for projection_token in projection_tokens + for token in projection_token.flatten() + ) + if has_wildcard_projection and not has_where: + raise ValueError("SELECT * without filtering is not allowed.") + if not has_where and self._projection_column_count(projection_tokens) > 4: + raise ValueError( + "Selecting more than 4 columns without filtering is not allowed." + ) + + def _contains_subquery(self, token_list: TokenList) -> bool: + for token in token_list.tokens: + if isinstance(token, Parenthesis): + if any( + nested.ttype is DML and nested.normalized == "SELECT" + for nested in token.flatten() + ): + return True + if isinstance(token, TokenList) and self._contains_subquery(token): + return True + return False + + def _projection_tokens(self, statement: Statement) -> List[Token]: + projection: List[Token] = [] + found_select = False + + for token in statement.tokens: + if token.is_whitespace or token.ttype in Comment: + continue + if not found_select: + if token.ttype is DML and token.normalized == "SELECT": + found_select = True + continue + if token.ttype in Keyword and token.normalized == "FROM": + break + projection.append(token) + return projection + + def _projection_column_count(self, projection_tokens: List[Token]) -> int: + identifier_list = next( + ( + token + for token in projection_tokens + if isinstance(token, IdentifierList) + ), + None, + ) + if identifier_list is not None: + return sum(1 for _ in identifier_list.get_identifiers()) + + count = 0 + has_current_expression = False + + for token in projection_tokens: + if token.is_whitespace or token.ttype in Comment: + continue + if token.ttype is Punctuation and token.value == ",": + if has_current_expression: + count += 1 + has_current_expression = False + continue + has_current_expression = True + + if has_current_expression: + count += 1 + return count diff --git a/tests/datafabric/test_entities_service.py b/tests/datafabric/test_entities_service.py new file mode 100644 index 000000000..ccd256577 --- /dev/null +++ b/tests/datafabric/test_entities_service.py @@ -0,0 +1,117 @@ +import re +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from uipath.platform import UiPathApiConfig, UiPathExecutionContext +from uipath.platform.entities._entities_service import EntitiesService + + +@pytest.fixture +def service() -> EntitiesService: + config = UiPathApiConfig(base_url="https://test.uipath.com/org/tenant", secret="secret") + execution_context = UiPathExecutionContext() + return EntitiesService(config=config, execution_context=execution_context) + + +@pytest.mark.parametrize( + "sql_query", + [ + "SELECT id FROM Customers WHERE id = 1", + "SELECT id, name FROM Customers LIMIT 10", + "SELECT * FROM Customers WHERE status = 'Active'", + "SELECT id, name, email, phone FROM Customers LIMIT 5", + "SELECT DISTINCT id FROM Customers WHERE id > 100", + ], +) +def test_validate_sql_query_allows_supported_select_queries( + sql_query: str, service: EntitiesService +) -> None: + service._validate_sql_query(sql_query) + + +@pytest.mark.parametrize( + "sql_query,error_message", + [ + ("", "SQL query cannot be empty."), + (" ", "SQL query cannot be empty."), + ("SELECT id FROM Customers; SELECT id FROM Orders", "Only a single SELECT statement is allowed."), + ("INSERT INTO Customers VALUES (1)", "Only SELECT statements are allowed."), + ( + "WITH cte AS (SELECT id FROM Customers) SELECT id FROM cte", + "SQL construct 'WITH' is not allowed in entity queries.", + ), + ("SELECT id FROM Customers UNION SELECT id FROM Orders", "SQL construct 'UNION' is not allowed in entity queries."), + ("SELECT id, SUM(amount) OVER (PARTITION BY id) FROM Orders LIMIT 10", "SQL construct 'OVER' is not allowed in entity queries."), + ("SELECT id FROM (SELECT id FROM Customers) c", "Subqueries are not allowed."), + ("SELECT id FROM Customers", "Queries without WHERE must include a LIMIT clause."), + ("SELECT * FROM Customers LIMIT 10", "SELECT * without filtering is not allowed."), + ( + "SELECT id, name, email, phone, address FROM Customers LIMIT 10", + "Selecting more than 4 columns without filtering is not allowed.", + ), + ], +) +def test_validate_sql_query_rejects_disallowed_queries( + sql_query: str, error_message: str, service: EntitiesService +) -> None: + with pytest.raises(ValueError, match=re.escape(error_message)): + service._validate_sql_query(sql_query) + + +def test_query_multiple_entities_rejects_invalid_sql_before_network_call( + service: EntitiesService, +) -> None: + service.request = MagicMock() # type: ignore[method-assign] + + with pytest.raises( + ValueError, match=re.escape("Only SELECT statements are allowed.") + ): + service.query_multiple_entities("UPDATE Customers SET name = 'X'") + + service.request.assert_not_called() # type: ignore[attr-defined] + + +def test_query_multiple_entities_calls_request_for_valid_sql( + service: EntitiesService, +) -> None: + response = MagicMock() + response.json.return_value = {"results": [{"id": 1}, {"id": 2}]} + + service.request = MagicMock(return_value=response) # type: ignore[method-assign] + + result = service.query_multiple_entities("SELECT id FROM Customers WHERE id > 0") + + assert result == [{"id": 1}, {"id": 2}] + service.request.assert_called_once() # type: ignore[attr-defined] + + +@pytest.mark.anyio +async def test_query_multiple_entities_async_rejects_invalid_sql_before_network_call( + service: EntitiesService, +) -> None: + service.request_async = AsyncMock() # type: ignore[method-assign] + + with pytest.raises(ValueError, match=re.escape("Subqueries are not allowed.")): + await service.query_multiple_entities_async( + "SELECT id FROM Customers WHERE id IN (SELECT id FROM Orders)" + ) + + service.request_async.assert_not_called() # type: ignore[attr-defined] + + +@pytest.mark.anyio +async def test_query_multiple_entities_async_calls_request_for_valid_sql( + service: EntitiesService, +) -> None: + response = MagicMock() + response.json.return_value = {"results": [{"id": "c1"}]} + + service.request_async = AsyncMock(return_value=response) # type: ignore[method-assign] + + result = await service.query_multiple_entities_async( + "SELECT id FROM Customers WHERE id = 'c1'" + ) + + assert result == [{"id": "c1"}] + service.request_async.assert_called_once() # type: ignore[attr-defined] diff --git a/tests/sdk/services/test_entities_service.py b/tests/sdk/services/test_entities_service.py deleted file mode 100644 index 4c6c85882..000000000 --- a/tests/sdk/services/test_entities_service.py +++ /dev/null @@ -1,262 +0,0 @@ -import uuid -from dataclasses import make_dataclass -from typing import Optional - -import pytest -from pytest_httpx import HTTPXMock - -from uipath.platform import UiPathApiConfig, UiPathExecutionContext -from uipath.platform.entities import Entity -from uipath.platform.entities._entities_service import EntitiesService - - -@pytest.fixture -def service( - config: UiPathApiConfig, - execution_context: UiPathExecutionContext, - monkeypatch: pytest.MonkeyPatch, -) -> EntitiesService: - return EntitiesService(config=config, execution_context=execution_context) - - -@pytest.fixture(params=[True, False], ids=["correct_schema", "incorrect_schema"]) -def record_schema(request): - is_correct = request.param - field_type = int if is_correct else str - schema_name = f"RecordSchema{'Correct' if is_correct else 'Incorrect'}" - - RecordSchema = make_dataclass( - schema_name, [("name", str), ("integer_field", field_type)] - ) - - return RecordSchema, is_correct - - -@pytest.fixture(params=[True, False], ids=["optional_field", "required_field"]) -def record_schema_optional(request): - is_optional = request.param - field_type = Optional[int] | None if is_optional else int - schema_name = f"RecordSchema{'Optional' if is_optional else 'Required'}" - - RecordSchemaOptional = make_dataclass( - schema_name, [("name", str), ("integer_field", field_type)] - ) - - return RecordSchemaOptional, is_optional - - -class TestEntitiesService: - def test_retrieve( - self, - httpx_mock: HTTPXMock, - service: EntitiesService, - base_url: str, - org: str, - tenant: str, - version: str, - ) -> None: - entity_key = uuid.uuid4() - httpx_mock.add_response( - url=f"{base_url}{org}{tenant}/datafabric_/api/Entity/{entity_key}", - status_code=200, - json={ - "name": "TestEntity", - "displayName": "TestEntity", - "entityType": "TestEntityType", - "description": "TestEntity Description", - "fields": [ - { - "id": "12345", - "name": "field_name", - "isPrimaryKey": True, - "isForeignKey": False, - "isExternalField": False, - "isHiddenField": True, - "isUnique": True, - "referenceType": "ManyToOne", - "sqlType": {"name": "VARCHAR", "LengthLimit": 100}, - "isRequired": True, - "displayName": "Field Display Name", - "description": "This is a brief description of the field.", - "isSystemField": False, - "isAttachment": False, - "isRbacEnabled": True, - } - ], - "isRbacEnabled": False, - "id": f"{entity_key}", - }, - ) - - entity = service.retrieve(entity_key=str(entity_key)) - - assert isinstance(entity, Entity) - assert entity.id == f"{entity_key}" - assert entity.name == "TestEntity" - assert entity.display_name == "TestEntity" - assert entity.entity_type == "TestEntityType" - assert entity.description == "TestEntity Description" - assert entity.fields is not None - assert entity.fields[0].id == "12345" - assert entity.fields[0].name == "field_name" - assert entity.fields[0].is_primary_key - assert not entity.fields[0].is_foreign_key - assert entity.fields[0].sql_type.name == "VARCHAR" - assert entity.fields[0].sql_type.length_limit == 100 - - sent_request = httpx_mock.get_request() - if sent_request is None: - raise Exception("No request was sent") - - assert sent_request.method == "GET" - assert ( - sent_request.url - == f"{base_url}{org}{tenant}/datafabric_/api/Entity/{entity_key}" - ) - - def test_retrieve_records_with_no_schema_succeeds( - self, - httpx_mock: HTTPXMock, - service: EntitiesService, - base_url: str, - org: str, - tenant: str, - version: str, - ) -> None: - entity_key = uuid.uuid4() - httpx_mock.add_response( - url=f"{base_url}{org}{tenant}/datafabric_/api/EntityService/entity/{str(entity_key)}/read?start=0&limit=1", - status_code=200, - json={ - "totalCount": 1, - "value": [ - {"Id": "12345", "name": "record_name", "integer_field": 10}, - {"Id": "12346", "name": "record_name2", "integer_field": 11}, - ], - }, - ) - - records = service.list_records(entity_key=str(entity_key), start=0, limit=1) - - sent_request = httpx_mock.get_request() - if sent_request is None: - raise Exception("No request was sent") - - assert isinstance(records, list) - assert len(records) == 2 - assert records[0].id == "12345" - assert records[0].name == "record_name" - assert records[0].integer_field == 10 - assert records[1].id == "12346" - assert records[1].name == "record_name2" - assert records[1].integer_field == 11 - - def test_retrieve_records_with_schema_succeeds( - self, - httpx_mock: HTTPXMock, - service: EntitiesService, - base_url: str, - org: str, - tenant: str, - version: str, - record_schema, - ) -> None: - entity_key = uuid.uuid4() - httpx_mock.add_response( - url=f"{base_url}{org}{tenant}/datafabric_/api/EntityService/entity/{str(entity_key)}/read?start=0&limit=1", - status_code=200, - json={ - "totalCount": 1, - "value": [ - {"Id": "12345", "name": "record_name", "integer_field": 10}, - {"Id": "12346", "name": "record_name2", "integer_field": 11}, - ], - }, - ) - - # Define the schema for the record. A wrong schema should make the validation fail - RecordSchema, is_schema_correct = record_schema - - if is_schema_correct: - records = service.list_records( - entity_key=str(entity_key), schema=RecordSchema, start=0, limit=1 - ) - - sent_request = httpx_mock.get_request() - if sent_request is None: - raise Exception("No request was sent") - - assert isinstance(records, list) - assert len(records) == 2 - assert records[0].id == "12345" - assert records[0].name == "record_name" - assert records[0].integer_field == 10 - assert records[1].id == "12346" - assert records[1].name == "record_name2" - assert records[1].integer_field == 11 - else: - # Validation should fail and raise an exception - with pytest.raises((ValueError, TypeError)): - service.list_records( - entity_key=str(entity_key), schema=RecordSchema, start=0, limit=1 - ) - - # Schema validation should take into account optional fields - def test_retrieve_records_with_optional_fields( - self, - httpx_mock: HTTPXMock, - service: EntitiesService, - base_url: str, - org: str, - tenant: str, - version: str, - record_schema_optional, - ) -> None: - entity_key = uuid.uuid4() - httpx_mock.add_response( - url=f"{base_url}{org}{tenant}/datafabric_/api/EntityService/entity/{str(entity_key)}/read?start=0&limit=1", - status_code=200, - json={ - "totalCount": 1, - "value": [ - { - "Id": "12345", - "name": "record_name", - }, - { - "Id": "12346", - "name": "record_name2", - }, - ], - }, - ) - - RecordSchemaOptional, is_field_optional = record_schema_optional - - if is_field_optional: - records = service.list_records( - entity_key=str(entity_key), - schema=RecordSchemaOptional, - start=0, - limit=1, - ) - - sent_request = httpx_mock.get_request() - if sent_request is None: - raise Exception("No request was sent") - - assert isinstance(records, list) - assert len(records) == 2 - assert records[0].id == "12345" - assert records[0].name == "record_name" - assert records[1].id == "12346" - assert records[1].name == "record_name2" - else: - # Validation should fail and raise an exception for missing required field - with pytest.raises((ValueError, TypeError)): - service.list_records( - entity_key=str(entity_key), - schema=RecordSchemaOptional, - start=0, - limit=1, - ) diff --git a/uv.lock b/uv.lock index 9426a50ef..6632bae9b 100644 --- a/uv.lock +++ b/uv.lock @@ -2368,6 +2368,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/46/2c/1462b1d0a634697ae9e55b3cecdcb64788e8b7d63f54d923fcd0bb140aed/soupsieve-2.8.3-py3-none-any.whl", hash = "sha256:ed64f2ba4eebeab06cc4962affce381647455978ffc1e36bb79a545b91f45a95", size = 37016, upload-time = "2026-01-20T04:27:01.012Z" }, ] +[[package]] +name = "sqlparse" +version = "0.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" }, +] + [[package]] name = "stevedore" version = "5.6.0" @@ -2548,6 +2557,7 @@ dependencies = [ { name = "python-dotenv" }, { name = "python-socketio" }, { name = "rich" }, + { name = "sqlparse" }, { name = "tenacity" }, { name = "truststore" }, { name = "uipath-core" }, @@ -2600,6 +2610,7 @@ requires-dist = [ { name = "python-dotenv", specifier = ">=1.0.1" }, { name = "python-socketio", specifier = ">=5.15.0,<6.0.0" }, { name = "rich", specifier = ">=14.2.0" }, + { name = "sqlparse", specifier = ">=0.4.4" }, { name = "tenacity", specifier = ">=9.0.0" }, { name = "truststore", specifier = ">=0.10.1" }, { name = "uipath-core", specifier = ">=0.3.0,<0.4.0" }, From ccac3cfd1fa5eb61f7902163b6b035dadf9b06b3 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Tue, 17 Feb 2026 10:11:01 +0530 Subject: [PATCH 4/9] Remvoed uipath from dev-deps --- pyproject.toml | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7d355f7dd..74ce889d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,8 +75,7 @@ dev = [ "mkdocs-llmstxt>=0.5.0", "inflection>=0.5.1", "types-toml>=0.10.8", - "pytest-timeout>=2.4.0", - "uipath", + "pytest-timeout>=2.4.0" ] [tool.hatch.build.targets.wheel] @@ -150,6 +149,3 @@ name = "testpypi" url = "https://test.pypi.org/simple/" publish-url = "https://test.pypi.org/legacy/" explicit = true - -[tool.uv.sources] -uipath = { workspace = true } From e32369f44f2bb05611f7b414b5581956760cb3b2 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Tue, 17 Feb 2026 10:13:06 +0530 Subject: [PATCH 5/9] Reverted an accidental removal --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 74ce889d0..326e5d740 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,7 @@ dev = [ "mkdocs-llmstxt>=0.5.0", "inflection>=0.5.1", "types-toml>=0.10.8", - "pytest-timeout>=2.4.0" + "pytest-timeout>=2.4.0", ] [tool.hatch.build.targets.wheel] From 5bc48d5d0aaa2c9c6fa4da4bdd43c4785ca57d0f Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Tue, 17 Feb 2026 10:18:01 +0530 Subject: [PATCH 6/9] Updated response to handler failure cases --- src/uipath/platform/entities/_entities_service.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/uipath/platform/entities/_entities_service.py b/src/uipath/platform/entities/_entities_service.py index 7b631719a..8b26f3d2b 100644 --- a/src/uipath/platform/entities/_entities_service.py +++ b/src/uipath/platform/entities/_entities_service.py @@ -465,8 +465,11 @@ def query_multiple_entities( full_url = f"{self._url.base_url}{spec.endpoint}" response = self.request(spec.method, full_url, json=spec.json, headers=headers) - records_data = response.json().get("results", []) - return records_data + if response.status_code == 200: + records_data = response.json().get("results", []) + return records_data + else: + response.raise_for_status() @traced(name="query_multiple_entities_async", run_type="uipath") @@ -510,8 +513,11 @@ async def query_multiple_entities_async( full_url = f"{self._url.base_url}{spec.endpoint}" response = await self.request_async(spec.method, full_url, json=spec.json, headers=headers) - records_data = response.json().get("results", []) - return records_data + if response.status_code == 200: + records_data = response.json().get("results", []) + return records_data + else: + response.raise_for_status() @traced(name="entity_record_insert_batch", run_type="uipath") def insert_records( From 2f9c140752b30b2cf1d87f493cf8feaf693d2af3 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Thu, 5 Mar 2026 12:05:16 +0530 Subject: [PATCH 7/9] feat: restore Data Fabric context changes after main merge --- .../platform/entities/_entities_service.py | 101 ++++++- packages/uipath-platform/uv.lock | 30 ++- .../uipath/src/uipath/agent/models/agent.py | 5 + .../uipath/tests/agent/models/test_agent.py | 247 +++--------------- packages/uipath/uv.lock | 60 ++++- 5 files changed, 219 insertions(+), 224 deletions(-) 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 90a6ca1cb..039c431ec 100644 --- a/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py +++ b/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py @@ -1,4 +1,4 @@ -from typing import Any, List, Optional, Type +from typing import Any, Dict, List, Optional, Type from httpx import Response from uipath.core.tracing import traced @@ -391,6 +391,94 @@ class CustomerRecord: EntityRecord.from_data(data=record, model=schema) for record in records_data ] + @traced(name="entity_query_records", run_type="uipath") + def query_entity_records( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = entities_service.query_entity_records( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = entities_service.query_entity_records( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + # Use absolute URL to bypass scoping since org/tenant are embedded in the path + full_url = f"{self._url.base_url}{spec.endpoint}" + response = self.request(spec.method, full_url, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + return records_data + + + @traced(name="query_entities_async", run_type="uipath") + async def query_entity_records_async( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Asynchronously query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = await entities_service.query_entity_records_async( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = await entities_service.query_entity_records_async( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + full_url = f"{self._url.base_url}{spec.endpoint}" + response = await self.request_async(spec.method, full_url, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + return records_data + @traced(name="entity_record_insert_batch", run_type="uipath") def insert_records( self, @@ -874,6 +962,17 @@ def _list_records_spec( params=({"start": start, "limit": limit}), ) + def _query_entity_records_spec( + self, + sql_query: str, + ) -> RequestSpec: + endpoint = f"/dataservice_/{self._url.org_name}/{self._url.tenant_name}/datafabric_/api/v1/query/execute" + return RequestSpec( + method="POST", + endpoint=Endpoint(endpoint), + json={"query": sql_query}, + ) + def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec: return RequestSpec( method="POST", diff --git a/packages/uipath-platform/uv.lock b/packages/uipath-platform/uv.lock index af0f80415..31af4ced8 100644 --- a/packages/uipath-platform/uv.lock +++ b/packages/uipath-platform/uv.lock @@ -1047,16 +1047,34 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.3" -source = { registry = "https://pypi.org/simple" } +version = "0.5.6" +source = { editable = "../uipath-core" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2f/76/568bbe81e2c502b0b3d34b35f0f2d7557ceed58fc9161820d186276b47ac/uipath_core-0.5.3.tar.gz", hash = "sha256:5ff386c9bf85006648f111496b74534925fab1de4b35d5d0c2f6dfdf81e6e103", size = 119096, upload-time = "2026-02-25T14:08:47.548Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1c/35/87a346abe7485c0a63802487050e3550723bfd97925f85cc8814d34bb2a3/uipath_core-0.5.3-py3-none-any.whl", hash = "sha256:2ad9670d3d8e62d7e4f5ed090dffeff00281b8d20d159fff67cac941889d6748", size = 42858, upload-time = "2026-02-25T14:08:46.037Z" }, + +[package.metadata] +requires-dist = [ + { name = "opentelemetry-instrumentation", specifier = ">=0.60b0,<1.0.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.39.0,<2.0.0" }, + { name = "pydantic", specifier = ">=2.12.5,<3.0.0" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.14.1" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]] @@ -1092,7 +1110,7 @@ requires-dist = [ { name = "pydantic-function-models", specifier = ">=0.1.11" }, { name = "tenacity", specifier = ">=9.0.0" }, { name = "truststore", specifier = ">=0.10.1" }, - { name = "uipath-core", specifier = ">=0.5.3,<0.6.0" }, + { name = "uipath-core", editable = "../uipath-core" }, ] [package.metadata.requires-dev] diff --git a/packages/uipath/src/uipath/agent/models/agent.py b/packages/uipath/src/uipath/agent/models/agent.py index bf89c6de0..c39d042c6 100644 --- a/packages/uipath/src/uipath/agent/models/agent.py +++ b/packages/uipath/src/uipath/agent/models/agent.py @@ -142,6 +142,7 @@ class AgentContextRetrievalMode(str, CaseInsensitiveEnum): STRUCTURED = "Structured" DEEP_RAG = "DeepRAG" BATCH_TRANSFORM = "BatchTransform" + DATA_FABRIC = "DataFabric" UNKNOWN = "Unknown" # fallback branch discriminator @@ -376,6 +377,9 @@ class AgentContextSettings(BaseCfg): output_columns: Optional[List[AgentContextOutputColumn]] = Field( None, alias="outputColumns" ) + entity_identifiers: Optional[List[str]] = Field( + None, alias="entityIdentifiers" + ) class AgentContextResourceConfig(BaseAgentResourceConfig): @@ -1199,6 +1203,7 @@ def _normalize_resources(v: Dict[str, Any]) -> None: "structured": "Structured", "deeprag": "DeepRAG", "batchtransform": "BatchTransform", + "datafabric": "DataFabric", "unknown": "Unknown", } diff --git a/packages/uipath/tests/agent/models/test_agent.py b/packages/uipath/tests/agent/models/test_agent.py index 7b6dababa..0eaf43e12 100644 --- a/packages/uipath/tests/agent/models/test_agent.py +++ b/packages/uipath/tests/agent/models/test_agent.py @@ -1,5 +1,3 @@ -from typing import Any - import pytest from pydantic import TypeAdapter @@ -2480,8 +2478,6 @@ def test_agent_with_ixp_vs_escalation(self): "outcomeMapping": None, "recipients": [], "type": "actionCenter", - "taskTitle": "Test IXP Escalation", - "priority": "High", "properties": { "appName": None, "appVersion": 1, @@ -3221,212 +3217,51 @@ def test_is_conversational_false_by_default(self): assert config.is_conversational is False -class TestAgentBuilderConfigResources: - """Tests for AgentDefinition resource configuration parsing.""" +class TestDataFabricContextConfig: + """Tests for Data Fabric context resource configuration.""" - def _agent_dict_with_resources(self, resources: list[Any]) -> dict[str, Any]: - """Helper method that returns an agent dict with default fields and provided resources.""" - return { - "version": "1.0.0", - "id": "test-agent-id", - "name": "Test Agent", - "metadata": {"isConversational": False, "storageVersion": "22.0.0"}, - "messages": [ - { - "role": "System", - "content": "You are a test agent.", - } - ], - "inputSchema": {"type": "object", "properties": {}}, - "outputSchema": {"type": "object", "properties": {}}, + def test_datafabric_retrieval_mode_exists(self): + """Test that DATA_FABRIC retrieval mode is defined.""" + assert AgentContextRetrievalMode.DATA_FABRIC == "DataFabric" + + def test_datafabric_context_config_parses(self): + """Test that Data Fabric context config parses correctly.""" + config = { + "$resourceType": "context", + "name": "Customer Data", + "description": "Query customer and order data", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", "settings": { - "model": "gpt-4o", - "maxTokens": 4096, - "temperature": 0, - "engine": "basic-v1", + "retrievalMode": "DataFabric", + "resultCount": 100, + "entityIdentifiers": ["customers-key", "orders-key"], }, - "resources": resources, } - def test_escalation_with_static_group_name_recipient_type(self): - """Test that escalation with StaticGroupName recipientType is parsed correctly.""" - resources = [ - { - "$resourceType": "escalation", - "id": "escalation-1", - "channels": [ - { - "name": "Test Channel", - "description": "Test channel description", - "type": "ActionCenter", - "inputSchema": { - "type": "object", - "properties": {"field": {"type": "string"}}, - }, - "outputSchema": {"type": "object", "properties": {}}, - "outcomeMapping": {"Approve": "continue"}, - "properties": { - "appName": "TestApp", - "appVersion": 1, - "folderName": "TestFolder", - "resourceKey": "test-key", - }, - "recipients": [ - { - "value": "TestGroup", - "type": "staticgroupname", - } - ], - "taskTitle": "Test Task", - "priority": "Medium", - } - ], - "isAgentMemoryEnabled": False, - "escalationType": 0, - "name": "Test Escalation", - "description": "Test escalation", - } - ] - - json_data = self._agent_dict_with_resources(resources) - config: AgentDefinition = TypeAdapter(AgentDefinition).validate_python( - json_data - ) - - escalation_resource = config.resources[0] - assert isinstance(escalation_resource, AgentEscalationResourceConfig) - recipient = escalation_resource.channels[0].recipients[0] - assert isinstance(recipient, StandardRecipient) - assert recipient.type == AgentEscalationRecipientType.GROUP_NAME - assert recipient.value == "TestGroup" - - def test_escalation_with_lowercase_userid_recipient_type(self): - """Test that escalation with lowercase userid recipientType is parsed correctly.""" - resources = [ - { - "$resourceType": "escalation", - "id": "escalation-2", - "channels": [ - { - "name": "Test Channel", - "description": "Test channel description", - "type": "ActionCenter", - "inputSchema": { - "type": "object", - "properties": {"field": {"type": "string"}}, - }, - "outputSchema": {"type": "object", "properties": {}}, - "outcomeMapping": {"Approve": "continue"}, - "properties": { - "appName": "TestApp", - "appVersion": 1, - "folderName": "TestFolder", - "resourceKey": "test-key", - }, - "recipients": [ - { - "value": "user-123", - "type": "userid", - } - ], - "taskTitle": "Test Task", - "priority": "Medium", - } - ], - "isAgentMemoryEnabled": False, - "escalationType": 0, - "name": "Test Escalation", - "description": "Test escalation", - } - ] - - json_data = self._agent_dict_with_resources(resources) - config: AgentDefinition = TypeAdapter(AgentDefinition).validate_python( - json_data - ) - - escalation_resource = config.resources[0] - assert isinstance(escalation_resource, AgentEscalationResourceConfig) - recipient = escalation_resource.channels[0].recipients[0] - assert isinstance(recipient, StandardRecipient) - assert recipient.type == AgentEscalationRecipientType.USER_ID - assert recipient.value == "user-123" - - def test_process_tool_missing_output_schema(self): - """Test that process tool without outputSchema is parsed correctly.""" - resources = [ - { - "$resourceType": "tool", - "type": "ProcessOrchestration", - "id": "process-tool-1", - "inputSchema": { - "type": "object", - "properties": {"input": {"type": "string"}}, - }, - "arguments": {}, - "settings": {"timeout": 0, "maxAttempts": 0, "retryDelay": 0}, - "properties": { - "processName": "TestProcess", - "folderPath": "TestFolder", - }, - "name": "Test Process", - "description": "Test process tool", - } - ] - - json_data = self._agent_dict_with_resources(resources) - config: AgentDefinition = TypeAdapter(AgentDefinition).validate_python( - json_data - ) - - tool_resource = config.resources[0] - assert isinstance(tool_resource, AgentProcessToolResourceConfig) - assert tool_resource.output_schema == {"type": "object", "properties": {}} - - def test_escalation_missing_escalation_type_defaults_to_zero(self): - """Test that missing escalationType defaults to 0.""" - resources = [ - { - "$resourceType": "escalation", - "id": "escalation-3", - "channels": [ - { - "name": "Test Channel", - "description": "Test channel description", - "type": "ActionCenter", - "inputSchema": { - "type": "object", - "properties": {"field": {"type": "string"}}, - }, - "outputSchema": {"type": "object", "properties": {}}, - "outcomeMapping": {"Approve": "continue"}, - "properties": { - "appName": "TestApp", - "appVersion": 1, - "folderName": "TestFolder", - "resourceKey": "test-key", - }, - "recipients": [ - { - "value": "user-123", - "type": "UserId", - } - ], - "taskTitle": "Test Task", - "priority": "Medium", - } - ], - "isAgentMemoryEnabled": False, - "name": "Test Escalation", - "description": "Test escalation", - } - ] + parsed = AgentContextResourceConfig.model_validate(config) + + assert parsed.name == "Customer Data" + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers == ["customers-key", "orders-key"] + + def test_datafabric_context_config_without_entity_identifiers(self): + """Test that entity_identifiers is optional.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "Test", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", + "settings": { + "retrievalMode": "DataFabric", + "resultCount": 10, + }, + } - json_data = self._agent_dict_with_resources(resources) - config: AgentDefinition = TypeAdapter(AgentDefinition).validate_python( - json_data - ) + parsed = AgentContextResourceConfig.model_validate(config) - escalation_resource = config.resources[0] - assert isinstance(escalation_resource, AgentEscalationResourceConfig) - assert escalation_resource.escalation_type == 0 + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers is None diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 3c772cef2..4fb270842 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2602,8 +2602,8 @@ requires-dist = [ { name = "rich", specifier = ">=14.2.0" }, { name = "tenacity", specifier = ">=9.0.0" }, { name = "truststore", specifier = ">=0.10.1" }, - { name = "uipath-core", specifier = ">=0.5.2,<0.6.0" }, - { name = "uipath-platform", specifier = ">=0.0.4,<0.1.0" }, + { name = "uipath-core", editable = "../uipath-core" }, + { name = "uipath-platform", editable = "../uipath-platform" }, { name = "uipath-runtime", specifier = ">=0.9.1,<0.10.0" }, ] @@ -2638,22 +2638,40 @@ dev = [ [[package]] name = "uipath-core" -version = "0.5.5" -source = { registry = "https://pypi.org/simple" } +version = "0.5.6" +source = { editable = "../uipath-core" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/de/01/4eac9fb79b0396ef10b099af7afc8c71d8081ab33538543bb7d7df579412/uipath_core-0.5.5.tar.gz", hash = "sha256:31fbba640b8b3e128e9cb4c1fd8c9cc492230cb9998397895f40f2b755590f96", size = 112378, upload-time = "2026-03-04T15:47:19.645Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/fa/1ce0c97947b47b50d3438edbec174f6d7e64596ff58f38fb4c1aa9663a55/uipath_core-0.5.5-py3-none-any.whl", hash = "sha256:6660b7c9c4ab15433d4df325a31a12d47418206b79e0d05640e73c3e74895f7a", size = 42026, upload-time = "2026-03-04T15:47:18.424Z" }, + +[package.metadata] +requires-dist = [ + { name = "opentelemetry-instrumentation", specifier = ">=0.60b0,<1.0.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.39.0,<2.0.0" }, + { name = "pydantic", specifier = ">=2.12.5,<3.0.0" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.14.1" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]] name = "uipath-platform" version = "0.0.11" -source = { registry = "https://pypi.org/simple" } +source = { editable = "../uipath-platform" } dependencies = [ { name = "httpx" }, { name = "pydantic-function-models" }, @@ -2661,9 +2679,29 @@ dependencies = [ { name = "truststore" }, { name = "uipath-core" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b4/b4/82ca80eb2ad17fa3ffcd140d59e5527efe884323fdf09f7adfccada73759/uipath_platform-0.0.11.tar.gz", hash = "sha256:5f4d20ab4407ae4e3846df5c3b938b7d32382362e160db459bc2038c21d8f798", size = 261785, upload-time = "2026-03-04T15:47:16.387Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/07/a1/f263ea9117d7aedb8d667055f8ffa6520db7762a29d18e5e56828cf6711b/uipath_platform-0.0.11-py3-none-any.whl", hash = "sha256:5d403341dd83b6eb432f4988861031c63cee95781a505c0b59a71bf75003efa0", size = 158043, upload-time = "2026-03-04T15:47:14.709Z" }, + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.28.1" }, + { name = "pydantic-function-models", specifier = ">=0.1.11" }, + { name = "tenacity", specifier = ">=9.0.0" }, + { name = "truststore", specifier = ">=0.10.1" }, + { name = "uipath-core", editable = "../uipath-core" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.19.0" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]] From edf56edb745470e7de7fc8413b79f1ea49288bfb Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Thu, 5 Mar 2026 12:09:30 +0530 Subject: [PATCH 8/9] fix: align entity SQL query APIs with merged feature branch tests --- .../platform/entities/_entities_service.py | 182 ++++++++++++------ 1 file changed, 120 insertions(+), 62 deletions(-) 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 039c431ec..7942f5e74 100644 --- a/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py +++ b/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py @@ -1,3 +1,4 @@ +import re from typing import Any, Dict, List, Optional, Type from httpx import Response @@ -13,6 +14,30 @@ EntityRecordsBatchResponse, ) +_FORBIDDEN_SQL_KEYWORDS = { + "INSERT", + "UPDATE", + "DELETE", + "MERGE", + "DROP", + "ALTER", + "CREATE", + "TRUNCATE", + "REPLACE", +} +_DISALLOWED_SQL_OPERATORS = [ + "WITH", + "UNION", + "INTERSECT", + "EXCEPT", + "OVER", + "ROLLUP", + "CUBE", + "GROUPING SETS", + "PARTITION BY", +] +_SQL_KEYWORD_PATTERN = re.compile(r"\b([A-Z]+(?:\s+BY|\s+SETS)?)\b") + class EntitiesService(BaseService): """Service for managing UiPath Data Service entities. @@ -397,87 +422,60 @@ def query_entity_records( sql_query: str, schema: Optional[Type[Any]] = None, ) -> List[Dict[str, Any]]: - """Query entity records using a SQL query. - - This method allows executing SQL queries directly against entity data - via the Data Fabric query endpoint. - - Args: - sql_query (str): The full SQL query to execute. Should be a valid - SELECT statement targeting the entity. - schema (Optional[Type[Any]]): Optional schema class for validation. + """Backward-compatible alias for query_multiple_entities.""" + return self.query_multiple_entities(sql_query, schema) - Returns: - List[Dict[str, Any]]: A list of record dictionaries matching the query. - - Examples: - Basic query:: - records = entities_service.query_entity_records( - "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" - ) - - Query with specific fields:: + @traced(name="query_entities_async", run_type="uipath") + async def query_entity_records_async( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Backward-compatible alias for query_multiple_entities_async.""" + return await self.query_multiple_entities_async(sql_query, schema) - records = entities_service.query_entity_records( - "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" - ) - """ - spec = self._query_entity_records_spec(sql_query) + @traced(name="query_multiple_entities", run_type="uipath") + def query_multiple_entities( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Query entity records using a validated SQL query.""" + self._validate_sql_query(sql_query) + spec = self._query_multiple_entities_spec(sql_query) headers = { "X-UiPath-Internal-TenantName": self._url.tenant_name, "X-UiPath-Internal-AccountName": self._url.org_name, } - # Use absolute URL to bypass scoping since org/tenant are embedded in the path full_url = f"{self._url.base_url}{spec.endpoint}" response = self.request(spec.method, full_url, json=spec.json, headers=headers) + status_code = getattr(response, "status_code", 200) + if isinstance(status_code, int) and status_code >= 400: + response.raise_for_status() + return response.json().get("results", []) - records_data = response.json().get("results", []) - return records_data - - - @traced(name="query_entities_async", run_type="uipath") - async def query_entity_records_async( + @traced(name="query_multiple_entities_async", run_type="uipath") + async def query_multiple_entities_async( self, sql_query: str, schema: Optional[Type[Any]] = None, ) -> List[Dict[str, Any]]: - """Asynchronously query entity records using a SQL query. - - This method allows executing SQL queries directly against entity data - via the Data Fabric query endpoint. - - Args: - sql_query (str): The full SQL query to execute. Should be a valid - SELECT statement targeting the entity. - schema (Optional[Type[Any]]): Optional schema class for validation. - - Returns: - List[Dict[str, Any]]: A list of record dictionaries matching the query. - - Examples: - Basic query:: - - records = await entities_service.query_entity_records_async( - "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" - ) - - Query with specific fields:: - - records = await entities_service.query_entity_records_async( - "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" - ) - """ - spec = self._query_entity_records_spec(sql_query) + """Asynchronously query entity records using a validated SQL query.""" + self._validate_sql_query(sql_query) + spec = self._query_multiple_entities_spec(sql_query) headers = { "X-UiPath-Internal-TenantName": self._url.tenant_name, "X-UiPath-Internal-AccountName": self._url.org_name, } full_url = f"{self._url.base_url}{spec.endpoint}" - response = await self.request_async(spec.method, full_url, json=spec.json, headers=headers) - - records_data = response.json().get("results", []) - return records_data + response = await self.request_async( + spec.method, full_url, json=spec.json, headers=headers + ) + status_code = getattr(response, "status_code", 200) + if isinstance(status_code, int) and status_code >= 400: + response.raise_for_status() + return response.json().get("results", []) @traced(name="entity_record_insert_batch", run_type="uipath") def insert_records( @@ -973,6 +971,12 @@ def _query_entity_records_spec( json={"query": sql_query}, ) + def _query_multiple_entities_spec( + self, + sql_query: str, + ) -> RequestSpec: + return self._query_entity_records_spec(sql_query) + def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec: return RequestSpec( method="POST", @@ -1001,3 +1005,57 @@ def _delete_batch_spec(self, entity_key: str, record_ids: List[str]) -> RequestS ), json=record_ids, ) + + def _validate_sql_query(self, sql_query: str) -> None: + query = sql_query.strip() + if not query: + raise ValueError("SQL query cannot be empty.") + + if ";" in query.rstrip(";"): + raise ValueError("Only a single SELECT statement is allowed.") + + normalized_query = re.sub(r"\s+", " ", query).strip() + normalized_upper = normalized_query.upper() + normalized_keywords = set(_SQL_KEYWORD_PATTERN.findall(normalized_upper)) + + # Keep legacy behavior: non-SELECT statements fail with this message, + # while CTE-based queries are reported as disallowed WITH usage. + if not normalized_upper.startswith("SELECT "): + if not normalized_upper.startswith("WITH "): + raise ValueError("Only SELECT statements are allowed.") + + for keyword in _FORBIDDEN_SQL_KEYWORDS: + if keyword in normalized_keywords: + raise ValueError(f"SQL keyword '{keyword}' is not allowed.") + + for operator in _DISALLOWED_SQL_OPERATORS: + if operator in normalized_keywords: + raise ValueError( + f"SQL construct '{operator}' is not allowed in entity queries." + ) + + if re.search(r"\(\s*SELECT\b", normalized_upper): + raise ValueError("Subqueries are not allowed.") + + has_where = bool(re.search(r"\bWHERE\b", normalized_upper)) + has_limit = bool(re.search(r"\bLIMIT\s+\d+\b", normalized_upper)) + if not has_where and not has_limit: + raise ValueError("Queries without WHERE must include a LIMIT clause.") + + projection = self._projection_segment(normalized_query) + if "*" in projection and not has_where: + raise ValueError("SELECT * without filtering is 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." + ) + + def _projection_segment(self, normalized_query: str) -> str: + match = re.match(r"(?is)\s*SELECT\s+(.*?)\s+FROM\s+", normalized_query) + return match.group(1) if match else "" + + def _projection_column_count(self, projection: str) -> int: + cleaned = projection.strip() + if not cleaned: + return 0 + return len([part for part in cleaned.split(",") if part.strip()]) From ed94fe141fb1d4d7f8925d030f63458425211bd1 Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Wed, 11 Mar 2026 16:10:01 +0530 Subject: [PATCH 9/9] Updated the entities_service implementation. Introduced a new context type called DATA_FABRIC --- .../platform/entities/_entities_service.py | 108 +++++++++--------- .../tests/services/test_entities_service.py | 61 ++++++---- packages/uipath-platform/uv.lock | 28 ++++- .../uipath/src/uipath/agent/models/agent.py | 60 +++++++--- .../uipath/tests/agent/models/test_agent.py | 108 +++++++++++++----- packages/uipath/uv.lock | 58 ++++++++-- 6 files changed, 283 insertions(+), 140 deletions(-) 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 7942f5e74..e87b8d8f4 100644 --- a/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py +++ b/packages/uipath-platform/src/uipath/platform/entities/_entities_service.py @@ -37,6 +37,7 @@ "PARTITION BY", ] _SQL_KEYWORD_PATTERN = re.compile(r"\b([A-Z]+(?:\s+BY|\s+SETS)?)\b") +_QUOTED_STRING_PATTERN = re.compile(r"'[^']*'") class EntitiesService(BaseService): @@ -420,61 +421,60 @@ class CustomerRecord: def query_entity_records( self, sql_query: str, - schema: Optional[Type[Any]] = None, ) -> List[Dict[str, Any]]: - """Backward-compatible alias for query_multiple_entities.""" - return self.query_multiple_entities(sql_query, schema) + """Query entity records using a validated SQL query. + + PREVIEW: This method is in preview and may change in future releases. + Args: + 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. + + Returns: + List[Dict[str, Any]]: A list of result records as dictionaries. - @traced(name="query_entities_async", run_type="uipath") + Raises: + 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) + + @traced(name="entity_query_records", run_type="uipath") async def query_entity_records_async( self, sql_query: str, - schema: Optional[Type[Any]] = None, ) -> List[Dict[str, Any]]: - """Backward-compatible alias for query_multiple_entities_async.""" - return await self.query_multiple_entities_async(sql_query, schema) + """Asynchronously query entity records using a validated SQL query. - @traced(name="query_multiple_entities", run_type="uipath") - def query_multiple_entities( - self, - sql_query: str, - schema: Optional[Type[Any]] = 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. + + Args: + 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. + + Returns: + List[Dict[str, Any]]: A list of result records as dictionaries. + + Raises: + 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) + + def _query_entities_for_records(self, sql_query: str) -> List[Dict[str, Any]]: self._validate_sql_query(sql_query) - spec = self._query_multiple_entities_spec(sql_query) - headers = { - "X-UiPath-Internal-TenantName": self._url.tenant_name, - "X-UiPath-Internal-AccountName": self._url.org_name, - } - full_url = f"{self._url.base_url}{spec.endpoint}" - response = self.request(spec.method, full_url, json=spec.json, headers=headers) - status_code = getattr(response, "status_code", 200) - if isinstance(status_code, int) and status_code >= 400: - response.raise_for_status() + spec = self._query_entity_records_spec(sql_query) + response = self.request(spec.method, spec.endpoint, json=spec.json) return response.json().get("results", []) - @traced(name="query_multiple_entities_async", run_type="uipath") - async def query_multiple_entities_async( - self, - sql_query: str, - schema: Optional[Type[Any]] = None, + async def _query_entities_for_records_async( + self, sql_query: str ) -> List[Dict[str, Any]]: - """Asynchronously query entity records using a validated SQL query.""" self._validate_sql_query(sql_query) - spec = self._query_multiple_entities_spec(sql_query) - headers = { - "X-UiPath-Internal-TenantName": self._url.tenant_name, - "X-UiPath-Internal-AccountName": self._url.org_name, - } - full_url = f"{self._url.base_url}{spec.endpoint}" - response = await self.request_async( - spec.method, full_url, json=spec.json, headers=headers - ) - status_code = getattr(response, "status_code", 200) - if isinstance(status_code, int) and status_code >= 400: - response.raise_for_status() + spec = self._query_entity_records_spec(sql_query) + response = await self.request_async(spec.method, spec.endpoint, json=spec.json) return response.json().get("results", []) @traced(name="entity_record_insert_batch", run_type="uipath") @@ -964,19 +964,12 @@ def _query_entity_records_spec( self, sql_query: str, ) -> RequestSpec: - endpoint = f"/dataservice_/{self._url.org_name}/{self._url.tenant_name}/datafabric_/api/v1/query/execute" return RequestSpec( method="POST", - endpoint=Endpoint(endpoint), + endpoint=Endpoint("datafabric_/api/v1/query/execute"), json={"query": sql_query}, ) - def _query_multiple_entities_spec( - self, - sql_query: str, - ) -> RequestSpec: - return self._query_entity_records_spec(sql_query) - def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec: return RequestSpec( method="POST", @@ -1007,29 +1000,30 @@ def _delete_batch_spec(self, entity_key: str, record_ids: List[str]) -> RequestS ) def _validate_sql_query(self, sql_query: str) -> None: - query = sql_query.strip() + query = sql_query.strip().rstrip(";").strip() if not query: raise ValueError("SQL query cannot be empty.") - if ";" in query.rstrip(";"): + # Strip quoted strings before checking for semicolons so that + # values like WHERE name = 'foo;bar' don't trigger a false positive. + unquoted = _QUOTED_STRING_PATTERN.sub("''", query) + if ";" in unquoted: raise ValueError("Only a single SELECT statement is allowed.") normalized_query = re.sub(r"\s+", " ", query).strip() normalized_upper = normalized_query.upper() - normalized_keywords = set(_SQL_KEYWORD_PATTERN.findall(normalized_upper)) + extracted_keywords = set(_SQL_KEYWORD_PATTERN.findall(normalized_upper)) - # Keep legacy behavior: non-SELECT statements fail with this message, - # while CTE-based queries are reported as disallowed WITH usage. if not normalized_upper.startswith("SELECT "): if not normalized_upper.startswith("WITH "): raise ValueError("Only SELECT statements are allowed.") for keyword in _FORBIDDEN_SQL_KEYWORDS: - if keyword in normalized_keywords: + if keyword in extracted_keywords: raise ValueError(f"SQL keyword '{keyword}' is not allowed.") for operator in _DISALLOWED_SQL_OPERATORS: - if operator in normalized_keywords: + if operator in extracted_keywords: raise ValueError( f"SQL construct '{operator}' is not allowed in entity queries." ) diff --git a/packages/uipath-platform/tests/services/test_entities_service.py b/packages/uipath-platform/tests/services/test_entities_service.py index feb80a161..b7173e510 100644 --- a/packages/uipath-platform/tests/services/test_entities_service.py +++ b/packages/uipath-platform/tests/services/test_entities_service.py @@ -1,7 +1,7 @@ +import re import uuid from dataclasses import make_dataclass from typing import Optional -import re from unittest.mock import AsyncMock, MagicMock import pytest @@ -271,31 +271,49 @@ def test_retrieve_records_with_optional_fields( "SELECT * FROM Customers WHERE status = 'Active'", "SELECT id, name, email, phone FROM Customers LIMIT 5", "SELECT DISTINCT id FROM Customers WHERE id > 100", + "SELECT id FROM Customers WHERE name = 'foo;bar'", + "SELECT id FROM Customers WHERE id = 1;", ], ) def test_validate_sql_query_allows_supported_select_queries( - self, - sql_query: str, service: EntitiesService + self, sql_query: str, service: EntitiesService ) -> None: service._validate_sql_query(sql_query) - @pytest.mark.parametrize( "sql_query,error_message", [ ("", "SQL query cannot be empty."), (" ", "SQL query cannot be empty."), - ("SELECT id FROM Customers; SELECT id FROM Orders", "Only a single SELECT statement is allowed."), + ( + "SELECT id FROM Customers; SELECT id FROM Orders", + "Only a single SELECT statement is allowed.", + ), ("INSERT INTO Customers VALUES (1)", "Only SELECT statements are allowed."), ( "WITH cte AS (SELECT id FROM Customers) SELECT id FROM cte", "SQL construct 'WITH' is not allowed in entity queries.", ), - ("SELECT id FROM Customers UNION SELECT id FROM Orders", "SQL construct 'UNION' is not allowed in entity queries."), - ("SELECT id, SUM(amount) OVER (PARTITION BY id) FROM Orders LIMIT 10", "SQL construct 'OVER' is not allowed in entity queries."), - ("SELECT id FROM (SELECT id FROM Customers) c", "Subqueries are not allowed."), - ("SELECT id FROM Customers", "Queries without WHERE must include a LIMIT clause."), - ("SELECT * FROM Customers LIMIT 10", "SELECT * without filtering is not allowed."), + ( + "SELECT id FROM Customers UNION SELECT id FROM Orders", + "SQL construct 'UNION' is not allowed in entity queries.", + ), + ( + "SELECT id, SUM(amount) OVER (PARTITION BY id) FROM Orders LIMIT 10", + "SQL construct 'OVER' is not allowed in entity queries.", + ), + ( + "SELECT id FROM (SELECT id FROM Customers) c", + "Subqueries are not allowed.", + ), + ( + "SELECT id FROM Customers", + "Queries without WHERE must include a LIMIT clause.", + ), + ( + "SELECT * FROM Customers LIMIT 10", + "SELECT * without filtering is not allowed.", + ), ( "SELECT id, name, email, phone, address FROM Customers LIMIT 10", "Selecting more than 4 columns without filtering is not allowed.", @@ -303,14 +321,12 @@ def test_validate_sql_query_allows_supported_select_queries( ], ) def test_validate_sql_query_rejects_disallowed_queries( - self, - sql_query: str, error_message: str, service: EntitiesService + self, sql_query: str, error_message: str, service: EntitiesService ) -> None: with pytest.raises(ValueError, match=re.escape(error_message)): service._validate_sql_query(sql_query) - - def test_query_multiple_entities_rejects_invalid_sql_before_network_call( + def test_query_entity_records_rejects_invalid_sql_before_network_call( self, service: EntitiesService, ) -> None: @@ -319,12 +335,11 @@ def test_query_multiple_entities_rejects_invalid_sql_before_network_call( with pytest.raises( ValueError, match=re.escape("Only SELECT statements are allowed.") ): - service.query_multiple_entities("UPDATE Customers SET name = 'X'") + service.query_entity_records("UPDATE Customers SET name = 'X'") service.request.assert_not_called() # type: ignore[attr-defined] - - def test_query_multiple_entities_calls_request_for_valid_sql( + def test_query_entity_records_calls_request_for_valid_sql( self, service: EntitiesService, ) -> None: @@ -333,29 +348,27 @@ def test_query_multiple_entities_calls_request_for_valid_sql( service.request = MagicMock(return_value=response) # type: ignore[method-assign] - result = service.query_multiple_entities("SELECT id FROM Customers WHERE id > 0") + result = service.query_entity_records("SELECT id FROM Customers WHERE id > 0") assert result == [{"id": 1}, {"id": 2}] service.request.assert_called_once() # type: ignore[attr-defined] - @pytest.mark.anyio - async def test_query_multiple_entities_async_rejects_invalid_sql_before_network_call( + async def test_query_entity_records_async_rejects_invalid_sql_before_network_call( self, service: EntitiesService, ) -> None: service.request_async = AsyncMock() # type: ignore[method-assign] with pytest.raises(ValueError, match=re.escape("Subqueries are not allowed.")): - await service.query_multiple_entities_async( + await service.query_entity_records_async( "SELECT id FROM Customers WHERE id IN (SELECT id FROM Orders)" ) service.request_async.assert_not_called() # type: ignore[attr-defined] - @pytest.mark.anyio - async def test_query_multiple_entities_async_calls_request_for_valid_sql( + async def test_query_entity_records_async_calls_request_for_valid_sql( self, service: EntitiesService, ) -> None: @@ -364,7 +377,7 @@ async def test_query_multiple_entities_async_calls_request_for_valid_sql( service.request_async = AsyncMock(return_value=response) # type: ignore[method-assign] - result = await service.query_multiple_entities_async( + result = await service.query_entity_records_async( "SELECT id FROM Customers WHERE id = 'c1'" ) diff --git a/packages/uipath-platform/uv.lock b/packages/uipath-platform/uv.lock index 695f52305..a1c53c23b 100644 --- a/packages/uipath-platform/uv.lock +++ b/packages/uipath-platform/uv.lock @@ -1048,15 +1048,33 @@ wheels = [ [[package]] name = "uipath-core" version = "0.5.6" -source = { registry = "https://pypi.org/simple" } +source = { editable = "../uipath-core" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ea/8a/d129d33a81865f99d9134391a52f8691f557d95a18a38df4d88917b3e235/uipath_core-0.5.6.tar.gz", hash = "sha256:bebaf2e62111e844739e4f4e4dc47c48bac93b7e6fce6754502a9f4979c41888", size = 112659, upload-time = "2026-03-04T18:04:42.963Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/8f/77ab712518aa2a8485a558a0de245ac425e07fd8b74cfa8951550f0aea63/uipath_core-0.5.6-py3-none-any.whl", hash = "sha256:4a741fc760605165b0541b3abb6ade728bfa386e000ace00054bc43995720e5b", size = 42047, upload-time = "2026-03-04T18:04:41.606Z" }, + +[package.metadata] +requires-dist = [ + { name = "opentelemetry-instrumentation", specifier = ">=0.60b0,<1.0.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.39.0,<2.0.0" }, + { name = "pydantic", specifier = ">=2.12.5,<3.0.0" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.14.1" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]] @@ -1092,7 +1110,7 @@ requires-dist = [ { name = "pydantic-function-models", specifier = ">=0.1.11" }, { name = "tenacity", specifier = ">=9.0.0" }, { name = "truststore", specifier = ">=0.10.1" }, - { name = "uipath-core", specifier = ">=0.5.4,<0.6.0" }, + { name = "uipath-core", editable = "../uipath-core" }, ] [package.metadata.requires-dev] diff --git a/packages/uipath/src/uipath/agent/models/agent.py b/packages/uipath/src/uipath/agent/models/agent.py index ed1af51f4..c2c09533c 100644 --- a/packages/uipath/src/uipath/agent/models/agent.py +++ b/packages/uipath/src/uipath/agent/models/agent.py @@ -146,6 +146,14 @@ class AgentContextRetrievalMode(str, CaseInsensitiveEnum): UNKNOWN = "Unknown" # fallback branch discriminator +class AgentContextType(str, CaseInsensitiveEnum): + """Agent context type enumeration.""" + + INDEX = "index" + ATTACHMENTS = "attachments" + DATA_FABRIC_ENTITY_SET = "datafabricentityset" + + class AgentMessageRole(str, CaseInsensitiveEnum): """Agent message role enumeration.""" @@ -377,9 +385,16 @@ class AgentContextSettings(BaseCfg): output_columns: Optional[List[AgentContextOutputColumn]] = Field( None, alias="outputColumns" ) - entity_identifiers: Optional[List[str]] = Field( - None, alias="entityIdentifiers" - ) + + +class DataFabricEntityItem(BaseCfg): + """A single Data Fabric entity reference.""" + + id: str + reference_key: Optional[str] = Field(None, alias="referenceKey") + name: str + folder_id: str = Field(alias="folderId") + description: Optional[str] = None class AgentContextResourceConfig(BaseAgentResourceConfig): @@ -388,13 +403,29 @@ class AgentContextResourceConfig(BaseAgentResourceConfig): resource_type: Literal[AgentResourceType.CONTEXT] = Field( alias="$resourceType", default=AgentResourceType.CONTEXT, frozen=True ) - folder_path: str = Field(alias="folderPath") - index_name: str = Field(alias="indexName") - settings: AgentContextSettings = Field(..., description="Context settings") + context_type: Optional[AgentContextType] = Field(None, alias="contextType") + folder_path: Optional[str] = Field(None, alias="folderPath") + index_name: Optional[str] = Field(None, alias="indexName") + settings: Optional[AgentContextSettings] = Field( + None, description="Context settings" + ) + entity_set: Optional[List[DataFabricEntityItem]] = Field(None, alias="entitySet") argument_properties: Dict[str, AgentToolArgumentProperties] = Field( {}, alias="argumentProperties" ) + @property + def is_datafabric(self) -> bool: + """Check if this context is a Data Fabric entity set resource.""" + return self.context_type == AgentContextType.DATA_FABRIC_ENTITY_SET + + @property + def datafabric_entity_identifiers(self) -> list[str]: + """Extract entity identifiers from entitySet.""" + if self.entity_set: + return [item.id for item in self.entity_set] + return [] + class AgentMcpTool(BaseCfg): """Agent MCP tool model.""" @@ -1236,14 +1267,15 @@ def _normalize_resources(v: Dict[str, Any]) -> None: ) if res["$resourceType"] == "context": - settings = res.get("settings", {}) - rm = settings.get("retrievalMode") or settings.get("retrieval_mode") - settings["retrievalMode"] = ( - CONTEXT_MODE_MAP.get(rm.lower(), "Unknown") - if isinstance(rm, str) - else "Unknown" - ) - res["settings"] = settings + settings = res.get("settings") + if settings is not None: + rm = settings.get("retrievalMode") or settings.get("retrieval_mode") + settings["retrievalMode"] = ( + CONTEXT_MODE_MAP.get(rm.lower(), "Unknown") + if isinstance(rm, str) + else "Unknown" + ) + res["settings"] = settings out.append(res) diff --git a/packages/uipath/tests/agent/models/test_agent.py b/packages/uipath/tests/agent/models/test_agent.py index 0eaf43e12..93d916acb 100644 --- a/packages/uipath/tests/agent/models/test_agent.py +++ b/packages/uipath/tests/agent/models/test_agent.py @@ -7,6 +7,7 @@ AgentBuiltInValidatorGuardrail, AgentContextResourceConfig, AgentContextRetrievalMode, + AgentContextType, AgentCustomGuardrail, AgentDefinition, AgentEscalationChannel, @@ -1637,6 +1638,7 @@ def test_agent_with_unknown_context_retrieval_mode(self): context_resource = config.resources[0] assert isinstance(context_resource, AgentContextResourceConfig) assert context_resource.resource_type == AgentResourceType.CONTEXT + assert context_resource.settings is not None assert ( context_resource.settings.retrieval_mode == AgentContextRetrievalMode.UNKNOWN @@ -3041,6 +3043,7 @@ def test_case_insensitive_enums(self): context = config.resources[1] assert isinstance(context, AgentContextResourceConfig) assert context.resource_type == AgentResourceType.CONTEXT + assert context.settings is not None assert context.settings.retrieval_mode == AgentContextRetrievalMode.SEMANTIC assert context.settings.query is not None assert ( @@ -3220,48 +3223,95 @@ def test_is_conversational_false_by_default(self): class TestDataFabricContextConfig: """Tests for Data Fabric context resource configuration.""" - def test_datafabric_retrieval_mode_exists(self): - """Test that DATA_FABRIC retrieval mode is defined.""" - assert AgentContextRetrievalMode.DATA_FABRIC == "DataFabric" - def test_datafabric_context_config_parses(self): - """Test that Data Fabric context config parses correctly.""" + """Test v48 format with contextType and entitySet.""" config = { "$resourceType": "context", - "name": "Customer Data", - "description": "Query customer and order data", - "isEnabled": True, - "folderPath": "Shared", - "indexName": "", - "settings": { - "retrievalMode": "DataFabric", - "resultCount": 100, - "entityIdentifiers": ["customers-key", "orders-key"], - }, + "name": "TestDataFabric", + "description": "", + "contextType": "datafabricentityset", + "entitySet": [ + { + "id": "abc-123", + "name": "Customers", + "folderId": "folder-1", + "description": "Customer records", + }, + { + "id": "def-456", + "referenceKey": "orders-ref", + "name": "Orders", + "folderId": "folder-2", + "description": None, + }, + ], } parsed = AgentContextResourceConfig.model_validate(config) - assert parsed.name == "Customer Data" - assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC - assert parsed.settings.entity_identifiers == ["customers-key", "orders-key"] + assert parsed.context_type == AgentContextType.DATA_FABRIC_ENTITY_SET + assert parsed.folder_path is None + assert parsed.index_name is None + assert parsed.settings is None + assert parsed.entity_set is not None + assert len(parsed.entity_set) == 2 + assert parsed.entity_set[0].id == "abc-123" + assert parsed.entity_set[0].name == "Customers" + assert parsed.entity_set[0].folder_id == "folder-1" + assert parsed.entity_set[0].description == "Customer records" + assert parsed.entity_set[0].reference_key is None + assert parsed.entity_set[1].reference_key == "orders-ref" + assert parsed.entity_set[1].description is None + + def test_is_datafabric(self): + """Test is_datafabric property with datafabricentityset contextType.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "", + "contextType": "datafabricentityset", + "entitySet": [ + {"id": "abc-123", "name": "E", "folderId": "", "description": None} + ], + } + parsed = AgentContextResourceConfig.model_validate(config) + assert parsed.is_datafabric is True - def test_datafabric_context_config_without_entity_identifiers(self): - """Test that entity_identifiers is optional.""" + def test_is_datafabric_false_for_semantic(self): + """Test is_datafabric is False for Semantic contexts.""" config = { "$resourceType": "context", "name": "Test", - "description": "Test", - "isEnabled": True, + "description": "", "folderPath": "Shared", - "indexName": "", - "settings": { - "retrievalMode": "DataFabric", - "resultCount": 10, - }, + "indexName": "my-index", + "settings": {"retrievalMode": "Semantic", "resultCount": 5}, } + parsed = AgentContextResourceConfig.model_validate(config) + assert parsed.is_datafabric is False + def test_datafabric_entity_identifiers(self): + """Test entity identifiers extracted from entitySet.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "", + "contextType": "datafabricentityset", + "entitySet": [ + {"id": "id-1", "name": "E1", "folderId": "", "description": None}, + {"id": "id-2", "name": "E2", "folderId": "", "description": None}, + ], + } parsed = AgentContextResourceConfig.model_validate(config) + assert parsed.datafabric_entity_identifiers == ["id-1", "id-2"] - assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC - assert parsed.settings.entity_identifiers is None + def test_datafabric_entity_identifiers_empty(self): + """Test empty entity identifiers when no entitySet provided.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "", + "contextType": "datafabricentityset", + } + parsed = AgentContextResourceConfig.model_validate(config) + assert parsed.datafabric_entity_identifiers == [] diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 643c97c30..67fdbf0fd 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2591,7 +2591,6 @@ dev = [ { name = "termynal" }, { name = "tomli-w" }, { name = "types-toml" }, - { name = "uipath" }, { name = "virtualenv" }, ] @@ -2645,28 +2644,45 @@ dev = [ { name = "termynal", specifier = ">=0.13.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, { name = "types-toml", specifier = ">=0.10.8" }, - { name = "uipath", editable = "." }, { name = "virtualenv", specifier = ">=20.36.1" }, ] [[package]] name = "uipath-core" version = "0.5.6" -source = { registry = "https://pypi.org/simple" } +source = { editable = "../uipath-core" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ea/8a/d129d33a81865f99d9134391a52f8691f557d95a18a38df4d88917b3e235/uipath_core-0.5.6.tar.gz", hash = "sha256:bebaf2e62111e844739e4f4e4dc47c48bac93b7e6fce6754502a9f4979c41888", size = 112659, upload-time = "2026-03-04T18:04:42.963Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/8f/77ab712518aa2a8485a558a0de245ac425e07fd8b74cfa8951550f0aea63/uipath_core-0.5.6-py3-none-any.whl", hash = "sha256:4a741fc760605165b0541b3abb6ade728bfa386e000ace00054bc43995720e5b", size = 42047, upload-time = "2026-03-04T18:04:41.606Z" }, + +[package.metadata] +requires-dist = [ + { name = "opentelemetry-instrumentation", specifier = ">=0.60b0,<1.0.0" }, + { name = "opentelemetry-sdk", specifier = ">=1.39.0,<2.0.0" }, + { name = "pydantic", specifier = ">=2.12.5,<3.0.0" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.14.1" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]] name = "uipath-platform" -version = "0.0.17" -source = { registry = "https://pypi.org/simple" } +version = "0.0.18" +source = { editable = "../uipath-platform" } dependencies = [ { name = "httpx" }, { name = "pydantic-function-models" }, @@ -2674,9 +2690,29 @@ dependencies = [ { name = "truststore" }, { name = "uipath-core" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1e/c9/e7568133f3a657af16b67444c4e090259941078acc62acb1e2c072903da4/uipath_platform-0.0.17.tar.gz", hash = "sha256:a2c228462d7e2642dcfc249547d9b8e94ba1c72b68f16ba673ee3e58204e9365", size = 264143, upload-time = "2026-03-06T20:34:22.23Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/36/d1/a1c357dbea16a8b5d5b8ae5311ff2353cc03a8b5dd15ff83b0b693687930/uipath_platform-0.0.17-py3-none-any.whl", hash = "sha256:7b88f2b4eb189877fb2f99d704fc0cbc6e4244f01dac59cf812fa8a03db95e36", size = 159073, upload-time = "2026-03-06T20:34:20.993Z" }, + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.28.1" }, + { name = "pydantic-function-models", specifier = ">=0.1.11" }, + { name = "tenacity", specifier = ">=9.0.0" }, + { name = "truststore", specifier = ">=0.10.1" }, + { name = "uipath-core", editable = "../uipath-core" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "bandit", specifier = ">=1.8.2" }, + { name = "mypy", specifier = ">=1.19.0" }, + { name = "pre-commit", specifier = ">=4.1.0" }, + { name = "pytest", specifier = ">=7.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=4.1.0" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, + { name = "pytest-mock", specifier = ">=3.11.1" }, + { name = "pytest-trio", specifier = ">=0.8.0" }, + { name = "ruff", specifier = ">=0.9.4" }, + { name = "rust-just", specifier = ">=1.39.0" }, ] [[package]]