From 39664957fe81df268c80deafe824476d1fdc4825 Mon Sep 17 00:00:00 2001 From: Zakaria Date: Sun, 25 Oct 2026 17:54:43 -0400 Subject: [PATCH 1/3] 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/")); + } +} -- 2.52.0 From ac96c6451e8edd11fe57f402d227a69e3dc96273 Mon Sep 17 00:00:00 2001 From: Zakaria Date: Sun, 25 Oct 2026 18:04:11 -0400 Subject: [PATCH 2/3] fix(athena): accept a query that ends in a semicolon Athena runs a statement with or without its closing semicolon, but the query engine wraps it in COPY (...) TO ..., where the semicolon is a syntax error, so SELECT * FROM orders; always failed. Trailing semicolons are now dropped before the statement is handed to the engine. --- .../floci/services/athena/AthenaService.java | 12 +++++++++++- .../athena/AthenaTrailingSemicolonTest.java | 15 +++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java 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 1b2a0ee..13251fb 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 @@ -59,6 +59,7 @@ public class AthenaService { public static final String DEFAULT_CATALOG = "AwsDataCatalog"; private static final String DEFAULT_OUTPUT_BUCKET = "floci-athena-results"; private static final String DEFAULT_WORKGROUP = "primary"; + private static final Pattern TRAILING_SEMICOLONS = Pattern.compile("[\\s;]+$"); private static final String DEFAULT_ENGINE_VERSION = "Athena engine version 3"; private static final long MIN_BYTES_SCANNED_CUTOFF_PER_QUERY = 10_000_000L; private static final String WORKGROUP_RESOURCE = "workgroup/"; @@ -191,7 +192,7 @@ public class AthenaService { if (outputLocation != null) { ensureOutputBucket(outputLocation); } - duckClient.execute(query, setupDdl, outputLocation); + duckClient.execute(withoutTrailingSemicolons(query), setupDdl, outputLocation); return null; }).onSuccess(v -> { markSucceeded(id, execution); @@ -555,6 +556,15 @@ public class AthenaService { return properties; } + /** + * Athena accepts one statement with or without a closing semicolon. The engine wraps the + * statement in {@code COPY (...) TO ...}, where a semicolon is a syntax error, so it is + * dropped before the query is handed over. + */ + static String withoutTrailingSemicolons(String query) { + return query == null ? null : TRAILING_SEMICOLONS.matcher(query).replaceFirst(""); + } + private static boolean producesResultRows(String query) { return RESULT_STATEMENT_PATTERN.matcher(statementText(query)).find(); } diff --git a/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java new file mode 100644 index 0000000..c19e3b8 --- /dev/null +++ b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java @@ -0,0 +1,15 @@ +package io.github.hectorvent.floci.services.athena; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class AthenaTrailingSemicolonTest { + + @Test + void dropsTheClosingSemicolonsAthenaAccepts() { + assertEquals("SELECT * FROM orders LIMIT 10", AthenaService.withoutTrailingSemicolons("SELECT * FROM orders LIMIT 10;")); + assertEquals("SELECT 1", AthenaService.withoutTrailingSemicolons("SELECT 1 ; ;\n")); + assertEquals("SELECT ';' AS s", AthenaService.withoutTrailingSemicolons("SELECT ';' AS s")); + } +} -- 2.52.0 From 57aa86614499bb43182f0c764fcf01a67890f287 Mon Sep 17 00:00:00 2001 From: Zakaria Date: Sun, 25 Oct 2026 18:07:31 -0400 Subject: [PATCH 3/3] fix(athena): report a failed query's error the way Athena does StateChangeReason carried the raw engine call failure, HTTP envelope and all: floci-duck execute returned HTTP 500: {"status":"error","message":...}, ending with the statement as the engine saw it, wrapped in COPY (...) TO. It is now the engine's own message alone, such as Catalog Error: Table with name nope does not exist! --- .../floci/services/athena/AthenaService.java | 29 ++++++++++++++++++- .../services/athena/AthenaQueryTextTest.java | 22 ++++++++++++++ .../athena/AthenaTrailingSemicolonTest.java | 15 ---------- 3 files changed, 50 insertions(+), 16 deletions(-) create mode 100644 src/test/java/io/github/hectorvent/floci/services/athena/AthenaQueryTextTest.java delete mode 100644 src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java 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 13251fb..872668f 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 @@ -59,6 +59,8 @@ public class AthenaService { public static final String DEFAULT_CATALOG = "AwsDataCatalog"; private static final String DEFAULT_OUTPUT_BUCKET = "floci-athena-results"; private static final String DEFAULT_WORKGROUP = "primary"; + private static final Pattern ENGINE_FAILURE = Pattern.compile("^floci-duck \\w+ returned HTTP \\d+: (\\{.*})$", Pattern.DOTALL); + private static final com.fasterxml.jackson.databind.ObjectMapper FAILURE_READER = new com.fasterxml.jackson.databind.ObjectMapper(); private static final Pattern TRAILING_SEMICOLONS = Pattern.compile("[\\s;]+$"); private static final String DEFAULT_ENGINE_VERSION = "Athena engine version 3"; private static final long MIN_BYTES_SCANNED_CUTOFF_PER_QUERY = 10_000_000L; @@ -610,10 +612,35 @@ public class AthenaService { private void markFailed(String id, QueryExecution execution, Throwable failure) { execution.getStatus().setState(QueryExecutionState.FAILED); - execution.getStatus().setStateChangeReason(failure.getMessage()); + execution.getStatus().setStateChangeReason(failureReason(failure.getMessage())); queryStore.put(id, execution); } + /** + * The engine's own explanation of a failed query, as Athena reports it: without the HTTP + * envelope of the engine call, and without the echoed statement, which shows the + * internal {@code COPY ... TO} wrapper rather than the query that was sent. + */ + static String failureReason(String message) { + if (message == null) { + return null; + } + Matcher envelope = ENGINE_FAILURE.matcher(message); + if (!envelope.matches()) { + return message; + } + try { + String engineMessage = FAILURE_READER.readTree(envelope.group(1)).path("message").asText(""); + if (engineMessage.isBlank()) { + return message; + } + int echo = engineMessage.indexOf("\n\nLINE "); + return (echo >= 0 ? engineMessage.substring(0, echo) : engineMessage).strip(); + } catch (Exception e) { + return message; + } + } + private record CreateDatabaseDdl(String name, boolean ifNotExists, String comment, String location, Map properties) { } diff --git a/src/test/java/io/github/hectorvent/floci/services/athena/AthenaQueryTextTest.java b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaQueryTextTest.java new file mode 100644 index 0000000..e68d548 --- /dev/null +++ b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaQueryTextTest.java @@ -0,0 +1,22 @@ +package io.github.hectorvent.floci.services.athena; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class AthenaQueryTextTest { + + @Test + void dropsTheClosingSemicolonsAthenaAccepts() { + assertEquals("SELECT * FROM orders LIMIT 10", AthenaService.withoutTrailingSemicolons("SELECT * FROM orders LIMIT 10;")); + assertEquals("SELECT 1", AthenaService.withoutTrailingSemicolons("SELECT 1 ; ;\n")); + assertEquals("SELECT ';' AS s", AthenaService.withoutTrailingSemicolons("SELECT ';' AS s")); + } + + @Test + void reportsTheEngineErrorWithoutItsEnvelopeOrTheWrappedStatement() { + String engine = "floci-duck execute returned HTTP 500: {\"status\":\"error\",\"message\":\"Catalog Error: Table with name nope does not exist!\\nDid you mean \\\"orders\\\"?\\n\\nLINE 1: COPY (SELECT * FROM nope) TO 's3://results/q.csv'\"}"; + assertEquals("Catalog Error: Table with name nope does not exist!\nDid you mean \"orders\"?", AthenaService.failureReason(engine)); + assertEquals("Failed to call floci-duck execute: timeout", AthenaService.failureReason("Failed to call floci-duck execute: timeout")); + } +} diff --git a/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java b/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java deleted file mode 100644 index c19e3b8..0000000 --- a/src/test/java/io/github/hectorvent/floci/services/athena/AthenaTrailingSemicolonTest.java +++ /dev/null @@ -1,15 +0,0 @@ -package io.github.hectorvent.floci.services.athena; - -import org.junit.jupiter.api.Test; - -import static org.junit.jupiter.api.Assertions.assertEquals; - -class AthenaTrailingSemicolonTest { - - @Test - void dropsTheClosingSemicolonsAthenaAccepts() { - assertEquals("SELECT * FROM orders LIMIT 10", AthenaService.withoutTrailingSemicolons("SELECT * FROM orders LIMIT 10;")); - assertEquals("SELECT 1", AthenaService.withoutTrailingSemicolons("SELECT 1 ; ;\n")); - assertEquals("SELECT ';' AS s", AthenaService.withoutTrailingSemicolons("SELECT ';' AS s")); - } -} -- 2.52.0