fix: add session retries to DDN dataflow client

cg/meta/archive/archive.py

@@ -3,6 +3,7 @@
 
 import rich_click as click
 from housekeeper.store.models import Archive, File
+from requests import RequestException
 
 from cg.apps.housekeeper.hk import HousekeeperAPI
 from cg.constants import SequencingFileTag
@@ -183,6 +184,8 @@
                 self.housekeeper_api.update_archive_retrieved_at(
                     old_retrieval_job_id=task_id, new_retrieval_job_id=None
                 )
+        except RequestException as error:
+            LOG.error(f"Failed to fetch status for job with id {task_id}: {error}")
 
     def sort_archival_ids_on_archive_location(
         self, archive_entries: list[Archive]

cg/meta/archive/ddn/ddn_data_flow_client.py

@@ -2,15 +2,17 @@
 
 import logging
 from datetime import datetime
+from http import HTTPStatus
 from pathlib import Path
 from urllib.parse import urljoin
 
 from housekeeper.store.models import File
 from pydantic import BaseModel
-from requests import Response
+from requests import Response, Session
+from requests.adapters import HTTPAdapter
+from urllib3 import Retry
 
 from cg.exc import ArchiveJobFailedError, DdnDataflowAuthenticationError, DdnDataflowDeleteFileError
-from cg.io.api import get, post
 from cg.meta.archive.ddn.constants import (
     DELETE_FILE_SUCCESSFUL_MESSAGE,
     DESTINATION_ATTRIBUTE,
@@ -55,8 +57,32 @@ def __init__(self, config: DataFlowConfig):
             "Content-Type": "application/json",
             "accept": "application/json",
         }
+        self.session = self._get_session()
+        self.no_retry_session = Session()
         self._set_auth_tokens()
 
+    def _get_session(self) -> Session:
+        session = Session()
+        self._configure_retries(session)
+        return session
+
+    @staticmethod
+    def _configure_retries(session: Session) -> None:
+        """Configures retries for the session."""
+        retry_strategy = Retry(
+            total=5,
+            status_forcelist=[
+                HTTPStatus.TOO_MANY_REQUESTS,
+                HTTPStatus.INTERNAL_SERVER_ERROR,
+                HTTPStatus.BAD_GATEWAY,
+                HTTPStatus.SERVICE_UNAVAILABLE,
+                HTTPStatus.GATEWAY_TIMEOUT,
+            ],
+            backoff_factor=2,
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("https://", adapter)
+
     def _set_auth_tokens(self) -> None:
         """Retrieves and sets auth and refresh token from the REST-API."""
         auth_token: AuthToken = self._get_auth_token()
@@ -66,7 +92,7 @@ def _set_auth_tokens(self) -> None:
 
     def _get_auth_token(self) -> AuthToken:
         """Retrieves auth and refresh token from the REST-API."""
-        response: Response = post(
+        response: Response = self.session.post(
             url=urljoin(base=self.url, url=DataflowEndpoints.GET_AUTH_TOKEN),
             headers=self.headers,
             json=AuthPayload(
@@ -87,7 +113,7 @@ def _refresh_auth_token(self) -> None:
         self.token_expiration: datetime = datetime.fromtimestamp(auth_token.expire)
 
     def _get_refreshed_auth_token(self) -> AuthToken:
-        response: Response = post(
+        response: Response = self.session.post(
             url=urljoin(base=self.url, url=DataflowEndpoints.REFRESH_AUTH_TOKEN),
             headers=self.headers,
             json=RefreshPayload(refresh=self.refresh_token).model_dump(),
@@ -204,6 +230,7 @@ def delete_file(self, file_and_sample: FileAndSample) -> None:
             endpoint=DataflowEndpoints.DELETE_FILE,
             headers=dict(self.headers, **self.auth_header),
             body=delete_file_payload,
+            session=self.no_retry_session,
         )
         delete_file_response = DeleteFileResponse.model_validate(response.json())
         if delete_file_response.message != DELETE_FILE_SUCCESSFUL_MESSAGE:
@@ -226,11 +253,15 @@ def _retrieve_files(self, body: BaseModel, headers: dict) -> RetrievalResponse:
         return RetrievalResponse.model_validate(response.json())
 
     def _post_request(
-        self, body: BaseModel, endpoint: DataflowEndpoints, headers: dict
+        self,
+        body: BaseModel,
+        endpoint: DataflowEndpoints,
+        headers: dict,
+        session: Session | None = None,
     ) -> Response:
         """Posts a request with the provided body and headers to the given endpoint."""
         LOG.info(get_request_log(body=body.model_dump(by_alias=True)))
-        response: Response = post(
+        response: Response = (session or self.session).post(
             url=urljoin(self.url, endpoint),
             headers=headers,
             json=body.model_dump(by_alias=True),
@@ -243,7 +274,7 @@ def _get_job_status(self, headers: dict, job_id: int) -> GetJobStatusResponse:
         """Gets the job status for the provided job_id."""
         LOG.info(f"Sending GET request for job {job_id}")
         url: str = urljoin(self.url, DataflowEndpoints.GET_JOB_STATUS + str(job_id))
-        response: Response = get(
+        response: Response = self.session.get(
             url=url,
             headers=headers,
             json={},

tests/meta/archive/conftest.py

@@ -7,7 +7,7 @@
 from click.testing import CliRunner
 from housekeeper.store.models import Bundle, File
 from pytest_mock import MockerFixture
-from requests import Response
+from requests import Response, Session
 
 from cg.apps.housekeeper.hk import HousekeeperAPI
 from cg.constants import SequencingFileTag
@@ -16,7 +16,6 @@
 from cg.constants.subject import Sex
 from cg.io.controller import WriteStream
 from cg.meta.archive.archive import SpringArchiveAPI
-from cg.meta.archive.ddn import ddn_data_flow_client
 from cg.meta.archive.ddn.constants import ROOT_TO_TRIM
 from cg.meta.archive.ddn.ddn_data_flow_client import DDNDataFlowClient
 from cg.meta.archive.ddn.models import AuthToken, MiriaObject, TransferPayload
@@ -130,7 +129,7 @@ def ddn_dataflow_client(
         },
     ).encode()
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=mock_ddn_auth_success_response,
     )

tests/meta/archive/test_archive_api.py

@@ -6,14 +6,13 @@
 import pytest
 from housekeeper.store.models import File, Version
 from pytest_mock import MockerFixture
-from requests import HTTPError, Response
+from requests import HTTPError, Response, Session
 
 from cg.apps.housekeeper.hk import HousekeeperAPI
 from cg.constants.archiving import ArchiveLocations
 from cg.constants.housekeeper_tags import SequencingFileTag
 from cg.exc import MissingFilesError, SampleFilesCurrentlyArchivingError
 from cg.meta.archive.archive import ARCHIVE_HANDLERS, FileAndSample, SpringArchiveAPI
-from cg.meta.archive.ddn import ddn_data_flow_client
 from cg.meta.archive.ddn.constants import FAILED_JOB_STATUSES, ONGOING_JOB_STATUSES, JobStatus
 from cg.meta.archive.ddn.ddn_data_flow_client import DDNDataFlowClient
 from cg.meta.archive.ddn.models import AuthToken, GetJobStatusResponse, MiriaObject
@@ -145,7 +144,7 @@ def test_archive_all_non_archived_spring_files(
         return_value=test_auth_token,
     )
     mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -307,7 +306,7 @@ def test_retrieve_case(
         return_value=test_auth_token,
     )
     mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -400,7 +399,7 @@ def test_retrieve_sample(
         return_value=test_auth_token,
     )
     mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -483,7 +482,7 @@ def test_retrieve_order(
         return_value=test_auth_token,
     )
     mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -536,7 +535,7 @@ def test_delete_file_raises_http_error(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=failed_delete_file_response,
     )
@@ -579,7 +578,7 @@ def test_delete_file_success(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_delete_file_response,
     )

tests/meta/archive/test_archive_cli.py

@@ -4,12 +4,11 @@
 from click.testing import CliRunner
 from housekeeper.store.models import Archive, File
 from pytest_mock import MockerFixture
-from requests import HTTPError, Response
+from requests import HTTPError, Response, Session
 
 from cg.cli.archive import archive_spring_files, delete_file, update_job_statuses
 from cg.constants import EXIT_SUCCESS, SequencingFileTag
 from cg.constants.archiving import ArchiveLocations
-from cg.meta.archive.ddn import ddn_data_flow_client
 from cg.meta.archive.ddn.constants import FAILED_JOB_STATUSES, ONGOING_JOB_STATUSES, JobStatus
 from cg.meta.archive.ddn.ddn_data_flow_client import DDNDataFlowClient
 from cg.meta.archive.ddn.models import ArchivalResponse, AuthToken, GetJobStatusResponse
@@ -64,7 +63,7 @@ def test_archive_spring_files_success(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -111,7 +110,7 @@ def test_get_archival_job_status(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -171,7 +170,7 @@ def test_get_retrieval_job_status(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -232,7 +231,7 @@ def test_delete_file_raises_http_error(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=failed_delete_file_response,
     )
@@ -278,7 +277,7 @@ def test_delete_file_success(
         return_value=test_auth_token,
     )
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_delete_file_response,
     )

tests/meta/archive/test_archiving.py

@@ -5,12 +5,11 @@
 
 import pytest
 from pytest_mock import MockerFixture
-from requests import Response
+from requests import Response, Session
 
 from cg.constants.constants import FileFormat
 from cg.exc import DdnDataflowAuthenticationError
 from cg.io.controller import WriteStream
-from cg.meta.archive.ddn import ddn_data_flow_client
 from cg.meta.archive.ddn.constants import (
     DESTINATION_ATTRIBUTE,
     OSTYPE,
@@ -120,7 +119,7 @@ def test_ddn_dataflow_client_initialization(
 
     # WHEN initializing the DDNDataFlowClient class with the valid DDNConfig object
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_response,
     )
@@ -143,7 +142,7 @@ def test_set_auth_tokens(
 
     # WHEN initializing the DDNDataFlowClient class with the valid DDNConfig object
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_response,
     )
@@ -167,7 +166,7 @@ def test_ddn_dataflow_client_initialization_invalid_credentials(
 
     # WHEN initializing the DDNDataFlowClient class with the invalid credentials
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=unauthorized_response,
     )
@@ -258,7 +257,7 @@ def test__refresh_auth_token(
 
     # WHEN refreshing the auth token
     mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_response,
     )
@@ -283,7 +282,7 @@ def test_archive_file(
     # GIVEN two paths that should be archived
     # WHEN running the archive method and providing two paths
     mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client,
+        Session,
         "post",
         return_value=ok_miria_response,
     )
@@ -324,9 +323,7 @@ def test_retrieve_files(
     # GIVEN a file and sample which is archived
 
     # WHEN running retrieve_files and providing a FileAndSample object
-    mock_request_submitter = mocker.patch.object(
-        ddn_data_flow_client, "post", return_value=ok_miria_response
-    )
+    mock_request_submitter = mocker.patch.object(Session, "post", return_value=ok_miria_response)
     job_id: int = ddn_dataflow_client.retrieve_files(files_and_samples=[file_and_sample])
 
     # THEN an integer should be returned