Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 18 additions & 15 deletions pyatlan/client/common/transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import httpx

from pyatlan.client.constants import BULK_UPDATE, INDEX_SEARCH
from pyatlan.errors import ErrorCode
from pyatlan.model.search import DSL, Bool, IndexSearchRequest, Term

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -91,8 +90,7 @@ def find_existing_policy(
"""
Search for an existing AuthPolicy by name and persona GUID (synchronous).

Raises:
ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the search call fails.
Returns None on failure so the retry loop can proceed normally.
"""
try:
search_request = build_policy_search_request(policy_name, persona_guid)
Expand All @@ -101,9 +99,14 @@ def find_existing_policy(
return raw_json["entities"][0]
return None
except Exception as e:
raise ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY.exception_with_parameters(
policy_name, persona_guid, str(e)
) from e
logger.warning(
"Duplicate policy search failed for '%s' (persona %s): %s. "
"Retry will proceed normally.",
policy_name,
persona_guid,
str(e),
)
return None


async def find_existing_policy_async(
Expand All @@ -112,8 +115,7 @@ async def find_existing_policy_async(
"""
Search for an existing AuthPolicy by name and persona GUID (asynchronous).

Raises:
ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the search call fails.
Returns None on failure so the retry loop can proceed normally.
"""
try:
search_request = build_policy_search_request(policy_name, persona_guid)
Expand All @@ -122,9 +124,14 @@ async def find_existing_policy_async(
return raw_json["entities"][0]
return None
except Exception as e:
raise ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY.exception_with_parameters(
policy_name, persona_guid, str(e)
) from e
logger.warning(
"Duplicate policy search failed for '%s' (persona %s): %s. "
"Retry will proceed normally.",
policy_name,
persona_guid,
str(e),
)
return None


def check_for_duplicate_policy(
Expand All @@ -137,8 +144,6 @@ def check_for_duplicate_policy(
Returns a mock response with the existing policy if a duplicate is found,
or None to let the retry proceed normally.

Raises:
ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the duplicate search fails.
"""
parsed = parse_auth_policy_entity(request)
if not parsed:
Expand All @@ -165,8 +170,6 @@ async def check_for_duplicate_policy_async(
Returns a mock response with the existing policy if a duplicate is found,
or None to let the retry proceed normally.

Raises:
ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY: if the duplicate search fails.
"""
parsed = parse_auth_policy_entity(request)
if not parsed:
Expand Down
22 changes: 14 additions & 8 deletions pyatlan/client/transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,13 @@ def _retry_operation(
"_retry_operation retrying response=%s retry=%s", response, retry
)

# ONLY during retry: check if this is a policy creation and if duplicate exists
retry = retry.increment()
retry.sleep(response)

# AFTER backoff: check if this is a policy creation and if duplicate exists.
# The sleep gives the index time to propagate the entity created by
# the previous request that may have succeeded server-side but timed
# out client-side.
if self._client:
duplicate_response = check_for_duplicate_policy(
self._client, request
Expand All @@ -119,9 +125,6 @@ def _retry_operation(
)
return duplicate_response

retry = retry.increment()
retry.sleep(response)

try:
response = send_method(request)
except httpx.HTTPError as e:
Expand Down Expand Up @@ -226,7 +229,13 @@ async def _retry_operation_async(
retry,
)

# ONLY during retry: check if this is a policy creation and if duplicate exists
retry = retry.increment()
await retry.asleep(response)

# AFTER backoff: check if this is a policy creation and if duplicate exists.
# The sleep gives the index time to propagate the entity created by
# the previous request that may have succeeded server-side but timed
# out client-side.
if self._client:
duplicate_response = await check_for_duplicate_policy_async(
self._client, request
Expand All @@ -238,9 +247,6 @@ async def _retry_operation_async(
)
return duplicate_response

retry = retry.increment()
await retry.asleep(response)

try:
response = await send_method(request)
except httpx.HTTPError as e:
Expand Down
26 changes: 10 additions & 16 deletions tests/unit/test_transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,12 +175,11 @@ def test_returns_none_when_raw_json_is_none(self):
result = find_existing_policy(client, POLICY_NAME, PERSONA_GUID)
assert result is None

def test_raises_error_code_on_exception(self):
def test_returns_none_on_exception(self):
client = MagicMock()
client._call_api.side_effect = Exception("search failed")
with pytest.raises(Exception) as exc_info:
find_existing_policy(client, POLICY_NAME, PERSONA_GUID)
assert "ATLAN-PYTHON-500-007" in str(exc_info.value)
result = find_existing_policy(client, POLICY_NAME, PERSONA_GUID)
assert result is None


# ---------------------------------------------------------------------------
Expand All @@ -204,12 +203,11 @@ async def test_returns_none_when_no_entities(self):
assert result is None

@pytest.mark.asyncio
async def test_raises_error_code_on_exception(self):
async def test_returns_none_on_exception(self):
client = MagicMock()
client._call_api = AsyncMock(side_effect=Exception("async search failed"))
with pytest.raises(Exception) as exc_info:
await find_existing_policy_async(client, POLICY_NAME, PERSONA_GUID)
assert "ATLAN-PYTHON-500-007" in str(exc_info.value)
result = await find_existing_policy_async(client, POLICY_NAME, PERSONA_GUID)
assert result is None


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -244,13 +242,11 @@ def test_returns_mock_response_when_duplicate_found(self):
body = resp.json()
assert body["guidAssignments"][TEMP_GUID] == EXISTING_GUID

def test_propagates_search_error(self):
def test_returns_none_on_search_error(self):
client = MagicMock()
client._call_api.side_effect = Exception("search failed")
req = _make_bulk_request()
with pytest.raises(Exception) as exc_info:
check_for_duplicate_policy(client, req)
assert "ATLAN-PYTHON-500-007" in str(exc_info.value)
assert check_for_duplicate_policy(client, req) is None


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -283,13 +279,11 @@ async def test_returns_mock_response_when_duplicate_found(self):
assert resp.json()["guidAssignments"][TEMP_GUID] == EXISTING_GUID

@pytest.mark.asyncio
async def test_propagates_search_error(self):
async def test_returns_none_on_search_error(self):
client = MagicMock()
client._call_api = AsyncMock(side_effect=Exception("async search failed"))
req = _make_bulk_request()
with pytest.raises(Exception) as exc_info:
await check_for_duplicate_policy_async(client, req)
assert "ATLAN-PYTHON-500-007" in str(exc_info.value)
assert await check_for_duplicate_policy_async(client, req) is None


# ---------------------------------------------------------------------------
Expand Down
Loading