Merge pull request #34 from PrefectHQ/feature/retries

Add retries and backoff to get_directory

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:

requirements.txt

@@ -1,2 +1,3 @@
 prefect>=2.13.5
 python-gitlab>=3.12.0
+tenacity>=8.2.3

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