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) {
-        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();

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