CAMEL-24243: camel-aws2-athena - do not relaunch a still-running query when waitTimeout expires

components/camel-aws/camel-aws2-athena/pom.xml

@@ -72,6 +72,11 @@
             <version>${mockito-version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
 
         <!-- test infra -->
        <dependency>

components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java

@@ -120,11 +120,32 @@ boolean shouldAttempt() {
             return false;
         }
 
+        if (isWaitTimeoutExceeded()) {
+            LOG.trace("AWS Athena start query execution exceeded the wait timeout of {} while the query was still"
+                      + " running, will not relaunch it",
+                    this.waitTimeout);
+            return false;
+        }
+
         // if this.isRetry, return true
 
         return true;
     }
 
+    /**
+     * Did the wait window elapse while the query was still running? Only a completed-and-retryable query is meant to
+     * consume another attempt, so a query that is merely slow must not be relaunched.
+     * <p>
+     * Called from {@link #shouldAttempt()} once success and failure have been ruled out, so {@code isRetry} is the only
+     * completion state left to exclude here.
+     */
+    private boolean isWaitTimeoutExceeded() {
+        if (this.attempts == 0 || this.isRetry) {
+            return false;
+        }
+        return now() - this.startMs >= this.waitTimeout;
+    }
+
     /**
      * Should there be a wait for the query to complete?
      */

components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java

@@ -28,6 +28,7 @@
 import software.amazon.awssdk.services.athena.model.QueryExecutionState;
 import software.amazon.awssdk.services.athena.model.QueryExecutionStatus;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -155,6 +156,60 @@ public void testStartQueryExecutionHappyPath() {
         assertFalse(helper.shouldAttempt());
     }
 
+    @Test
+    void aStillRunningQueryIsNotRelaunchedWhenTheWaitTimeoutExpires() {
+        Athena2Configuration configuration = new Athena2Configuration();
+        configuration.setMaxAttempts(3);
+        configuration.setWaitTimeout(1);
+        configuration.setInitialDelay(1);
+        configuration.setDelay(1);
+        configuration.setRetry("always");
+
+        Athena2QueryHelper helper = new Athena2QueryHelper(
+                new DefaultExchange(new DefaultCamelContext()),
+                configuration);
+
+        assertThat(helper.shouldAttempt()).isTrue();
+        helper.markAttempt();
+
+        // the query is still running when the wait window elapses, so no completion state gets set
+        helper.doWait();
+        helper.setStatusFrom(newGetQueryExecutionResponse(QueryExecutionState.RUNNING));
+
+        assertThat(helper.shouldWait()).isFalse();
+        assertThat(helper.isSuccess()).isFalse();
+        assertThat(helper.isFailure()).isFalse();
+        assertThat(helper.isRetry()).isFalse();
+
+        // attempts are reserved for retrying completed-and-retryable queries; relaunching here would
+        // orphan the running query and bill the same SQL twice
+        assertThat(helper.shouldAttempt()).isFalse();
+        assertThat(helper.getAttempts()).isEqualTo(1);
+    }
+
+    @Test
+    void aRetryableFailureStillConsumesAnotherAttemptAfterTheWaitTimeout() {
+        Athena2Configuration configuration = new Athena2Configuration();
+        configuration.setMaxAttempts(3);
+        configuration.setWaitTimeout(1);
+        configuration.setInitialDelay(1);
+        configuration.setDelay(1);
+        configuration.setRetry("always");
+
+        Athena2QueryHelper helper = new Athena2QueryHelper(
+                new DefaultExchange(new DefaultCamelContext()),
+                configuration);
+
+        assertThat(helper.shouldAttempt()).isTrue();
+        helper.markAttempt();
+
+        helper.doWait();
+        helper.setStatusFrom(newGetQueryExecutionResponse(QueryExecutionState.FAILED, "GENERIC_INTERNAL_ERROR"));
+
+        assertThat(helper.isRetry()).isTrue();
+        assertThat(helper.shouldAttempt()).isTrue();
+    }
+
     @Test
     public void isComplete() {
         Athena2QueryHelper helper = defaultAthena2QueryHelper();