From 39664957fe81df268c80deafe824476d1fdc4825 Mon Sep 17 00:00:00 2001 From: Zakaria Date: Sun, 25 Oct 2026 17:54:43 -0400 Subject: [PATCH] fix(athena): run and list queries in their workgroup StartQueryExecution ignored the workgroup: results always went to the default bucket, even when the workgroup has a result location (or enforces it), and an unknown workgroup was accepted. ListQueryExecutions ignored the WorkGroup filter and returned every query. Now, as in AWS, a query without a result location uses its workgroup's, an enforcing workgroup's location always wins, an unknown workgroup is an InvalidRequestException, and ListQueryExecutions returns the named workgroup's queries (the primary workgroup's when none is named), newest first. --- .../services/athena/AthenaJsonHandler.java | 4 +- .../floci/services/athena/AthenaService.java | 47 +++++++- ...AthenaWorkGroupQueriesIntegrationTest.java | 109 ++++++++++++++++++ 3 files changed, 155 insertions(+), 5 deletions(-) create mode 100644 src/test/java/io/github/hectorvent/floci/services/athena/AthenaWorkGroupQueriesIntegrationTest.java diff --git a/src/main/java/io/github/hectorvent/floci/services/athena/AthenaJsonHandler.java b/src/main/java/io/github/hectorvent/floci/services/athena/AthenaJsonHandler.java index 3fd01c1..ead1a7f 100644 --- a/src/main/java/io/github/hectorvent/floci/services/athena/AthenaJsonHandler.java +++ b/src/main/java/io/github/hectorvent/floci/services/athena/AthenaJsonHandler.java @@ -47,7 +47,7 @@ public class AthenaJsonHandler { resultConfiguration = mapper.treeToValue(request.get("ResultConfiguration"), ResultConfiguration.class); } - String id = athenaService.startQueryExecution(query, workGroup, context, resultConfiguration); + String id = athenaService.startQueryExecution(query, workGroup, context, resultConfiguration, region); yield Response.ok(Map.of("QueryExecutionId", id)).build(); } case "GetQueryExecution" -> { @@ -62,7 +62,7 @@ public class AthenaJsonHandler { } case "ListQueryExecutions" -> { yield Response.ok(Map.of("QueryExecutionIds", - athenaService.listQueryExecutions().stream() + athenaService.listQueryExecutions(request.has("WorkGroup") ? request.get("WorkGroup").asText() : null).stream() .map(QueryExecution::getQueryExecutionId).toList())).build(); } case "StopQueryExecution" -> { diff --git a/src/main/java/io/github/hectorvent/floci/services/athena/AthenaService.java b/src/main/java/io/github/hectorvent/floci/services/athena/AthenaService.java index 3f88b10..1b2a0ee 100644 --- a/src/main/java/io/github/hectorvent/floci/services/athena/AthenaService.java +++ b/src/main/java/io/github/hectorvent/floci/services/athena/AthenaService.java @@ -122,6 +122,20 @@ public class AthenaService { String workGroup, QueryExecutionContext context, ResultConfiguration resultConfiguration) { + return startQueryExecution(query, workGroup, context, resultConfiguration, null); + } + + /** + * Runs a query in a workgroup. As in AWS, the workgroup's result location applies when the + * request names none, and always when the workgroup enforces its configuration; an unknown + * workgroup is refused. A null region skips the workgroup lookup. + */ + public String startQueryExecution(String query, + String workGroup, + QueryExecutionContext context, + ResultConfiguration resultConfiguration, + String region) { + WorkGroupConfiguration workGroupConfiguration = region == null ? null : workGroupFor(region, workGroup).getConfiguration(); String id = UUID.randomUUID().toString(); String database = context != null && context.getDatabase() != null ? context.getDatabase() : "default"; QueryExecutionContext resolvedContext = context != null ? context : new QueryExecutionContext(); @@ -135,7 +149,7 @@ public class AthenaService { // directory prefix — the same key is written and returned to the client. // Statements that do not return rows must not be wrapped in a result COPY. String outputLocation = producesResultRows(query) - ? resolveOutputLocation(resultConfiguration, id) + ? resolveOutputLocation(resultConfiguration, workGroupConfiguration, id) : null; ResultConfiguration resolvedResult = outputLocation != null ? new ResultConfiguration(outputLocation) @@ -214,6 +228,16 @@ public class AthenaService { return queryStore.scan(k -> true); } + /** A workgroup's queries, most recent first, as ListQueryExecutions returns them. */ + public List listQueryExecutions(String workGroup) { + String name = workGroup == null || workGroup.isBlank() ? DEFAULT_WORKGROUP : workGroup; + return queryStore.scan(k -> true).stream() + .filter(execution -> name.equals(execution.getWorkGroup() == null ? DEFAULT_WORKGROUP : execution.getWorkGroup())) + .sorted(Comparator.comparing((QueryExecution execution) -> execution.getStatus().getSubmissionDateTime(), + Comparator.nullsLast(Comparator.naturalOrder())).reversed()) + .toList(); + } + public void stopQueryExecution(String id) { QueryExecution execution = getQueryExecution(id); execution.getStatus().setState(QueryExecutionState.CANCELLED); @@ -584,9 +608,20 @@ public class AthenaService { String location, Map properties) { } - private String resolveOutputLocation(ResultConfiguration rc, String queryId) { - String base = (rc != null && rc.getOutputLocation() != null && !rc.getOutputLocation().isBlank()) + private String resolveOutputLocation(ResultConfiguration rc, WorkGroupConfiguration workGroupConfiguration, String queryId) { + String requested = rc != null && rc.getOutputLocation() != null && !rc.getOutputLocation().isBlank() ? rc.getOutputLocation() + : null; + String workGroupLocation = workGroupConfiguration != null + && workGroupConfiguration.getResultConfiguration() != null + && workGroupConfiguration.getResultConfiguration().getOutputLocation() != null + && !workGroupConfiguration.getResultConfiguration().getOutputLocation().isBlank() + ? workGroupConfiguration.getResultConfiguration().getOutputLocation() + : null; + boolean enforced = workGroupConfiguration != null && Boolean.TRUE.equals(workGroupConfiguration.getEnforceWorkGroupConfiguration()); + String base = enforced && workGroupLocation != null ? workGroupLocation + : requested != null ? requested + : workGroupLocation != null ? workGroupLocation : "s3://" + DEFAULT_OUTPUT_BUCKET + "/results/"; return base.endsWith("/") ? base + queryId + ".csv" : base + "/" + queryId + ".csv"; } @@ -692,6 +727,12 @@ public class AthenaService { return engineVersion; } + private WorkGroup workGroupFor(String region, String name) { + return name == null || name.isBlank() || DEFAULT_WORKGROUP.equals(name) + ? primaryWorkGroup(region) + : requireWorkGroup(region, name); + } + private WorkGroup primaryWorkGroup(String region) { return workGroupStore.get(workGroupKey(region, DEFAULT_WORKGROUP)).orElseGet(() -> { WorkGroup workGroup = new WorkGroup(); diff --git a/src/test/java/io/github/hectorvent/floci/services/athena/AthenaWorkGroupQueriesIntegrationTest.java b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaWorkGroupQueriesIntegrationTest.java new file mode 100644 index 0000000..bdb4b31 --- /dev/null +++ b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaWorkGroupQueriesIntegrationTest.java @@ -0,0 +1,109 @@ +package io.github.hectorvent.floci.services.athena; + +import io.github.hectorvent.floci.testing.RestAssuredJsonUtils; +import io.quarkus.test.junit.QuarkusTest; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.UUID; + +import static io.restassured.RestAssured.given; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.startsWith; +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** Queries belong to their workgroup: they use its result location and are listed under it. */ +@QuarkusTest +class AthenaWorkGroupQueriesIntegrationTest { + + private static final String CONTENT_TYPE = "application/x-amz-json-1.1"; + + @BeforeAll + static void configureRestAssured() { + RestAssuredJsonUtils.configureAwsContentTypes(); + } + + private static io.restassured.response.ValidatableResponse athena(String action, String body) { + return given() + .header("X-Amz-Target", "AmazonAthena." + action) + .contentType(CONTENT_TYPE) + .body(body) + .when().post("/") + .then(); + } + + private static String workGroup(String location, boolean enforce) { + String name = "wg-" + UUID.randomUUID().toString().substring(0, 8); + athena("CreateWorkGroup", """ + {"Name": "%s", "Configuration": {"ResultConfiguration": {"OutputLocation": "%s"}, "EnforceWorkGroupConfiguration": %s}} + """.formatted(name, location, enforce)).statusCode(200); + return name; + } + + private static String start(String workGroup, String outputLocation) { + String result = outputLocation == null ? "" : ", \"ResultConfiguration\": {\"OutputLocation\": \"" + outputLocation + "\"}"; + return athena("StartQueryExecution", "{\"QueryString\": \"SELECT 1\", \"WorkGroup\": \"" + workGroup + "\"" + result + "}") + .statusCode(200) + .extract().jsonPath().getString("QueryExecutionId"); + } + + private static String outputLocation(String queryId) { + return athena("GetQueryExecution", "{\"QueryExecutionId\": \"" + queryId + "\"}") + .statusCode(200) + .body("QueryExecution.WorkGroup", not(equalTo(null))) + .extract().jsonPath().getString("QueryExecution.ResultConfiguration.OutputLocation"); + } + + @Test + void usesTheWorkGroupResultLocationWhenTheQueryNamesNone() { + String name = workGroup("s3://team-results/athena/", false); + String queryId = start(name, null); + + assertEquals("s3://team-results/athena/" + queryId + ".csv", outputLocation(queryId)); + } + + @Test + void aRequestedLocationWinsUnlessTheWorkGroupEnforcesItsOwn() { + String open = workGroup("s3://team-results/open/", false); + String openQuery = start(open, "s3://mine/results/"); + assertEquals("s3://mine/results/" + openQuery + ".csv", outputLocation(openQuery)); + + String enforced = workGroup("s3://team-results/enforced/", true); + String enforcedQuery = start(enforced, "s3://mine/results/"); + assertEquals("s3://team-results/enforced/" + enforcedQuery + ".csv", outputLocation(enforcedQuery)); + } + + @Test + void refusesAWorkGroupThatDoesNotExist() { + athena("StartQueryExecution", "{\"QueryString\": \"SELECT 1\", \"WorkGroup\": \"no-such-group\"}") + .statusCode(400) + .body("__type", equalTo("InvalidRequestException")); + } + + @Test + void listsOnlyTheWorkGroupQueriesNewestFirst() { + String name = workGroup("s3://team-results/listed/", false); + String first = start(name, null); + String second = start(name, null); + String elsewhere = start("primary", null); + + List listed = athena("ListQueryExecutions", "{\"WorkGroup\": \"" + name + "\"}") + .statusCode(200) + .extract().jsonPath().getList("QueryExecutionIds", String.class); + assertEquals(List.of(second, first), listed); + + athena("ListQueryExecutions", "{}") + .statusCode(200) + .body("QueryExecutionIds", hasItem(elsewhere)) + .body("QueryExecutionIds", not(hasItem(first))); + } + + @Test + void thePrimaryWorkGroupKeepsItsDefaultLocation() { + String queryId = start("primary", null); + org.hamcrest.MatcherAssert.assertThat(outputLocation(queryId), startsWith("s3://floci-athena-results/results/")); + } +}