From f5b39eb217ca06d7e2ae5927c25f9a9fc50b9739 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patrick=20Jos=C3=A9=20Pereira?= <patrickelectric@gmail.com>
Date: Wed, 30 Aug 2023 14:08:16 -0300
Subject: [PATCH] registry: Add aiohttp_retry to avoid connection problems
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Patrick José Pereira <patrickelectric@gmail.com>
---
 blueos_repository/registry.py | 11 +++++++----
 pyproject.toml                |  1 +
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/blueos_repository/registry.py b/blueos_repository/registry.py
index b87190c..2ecd30c 100644
--- a/blueos_repository/registry.py
+++ b/blueos_repository/registry.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, List, Optional
 
 import aiohttp
+import aiohttp_retry
 
 
 class Registry:
@@ -13,6 +14,8 @@ class Registry:
     docker_url: str = "https://hub.docker.com/"
     token: Optional[str] = None
 
+    retry_options = aiohttp_retry.ExponentialRetry(attempts=5)
+
     async def _get_token(self, repo: str) -> str:
         """[summary]
         Gets a token for dockerhub.com
@@ -32,7 +35,7 @@ async def _get_token(self, repo: str) -> str:
         }
 
         auth_url = f"https://auth.docker.io/token?service=registry.docker.io&scope=repository:{repo}:pull"
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp_retry.RetryClient(retry_options=self.retry_options) as session:
             async with session.get(auth_url + "/token", params=payload) as resp:
                 if resp.status != 200:
                     print(f"Error status {resp.status}")
@@ -43,7 +46,7 @@ async def fetch_remote_tags(self, repository: str) -> List[str]:
         """Fetches the tags available for an image in DockerHub"""
         print(f"fetching tags in {repository}")
         self.token = await self._get_token(repository)
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp_retry.RetryClient(retry_options=self.retry_options) as session:
             async with session.get(
                 f"{self.docker_url}/v2/repositories/{repository}/tags/?page_size=25&page=1&ordering=last_updated"
             ) as resp:
@@ -79,7 +82,7 @@ async def fetch_digest(self, digest: str, repository: str) -> str:
             "Authorization": f"Bearer {self.token}",
             "Accept": "application/vnd.docker.distribution.manifest.v2+json,application/vnd.oci.image.manifest.v1+json,application/vnd.oci.image.index.v1+json",
         }
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp_retry.RetryClient(retry_options=self.retry_options) as session:
             url = f"{self.index_url}/v2/{repository}/manifests/{digest}"
             async with session.get(url, headers=header) as resp:
                 if resp.status != 200:
@@ -97,7 +100,7 @@ async def fetch_labels(self, repo: str) -> Dict[str, Any]:
         }
         repository, tag = repo.split(":")
         print(f"fetching labels for {repository}:{tag}")
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp_retry.RetryClient(retry_options=self.retry_options) as session:
             url = f"{self.index_url}/v2/{repository}/manifests/{tag}"
             async with session.get(url, headers=header) as resp:
                 if resp.status != 200:
diff --git a/pyproject.toml b/pyproject.toml
index 60c30b0..08abf6e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,6 +10,7 @@ packages = [{include = "blueos_repository"}]
 
 [tool.poetry.dependencies]
 aiohttp = "3.7.4"
+aiohttp-retry = "2.8.3"
 json-five = "1.1.1"
 python = "^3.10"
 semver = "^2.13.0"