[changes] Fixed AtlanClient event stream handling

Author: Aryamanz29

pyatlan/client/asset.py

@@ -25,15 +25,20 @@
 )
 from warnings import warn
 
-import httpx
 from pydantic.v1 import (
     StrictStr,
     ValidationError,
     constr,
     parse_obj_as,
     validate_arguments,
 )
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
+from tenacity import (
+    RetryError,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_fixed,
+)
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (
@@ -844,7 +849,10 @@ def delete_by_guid(self, guid: Union[str, List[str]]) -> AssetMutationResponse:
         )
         response = AssetMutationResponse(**raw_json)
         for asset in response.assets_deleted(asset_type=Asset):
-            self._wait_till_deleted(asset)
+            try:
+                self._wait_till_deleted(asset)
+            except RetryError as err:
+                raise ErrorCode.RETRY_OVERRUN.exception_with_parameters() from err
         return response
 
     @retry(
@@ -854,12 +862,9 @@ def delete_by_guid(self, guid: Union[str, List[str]]) -> AssetMutationResponse:
         wait=wait_fixed(1),
     )
     def _wait_till_deleted(self, asset: Asset):
-        try:
-            asset = self.retrieve_minimal(guid=asset.guid, asset_type=Asset)
-            if asset.status == EntityStatus.DELETED:
-                return
-        except httpx.TransportError as err:
-            raise ErrorCode.RETRY_OVERRUN.exception_with_parameters() from err
+        asset = self.retrieve_minimal(guid=asset.guid, asset_type=Asset)
+        if asset.status == EntityStatus.DELETED:
+            return
 
     @validate_arguments
     def restore(self, asset_type: Type[A], qualified_name: str) -> bool:

pyatlan/client/atlan.py

@@ -7,7 +7,6 @@
 import copy
 import json
 import logging
-import shutil
 import uuid
 from contextvars import ContextVar
 from http import HTTPStatus
@@ -343,7 +342,8 @@ def update_headers(self, header: Dict[str, str]):
     def _handle_file_download(self, raw_response: Any, file_path: str) -> str:
         try:
             download_file = open(file_path, "wb")
-            shutil.copyfileobj(raw_response, download_file)
+            for chunk in raw_response:
+                download_file.write(chunk)
         except Exception as err:
             raise ErrorCode.UNABLE_TO_DOWNLOAD_FILE.exception_with_parameters(
                 str((hasattr(err, "strerror") and err.strerror) or err), file_path
@@ -374,15 +374,49 @@ def _call_api_internal(
                     timeout=timeout,
                 )
             elif api.consumes == EVENT_STREAM and api.produces == EVENT_STREAM:
-                response = self._session.request(
+                with self._session.stream(
                     api.method.value,
                     path,
                     **params,
-                    stream=True,
                     timeout=timeout,
-                )
-                if download_file_path:
-                    return self._handle_file_download(response.raw, download_file_path)
+                ) as stream_response:
+                    if download_file_path:
+                        return self._handle_file_download(
+                            stream_response.iter_raw(), download_file_path
+                        )
+
+                    # For event streams, we need to read the content while the stream is open
+                    # Store the response data and create a mock response object for common processing
+                    content = stream_response.read()
+                    text = content.decode("utf-8") if content else ""
+                    lines = []
+
+                    # Only process lines for successful responses to avoid errors on error responses
+                    if stream_response.status_code == api.expected_status:
+                        # Reset stream position and get lines
+                        lines = text.splitlines() if text else []
+
+                    response_data = {
+                        "status_code": stream_response.status_code,
+                        "headers": stream_response.headers,
+                        "text": text,
+                        "content": content,
+                        "lines": lines,
+                    }
+
+                    # Create a simple namespace object to mimic the response interface
+                    response = SimpleNamespace(
+                        status_code=response_data["status_code"],
+                        headers=response_data["headers"],
+                        text=response_data["text"],
+                        content=response_data["content"],
+                        _stream_lines=response_data[
+                            "lines"
+                        ],  # Store lines for event processing
+                        json=lambda: json.loads(response_data["text"])
+                        if response_data["text"]
+                        else {},
+                    )
             else:
                 response = self._session.request(
                     api.method.value,
@@ -429,14 +463,16 @@ def _call_api_internal(
                             response,
                         )
                     if api.consumes == EVENT_STREAM and api.produces == EVENT_STREAM:
-                        for line in response.iter_lines(decode_unicode=True):
-                            if not line:
-                                continue
-                            if not line.startswith("data: "):
-                                raise ErrorCode.UNABLE_TO_DESERIALIZE.exception_with_parameters(
-                                    line
-                                )
-                            events.append(json.loads(line.split("data: ")[1]))
+                        # Process event stream using stored lines from the streaming response
+                        if hasattr(response, "_stream_lines"):
+                            for line in response._stream_lines:
+                                if not line:
+                                    continue
+                                if not line.startswith("data: "):
+                                    raise ErrorCode.UNABLE_TO_DESERIALIZE.exception_with_parameters(
+                                        line
+                                    )
+                                events.append(json.loads(line.split("data: ")[1]))
                     if text_response:
                         response_ = response.text
                     else:

tests/integration/purpose_test.py

@@ -271,7 +271,7 @@ def test_retrieve_purpose(
             assert full.policy_mask_type == DataMaskingType.REDACT
 
 
-@pytest.mark.skip(reason="Test failing with HekaException")
+# @pytest.mark.skip(reason="Test failing with HekaException")
 @pytest.mark.order(after="test_retrieve_purpose")
 def test_run_query_without_policy(client: AtlanClient, assign_tag_to_asset, query):
     response = client.queries.stream(request=query)