Entities moved to public-preview with warning. New Query API introduced. (#1432)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: UIPath-Harshit

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.0.18"
+version = "0.0.19"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
@@ -10,6 +10,7 @@ dependencies = [
   "truststore>=0.10.1",
   "uipath-core>=0.5.4, <0.6.0",
   "pydantic-function-models>=0.1.11",
+  "sqlparse>=0.5.5",
 ]
 classifiers = [
   "Intended Audience :: Developers",

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

@@ -1,6 +1,9 @@
-from typing import Any, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
+import sqlparse
 from httpx import Response
+from sqlparse.sql import Parenthesis, Where
+from sqlparse.tokens import DML, Keyword, Wildcard
 from uipath.core.tracing import traced
 
 from ..common._base_service import BaseService
@@ -13,6 +16,20 @@
     EntityRecordsBatchResponse,
 )
 
+_FORBIDDEN_DML = {"INSERT", "UPDATE", "DELETE", "MERGE", "REPLACE"}
+_FORBIDDEN_DDL = {"DROP", "ALTER", "CREATE", "TRUNCATE"}
+_DISALLOWED_KEYWORDS = [
+    "WITH",
+    "UNION",
+    "INTERSECT",
+    "EXCEPT",
+    "OVER",
+    "ROLLUP",
+    "CUBE",
+    "GROUPING",
+    "PARTITION",
+]
+
 
 class EntitiesService(BaseService):
     """Service for managing UiPath Data Service entities.
@@ -22,6 +39,10 @@ class EntitiesService(BaseService):
 
     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 are subject to change in future versions.
     """
 
     def __init__(
@@ -391,6 +412,66 @@ 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,
+    ) -> 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 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,
+    ) -> 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.
+
+        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_entity_records_spec(sql_query)
+        response = self.request(spec.method, spec.endpoint, json=spec.json)
+        return response.json().get("results", [])
+
+    async def _query_entities_for_records_async(
+        self, sql_query: str
+    ) -> List[Dict[str, Any]]:
+        self._validate_sql_query(sql_query)
+        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")
     def insert_records(
         self,
@@ -874,6 +955,16 @@ def _list_records_spec(
             params=({"start": start, "limit": limit}),
         )
 
+    def _query_entity_records_spec(
+        self,
+        sql_query: str,
+    ) -> RequestSpec:
+        return RequestSpec(
+            method="POST",
+            endpoint=Endpoint("datafabric_/api/v1/query/execute"),
+            json={"query": sql_query},
+        )
+
     def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec:
         return RequestSpec(
             method="POST",
@@ -902,3 +993,100 @@ 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().rstrip(";").strip()
+        if not query:
+            raise ValueError("SQL query cannot be empty.")
+
+        statements = sqlparse.parse(query)
+        if len(statements) != 1 or not statements[0].tokens:
+            raise ValueError("Only a single SELECT statement is allowed.")
+
+        stmt = statements[0]
+        stmt_type = stmt.get_type()
+
+        if stmt_type != "SELECT":
+            raise ValueError("Only SELECT statements are allowed.")
+
+        keywords = set()
+        for token in stmt.flatten():
+            if token.ttype in Keyword:
+                keywords.add(token.normalized)
+
+        for kw in _FORBIDDEN_DML:
+            if kw in keywords:
+                raise ValueError(f"SQL keyword '{kw}' is not allowed.")
+
+        for kw in _FORBIDDEN_DDL:
+            if kw in keywords:
+                raise ValueError(f"SQL keyword '{kw}' is not allowed.")
+
+        for kw in _DISALLOWED_KEYWORDS:
+            if kw in keywords:
+                raise ValueError(
+                    f"SQL construct '{kw}' is not allowed in entity queries."
+                )
+
+        if self._has_subquery(stmt):
+            raise ValueError("Subqueries are not allowed.")
+
+        has_where = any(isinstance(t, Where) for t in stmt.tokens)
+        has_limit = "LIMIT" in keywords
+        if not has_where and not has_limit:
+            raise ValueError("Queries without WHERE must include a LIMIT clause.")
+
+        projection = self._projection_tokens(stmt)
+        has_wildcard = any(t.ttype is Wildcard for t in projection)
+        if has_wildcard 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."
+            )
+
+    @staticmethod
+    def _has_subquery(stmt: sqlparse.sql.Statement) -> bool:
+        """Recursively walk the AST looking for SELECT inside parentheses."""
+
+        def _walk(token: sqlparse.sql.Token) -> bool:
+            if isinstance(token, Parenthesis):
+                for child in token.flatten():
+                    if child.ttype is DML and child.normalized == "SELECT":
+                        return True
+            if hasattr(token, "tokens"):
+                for child in token.tokens:
+                    if _walk(child):
+                        return True
+            return False
+
+        for token in stmt.tokens:
+            if _walk(token):
+                return True
+        return False
+
+    @staticmethod
+    def _projection_tokens(
+        stmt: sqlparse.sql.Statement,
+    ) -> list[sqlparse.sql.Token]:
+        """Extract tokens between the first SELECT and FROM."""
+        tokens: list[sqlparse.sql.Token] = []
+        collecting = False
+        for token in stmt.flatten():
+            if token.ttype is DML and token.normalized == "SELECT":
+                collecting = True
+                continue
+            if token.ttype is Keyword and token.normalized == "FROM":
+                break
+            if collecting:
+                tokens.append(token)
+        return tokens
+
+    @staticmethod
+    def _projection_column_count(
+        projection: list[sqlparse.sql.Token],
+    ) -> int:
+        text = "".join(t.value for t in projection).strip()
+        if not text:
+            return 0
+        return len([part for part in text.split(",") if part.strip()])

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

@@ -1,6 +1,8 @@
+import re
 import uuid
 from dataclasses import make_dataclass
 from typing import Optional
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from pytest_httpx import HTTPXMock
@@ -260,3 +262,130 @@ def test_retrieve_records_with_optional_fields(
                     start=0,
                     limit=1,
                 )
+
+    @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",
+            "SELECT id FROM Customers WHERE name = 'foo;bar'",
+            "SELECT id FROM Customers WHERE id = 1;",
+            "SELECT id FROM Customers WHERE name = 'DELETE'",
+            "SELECT id FROM Customers WHERE status = 'UPDATE me'",
+        ],
+    )
+    def test_validate_sql_query_allows_supported_select_queries(
+        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.",
+            ),
+            ("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 COALESCE((SELECT max(id) FROM Orders), 0) FROM Customers WHERE id = 1",
+                "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(
+        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_entity_records_rejects_invalid_sql_before_network_call(
+        self,
+        service: EntitiesService,
+    ) -> None:
+        service.request = MagicMock()  # type: ignore[method-assign]
+
+        with pytest.raises(
+            ValueError, match=re.escape("Only SELECT statements are allowed.")
+        ):
+            service.query_entity_records("UPDATE Customers SET name = 'X'")
+
+        service.request.assert_not_called()
+
+    def test_query_entity_records_calls_request_for_valid_sql(
+        self,
+        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_entity_records("SELECT id FROM Customers WHERE id > 0")
+
+        assert result == [{"id": 1}, {"id": 2}]
+        service.request.assert_called_once()
+
+    @pytest.mark.anyio
+    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_entity_records_async(
+                "SELECT id FROM Customers WHERE id IN (SELECT id FROM Orders)"
+            )
+
+        service.request_async.assert_not_called()
+
+    @pytest.mark.anyio
+    async def test_query_entity_records_async_calls_request_for_valid_sql(
+        self,
+        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_entity_records_async(
+            "SELECT id FROM Customers WHERE id = 'c1'"
+        )
+
+        assert result == [{"id": "c1"}]
+        service.request_async.assert_called_once()