add retry on database connection (#59)

* add retry on database connection

* qa

* qa

Author: francesconazzaro

.github/workflows/main.yaml

@@ -4,12 +4,12 @@ on:
   workflow_dispatch:
   push:
     branches:
-      - main
+    - main
     tags:
-      - "v*"
+    - "v*"
   pull_request:
     branches:
-      - main
+    - main
 
 
 jobs:

cads_worker/worker.py

@@ -4,6 +4,7 @@
 import os
 import random
 import socket
+import time
 from typing import Any
 
 import cacholote
@@ -23,6 +24,8 @@
 
 LEVELS_MAPPING = logging.getLevelNamesMapping()
 
+DB_CONNECTION_RETRIES = int(os.getenv("WORKER_DB_CONNECTION_RETRIES", 3))
+
 
 @functools.lru_cache
 def create_session_maker() -> cads_broker.database.sa.orm.sessionmaker:
@@ -32,13 +35,34 @@ def create_session_maker() -> cads_broker.database.sa.orm.sessionmaker:
 def ensure_session(func):
     @functools.wraps(func)
     def wrapper(self, *args, session=None, **kwargs):
-        close_session = False
-        if session is None:
-            session = create_session_maker()()
-            close_session = True
-        func(self, *args, session=session, **kwargs)
-        if close_session:
-            session.close()
+        retries = 1
+        while retries <= DB_CONNECTION_RETRIES:
+            try:
+                close_session = False
+                # create a new session if not provided
+                if session is None:
+                    session = create_session_maker()()
+                    close_session = True
+                # run the function
+                result = func(self, *args, session=session, **kwargs)
+                # close the session if we created it
+                if close_session:
+                    session.close()
+                return result
+            except cads_broker.database.sa.exc.OperationalError as e:
+                exception = e
+                retries += 1
+                self.logger.warning(
+                    f"Database operation failed. Retrying {retries}/{DB_CONNECTION_RETRIES}...",
+                    error=str(e),
+                )
+                # close the session anyway because it could be broken
+                session.close()
+                session = None
+                time.sleep(os.getenv("WORKER_DB_CONNECTION_RETRY_SLEEP", 2))
+
+        self.logger.error("Max retries reached. Aborting operation.")
+        raise exception
 
     return wrapper
 

tests/test_40_worker.py

@@ -0,0 +1,104 @@
+from unittest.mock import patch
+
+import pytest
+from sqlalchemy.orm import Session
+
+from cads_worker import worker
+
+
+def mock_session_maker():
+    return Session()
+
+
+def test_ensure_session_decorator():
+    # Test case 1: when session is None
+    @worker.ensure_session
+    def sample_function(self, session=None):
+        assert isinstance(session, Session)
+        return session
+
+    with patch(
+        "cads_worker.worker.create_session_maker", return_value=mock_session_maker
+    ):
+        result = sample_function(self=None)
+        assert isinstance(result, Session)
+
+    # Test case 2: when session is provided
+    mock_session = Session()
+    result = sample_function(self=None, session=mock_session)
+    assert result is mock_session
+
+    # Clean up
+    mock_session.close()
+
+
+def test_ensure_session_decorator_nested():
+    # Test nested function calls with session
+    @worker.ensure_session
+    def outer_function(self, session=None):
+        @worker.ensure_session
+        def inner_function(self, session=None):
+            return session
+
+        return inner_function(self=None, session=session)
+
+    with patch(
+        "cads_worker.worker.create_session_maker", return_value=mock_session_maker
+    ):
+        result = outer_function(self=None)
+    assert isinstance(result, Session)
+    result.close()
+
+
+def test_ensure_session_decorator_error():
+    # Test error handling
+    @worker.ensure_session
+    def failing_function(self, session=None):
+        raise ValueError("Test error")
+
+    with pytest.raises(ValueError):
+        failing_function(self=None)
+
+
+call_count = 1
+
+
+def test_ensure_session_retry():
+    # Test retries
+
+    context = worker.Context()
+
+    @worker.ensure_session
+    def failing_function(self, session=None):
+        print("failing function called")
+        raise worker.cads_broker.database.sa.exc.OperationalError(
+            "Simulated DB error", None, None
+        )
+
+    with patch(
+        "cads_worker.worker.create_session_maker", return_value=mock_session_maker
+    ):
+        with pytest.raises(worker.cads_broker.database.sa.exc.OperationalError):
+            # This should raise after max retries
+            failing_function(self=context)
+
+    global call_count
+
+    @worker.ensure_session
+    def successful_function(self, session=None):
+        global call_count
+        if call_count < 3:
+            print("failing function called - retries:", call_count)
+            call_count += 1
+            raise worker.cads_broker.database.sa.exc.OperationalError(
+                "Simulated DB error", None, None
+            )
+        else:
+            print("successful function called - retries:", call_count)
+            return session
+
+    with patch(
+        "cads_worker.worker.create_session_maker", return_value=mock_session_maker
+    ):
+        result = successful_function(self=context)
+    assert isinstance(result, Session)