1- import logging
21from typing import Any , Dict , List , Optional , Type
32
43import sqlparse
87from uipath .core .tracing import traced
98
109from ..common ._base_service import BaseService
11- from ..common ._bindings import EntityResourceOverwrite , _resource_overwrites
1210from ..common ._config import UiPathApiConfig
1311from ..common ._execution_context import UiPathExecutionContext
1412from ..common ._models import Endpoint , RequestSpec
15- from ..orchestrator ._folder_service import FolderService
1613from .entities import (
1714 Entity ,
1815 EntityRecord ,
1916 EntityRecordsBatchResponse ,
20- EntityRouting ,
2117 QueryRoutingOverrideContext ,
2218)
2319
24- logger = logging .getLogger (__name__ )
25-
2620_FORBIDDEN_DML = {"INSERT" , "UPDATE" , "DELETE" , "MERGE" , "REPLACE" }
2721_FORBIDDEN_DDL = {"DROP" , "ALTER" , "CREATE" , "TRUNCATE" }
2822_DISALLOWED_KEYWORDS = [
@@ -53,32 +47,9 @@ class EntitiesService(BaseService):
5347 """
5448
5549 def __init__ (
56- self ,
57- config : UiPathApiConfig ,
58- execution_context : UiPathExecutionContext ,
59- folders_service : Optional [FolderService ] = None ,
60- folders_map : Optional [Dict [str , str ]] = None ,
50+ self , config : UiPathApiConfig , execution_context : UiPathExecutionContext
6151 ) -> None :
6252 super ().__init__ (config = config , execution_context = execution_context )
63- self ._folders_service = folders_service
64- self ._folders_map = folders_map or {}
65-
66- def with_folders_map (self , folders_map : Dict [str , str ]) -> "EntitiesService" :
67- """Return a new EntitiesService configured with the given folders map.
68-
69- The map is used to build a routing context automatically when
70- ``query_entity_records`` is called without an explicit routing context.
71- Folder paths in the map are resolved to folder keys via ``FolderService``.
72-
73- Args:
74- folders_map: Mapping of entity name to folder path.
75- """
76- return EntitiesService (
77- config = self ._config ,
78- execution_context = self ._execution_context ,
79- folders_service = self ._folders_service ,
80- folders_map = folders_map ,
81- )
8253
8354 @traced (name = "entity_retrieve" , run_type = "uipath" )
8455 def retrieve (self , entity_key : str ) -> Entity :
@@ -446,6 +417,7 @@ class CustomerRecord:
446417 def query_entity_records (
447418 self ,
448419 sql_query : str ,
420+ routing_context : Optional [QueryRoutingOverrideContext ] = None ,
449421 ) -> List [Dict [str , Any ]]:
450422 """Query entity records using a validated SQL query.
451423
@@ -455,10 +427,9 @@ def query_entity_records(
455427 sql_query (str): A SQL SELECT query to execute against Data Service entities.
456428 Only SELECT statements are allowed. Queries without WHERE must include
457429 a LIMIT clause. Subqueries and multi-statement queries are not permitted.
458-
459- Notes:
460- A routing context is always derived from the configured ``folders_map``
461- when present and included in the request body.
430+ routing_context (Optional[QueryRoutingOverrideContext]): Per-entity routing context
431+ for multi-folder queries. When present, included in the request body
432+ and takes precedence over the folder header on the backend.
462433
463434 Returns:
464435 List[Dict[str, Any]]: A list of result records as dictionaries.
@@ -467,12 +438,15 @@ def query_entity_records(
467438 ValueError: If the SQL query fails validation (e.g., non-SELECT, missing
468439 WHERE/LIMIT, forbidden keywords, subqueries).
469440 """
470- return self ._query_entities_for_records (sql_query )
441+ return self ._query_entities_for_records (
442+ sql_query , routing_context = routing_context
443+ )
471444
472445 @traced (name = "entity_query_records" , run_type = "uipath" )
473446 async def query_entity_records_async (
474447 self ,
475448 sql_query : str ,
449+ routing_context : Optional [QueryRoutingOverrideContext ] = None ,
476450 ) -> List [Dict [str , Any ]]:
477451 """Asynchronously query entity records using a validated SQL query.
478452
@@ -482,10 +456,9 @@ async def query_entity_records_async(
482456 sql_query (str): A SQL SELECT query to execute against Data Service entities.
483457 Only SELECT statements are allowed. Queries without WHERE must include
484458 a LIMIT clause. Subqueries and multi-statement queries are not permitted.
485-
486- Notes:
487- A routing context is always derived from the configured ``folders_map``
488- when present and included in the request body.
459+ routing_context (Optional[QueryRoutingOverrideContext]): Per-entity routing context
460+ for multi-folder queries. When present, included in the request body
461+ and takes precedence over the folder header on the backend.
489462
490463 Returns:
491464 List[Dict[str, Any]]: A list of result records as dictionaries.
@@ -494,24 +467,28 @@ async def query_entity_records_async(
494467 ValueError: If the SQL query fails validation (e.g., non-SELECT, missing
495468 WHERE/LIMIT, forbidden keywords, subqueries).
496469 """
497- return await self ._query_entities_for_records_async (sql_query )
470+ return await self ._query_entities_for_records_async (
471+ sql_query , routing_context = routing_context
472+ )
498473
499474 def _query_entities_for_records (
500475 self ,
501476 sql_query : str ,
477+ * ,
478+ routing_context : Optional [QueryRoutingOverrideContext ] = None ,
502479 ) -> List [Dict [str , Any ]]:
503480 self ._validate_sql_query (sql_query )
504- routing_context = self ._build_routing_context_from_map ()
505481 spec = self ._query_entity_records_spec (sql_query , routing_context )
506482 response = self .request (spec .method , spec .endpoint , json = spec .json )
507483 return response .json ().get ("results" , [])
508484
509485 async def _query_entities_for_records_async (
510486 self ,
511487 sql_query : str ,
488+ * ,
489+ routing_context : Optional [QueryRoutingOverrideContext ] = None ,
512490 ) -> List [Dict [str , Any ]]:
513491 self ._validate_sql_query (sql_query )
514- routing_context = await self ._build_routing_context_from_map_async ()
515492 spec = self ._query_entity_records_spec (sql_query , routing_context )
516493 response = await self .request_async (spec .method , spec .endpoint , json = spec .json )
517494 return response .json ().get ("results" , [])
@@ -1015,131 +992,6 @@ def _query_entity_records_spec(
1015992 json = body ,
1016993 )
1017994
1018- def _build_routing_context_from_map (
1019- self ,
1020- ) -> Optional [QueryRoutingOverrideContext ]:
1021- """Build a routing context from the configured folders_map and context overwrites.
1022-
1023- Folder paths in the map are resolved to folder keys via FolderService.
1024- Entity overwrites from the active ``ResourceOverwritesContext`` are
1025- merged in, supplying ``override_entity_name`` when the overwrite
1026- provides a different entity name.
1027-
1028- Returns:
1029- A QueryRoutingOverrideContext if routing entries exist,
1030- None otherwise.
1031- """
1032- resolved = self ._resolve_folder_paths_to_ids ()
1033- return self ._build_routing_context_from_resolved_map (resolved )
1034-
1035- async def _build_routing_context_from_map_async (
1036- self ,
1037- ) -> Optional [QueryRoutingOverrideContext ]:
1038- """Async version of _build_routing_context_from_map."""
1039- resolved = await self ._resolve_folder_paths_to_ids_async ()
1040- return self ._build_routing_context_from_resolved_map (resolved )
1041-
1042- def _resolve_folder_paths_to_ids (self ) -> Optional [dict [str , str ]]:
1043- if not self ._folders_map :
1044- return None
1045-
1046- resolved : dict [str , str ] = {}
1047- for folder_path in set (self ._folders_map .values ()):
1048- if self ._folders_service is not None :
1049- folder_key = self ._folders_service .retrieve_folder_key (folder_path )
1050- if folder_key is not None :
1051- resolved [folder_path ] = folder_key
1052- continue
1053- resolved [folder_path ] = folder_path
1054-
1055- return resolved
1056-
1057- async def _resolve_folder_paths_to_ids_async (self ) -> Optional [dict [str , str ]]:
1058- if not self ._folders_map :
1059- return None
1060-
1061- resolved : dict [str , str ] = {}
1062- for folder_path in set (self ._folders_map .values ()):
1063- if self ._folders_service is not None :
1064- folder_key = await self ._folders_service .retrieve_folder_key_async (
1065- folder_path
1066- )
1067- if folder_key is not None :
1068- resolved [folder_path ] = folder_key
1069- continue
1070- resolved [folder_path ] = folder_path
1071-
1072- return resolved
1073-
1074- @staticmethod
1075- def _get_entity_overwrites_from_context () -> Dict [str , EntityResourceOverwrite ]:
1076- """Extract entity overwrites from the active ResourceOverwritesContext.
1077-
1078- Returns:
1079- A dict mapping original entity name to its EntityResourceOverwrite.
1080- """
1081- context_overwrites = _resource_overwrites .get ()
1082- if not context_overwrites :
1083- return {}
1084-
1085- result : Dict [str , EntityResourceOverwrite ] = {}
1086- for key , overwrite in context_overwrites .items ():
1087- if isinstance (overwrite , EntityResourceOverwrite ):
1088- # Key format is "entity.<original_name>"
1089- original_name = key .split ("." , 1 )[1 ] if "." in key else key
1090- result [original_name ] = overwrite
1091- return result
1092-
1093- def _build_routing_context_from_resolved_map (
1094- self ,
1095- resolved : Optional [dict [str , str ]],
1096- ) -> Optional [QueryRoutingOverrideContext ]:
1097- entity_overwrites = self ._get_entity_overwrites_from_context ()
1098-
1099- routings : List [EntityRouting ] = []
1100-
1101- # Add routings from folders_map
1102- if self ._folders_map and resolved is not None :
1103- for name , folder_path in self ._folders_map .items ():
1104- overwrite = entity_overwrites .pop (name , None )
1105- override_name = (
1106- overwrite .resource_identifier
1107- if overwrite and overwrite .resource_identifier != name
1108- else None
1109- )
1110- folder_id = (
1111- overwrite .folder_identifier
1112- if overwrite
1113- else resolved .get (folder_path , folder_path )
1114- )
1115- routings .append (
1116- EntityRouting (
1117- entity_name = name ,
1118- folder_id = folder_id ,
1119- override_entity_name = override_name ,
1120- )
1121- )
1122-
1123- # Add routings from context overwrites not already in folders_map
1124- for original_name , overwrite in entity_overwrites .items ():
1125- override_name = (
1126- overwrite .resource_identifier
1127- if overwrite .resource_identifier != original_name
1128- else None
1129- )
1130- routings .append (
1131- EntityRouting (
1132- entity_name = original_name ,
1133- folder_id = overwrite .folder_identifier ,
1134- override_entity_name = override_name ,
1135- )
1136- )
1137-
1138- if not routings :
1139- return None
1140-
1141- return QueryRoutingOverrideContext (entity_routings = routings )
1142-
1143995 def _insert_batch_spec (self , entity_key : str , records : List [Any ]) -> RequestSpec :
1144996 return RequestSpec (
1145997 method = "POST" ,
0 commit comments