Merge pull request #878 from atlanhq/GOV-667

GOV-667 | feat: Prevent duplicate AuthPolicy creation on timeout retries

Author: Aryamanz29

pyatlan/client/common/transport.py

@@ -17,20 +17,26 @@
 
 from pyatlan.client.constants import BULK_UPDATE, INDEX_SEARCH
 from pyatlan.errors import ErrorCode
-from pyatlan.model.search import DSL, Bool, IndexSearchRequest, Term
+from pyatlan.model.search import Bool, DSL, IndexSearchRequest, Prefix, Term
 
 logger = logging.getLogger(__name__)
 
 
 def build_policy_search_request(
-    policy_name: str, persona_guid: str
+    policy_name: str, persona_qualified_name: str
 ) -> IndexSearchRequest:
-    """Build an IndexSearchRequest to find an existing AuthPolicy by name and persona."""
+    """
+    Build an IndexSearchRequest to find an existing AuthPolicy by name and persona.
+    Using persona GUID directly returns associated assets, not policies.
+    Policies require a hierarchical (prefix) match to be correctly retrieved.
+    """
     query = Bool(
         filter=[
+            Term(field="__state", value="ACTIVE"),
             Term(field="__typeName.keyword", value="AuthPolicy"),
+            Term(field="policyCategory", value="persona"),
             Term(field="name.keyword", value=policy_name),
-            Term(field="__persona", value=persona_guid),
+            Prefix(field="qualifiedName", value=persona_qualified_name),
         ]
     )
     return IndexSearchRequest(
@@ -57,7 +63,10 @@ def create_mock_response(
 def parse_auth_policy_entity(request: httpx.Request) -> Optional[tuple[str, str, str]]:
     """
     Parse the request body and return (policy_name, persona_guid, temp_guid)
-    if the request is a bulk POST containing an AuthPolicy, else None.
+    if the request is a bulk POST containing a NEW AuthPolicy creation, else None.
+
+    Only matches policy CREATES (temp GUIDs starting with "-"), not UPDATES
+    (real GUIDs), to prevent suppressing legitimate policy modifications.
     """
     if request.method != "POST" or BULK_UPDATE.path not in str(request.url):
         return None
@@ -75,27 +84,107 @@ def parse_auth_policy_entity(request: httpx.Request) -> Optional[tuple[str, str,
     for entity in body.get("entities", []):
         if entity.get("typeName") != "AuthPolicy":
             continue
+
+        entity_guid = entity.get("guid", "-1")
+        # Only match policy CREATES (temp GUIDs like "-1", "-2", etc.)
+        # Skip policy UPDATES (real GUIDs) to avoid suppressing modifications
+        if not isinstance(entity_guid, str) or not entity_guid.startswith("-"):
+            logger.debug(
+                "parse_auth_policy_entity: skipping duplicate check for policy with GUID %s (likely an update or invalid type)",
+                entity_guid,
+            )
+            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 policy_name and persona_guid:
-            return policy_name, persona_guid, entity.get("guid", "-1")
+            return policy_name, persona_guid, entity_guid
     return None
 
 
+def get_persona_qualified_name(client: Any, persona_guid: str) -> Optional[str]:
+    """
+    Fetch the qualifiedName of a Persona by its GUID via IndexSearch (synchronous).
+    """
+    try:
+        query = Bool(
+            filter=[
+                Term(field="__typeName.keyword", value="Persona"),
+                Term(field="__guid", value=persona_guid),
+            ]
+        )
+        search = IndexSearchRequest(
+            dsl=DSL(query=query, size=1, from_=0),
+            attributes=["qualifiedName"],
+        )
+        raw_json = client._call_api(INDEX_SEARCH, request_obj=search)
+        if raw_json and raw_json.get("entities"):
+            return raw_json["entities"][0].get("attributes", {}).get("qualifiedName")
+        return None
+    except Exception as e:
+        logger.debug(
+            "get_persona_qualified_name: could not fetch qualifiedName for persona %s: %s",
+            persona_guid,
+            e,
+        )
+        return None
+
+
+async def get_persona_qualified_name_async(
+    client: Any, persona_guid: str
+) -> Optional[str]:
+    """
+    Fetch the qualifiedName of a Persona by its GUID via IndexSearch (asynchronous).
+    """
+    try:
+        query = Bool(
+            filter=[
+                Term(field="__typeName.keyword", value="Persona"),
+                Term(field="__guid", value=persona_guid),
+            ]
+        )
+        search = IndexSearchRequest(
+            dsl=DSL(query=query, size=1, from_=0),
+            attributes=["qualifiedName"],
+        )
+        raw_json = await client._call_api(INDEX_SEARCH, request_obj=search)
+        if raw_json and raw_json.get("entities"):
+            return raw_json["entities"][0].get("attributes", {}).get("qualifiedName")
+        return None
+    except Exception as e:
+        logger.debug(
+            "get_persona_qualified_name_async: could not fetch qualifiedName for persona %s: %s",
+            persona_guid,
+            e,
+        )
+        return None
+
+
 def find_existing_policy(
     client: Any, policy_name: str, persona_guid: str
 ) -> Optional[dict]:
     """
     Search for an existing AuthPolicy by name and persona GUID (synchronous).
 
+    First resolves the persona GUID to its qualifiedName, then uses a qualifiedName
+    prefix query to scope the search to that persona.
+
     Raises:
-        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the search call fails.
+        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the policy search call fails.
     """
+    persona_qualified_name = get_persona_qualified_name(client, persona_guid)
+    if not persona_qualified_name:
+        raise ErrorCode.UNABLE_TO_RESOLVE_PERSONA_QUALIFIED_NAME.exception_with_parameters(
+            persona_guid
+        )
+
     try:
-        search_request = build_policy_search_request(policy_name, persona_guid)
+        search_request = build_policy_search_request(
+            policy_name, persona_qualified_name
+        )
         raw_json = client._call_api(INDEX_SEARCH, request_obj=search_request)
         if raw_json and raw_json.get("entities"):
             return raw_json["entities"][0]
@@ -112,11 +201,24 @@ async def find_existing_policy_async(
     """
     Search for an existing AuthPolicy by name and persona GUID (asynchronous).
 
+    First resolves the persona GUID to its qualifiedName, then uses a qualifiedName
+    prefix query to scope the search to that persona.
+
     Raises:
-        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the search call fails.
+        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the policy search call fails.
     """
+    persona_qualified_name = await get_persona_qualified_name_async(
+        client, persona_guid
+    )
+    if not persona_qualified_name:
+        raise ErrorCode.UNABLE_TO_RESOLVE_PERSONA_QUALIFIED_NAME.exception_with_parameters(
+            persona_guid
+        )
+
     try:
-        search_request = build_policy_search_request(policy_name, persona_guid)
+        search_request = build_policy_search_request(
+            policy_name, persona_qualified_name
+        )
         raw_json = await client._call_api(INDEX_SEARCH, request_obj=search_request)
         if raw_json and raw_json.get("entities"):
             return raw_json["entities"][0]
@@ -132,24 +234,36 @@ def check_for_duplicate_policy(
 ) -> Optional[httpx.Response]:
     """
     Check whether a bulk POST is creating an AuthPolicy that already exists (synchronous).
-    Only called during retry attempts, never on the first request.
+    Called before every attempt — including the first — so that repeated automation
+    runs that don't pre-check for an existing policy are handled transparently.
 
     Returns a mock response with the existing policy if a duplicate is found,
-    or None to let the retry proceed normally.
+    or None to let the request proceed normally.
 
-    Raises:
-        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the duplicate search fails.
+    Never raises: search failures are logged and treated as "not found" so that
+    a degraded index cannot block policy creation entirely.
     """
     parsed = parse_auth_policy_entity(request)
     if not parsed:
         return None
 
     policy_name, persona_guid, temp_guid = parsed
-    existing_policy = find_existing_policy(client, policy_name, persona_guid)
+    try:
+        existing_policy = find_existing_policy(client, policy_name, persona_guid)
+    except Exception as e:
+        logger.warning(
+            "Duplicate policy search failed for '%s' (persona %s): %s. "
+            "Proceeding with request.",
+            policy_name,
+            persona_guid,
+            str(e),
+        )
+        return None
     if existing_policy:
         logger.info(
-            f"Found existing policy '{policy_name}' with guid "
-            f"{existing_policy.get('guid')} during retry check"
+            "Found existing policy '%s' with guid %s — returning it instead of creating a duplicate.",
+            policy_name,
+            existing_policy.get("guid"),
         )
         return create_mock_response(existing_policy, temp_guid)
     return None
@@ -160,26 +274,38 @@ async def check_for_duplicate_policy_async(
 ) -> Optional[httpx.Response]:
     """
     Check whether a bulk POST is creating an AuthPolicy that already exists (asynchronous).
-    Only called during retry attempts, never on the first request.
+    Called before every attempt — including the first — so that repeated automation
+    runs that don't pre-check for an existing policy are handled transparently.
 
     Returns a mock response with the existing policy if a duplicate is found,
-    or None to let the retry proceed normally.
+    or None to let the request proceed normally.
 
-    Raises:
-        ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the duplicate search fails.
+    Never raises: search failures are logged and treated as "not found" so that
+    a degraded index cannot block policy creation entirely.
     """
     parsed = parse_auth_policy_entity(request)
     if not parsed:
         return None
 
     policy_name, persona_guid, temp_guid = parsed
-    existing_policy = await find_existing_policy_async(
-        client, policy_name, persona_guid
-    )
+    try:
+        existing_policy = await find_existing_policy_async(
+            client, policy_name, persona_guid
+        )
+    except Exception as e:
+        logger.warning(
+            "Duplicate policy search failed for '%s' (persona %s): %s. "
+            "Proceeding with request.",
+            policy_name,
+            persona_guid,
+            str(e),
+        )
+        return None
     if existing_policy:
         logger.info(
-            f"Found existing policy '{policy_name}' with guid "
-            f"{existing_policy.get('guid')} during retry check"
+            "Found existing policy '%s' with guid %s — returning it instead of creating a duplicate.",
+            policy_name,
+            existing_policy.get("guid"),
         )
         return create_mock_response(existing_policy, temp_guid)
     return None

pyatlan/client/transport.py

@@ -106,22 +106,23 @@ 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 = check_for_duplicate_policy(
-                        self._client, 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)
 
+            # Check before EVERY attempt (first + retries).
+            # On the first attempt this catches automation re-runs where the policy
+            # was already created by a previous run.
+            # On retries, the preceding sleep gives the index time to propagate an
+            # entity that was committed server-side before a gateway timeout.
+            if self._client:
+                duplicate_response = check_for_duplicate_policy(self._client, request)
+                if duplicate_response:
+                    logger.warning(
+                        "DUPLICATE PREVENTED: Policy already exists. "
+                        "Returning existing policy instead of creating a duplicate."
+                    )
+                    return duplicate_response
+
             try:
                 response = send_method(request)
             except httpx.HTTPError as e:
@@ -225,22 +226,25 @@ 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 = await check_for_duplicate_policy_async(
-                        self._client, 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)
 
+            # Check before EVERY attempt (first + retries).
+            # On the first attempt this catches automation re-runs where the policy
+            # was already created by a previous run.
+            # On retries, the preceding sleep gives the index time to propagate an
+            # entity that was committed server-side before a gateway timeout.
+            if self._client:
+                duplicate_response = await check_for_duplicate_policy_async(
+                    self._client, request
+                )
+                if duplicate_response:
+                    logger.warning(
+                        "DUPLICATE PREVENTED: Policy already exists. "
+                        "Returning existing policy instead of creating a duplicate."
+                    )
+                    return duplicate_response
+
             try:
                 response = await send_method(request)
             except httpx.HTTPError as e:

pyatlan/errors.py

@@ -1061,6 +1061,13 @@ class ErrorCode(Enum):
         "Check your backend connectivity and ensure the Atlan search service is accessible.",
         ApiError,
     )
+    UNABLE_TO_RESOLVE_PERSONA_QUALIFIED_NAME = (
+        500,
+        "ATLAN-PYTHON-500-008",
+        "Unable to resolve the qualifiedName for persona with GUID '{0}'.",
+        "Verify the persona exists and the Atlan search service is accessible.",
+        ApiError,
+    )
 
     def __init__(
         self,