diff --git a/prefect_gitlab/repositories.py b/prefect_gitlab/repositories.py index a2b0f1d..cb6673a 100644 --- a/prefect_gitlab/repositories.py +++ b/prefect_gitlab/repositories.py @@ -52,6 +52,7 @@ from prefect.utilities.asyncutils import sync_compatible from prefect.utilities.processutils import run_process from pydantic import VERSION as PYDANTIC_VERSION +from tenacity import retry, stop_after_attempt, wait_fixed, wait_random if PYDANTIC_VERSION.startswith("2."): from pydantic.v1 import Field, HttpUrl, validator @@ -60,6 +61,13 @@ from prefect_gitlab.credentials import GitLabCredentials +# Create get_directory retry settings + +MAX_CLONE_ATTEMPTS = 3 +CLONE_RETRY_MIN_DELAY_SECONDS = 1 +CLONE_RETRY_MIN_DELAY_JITTER_SECONDS = 0 +CLONE_RETRY_MAX_DELAY_JITTER_SECONDS = 3 + class GitLabRepository(ReadableDeploymentStorage): """ @@ -153,6 +161,15 @@ def _get_paths( return str(content_source), str(content_destination) @sync_compatible + @retry( + stop=stop_after_attempt(MAX_CLONE_ATTEMPTS), + wait=wait_fixed(CLONE_RETRY_MIN_DELAY_SECONDS) + + wait_random( + CLONE_RETRY_MIN_DELAY_JITTER_SECONDS, + CLONE_RETRY_MAX_DELAY_JITTER_SECONDS, + ), + reraise=True, + ) async def get_directory( self, from_path: Optional[str] = None, local_path: Optional[str] = None ) -> None: diff --git a/requirements.txt b/requirements.txt index 29e7b62..ce12fa3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ prefect>=2.13.5 python-gitlab>=3.12.0 +tenacity>=8.2.3 \ No newline at end of file diff --git a/tests/test_repositories.py b/tests/test_repositories.py index 7dc369b..7ed921d 100644 --- a/tests/test_repositories.py +++ b/tests/test_repositories.py @@ -15,7 +15,7 @@ import prefect_gitlab from prefect_gitlab.credentials import GitLabCredentials -from prefect_gitlab.repositories import GitLabRepository +from prefect_gitlab.repositories import GitLabRepository # noqa: E402 class TestGitLab: @@ -243,3 +243,22 @@ class p: assert set(os.listdir(tmp_dst)) == set([sub_dir_name]) assert set(os.listdir(Path(tmp_dst) / sub_dir_name)) == child_contents + + async def test_get_directory_retries(self, monkeypatch): + # Constants for the retry decorator + MAX_CLONE_ATTEMPTS = 3 + + # Create an instance of GitLabRepository + g = GitLabRepository(repository="https://gitlab.com/prefectHQ/prefect.git") + + # Prepare a MagicMock to simulate the process call within get_directory + mock = AsyncMock() + mock.return_value = AsyncMock(returncode=1) # Simulate failure + monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock) + + # Call get_directory and expect it to raise a RetryError after maximum attempts + with pytest.raises(OSError): + await g.get_directory() + print(mock.call_count) + # Verify that the function retried the expected number of times + assert mock.call_count == MAX_CLONE_ATTEMPTS