Compare commits

...

4 Commits

Author SHA1 Message Date
Zakaria 57aa866144 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!
2026-10-25 18:07:31 -04:00
Zakaria ac96c6451e 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.
2026-10-25 18:04:11 -04:00
Zakaria 39664957fe 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.
2026-10-25 17:54:43 -04:00
Zakaria 00b3af3813 Merge pull request 'fix(acm): keep the Export option on RequestCertificate' (#12) from fix/acm-request-export-option into main
Reviewed-on: #12
2026-10-25 17:21:08 +00:00
4 changed files with 216 additions and 7 deletions
@@ -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" -> {
@@ -59,6 +59,9 @@ 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;
private static final String WORKGROUP_RESOURCE = "workgroup/";
@@ -122,6 +125,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 +152,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)
@@ -177,7 +194,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);
@@ -214,6 +231,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);
@@ -531,6 +558,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();
}
@@ -576,17 +612,53 @@ 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<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 +764,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();
@@ -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"));
}
}
@@ -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/"));
}
}