From 9348379e7ac6acbe1644a9d36c7d8ff983e3bfad Mon Sep 17 00:00:00 2001
From: ahdamin <ahmed.amin@scilifelab.se>
Date: Tue, 8 Oct 2024 11:24:54 +0200
Subject: [PATCH] Add retry logic with tenacity

---
 genotype_api/database/base_handler.py |  9 +++++++++
 genotype_api/database/store.py        |  9 +++++++++
 poetry.lock                           | 17 ++++++++++++++++-
 pyproject.toml                        |  1 +
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/genotype_api/database/base_handler.py b/genotype_api/database/base_handler.py
index 3e3d295..7c0e48f 100644
--- a/genotype_api/database/base_handler.py
+++ b/genotype_api/database/base_handler.py
@@ -1,10 +1,13 @@
 from dataclasses import dataclass
 from typing import Type
 
+from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
 from sqlalchemy.orm import DeclarativeBase, Query
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
+from genotype_api.config import settings
 from genotype_api.database.models import Analysis, Sample
 
 
@@ -15,6 +18,12 @@ class BaseHandler:
     def __init__(self, session: AsyncSession):
         self.session = session
 
+    @retry(
+        stop=stop_after_attempt(settings.max_retries),
+        wait=wait_fixed(settings.retry_delay),
+        retry=retry_if_exception_type(OperationalError),
+        reraise=True,
+    )
     def _get_query(self, table: Type[DeclarativeBase]) -> Query:
         """Return a query for the given table."""
         return select(table)
diff --git a/genotype_api/database/store.py b/genotype_api/database/store.py
index 379e5f6..23eaf70 100644
--- a/genotype_api/database/store.py
+++ b/genotype_api/database/store.py
@@ -1,7 +1,10 @@
 """Module for the store handler."""
 
+from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
+from genotype_api.config import settings
 from genotype_api.database.crud.create import CreateHandler
 from genotype_api.database.crud.delete import DeleteHandler
 from genotype_api.database.crud.read import ReadHandler
@@ -24,6 +27,12 @@ def __init__(self, session: AsyncSession):
         UpdateHandler.__init__(self, session)
 
     @classmethod
+    @retry(
+        stop=stop_after_attempt(settings.max_retries),
+        wait=wait_fixed(settings.retry_delay),
+        retry=retry_if_exception_type(OperationalError),
+        reraise=True,
+    )
     async def create(cls) -> "Store":
         """Asynchronously create and return a Store instance with a session."""
         async with get_session() as session:  # Correctly use async context manager
diff --git a/poetry.lock b/poetry.lock
index 4435757..238f05f 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1653,6 +1653,21 @@ anyio = ">=3.4.0,<5"
 [package.extras]
 full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"]
 
+[[package]]
+name = "tenacity"
+version = "9.0.0"
+description = "Retry code until it succeeds"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "tenacity-9.0.0-py3-none-any.whl", hash = "sha256:93de0c98785b27fcf659856aa9f54bfbd399e29969b0621bc7f762bd441b4539"},
+    {file = "tenacity-9.0.0.tar.gz", hash = "sha256:807f37ca97d62aa361264d497b0e31e92b8027044942bfa756160d908320d73b"},
+]
+
+[package.extras]
+doc = ["reno", "sphinx"]
+test = ["pytest", "tornado (>=4.5)", "typeguard"]
+
 [[package]]
 name = "tomlkit"
 version = "0.13.2"
@@ -1777,4 +1792,4 @@ test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "4137cf8d00a4d14d8c728d07f899a30edda86c4c33c7b6252b6e94085bb04b6f"
+content-hash = "78c8f76d8fe31061c6c1261dcb06819e4534eaae54d6d524ef339d5e188c500a"
diff --git a/pyproject.toml b/pyproject.toml
index 3a9d2cf..4963809 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,6 +35,7 @@ uvicorn = "^0.29.0"
 uvloop = "^0.19.0"
 aiomysql = "^0.2.0"
 pytest-asyncio = "^0.24.0"
+tenacity = "^9.0.0"
 
 [tool.poetry.group.dev.dependencies]
 bump2version = "^1.0.1"