GOV-667: Check if policy was created during retry

Author: ankitpatnaik-atlan

pyatlan/client/atlan.py

@@ -211,8 +211,11 @@ def __init__(self, **data):
         # Configure httpx client with custom transport that supports retry and proxy
         # Note: We pass proxy/SSL config to the transport, not the client,
         # so that retry logic properly respects these settings
+        # Pass self reference to transport for duplicate checking during retries
         self._session = httpx.Client(
-            transport=PyatlanSyncTransport(retry=self.retry, **transport_kwargs),
+            transport=PyatlanSyncTransport(
+                retry=self.retry, client=self, **transport_kwargs
+            ),
             headers={
                 "x-atlan-agent": "sdk",
                 "x-atlan-agent-id": "python",
@@ -2005,7 +2008,9 @@ def max_retries(  # type: ignore[misc]
         if self.verify is not None:
             transport_kwargs["verify"] = self.verify
 
-        new_transport = PyatlanSyncTransport(retry=max_retries, **transport_kwargs)
+        new_transport = PyatlanSyncTransport(
+            retry=max_retries, client=self, **transport_kwargs
+        )
         self._session._transport = new_transport
 
         LOGGER.debug(

pyatlan/client/transport.py

@@ -6,13 +6,17 @@
 with httpx's HTTPTransport while respecting proxy and SSL configurations.
 """
 
+import json
 import logging
 from functools import partial
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import httpx
 from httpx_retries import Retry
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,9 +47,11 @@ class PyatlanSyncTransport(httpx.BaseTransport):
     def __init__(
         self,
         retry: Optional[Retry] = None,
+        client: Optional["AtlanClient"] = None,
         **kwargs: Any,
     ) -> None:
         self.retry = retry or Retry()
+        self._client = client  # Reference to AtlanClient for duplicate checking
         # Ensure trust_env is True by default to respect environment variables
         # unless explicitly overridden
         if "trust_env" not in kwargs:
@@ -97,6 +103,17 @@ def _retry_operation(
                 logger.debug(
                     "_retry_operation retrying response=%s retry=%s", response, retry
                 )
+
+                # 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)
+                    if duplicate_response:
+                        logger.warning(
+                            "RETRY PREVENTED: Policy already exists (likely from previous "
+                            "request that timed out but succeeded). Returning existing policy."
+                        )
+                        return duplicate_response
+
                 retry = retry.increment()
                 retry.sleep(response)
 
@@ -113,6 +130,128 @@ 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()
@@ -145,9 +284,11 @@ class PyatlanAsyncTransport(httpx.AsyncBaseTransport):
     def __init__(
         self,
         retry: Optional[Retry] = None,
+        client: Optional["AtlanClient"] = None,
         **kwargs: Any,
     ) -> None:
         self.retry = retry or Retry()
+        self._client = client  # Reference to AtlanClient for duplicate checking
         # Ensure trust_env is True by default to respect environment variables
         # unless explicitly overridden
         if "trust_env" not in kwargs:
@@ -201,6 +342,17 @@ async def _retry_operation_async(
                     response,
                     retry,
                 )
+
+                # 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)
+                    if duplicate_response:
+                        logger.warning(
+                            "RETRY PREVENTED: Policy already exists (likely from previous "
+                            "request that timed out but succeeded). Returning existing policy."
+                        )
+                        return duplicate_response
+
                 retry = retry.increment()
                 await retry.asleep(response)