Add trace id to dwb import error

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,
+            }

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()