Fix database concurrency conflicts with retry mechanism and enhanced error handling

- Add retry_on_database_conflict decorator with exponential backoff for database operations
- Implement transaction-safe database_transaction context manager with proper rollback
- Apply @retry_on_database_conflict to all WorkflowTracker event sourcing operations
- Add retry protection to view update operations in JobAccounting, WorkflowProgress, and JobProgress
- Enhance EngineManager with robust commit/rollback error handling and safe_commit method
- Add support for EventSourcing-specific exceptions (IntegrityError, OperationalError)
- Implement defensive checks for unknown workflows/tasks in event handlers
- Upgrade error logging from ERROR to WARNING level for retryable conflicts
- Add comprehensive test coverage for retry mechanisms and concurrent access scenarios

This resolves production issues where database unique constraint violations and
deadlocks were causing thread failures and incomplete workflow tracking in
high-concurrency scenarios like batch processing workflows.

Author: TorecLuik

biomero/database.py

@@ -24,10 +24,144 @@
 )
 from sqlalchemy.orm import sessionmaker, declarative_base, scoped_session
 from sqlalchemy.dialects.postgresql import UUID as PGUUID
+from sqlalchemy.exc import OperationalError, IntegrityError
 import os
+import time
+import random
+from contextlib import contextmanager
+from functools import wraps
+
+# Import EventSourcing exceptions
+try:
+    from eventsourcing.persistence import IntegrityError as EventSourcingIntegrityError
+    from eventsourcing.persistence import OperationalError as EventSourcingOperationalError
+except ImportError:
+    # Fallback for older versions
+    EventSourcingIntegrityError = None
+    EventSourcingOperationalError = None
 
 logger = logging.getLogger(__name__)
 
+# --------------------- CONCURRENCY HELPERS ---------------------------- #
+
+def retry_on_database_conflict(max_retries=10, base_delay=0.1, max_delay=5.0):
+    """Decorator to retry database operations on concurrency conflicts.
+    
+    Args:
+        max_retries: Maximum number of retry attempts
+        base_delay: Base delay between retries in seconds
+        max_delay: Maximum delay between retries in seconds
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            last_exception = None
+            
+            # Build list of exception types to catch
+            exception_types = [IntegrityError, OperationalError]
+            if EventSourcingIntegrityError:
+                exception_types.append(EventSourcingIntegrityError)
+            if EventSourcingOperationalError:
+                exception_types.append(EventSourcingOperationalError)
+            exception_types = tuple(exception_types)
+            
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except exception_types as e:
+                    last_exception = e
+                    
+                    # Don't retry on the last attempt
+                    if attempt == max_retries:
+                        break
+                    
+                    # Check if it's a conflict we can retry
+                    error_msg = str(e).lower()
+                    
+                    # For EventSourcing exceptions, also check the underlying cause
+                    if hasattr(e, '__cause__') and e.__cause__:
+                        error_msg += " " + str(e.__cause__).lower()
+                    
+                    # These are retryable concurrency conflicts
+                    retryable_conflicts = [
+                        'unique constraint',
+                        'duplicate key',
+                        'concurrent update',
+                        'transaction is aborted',
+                        'deadlock detected'
+                    ]
+                    
+                    # These are permanent data validation errors - do NOT retry
+                    non_retryable_errors = [
+                        'not null constraint',
+                        'null value in column',
+                        'cannot be null'
+                    ]
+                    
+                    # Check for non-retryable errors first
+                    if any(error in error_msg for error in non_retryable_errors):
+                        logger.debug(f"Non-retryable data validation error: {e}")
+                        break
+                    
+                    # Check for retryable conflicts
+                    if any(conflict in error_msg for conflict in retryable_conflicts):
+                        # Calculate exponential backoff with jitter
+                        delay = min(
+                            base_delay * (2 ** attempt) + random.uniform(0, 0.1),
+                            max_delay
+                        )
+                        logger.warning(
+                            f"Database conflict on attempt {attempt + 1}/{max_retries + 1}: {e}. "
+                            f"Retrying in {delay:.2f}s..."
+                        )
+                        
+                        time.sleep(delay)
+                        
+                        # Rollback current transaction
+                        try:
+                            EngineManager.rollback()
+                        except:
+                            pass
+                    else:
+                        # Not a retryable error
+                        break
+                        
+            # All retries exhausted, raise the last exception
+            logger.error(f"Database operation failed after {max_retries + 1} attempts: {last_exception}")
+            raise last_exception
+            
+        return wrapper
+    return decorator
+
+
+@contextmanager
+def database_transaction(isolation_level=None):
+    """Context manager for database transactions with proper error handling.
+    
+    Args:
+        isolation_level: SQLAlchemy isolation level (e.g., 'READ_COMMITTED', 'SERIALIZABLE')
+    """
+    # Don't manage session lifecycle here - let the scoped session handle it
+    # Just provide transaction boundaries and error handling
+    try:
+        # Set isolation level if specified
+        if isolation_level:
+            session = EngineManager.get_session()
+            session.connection(execution_options={'isolation_level': isolation_level})
+            
+        yield
+        
+        # Commit if we reach this point
+        EngineManager.commit()
+        
+    except Exception as e:
+        try:
+            EngineManager.rollback()
+            logger.debug(f"Transaction rolled back due to: {e}")
+        except Exception as rollback_error:
+            logger.error(f"Error during rollback: {rollback_error}")
+        raise
+
 # --------------------- VIEWS DB tables/classes ---------------------------- #
 
 # Base class for declarative class definitions
@@ -247,7 +381,40 @@ def commit(cls):
         """
         Commits the current transaction in the scoped session.
         """
-        cls._session.commit()
+        try:
+            cls._session.commit()
+        except Exception as e:
+            logger.warning(f"Database commit failed (will be retried if applicable): {e}")
+            cls.rollback()
+            raise
+    
+    @classmethod
+    def rollback(cls):
+        """
+        Rolls back the current transaction in the scoped session.
+        """
+        try:
+            cls._session.rollback()
+        except Exception as e:
+            logger.warning(f"Database rollback failed (forcing session removal): {e}")
+            # Force session removal on rollback failure
+            cls._session.remove()
+            raise
+    
+    @classmethod
+    def remove_session(cls):
+        """
+        Removes the current session from the scoped session registry.
+        """
+        cls._session.remove()
+    
+    @classmethod
+    @retry_on_database_conflict(max_retries=10)
+    def safe_commit(cls):
+        """
+        Commits the transaction with automatic retry on conflicts.
+        """
+        cls.commit()
 
     @classmethod
     def close_engine(cls):

biomero/eventsourcing.py

@@ -18,7 +18,7 @@
 from typing import Any, Dict, List
 from fabric import Result
 import logging
-from biomero.database import EngineManager
+from biomero.database import EngineManager, retry_on_database_conflict
 
 
 # Create a logger for this module
@@ -346,6 +346,7 @@ class WorkflowTracker(Application):
         update_task_progress: Updates the progress of a task.
     """
 
+    @retry_on_database_conflict(max_retries=3)
     def initiate_workflow(self,
                           name: str,
                           description: str,
@@ -366,9 +367,10 @@ def initiate_workflow(self,
         logger.debug(f"[WFT] Initiating workflow: name={name}, description={description}, user={user}, group={group}")
         workflow = WorkflowRun(name, description, user, group)
         self.save(workflow)
-        EngineManager.commit()
+        EngineManager.safe_commit()
         return workflow.id
 
+    @retry_on_database_conflict(max_retries=3)
     def add_task_to_workflow(self,
                              workflow_id: UUID,
                              task_name: str,
@@ -397,13 +399,14 @@ def add_task_to_workflow(self,
                     input_data,
                     kwargs)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
         workflow: WorkflowRun = self.repository.get(workflow_id)
         workflow.add_task(task.id)
         self.save(workflow)
-        EngineManager.commit()
+        EngineManager.safe_commit()
         return task.id
 
+    @retry_on_database_conflict(max_retries=3)
     def start_workflow(self, workflow_id: UUID):
         """
         Starts the workflow with the given UUID.
@@ -416,8 +419,9 @@ def start_workflow(self, workflow_id: UUID):
         workflow: WorkflowRun = self.repository.get(workflow_id)
         workflow.start_workflow()
         self.save(workflow)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def complete_workflow(self, workflow_id: UUID):
         """
         Marks the workflow with the given UUID as completed.
@@ -430,8 +434,9 @@ def complete_workflow(self, workflow_id: UUID):
         workflow: WorkflowRun = self.repository.get(workflow_id)
         workflow.complete_workflow()
         self.save(workflow)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def fail_workflow(self, workflow_id: UUID, error_message: str):
         """
         Marks the workflow with the given UUID as failed with an error message.
@@ -445,8 +450,9 @@ def fail_workflow(self, workflow_id: UUID, error_message: str):
         workflow: WorkflowRun = self.repository.get(workflow_id)
         workflow.fail_workflow(error_message)
         self.save(workflow)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def start_task(self, task_id: UUID):
         """
         Starts the task with the given UUID.
@@ -459,8 +465,9 @@ def start_task(self, task_id: UUID):
         task: Task = self.repository.get(task_id)
         task.start_task()
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def complete_task(self, task_id: UUID, message: str):
         """
         Marks the task with the given UUID as completed with a result message.
@@ -474,8 +481,9 @@ def complete_task(self, task_id: UUID, message: str):
         task: Task = self.repository.get(task_id)
         task.complete_task(message)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def fail_task(self, task_id: UUID, error_message: str):
         """
         Marks the task with the given UUID as failed with an error message.
@@ -489,8 +497,9 @@ def fail_task(self, task_id: UUID, error_message: str):
         task: Task = self.repository.get(task_id)
         task.fail_task(error_message)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def add_job_id(self, task_id, slurm_job_id):
         """
         Adds a Slurm job ID to the task with the given UUID.
@@ -504,8 +513,9 @@ def add_job_id(self, task_id, slurm_job_id):
         task: Task = self.repository.get(task_id)
         task.add_job_id(slurm_job_id)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def add_result(self, task_id, result):
         """
         Adds a result to the task with the given UUID.
@@ -519,8 +529,9 @@ def add_result(self, task_id, result):
         task: Task = self.repository.get(task_id)
         task.add_result(result)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def update_task_status(self, task_id, status):
         """
         Updates the status of the task with the given UUID.
@@ -534,8 +545,9 @@ def update_task_status(self, task_id, status):
         task: Task = self.repository.get(task_id)
         task.update_task_status(status)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()
 
+    @retry_on_database_conflict(max_retries=3)
     def update_task_progress(self, task_id, progress):
         """
         Updates the progress of the task with the given UUID.
@@ -549,7 +561,7 @@ def update_task_progress(self, task_id, progress):
         task: Task = self.repository.get(task_id)
         task.update_task_progress(progress)
         self.save(task)
-        EngineManager.commit()
+        EngineManager.safe_commit()