fix(platform): support entity resource overwrites with folder resolution

Author: UIPath-Harshit

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.21"
+version = "0.1.22"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

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

@@ -8,17 +8,18 @@
 from uipath.core.tracing import traced
 
 from ..common._base_service import BaseService
-from ..common._bindings import EntityResourceOverwrite, _resource_overwrites
+from ..common._bindings import _resource_overwrites
 from ..common._config import UiPathApiConfig
 from ..common._execution_context import UiPathExecutionContext
 from ..common._models import Endpoint, RequestSpec
 from ..common.constants import HEADER_FOLDER_KEY
 from ..orchestrator._folder_service import FolderService
 from ._entity_resolution import (
-    build_resolution_routing_context,
+    RoutingStrategy,
     build_resolution_service,
     create_resolution_plan,
     create_resolution_plan_async,
+    create_routing_strategy,
     fetch_resolved_entities,
     fetch_resolved_entities_async,
 )
@@ -27,7 +28,6 @@
     Entity,
     EntityRecord,
     EntityRecordsBatchResponse,
-    EntityRouting,
     EntitySetResolution,
     QueryRoutingOverrideContext,
 )
@@ -74,32 +74,11 @@ def __init__(
     ) -> None:
         super().__init__(config=config, execution_context=execution_context)
         self._folders_service = folders_service
-        self._folders_map = folders_map or {}
-        self._effective_entity_names = entity_name_overrides or {}
-        self._routing_context = routing_context
-
-    def with_folders_map(
-        self,
-        folders_map: Dict[str, str],
-        entity_name_overrides: Optional[Dict[str, str]] = None,
-    ) -> "EntitiesService":
-        """Return a new EntitiesService configured with the given folders map.
-
-        The map is used to build a routing context automatically when
-        ``query_entity_records`` is called without an explicit routing context.
-        Folder paths in the map are resolved to folder keys via ``FolderService``.
-
-        Args:
-            folders_map: Mapping of entity name to folder path.
-            entity_name_overrides: Mapping of original entity name to
-                overridden entity name.
-        """
-        return EntitiesService(
-            config=self._config,
-            execution_context=self._execution_context,
-            folders_service=self._folders_service,
+        self._routing_strategy: RoutingStrategy = create_routing_strategy(
             folders_map=folders_map,
-            entity_name_overrides=entity_name_overrides,
+            effective_entity_names=entity_name_overrides,
+            routing_context=routing_context,
+            folders_service=folders_service,
         )
 
     @traced(name="entity_retrieve", run_type="uipath")
@@ -629,7 +608,7 @@ def _query_entities_for_records(
         sql_query: str,
     ) -> List[Dict[str, Any]]:
         self._validate_sql_query(sql_query)
-        routing_context = self._build_routing_context_from_map()
+        routing_context = self._routing_strategy.resolve()
         spec = self._query_entity_records_spec(sql_query, routing_context)
         response = self.request(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
@@ -639,7 +618,7 @@ async def _query_entities_for_records_async(
         sql_query: str,
     ) -> List[Dict[str, Any]]:
         self._validate_sql_query(sql_query)
-        routing_context = await self._build_routing_context_from_map_async()
+        routing_context = await self._routing_strategy.resolve_async()
         spec = self._query_entity_records_spec(sql_query, routing_context)
         response = await self.request_async(spec.method, spec.endpoint, json=spec.json)
         return response.json().get("results", [])
@@ -1158,152 +1137,6 @@ def _query_entity_records_spec(
             json=body,
         )
 
-    def _build_routing_context_from_map(
-        self,
-    ) -> Optional[QueryRoutingOverrideContext]:
-        """Build a routing context from the configured folders_map and context overwrites.
-
-        If a pre-built routing context was provided (e.g. by
-        ``resolve_entity_set_async``), it is returned directly without
-        re-resolving folder paths.
-
-        Otherwise, folder paths in the map are resolved to folder keys via
-        FolderService and entity overwrites from the active
-        ``ResourceOverwritesContext`` are merged in.
-
-        Returns:
-            A QueryRoutingOverrideContext if routing entries exist,
-            None otherwise.
-        """
-        if self._routing_context is not None:
-            return self._routing_context
-        resolved = self._resolve_folder_paths_to_ids()
-        return self._build_routing_context_from_resolved_map(resolved)
-
-    async def _build_routing_context_from_map_async(
-        self,
-    ) -> Optional[QueryRoutingOverrideContext]:
-        """Async version of _build_routing_context_from_map."""
-        if self._routing_context is not None:
-            return self._routing_context
-        resolved = await self._resolve_folder_paths_to_ids_async()
-        return self._build_routing_context_from_resolved_map(resolved)
-
-    def _resolve_folder_paths_to_ids(self) -> Optional[dict[str, str]]:
-        entity_overwrites = self._get_entity_overwrites_from_context()
-        folder_paths = set(self._folders_map.values())
-        for overwrite in entity_overwrites.values():
-            if overwrite.folder_path:
-                folder_paths.add(overwrite.folder_path)
-
-        if not folder_paths:
-            return None
-
-        resolved: dict[str, str] = {}
-        for folder_path in folder_paths:
-            if self._folders_service is not None:
-                folder_key = self._folders_service.retrieve_folder_key(folder_path)
-                if folder_key is not None:
-                    resolved[folder_path] = folder_key
-                    continue
-            resolved[folder_path] = folder_path
-
-        return resolved
-
-    async def _resolve_folder_paths_to_ids_async(self) -> Optional[dict[str, str]]:
-        entity_overwrites = self._get_entity_overwrites_from_context()
-        folder_paths = set(self._folders_map.values())
-        for overwrite in entity_overwrites.values():
-            if overwrite.folder_path:
-                folder_paths.add(overwrite.folder_path)
-
-        if not folder_paths:
-            return None
-
-        resolved: dict[str, str] = {}
-        for folder_path in folder_paths:
-            if self._folders_service is not None:
-                folder_key = await self._folders_service.retrieve_folder_key_async(
-                    folder_path
-                )
-                if folder_key is not None:
-                    resolved[folder_path] = folder_key
-                    continue
-            resolved[folder_path] = folder_path
-
-        return resolved
-
-    @staticmethod
-    def _get_entity_overwrites_from_context() -> Dict[str, EntityResourceOverwrite]:
-        """Extract entity overwrites from the active ResourceOverwritesContext.
-
-        Returns:
-            A dict mapping original entity name to its EntityResourceOverwrite.
-        """
-        context_overwrites = _resource_overwrites.get()
-        if not context_overwrites:
-            return {}
-
-        result: Dict[str, EntityResourceOverwrite] = {}
-        for key, overwrite in context_overwrites.items():
-            if isinstance(overwrite, EntityResourceOverwrite):
-                # Key format is "entity.<original_id>"
-                original_name = key.split(".", 1)[1] if "." in key else key
-                result[original_name] = overwrite
-        return result
-
-    def _build_routing_context_from_resolved_map(
-        self,
-        resolved: Optional[dict[str, str]],
-    ) -> Optional[QueryRoutingOverrideContext]:
-        if self._folders_map:
-            return build_resolution_routing_context(
-                {
-                    name: (resolved or {}).get(folder_path, folder_path)
-                    for name, folder_path in self._folders_map.items()
-                },
-                self._effective_entity_names,
-            )
-
-        routings: List[EntityRouting] = []
-        if not self._folders_map:
-            # Fallback for direct SDK usage (no folders_map)
-            entity_overwrites = self._get_entity_overwrites_from_context()
-            for original_name, overwrite in entity_overwrites.items():
-                override_name = (
-                    overwrite.resource_identifier
-                    if overwrite.resource_identifier != original_name
-                    else None
-                )
-                routings.append(
-                    EntityRouting(
-                        entity_name=original_name,
-                        folder_id=self._resolve_overwrite_folder(overwrite, resolved),
-                        override_entity_name=override_name,
-                    )
-                )
-
-        if not routings:
-            return None
-
-        return QueryRoutingOverrideContext(entity_routings=routings)
-
-    @staticmethod
-    def _resolve_overwrite_folder(
-        overwrite: EntityResourceOverwrite,
-        resolved: Optional[dict[str, str]],
-    ) -> str:
-        """Return the folder key for an entity overwrite.
-
-        Uses folder_id directly when present (already a key).
-        Falls back to resolving folder_path through the resolved map.
-        """
-        if overwrite.folder_id:
-            return overwrite.folder_id
-        if overwrite.folder_path and resolved:
-            return resolved.get(overwrite.folder_path, overwrite.folder_path)
-        return overwrite.folder_identifier
-
     def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec:
         return RequestSpec(
             method="POST",