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