fix(ingest/mode): Handle 204 response and invalid json

Author: asikowitz

metadata-ingestion/src/datahub/ingestion/source/mode.py

@@ -2,6 +2,8 @@
 import logging
 import re
 import time
+from json import JSONDecodeError
+
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import lru_cache
@@ -15,6 +17,7 @@
 import yaml
 from liquid import Template, Undefined
 from pydantic import Field, validator
+from requests import Response
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import ConnectionError
 from requests.models import HTTPBasicAuth, HTTPError
@@ -1466,10 +1469,13 @@ def _get_request_json(self, url: str) -> Dict:
 
         @r.wraps
         def get_request():
+            response: Optional[Response] = None
             try:
                 response = self.session.get(
                     url, timeout=self.config.api_options.timeout
                 )
+                if response.status_code == 204:
+                    return {}
                 return response.json()
             except HTTPError as http_error:
                 error_response = http_error.response
@@ -1481,6 +1487,10 @@ def get_request():
                     raise HTTPError429
 
                 raise http_error
+            except JSONDecodeError as json_error:
+                raise HTTPError(
+                    f"{json_error.__class__.__name__}: {json_error}"
+                ) from json_error
 
         return get_request()
 

metadata-ingestion/tests/integration/mode/test_mode.py

@@ -1,11 +1,14 @@
 import json
 import pathlib
+from typing import Sequence
 from unittest.mock import patch
 
+import pytest
 from freezegun import freeze_time
 from requests.models import HTTPError
 
 from datahub.configuration.common import PipelineExecutionError
+from datahub.ingestion.api.source import StructuredLogEntry
 from datahub.ingestion.run.pipeline import Pipeline
 from tests.test_helpers import mce_helpers
 
@@ -28,13 +31,14 @@
     "https://app.mode.com/api/acryl/reports/24f66e1701b6/queries": "dataset_queries_24f66e1701b6.json",
 }
 
-RESPONSE_ERROR_LIST = ["https://app.mode.com/api/acryl/spaces/75737b70402e/reports"]
+ERROR_URL = "https://app.mode.com/api/acryl/spaces/75737b70402e/reports"
 
 test_resources_dir = pathlib.Path(__file__).parent
 
 
 class MockResponse:
     def __init__(self, error_list, status_code):
+        self.text = None
         self.json_data = None
         self.error_list = error_list
         self.status_code = status_code
@@ -49,6 +53,14 @@ def mount(self, prefix, adaptor):
         return self
 
     def get(self, url, timeout=40):
+        if self.error_list is not None and self.url in self.error_list:
+            http_error_msg = "{} Client Error: {} for url: {}".format(
+                400,
+                "Simulate error",
+                self.url,
+            )
+            raise HTTPError(http_error_msg, response=self)
+
         self.url = url
         self.timeout = timeout
         response_json_path = f"{test_resources_dir}/setup/{JSON_RESPONSE_MAP.get(url)}"
@@ -57,29 +69,46 @@ def get(self, url, timeout=40):
             self.json_data = data
         return self
 
-    def raise_for_status(self):
-        if self.error_list is not None and self.url in self.error_list:
-            http_error_msg = "{} Client Error: {} for url: {}".format(
-                400,
-                "Simulate error",
-                self.url,
-            )
-            raise HTTPError(http_error_msg, response=self)
+
+class MockResponseJson(MockResponse):
+    def __init__(
+        self,
+        status_code=200,
+        *,
+        json_empty_list: Sequence[str] = (),
+        json_error_list: Sequence[str] = (),
+    ):
+        super().__init__(None, status_code)
+        self.json_empty_list = json_empty_list
+        self.json_error_list = json_error_list
+
+    def json(self):
+        if self.url in self.json_empty_list:
+            return json.loads("")  # Shouldn't be called
+        if self.url in self.json_error_list:
+            return json.loads("{")
+        return super().json()
+
+    def get(self, url, timeout=40):
+        response = super().get(url, timeout)
+        if self.url in self.json_empty_list:
+            response.status_code = 204
+        return response
 
 
-def mocked_requests_sucess(*args, **kwargs):
+def mocked_requests_success(*args, **kwargs):
     return MockResponse(None, 200)
 
 
 def mocked_requests_failure(*args, **kwargs):
-    return MockResponse(RESPONSE_ERROR_LIST, 200)
+    return MockResponse([ERROR_URL], 200)
 
 
 @freeze_time(FROZEN_TIME)
 def test_mode_ingest_success(pytestconfig, tmp_path):
     with patch(
         "datahub.ingestion.source.mode.requests.Session",
-        side_effect=mocked_requests_sucess,
+        side_effect=mocked_requests_success,
     ):
         pipeline = Pipeline.create(
             {
@@ -142,8 +171,89 @@ def test_mode_ingest_failure(pytestconfig, tmp_path):
             }
         )
         pipeline.run()
-        try:
+        with pytest.raises(PipelineExecutionError) as exec_error:
             pipeline.raise_from_status()
-        except PipelineExecutionError as exec_error:
-            assert exec_error.args[0] == "Source reported errors"
-            assert len(exec_error.args[1].failures) == 1
+        assert exec_error.value.args[0] == "Source reported errors"
+        assert len(exec_error.value.args[1].failures) == 1
+        error_dict: StructuredLogEntry
+        _level, error_dict = exec_error.value.args[1].failures[0]
+        error = next(iter(error_dict.context))
+        assert "Simulate error" in error
+        assert ERROR_URL in error
+
+
+@freeze_time(FROZEN_TIME)
+def test_mode_ingest_json_empty(pytestconfig, tmp_path):
+    with patch(
+        "datahub.ingestion.source.mode.requests.Session",
+        side_effect=lambda *args, **kwargs: MockResponseJson(
+            json_empty_list=["https://app.mode.com/api/modeuser"]
+        ),
+    ):
+        global test_resources_dir
+        test_resources_dir = pytestconfig.rootpath / "tests/integration/mode"
+
+        pipeline = Pipeline.create(
+            {
+                "run_id": "mode-test",
+                "source": {
+                    "type": "mode",
+                    "config": {
+                        "token": "xxxx",
+                        "password": "xxxx",
+                        "connect_uri": "https://app.mode.com/",
+                        "workspace": "acryl",
+                    },
+                },
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/mode_mces.json",
+                    },
+                },
+            }
+        )
+        pipeline.run()
+        pipeline.raise_from_status(raise_warnings=True)
+
+
+@freeze_time(FROZEN_TIME)
+def test_mode_ingest_json_failure(pytestconfig, tmp_path):
+    with patch(
+        "datahub.ingestion.source.mode.requests.Session",
+        side_effect=lambda *args, **kwargs: MockResponseJson(
+            json_error_list=["https://app.mode.com/api/modeuser"]
+        ),
+    ):
+        global test_resources_dir
+        test_resources_dir = pytestconfig.rootpath / "tests/integration/mode"
+
+        pipeline = Pipeline.create(
+            {
+                "run_id": "mode-test",
+                "source": {
+                    "type": "mode",
+                    "config": {
+                        "token": "xxxx",
+                        "password": "xxxx",
+                        "connect_uri": "https://app.mode.com/",
+                        "workspace": "acryl",
+                    },
+                },
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/mode_mces.json",
+                    },
+                },
+            }
+        )
+        pipeline.run()
+        pipeline.raise_from_status(raise_warnings=False)
+        with pytest.raises(PipelineExecutionError) as exec_error:
+            pipeline.raise_from_status(raise_warnings=True)
+        assert len(exec_error.value.args[1].warnings) > 0
+        error_dict: StructuredLogEntry
+        _level, error_dict = exec_error.value.args[1].warnings[0]
+        error = next(iter(error_dict.context))
+        assert "JSONDecodeError" in error