diff --git a/components/camel-aws/camel-aws2-athena/pom.xml b/components/camel-aws/camel-aws2-athena/pom.xml index 6d9fb8f9b961a..7dd4317270fd6 100644 --- a/components/camel-aws/camel-aws2-athena/pom.xml +++ b/components/camel-aws/camel-aws2-athena/pom.xml @@ -72,6 +72,11 @@ ${mockito-version} test + + org.assertj + assertj-core + test + diff --git a/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java b/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java index 97fe349a05d06..467cd76735334 100644 --- a/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java +++ b/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. + *

+ * 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? */ diff --git a/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java b/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java index db81f368106c0..76688f21a3dc6 100644 --- a/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java +++ b/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();