GOV-667: Fix review comments

Author: ankitpatnaik-atlan

pyatlan/client/transport.py

@@ -19,6 +19,90 @@
 
 logger = logging.getLogger(__name__)
 
+def _find_existing_policy(
+    client: "AtlanClient", policy_name: str, persona_guid: str
+) -> Optional[dict]:
+    """Search for an existing AuthPolicy by name and persona GUID."""
+    try:
+        from pyatlan.client.constants import INDEX_SEARCH
+        from pyatlan.model.search import Bool, DSL, IndexSearchRequest, Term
+
+        query = Bool(
+            filter=[
+                Term(field="__typeName.keyword", value="AuthPolicy"),
+                Term(field="name.keyword", value=policy_name),
+                Term(field="__persona", value=persona_guid),
+            ]
+        )
+        search_request = IndexSearchRequest(
+            dsl=DSL(query=query, size=1, from_=0),
+            attributes=["name", "qualifiedName"],
+        )
+        raw_json = client._call_api(INDEX_SEARCH, request_obj=search_request)
+        if raw_json and raw_json.get("entities"):
+            return raw_json["entities"][0]
+        return None
+    except Exception as e:
+        logger.debug(f"Error searching for existing policy: {e}")
+        return None
+
+
+def _create_mock_response(
+    existing_policy: dict, temp_guid: str = "-1"
+) -> httpx.Response:
+    """Build a mock bulk-entity response containing an already-created policy."""
+    response_body = {
+        "mutatedEntities": {"CREATE": [existing_policy]},
+        "guidAssignments": {temp_guid: existing_policy.get("guid")},
+    }
+    return httpx.Response(
+        status_code=200,
+        json=response_body,
+        request=httpx.Request("POST", "http://mock"),
+    )
+
+
+def _check_for_duplicate_policy(
+    client: "AtlanClient", request: httpx.Request
+) -> Optional[httpx.Response]:
+    """
+    Check whether a bulk POST is creating an AuthPolicy that already exists.
+    Only called during retry attempts, never on the first request.
+
+    Returns a mock response with the existing policy if a duplicate is found,
+    or None to let the retry proceed normally.
+    """
+    try:
+        if request.method != "POST" or "/api/meta/entity/bulk" not in str(request.url):
+            return None
+        if not request.content:
+            return None
+
+        body = json.loads(request.content.decode("utf-8"))
+        for entity in body.get("entities", []):
+            if entity.get("typeName") != "AuthPolicy":
+                continue
+            policy_name = entity.get("attributes", {}).get("name")
+            access_control = entity.get("attributes", {}).get("accessControl")
+            persona_guid = (
+                access_control.get("guid")
+                if isinstance(access_control, dict)
+                else None
+            )
+            if not (policy_name and persona_guid):
+                continue
+            existing_policy = _find_existing_policy(client, policy_name, persona_guid)
+            if existing_policy:
+                logger.info(
+                    f"Found existing policy '{policy_name}' with guid "
+                    f"{existing_policy.get('guid')} during retry check"
+                )
+                return _create_mock_response(existing_policy, entity.get("guid", "-1"))
+        return None
+    except Exception as e:
+        logger.debug(f"Duplicate policy check failed (will proceed with retry): {e}")
+        return None
+
 
 class PyatlanSyncTransport(httpx.BaseTransport):
     """
@@ -106,7 +190,7 @@ def _retry_operation(
 
                 # ONLY during retry: check if this is a policy creation and if duplicate exists
                 if self._client:
-                    duplicate_response = self._check_for_duplicate_policy(request)
+                    duplicate_response = _check_for_duplicate_policy(self._client, request)
                     if duplicate_response:
                         logger.warning(
                             "RETRY PREVENTED: Policy already exists (likely from previous "
@@ -130,128 +214,6 @@ def _retry_operation(
             ):
                 return response
 
-    def _check_for_duplicate_policy(
-        self, request: httpx.Request
-    ) -> Optional[httpx.Response]:
-        """
-        Check if request is creating an AuthPolicy that already exists.
-        This is ONLY called during retry attempts, not on first request.
-
-        Returns:
-            A mock response with existing policy if duplicate found, None otherwise.
-        """
-        try:
-            # Only check POST requests to entity/bulk endpoint
-            if request.method != "POST" or "/api/meta/entity/bulk" not in str(
-                request.url
-            ):
-                return None
-
-            if not request.content:
-                return None
-
-            body = json.loads(request.content.decode("utf-8"))
-            entities = body.get("entities", [])
-
-            # Check if any entity is an AuthPolicy
-            for entity in entities:
-                if entity.get("typeName") == "AuthPolicy":
-                    policy_name = entity.get("attributes", {}).get("name")
-                    access_control = entity.get("attributes", {}).get("accessControl")
-
-                    # Extract persona GUID from accessControl relationship
-                    persona_guid = None
-                    if isinstance(access_control, dict):
-                        persona_guid = access_control.get("guid")
-
-                    if policy_name and persona_guid:
-                        # Check if policy already exists
-                        existing_policy = self._find_existing_policy(
-                            policy_name, persona_guid
-                        )
-                        if existing_policy:
-                            logger.info(
-                                f"Found existing policy '{policy_name}' with guid "
-                                f"{existing_policy.get('guid')} during retry check"
-                            )
-                            # Create mock response with existing policy
-                            return self._create_mock_response(existing_policy)
-
-            return None
-
-        except Exception as e:
-            # If duplicate check fails, log and continue with retry
-            logger.debug(
-                f"Duplicate policy check failed (will proceed with retry): {e}"
-            )
-            return None
-
-    def _find_existing_policy(
-        self, policy_name: str, persona_guid: str
-    ) -> Optional[dict]:
-        """
-        Search for existing policy by name and persona.
-
-        Returns:
-            Policy data if found, None otherwise.
-        """
-        try:
-            # Import here to avoid circular dependency
-            from pyatlan.client.constants import INDEX_SEARCH
-            from pyatlan.model.search import Bool, DSL, IndexSearchRequest, Term
-
-            # Build search request to find policy by name and persona
-            query = Bool(
-                filter=[
-                    Term(field="__typeName.keyword", value="AuthPolicy"),
-                    Term(field="name.keyword", value=policy_name),
-                    Term(field="__persona", value=persona_guid),
-                ]
-            )
-
-            search_request = IndexSearchRequest(
-                dsl=DSL(query=query, size=1, from_=0),
-                attributes=["name", "qualifiedName"],
-            )
-
-            # Execute search directly using INDEX_SEARCH endpoint
-            raw_json = self._client._call_api(
-                INDEX_SEARCH,
-                request_obj=search_request,
-            )
-
-            # Check if policy found
-            if (
-                raw_json
-                and raw_json.get("entities")
-                and len(raw_json["entities"]) > 0
-            ):
-                return raw_json["entities"][0]
-
-            return None
-
-        except Exception as e:
-            logger.debug(f"Error searching for existing policy: {e}")
-            return None
-
-    def _create_mock_response(self, existing_policy: dict) -> httpx.Response:
-        """
-        Create a mock HTTP response containing the existing policy.
-
-        This response mimics what the bulk entity creation endpoint would return.
-        """
-        # Create response body matching AssetMutationResponse format
-        response_body = {
-            "mutatedEntities": {"CREATE": [existing_policy]},
-            "guidAssignments": {"-1": existing_policy.get("guid")},
-        }
-
-        return httpx.Response(
-            status_code=200,
-            json=response_body,
-            request=httpx.Request("POST", "http://mock"),
-        )
-
     def close(self) -> None:
         """Close the underlying transport."""
         self._transport.close()
@@ -345,7 +307,7 @@ async def _retry_operation_async(
 
                 # ONLY during retry: check if this is a policy creation and if duplicate exists
                 if self._client:
-                    duplicate_response = self._check_for_duplicate_policy(request)
+                    duplicate_response = _check_for_duplicate_policy(self._client, request)
                     if duplicate_response:
                         logger.warning(
                             "RETRY PREVENTED: Policy already exists (likely from previous "