Add configuration options to skip ingestion errors

openverse_catalog/dags/providers/provider_api_scripts/provider_data_ingester.py

@@ -28,13 +28,11 @@ def __init__(self, error, traceback, query_params):
 
     def __str__(self):
         # Append query_param info to error message
-        return f"""{self.error}
-    query_params: {self.query_params}"""
+        return f"{self.error}\nquery_params: {self.query_params}"
 
-    def print_with_traceback(self):
+    def repr_with_traceback(self):
         # Append traceback
-        return f"""{str(self)}
-    {self.traceback}"""
+        return f"{str(self)}\n{self.traceback}"
 
 
 class ProviderDataIngester(ABC):
@@ -116,6 +114,7 @@ def __init__(self, conf: dict = None, date: str = None):
         # just these sets of params.
         self.override_query_params = None
         if query_params_list := conf.get("query_params_list"):
+            # Create a generator to facilitate fetching the next set of query_params.
             self.override_query_params = (qp for qp in query_params_list)
 
     def init_media_stores(self) -> dict[str, MediaStore]:
@@ -148,7 +147,7 @@ def ingest_records(self, **kwargs) -> None:
             query_params = self.get_query_params(query_params, **kwargs)
             if query_params is None:
                 # Break out of ingestion if no query_params are supplied. This can
-                # happen when the final `manual_query_params` is processed.
+                # happen when the final `override_query_params` is processed.
                 break
 
             try:
@@ -161,6 +160,16 @@ def ingest_records(self, **kwargs) -> None:
                     logger.info("Batch complete.")
                     should_continue = False
 
+            except AirflowException as error:
+                # AirflowExceptions should not be caught, as execution should not
+                # continue when the task is being stopped by Airflow.
+
+                # If errors have already been caught during processing, raise them
+                # as well.
+                if error_summary := self.get_ingestion_errors():
+                    raise error_summary from error
+                raise
+
             except Exception as error:
                 ingestion_error = IngestionError(
                     error, traceback.format_exc(), query_params
@@ -185,6 +194,16 @@ def ingest_records(self, **kwargs) -> None:
         self.commit_records()
 
         # If errors were caught during processing, raise them now
+        if error_summary := self.get_ingestion_errors():
+            raise error_summary
+
+    def get_ingestion_errors(self) -> AirflowException | None:
+        """ "
+        If any errors were skipped during ingestion, returns an AirflowException
+        with a summary of all errors.
+
+        It there are no errors to report, returns None.
+        """
         if self.ingestion_errors:
             # Log the affected query_params
             bad_query_params = ", \n".join(
@@ -195,11 +214,12 @@ def ingest_records(self, **kwargs) -> None:
                 f"{bad_query_params}"
             )
             errors_str = "\n".join(
-                e.print_with_traceback() for e in self.ingestion_errors
+                e.repr_with_traceback() for e in self.ingestion_errors
             )
-            raise AirflowException(
+            return AirflowException(
                 f"The following errors were encountered during ingestion:\n{errors_str}"
             )
+        return None
 
     def get_query_params(
         self, prev_query_params: Optional[Dict], **kwargs

tests/dags/providers/provider_api_scripts/test_provider_data_ingester.py

@@ -196,13 +196,11 @@ def test_ingest_records_stops_after_reaching_limit():
             assert process_batch_mock.call_count == 1
 
 
-def test_ingest_records_commits_on_exception(self):
+def test_ingest_records_commits_on_exception():
     with (
-        patch.object(self.ingester, "get_batch") as get_batch_mock,
-        patch.object(
-            self.ingester, "process_batch", return_value=3
-        ) as process_batch_mock,
-        patch.object(self.ingester, "commit_records") as commit_mock,
+        patch.object(ingester, "get_batch") as get_batch_mock,
+        patch.object(ingester, "process_batch", return_value=3) as process_batch_mock,
+        patch.object(ingester, "commit_records") as commit_mock,
     ):
         get_batch_mock.side_effect = [
             (EXPECTED_BATCH_DATA, True),  # First batch
@@ -212,7 +210,7 @@ def test_ingest_records_commits_on_exception(self):
         ]
 
         with pytest.raises(ValueError, match="Whoops :C"):
-            self.ingester.ingest_records()
+            ingester.ingest_records()
 
         # Check that get batch was only called thrice
         assert get_batch_mock.call_count == 3
@@ -286,41 +284,72 @@ def test_ingest_records_raises_IngestionError():
             (EXPECTED_BATCH_DATA, True),  # Second batch should not be reached
         ]
 
-        with pytest.raises(Exception) as error:
+        with pytest.raises(Exception, match="Mock exception message"):
             ingester.ingest_records()
 
         # By default, `skip_ingestion_errors` is False and get_batch_data
         # is no longer called after encountering an error
         assert get_batch_mock.call_count == 1
 
-        assert str(error.value) == "Mock exception message"
-
 
-def test_ingest_records_with_skip_ingestion_errors():
+@pytest.mark.parametrize(
+    "batches, expected_call_count, error_messages",
+    [
+        # Multiple errors are skipped
+        (
+            [
+                Exception("Mock exception 1"),
+                (EXPECTED_BATCH_DATA, True),  # First error
+                Exception("Mock exception 2"),
+                (EXPECTED_BATCH_DATA, False),  # Second error, `should_continue` False
+            ],
+            4,  # get_batch is called until `should_continue` is False, ignoring errors
+            ["Mock exception 1", "Mock exception 2"],
+        ),
+        # An AirflowException should not be skipped
+        (
+            [
+                (EXPECTED_BATCH_DATA, True),
+                AirflowException("An Airflow exception"),  # Second batch, should raise
+                (EXPECTED_BATCH_DATA, True),  # This batch should not be reached
+            ],
+            2,  # The final batch should not be reached
+            ["An Airflow exception"],
+        ),
+        # An AirflowException is raised, but there were already other ingestion errors
+        (
+            [
+                Exception("Some other exception"),  # First batch, should be skipped
+                AirflowException("An Airflow exception"),  # Second batch, should raise
+                (EXPECTED_BATCH_DATA, True),  # This batch should not be reached
+            ],
+            2,  # The final batch should not be reached
+            ["Some other exception"],  # Ingestion errors reported
+        ),
+    ],
+)
+def test_ingest_records_with_skip_ingestion_errors(
+    batches, expected_call_count, error_messages
+):
     ingester = MockProviderDataIngester({"skip_ingestion_errors": True})
 
     with (
         patch.object(ingester, "get_batch") as get_batch_mock,
         patch.object(ingester, "process_batch", return_value=10),
     ):
-        get_batch_mock.side_effect = [
-            Exception("Mock exception 1"),  # First batch
-            (EXPECTED_BATCH_DATA, True),  # Second batch
-            Exception("Mock exception 2"),  # Third batch
-            (EXPECTED_BATCH_DATA, False),  # Final batch
-        ]
+        get_batch_mock.side_effect = batches
 
         # ingest_records ultimately raises an exception
         with pytest.raises(AirflowException) as error:
             ingester.ingest_records()
 
         # get_batch was called four times before the exception was thrown,
         # despite errors being raised
-        assert get_batch_mock.call_count == 4
+        assert get_batch_mock.call_count == expected_call_count
 
         # All errors are summarized in the exception thrown at the end
-        assert "Mock exception 1" in str(error)
-        assert "Mock exception 2" in str(error)
+        for error_message in error_messages:
+            assert error_message in str(error)
 
 
 def test_commit_commits_all_stores():