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.

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" -> {

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<QueryExecution> 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<String, String> properties) {
     }
 
-    private String resolveOutputLocation(ResultConfiguration rc, String queryId) {
+    private String resolveOutputLocation(ResultConfiguration rc, WorkGroupConfiguration workGroupConfiguration, String queryId) {
-        String base = (rc != null && rc.getOutputLocation() != null && !rc.getOutputLocation().isBlank())
+        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();

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<String> 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/"));
+    }
+}