diff --git a/src/dataworkbench/gateway.py b/src/dataworkbench/gateway.py index 88a06b6..dff246d 100644 --- a/src/dataworkbench/gateway.py +++ b/src/dataworkbench/gateway.py @@ -1,6 +1,7 @@ from typing import Any import requests from uuid import UUID +import json from dataworkbench.utils import get_secret from dataworkbench.auth import TokenManager @@ -10,6 +11,11 @@ logger = setup_logger(__name__) +def _get_trace_id_from_response(response: requests.Response) -> str | None: + response_dict = json.loads(response.text) + return response_dict.get("traceId") + + class Gateway: """ Handles interactions with the remote data gateway API. @@ -135,7 +141,19 @@ def import_dataset( logger.info(f"Sending request to import dataset: {dataset_name}") try: - return self.__send_request(url, payload) + response = self.__send_request(url, payload) + logger.info( + f"Successfully imported dataset {dataset_name} to Data Workbench" + ) + return response except requests.exceptions.RequestException as e: + trace_id = ( + _get_trace_id_from_response(e.response) + if e.response is not None + else None + ) logger.error(f"Error creating data catalog entry: {e}") - raise + return { + "error": "Failed to create data catalog entry.", + "correlation-id": trace_id, + } diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 9428f13..64707a9 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -2,7 +2,7 @@ import requests from unittest.mock import patch, MagicMock from dataworkbench.gateway import Gateway -from requests.exceptions import RequestException +import json @pytest.fixture def mock_gateway(): @@ -34,9 +34,20 @@ def test_import_dataset_success(mock_gateway, mock_post): def test_import_dataset_failure(mock_gateway, mock_post): """Test dataset import failure.""" - mock_post.side_effect = requests.exceptions.RequestException("Request failed") - with pytest.raises(RequestException): - mock_gateway.import_dataset("dataset_name", "dataset_description", "schema_id", {"tag": "value"}, "folder_id") + response_body = {"type":"BusinessError","traceId":"8b01e7eb14484611add6138618daf112"} + mock_response = MagicMock() + mock_response.return_value.status_code = 400 + mock_response.text = json.dumps(response_body) + + http_error = requests.exceptions.HTTPError() + http_error.response = mock_response + + mock_response.raise_for_status.side_effect = http_error + mock_post.return_value = mock_response + + result = mock_gateway.import_dataset("dataset_name", "dataset_description", "schema_id", {"tag": "value"}, "folder_id") + + assert result == {"error": "Failed to create data catalog entry.", "correlation-id": response_body["traceId"]} mock_post.assert_called_once()