content: package org.pmiops.workbench.db;
  
  import com.google.common.base.Preconditions;
  import java.time.Duration;
  import java.util.concurrent.TimeUnit;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  import org.apache.commons.lang3.exception.ExceptionUtils;
  import org.springframework.dao.DataAccessException;
  import org.springframework.dao.DataRetrievalFailureException;
  import org.springframework.dao.RecoverableDataAccessException;
  import org.springframework.dao.TransientDataAccessException;
  
  /** Utilities to execute database operations with retry support. */
  public final class DbRetryUtils {
    private static final Logger logger = Logger.getLogger(DbRetryUtils.class.getName());
  
    private DbRetryUtils() {}
  
    /**
     * Executes a database operation and retries if retryable
     *
     * @param execute database operation to execute
     * @param retrySleep fixed retry sleep interval
     * @param maxNumAttempts maximum times this operation will be run
     * @param <T> database operation class
     * @return database operation class
     * @throws InterruptedException on thread interruption
     * @throws DataAccessException throws the last DB operation error if maxNumAttempts is exceeded.
     */
    public static <T> T executeAndRetry(
        DatabaseOperation<T> execute, Duration retrySleep, int maxNumAttempts)
        throws InterruptedException {
      Preconditions.checkArgument(maxNumAttempts > 0, "maxNumAttempts must be at least 1");
      int numAttempts = 1;
      while (numAttempts <= maxNumAttempts) {
        try {
          return execute.execute();
        } catch (DataAccessException e) {
          if (!shouldRetryQuery(e) || (numAttempts == maxNumAttempts)) {
            throw e;
          }
          logger.log(
              Level.WARNING,
              String.format(
                  "Caught exception, retrying DB operation. Attempts so far: %d", numAttempts),
              e);
        }
        ++numAttempts;
        TimeUnit.MILLISECONDS.sleep(retrySleep.toMillis());
      }
      throw new IllegalStateException(
          "Exceeded maximum number of retries without throwing an exception. This should never happen.");
    }
  
    /**
     * Tests an exception to see if it is retryable.
     *
     * @param dataAccessException execption to test
     * @return {@code true} if that is retryable {@link DataAccessException}.
     */
    public static boolean shouldRetryQuery(DataAccessException dataAccessException) {
      return ExceptionUtils.hasCause(dataAccessException, RecoverableDataAccessException.class)
          || ExceptionUtils.hasCause(dataAccessException, TransientDataAccessException.class)
          || ExceptionUtils.hasCause(dataAccessException, DataRetrievalFailureException.class);
    }
  
    /** How to execute this database operation. */
    @FunctionalInterface
    public interface DatabaseOperation<T> {
      T execute();
    }
  }
  