add more error handling for API responses

Author: ialarmedalien

.dlt/config.toml

@@ -6,6 +6,16 @@ http_show_error_body = "true"
 log_level = "INFO"
 log_format = "JSON"
 
+# REST Client retries
+# default: 60
+request_timeout = 60
+# default: 5
+request_max_attempts = 2
+# default: 1
+request_backoff_factor = 1.5
+# default: 300
+request_max_retry_delay = 600
+
 # do not raise an error if a job fails
 load.raise_on_failed_jobs = "false"
 

src/cdm_data_loaders/pipelines/ncbi_rest_api.py

@@ -19,6 +19,7 @@
 )
 from pydantic import AliasChoices, Field, model_validator
 from pydantic_settings import CliSuppress, SettingsConfigDict
+from requests.exceptions import HTTPError
 
 from cdm_data_loaders.pipelines.core import (
     run_cli,
@@ -38,6 +39,7 @@
 
 DATASET = "dataset"
 ANNOTATION = "annotation"
+ERROR = "error"
 
 dlt_logger = logging.getLogger("dlt")
 
@@ -143,6 +145,49 @@ def save_raw_response(config: Settings, response: Response, *_: Any, **__: Any)
 )
 
 
+def add_error(
+    error_list: list[dict[str, Any]],
+    error: Exception,
+    error_from: str,
+    assembly_id: str | None = None,
+    assembly_id_list: list[str] | None = None,
+) -> None:
+    """Add an error to the list of output errors.
+
+    :param error_list: running list of errors
+    :type error_list: list[dict[str, Any]]
+    :param error: the error object from the exception handler
+    :type error: Exception
+    :param error_from: what type of request was being made when the error occurred
+    :type error_from: str
+    :param assembly_id: ID of the assembly being fetched when the error occurred, defaults to None
+    :type assembly_id: str | None, optional
+    :param assembly_id_list: list of IDs being fetched when the error occurred, defaults to None
+    :type assembly_id_list: list[str] | None, optional
+    """
+    err_args = {
+        "assembly_id": assembly_id or None,
+        "assembly_id_list": assembly_id_list or None,
+        "error_class": type(error).__name__,
+        "error_from": error_from,
+        "message": str(error),
+        "request_url": None,
+        "status": None,
+        "reason": None,
+    }
+
+    if isinstance(error, HTTPError):
+        # save the URL, status code, and error message
+        err_args = {
+            **err_args,
+            "request_url": error.request.url,
+            "status": error.response.status_code,
+            "reason": error.response.reason,
+        }
+
+    error_list.append(err_args)
+
+
 def get_assembly_reports(assembly_id_list: list[str]) -> dict[str, Any]:
     """Retrieve dataset and annotation reports for a list of IDs from the NCBI datasets API.
 
@@ -154,14 +199,27 @@ def get_assembly_reports(assembly_id_list: list[str]) -> dict[str, Any]:
     if not assembly_id_list:
         return {}
 
+    errors = []
+
     # N.b. invalid IDs will not be present in dataset_reports
-    dataset_reports = get_dataset_reports(assembly_id_list)
-    annotation_reports = {assembly_id: get_annotation_report(assembly_id) for assembly_id in assembly_id_list}
+    dataset_reports = {}
+    try:
+        dataset_reports = get_dataset_reports(assembly_id_list)
+    except Exception as e:  # noqa: BLE001
+        add_error(errors, e, "dataset_report", assembly_id_list=assembly_id_list)
+
+    annotation_reports: dict[str, Any] = {}
+    for assembly_id in assembly_id_list:
+        try:
+            annotation_reports[assembly_id] = get_annotation_report(assembly_id)
+        except Exception as e:  # noqa: BLE001
+            add_error(errors, e, "annotation_report", assembly_id=assembly_id)
 
     # ensure every assembly_id in the list has either the downloaded dataset_report or None
     return {
         DATASET: {assembly_id: dataset_reports.get(assembly_id) for assembly_id in assembly_id_list},
         ANNOTATION: {assembly_id: annotation_reports.get(assembly_id) for assembly_id in assembly_id_list},
+        ERROR: errors,
     }
 
 
@@ -172,6 +230,7 @@ def get_dataset_reports(assembly_id_list: list[str]) -> dict[str, None | dict[st
 
     dlt_logger.info("fetching dataset reports for:\n%s", ", ".join(sorted(assembly_id_list)))
     assembly_dataset_reports = []
+
     for page in ncbi_genome_client.paginate(
         f"{'%2C'.join(assembly_id_list)}/dataset_report",
         params={
@@ -182,7 +241,7 @@ def get_dataset_reports(assembly_id_list: list[str]) -> dict[str, None | dict[st
         assembly_dataset_reports.extend(page)
 
     # return dataset reports, indexed by assembly_id
-    # invalid IDs are silently dropped by the API
+    # invalid IDs are silently dropped by the NCBI REST API
     datasets = {report.get("accession"): report for report in assembly_dataset_reports}
     # fill in the missing gaps in assembly_id_list with None
     return {assembly_id: datasets.get(assembly_id) for assembly_id in assembly_id_list}
@@ -192,6 +251,7 @@ def get_annotation_report(assembly_id: str) -> list[dict[str, Any]] | None:
     """Fetch the annotation report for an assembly from the NCBI datasets REST API."""
     dlt_logger.info("fetching annotation report for %s", assembly_id)
     page_data = []
+
     for page in ncbi_genome_client.paginate(
         f"{assembly_id}/annotation_report",
         params={
@@ -200,6 +260,7 @@ def get_annotation_report(assembly_id: str) -> list[dict[str, Any]] | None:
         hooks=REST_CLIENT_HOOKS,  # type: ignore[reportArgumentType]
     ):
         page_data.extend(page)
+
     # page_data is empty if the ID is invalid
     return page_data or None
 
@@ -239,6 +300,11 @@ def assembly_report_parser(
 
     dataset_reports: dict[str, dict[str, Any]] = assembly_reports.get(DATASET)  # type: ignore[reportAssignmentType]
     annotation_reports: dict[str, list[dict[str, Any]]] = assembly_reports.get(ANNOTATION)  # type: ignore[reportAssignmentType]
+    error_reports: list[dict[str, Any]] = assembly_reports.get(ERROR)  # type: ignore[reportAssignmentType]
+
+    if error_reports:
+        yield dlt.mark.with_table_name(error_reports, "ncbi_import_error")
+
     # yield the raw data to save as tables
     yield dlt.mark.with_table_name(
         [{"assembly_id": assembly_id, **(report or {})} for assembly_id, report in dataset_reports.items()],

tests/pipelines/test_ncbi_rest_api.py

@@ -2,18 +2,21 @@
 
 from pathlib import Path
 from typing import Any
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
 import pytest
 from pydantic import ValidationError
 from pydantic_settings import CliApp
+from requests import HTTPError
 
 from cdm_data_loaders.pipelines import ncbi_rest_api as ncbi_module
 from cdm_data_loaders.pipelines.cts_defaults import INPUT_MOUNT, OUTPUT_MOUNT, VALID_DESTINATIONS
 from cdm_data_loaders.pipelines.ncbi_rest_api import (
     ANNOTATION,
     DATASET,
     DATASET_NAME,
+    ERROR,
     Settings,
     assembly_list,
     cli,
@@ -148,7 +151,7 @@ def test_invalid_destination_raises(bad: str) -> None:
 @pytest.mark.parametrize("output", ["-o", "--output"])
 @pytest.mark.parametrize("dev_mode_flag", ["--dev", "--dev-mode", "--dev_mode"])
 @pytest.mark.parametrize("batch_size", ["-b", "--batch-size", "--batch_size"])
-def test_cli_all_variants(
+def test_cli_all_variants(  # noqa: PLR0913
     input_dir: str,
     destination: str,
     use_output_dir_for_pipeline_metadata: str,
@@ -378,6 +381,14 @@ def test_get_annotation_report_multi_page() -> None:
     check_annotation_report(annotation_report, ID_WITH_2K_ANNOTS)
 
 
+@mock.patch("tenacity.nap.time.sleep", MagicMock())
+@pytest.mark.vcr
+def test_get_annotation_report_multi_page_err() -> None:
+    """An error in the middle of a multi-page retrieval should stop the whole retrieval process."""
+    with pytest.raises(HTTPError, match="500 Server Error: Internal Server Error for url"):
+        get_annotation_report(ID_WITH_2K_ANNOTS)
+
+
 @pytest.mark.default_cassette("test_get_assembly_reports.yaml")
 @pytest.mark.vcr
 def test_get_annotation_report_invalid_id() -> None:
@@ -394,13 +405,115 @@ def test_get_assembly_reports_empty_id_list() -> None:
 def test_get_assembly_reports() -> None:
     """Test the retrieval of annotation and dataset reports."""
     assembly_reports = get_assembly_reports(ALL_IDS)
-    assert set(assembly_reports) == {DATASET, ANNOTATION}
+    assert set(assembly_reports) == {DATASET, ANNOTATION, ERROR}
     for datatype in [DATASET, ANNOTATION]:
         assert set(assembly_reports[datatype]) == set(ALL_IDS)
         assert assembly_reports[datatype][INVALID_ID] is None
     for assembly_id in [ID_WITH_2K_ANNOTS, ID_WITH_500_ANNOTS]:
         check_annotation_report(assembly_reports[ANNOTATION][assembly_id], assembly_id)
         check_dataset_report(assembly_reports[DATASET][assembly_id], assembly_id)
+    assert assembly_reports[ERROR] == []
+
+
+@mock.patch("tenacity.nap.time.sleep", MagicMock())
+@pytest.mark.default_cassette("test_get_assembly_reports_annotation_report_errors.yaml")
+@pytest.mark.vcr
+def test_get_assembly_reports_annotation_report_errors() -> None:
+    """Test the retrieval of assembly data when errors occur fetching annotation reports."""
+    assembly_reports = get_assembly_reports(ALL_IDS)
+    assert set(assembly_reports) == {DATASET, ANNOTATION, ERROR}
+    for datatype in [DATASET, ANNOTATION]:
+        assert set(assembly_reports[datatype]) == set(ALL_IDS)
+        assert assembly_reports[datatype][INVALID_ID] is None
+    for assembly_id in [ID_WITH_2K_ANNOTS, ID_WITH_500_ANNOTS]:
+        check_dataset_report(assembly_reports[DATASET][assembly_id], assembly_id)
+    # ID_WITH_500 succeeds, ID_WITH_2K does not
+    check_annotation_report(assembly_reports[ANNOTATION][ID_WITH_500_ANNOTS], ID_WITH_500_ANNOTS)
+    assert assembly_reports[ANNOTATION][ID_WITH_2K_ANNOTS] is None
+
+    assert assembly_reports[ERROR] == [
+        {
+            "assembly_id": ID_WITH_2K_ANNOTS,
+            "assembly_id_list": None,
+            "error_class": "HTTPError",
+            "error_from": "annotation_report",
+            "message": '500 Server Error: Internal Server Error for url: https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000003135.1/annotation_report?page_size=1000&page_token=eNrjYos2NDAwjAUABagBiw\nResponse: {"error":"Internal Server Error","code":500,"message":"Internal Server Error (For more help, see the NCBI Datasets Documentation at https://www.ncbi.nlm.nih.gov/datasets/docs/) (1D3311AB4E92F955000055F41870A1E3.1.1)"}\n',
+            "request_url": "https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000003135.1/annotation_report?page_size=1000&page_token=eNrjYos2NDAwjAUABagBiw",
+            "status": 500,
+            "reason": "Internal Server Error",
+        }
+    ]
+
+
+@mock.patch("tenacity.nap.time.sleep", MagicMock())
+@pytest.mark.vcr
+def test_get_assembly_reports_dataset_report_errors() -> None:
+    """Test the retrieval of assembly data when an error occurs fetching dataset reports."""
+    assembly_reports = get_assembly_reports(ALL_IDS)
+    assert set(assembly_reports) == {DATASET, ANNOTATION, ERROR}
+    for datatype in [DATASET, ANNOTATION]:
+        assert set(assembly_reports[datatype]) == set(ALL_IDS)
+        assert assembly_reports[datatype][INVALID_ID] is None
+    for assembly_id in [ID_WITH_2K_ANNOTS, ID_WITH_500_ANNOTS]:
+        check_annotation_report(assembly_reports[ANNOTATION][assembly_id], assembly_id)
+        assert assembly_reports[DATASET][assembly_id] is None
+
+    assert assembly_reports[ERROR] == [
+        {
+            "assembly_id": None,
+            "assembly_id_list": ALL_IDS,
+            "error_class": "HTTPError",
+            "error_from": "dataset_report",
+            "message": '404 Client Error: Not Found for url: https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1%2CGCF_000003135.1%2Cinvalid_id/dataset_report?page_size=1000\nResponse: {"error":"Not Found","code":404,"message":"Your request is invalid. (For more help, see the NCBI Datasets Documentation at https://www.ncbi.nlm.nih.gov/datasets/docs/)"}\n',
+            "request_url": "https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1%2CGCF_000003135.1%2Cinvalid_id/dataset_report?page_size=1000",
+            "status": 404,
+            "reason": "Not Found",
+        }
+    ]
+
+
+@mock.patch("tenacity.nap.time.sleep", MagicMock())
+@pytest.mark.vcr
+def test_get_assembly_reports_total_wipeout() -> None:
+    """Test the retrieval of assembly data when all queries fail."""
+    # TODO: another type of error?
+    output = get_assembly_reports(ALL_IDS)
+    assert output == {
+        DATASET: dict.fromkeys(ALL_IDS),
+        ANNOTATION: dict.fromkeys(ALL_IDS),
+        ERROR: [
+            {
+                "assembly_id": None,
+                "assembly_id_list": ALL_IDS,
+                "error_class": "HTTPError",
+                "error_from": "dataset_report",
+                "message": '404 Client Error: Not Found for url: https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1%2CGCF_000003135.1%2Cinvalid_id/dataset_report?page_size=1000\nResponse: {"error":"Not Found","code":404,"message":"Your request is invalid. (For more help, see the NCBI Datasets Documentation at https://www.ncbi.nlm.nih.gov/datasets/docs/)"}\n',
+                "request_url": "https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1%2CGCF_000003135.1%2Cinvalid_id/dataset_report?page_size=1000",
+                "status": 404,
+                "reason": "Not Found",
+            },
+            {
+                "assembly_id": ID_WITH_500_ANNOTS,
+                "assembly_id_list": None,
+                "error_class": "HTTPError",
+                "error_from": "annotation_report",
+                "message": '404 Client Error: Not Found for url: https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1/annotation_report?page_size=1000\nResponse: {"error":"Not Found","code":404,"message":"Your request is invalid. (For more help, see the NCBI Datasets Documentation at https://www.ncbi.nlm.nih.gov/datasets/docs/)"}\n',
+                "request_url": "https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000007725.1/annotation_report?page_size=1000",
+                "status": 404,
+                "reason": "Not Found",
+            },
+            {
+                "assembly_id": ID_WITH_2K_ANNOTS,
+                "assembly_id_list": None,
+                "error_class": "HTTPError",
+                "error_from": "annotation_report",
+                "message": '500 Server Error: Internal Server Error for url: https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000003135.1/annotation_report?page_size=1000&page_token=eNrjYos2NDAwjAUABagBiw\nResponse: {"error":"Internal Server Error","code":500,"message":"Internal Server Error (For more help, see the NCBI Datasets Documentation at https://www.ncbi.nlm.nih.gov/datasets/docs/) (1D32DF2AAA9FBB1500003B3C46C243D3.1.1)"}\n',
+                "request_url": "https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/GCF_000003135.1/annotation_report?page_size=1000&page_token=eNrjYos2NDAwjAUABagBiw",
+                "status": 500,
+                "reason": "Internal Server Error",
+            },
+        ],
+    }
 
 
 @pytest.mark.skip("FIXME: not working, possibly due to parallelization?")