Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 57aa866144 | |||
| ac96c6451e | |||
| 39664957fe | |||
| 00b3af3813 | |||
| d8d7dd021b | |||
| eaf5ca6ade |
@@ -563,8 +563,10 @@ public class AcmJsonHandler {
|
||||
|
||||
private CertificateOptions parseOptions(JsonNode optionsNode) {
|
||||
if (optionsNode.isMissingNode()) return null;
|
||||
String ctPref = optionsNode.path("CertificateTransparencyLoggingPreference").asText("ENABLED");
|
||||
return new CertificateOptions(ctPref, "DISABLED");
|
||||
String ctPref = parseOptionValue(optionsNode, "CertificateTransparencyLoggingPreference");
|
||||
String export = parseOptionValue(optionsNode, "Export");
|
||||
// An exportable public certificate is requested with Export ENABLED, and keeps it for life.
|
||||
return new CertificateOptions(ctPref != null ? ctPref : "ENABLED", export != null ? export : "DISABLED");
|
||||
}
|
||||
|
||||
private List<CertificateStatus> parseCertificateStatuses(JsonNode node) {
|
||||
|
||||
@@ -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,70 @@
|
||||
package io.github.hectorvent.floci.services.acm;
|
||||
|
||||
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 static io.restassured.RestAssured.given;
|
||||
import static org.hamcrest.Matchers.equalTo;
|
||||
import static org.hamcrest.Matchers.notNullValue;
|
||||
|
||||
@QuarkusTest
|
||||
class AcmExportableRequestTest {
|
||||
|
||||
private static final String ACM_CONTENT_TYPE = "application/x-amz-json-1.1";
|
||||
|
||||
@BeforeAll
|
||||
static void configureRestAssured() {
|
||||
RestAssuredJsonUtils.configureAwsContentTypes();
|
||||
}
|
||||
|
||||
@Test
|
||||
void requestCertificateKeepsTheExportOption() {
|
||||
String exportable = request("{\"DomainName\":\"exportable.example.com\",\"Options\":{\"Export\":\"ENABLED\"}}");
|
||||
|
||||
call("DescribeCertificate", "{\"CertificateArn\":\"" + exportable + "\"}")
|
||||
.statusCode(200)
|
||||
.body("Certificate.Options.Export", equalTo("ENABLED"))
|
||||
.body("Certificate.Options.CertificateTransparencyLoggingPreference", equalTo("ENABLED"));
|
||||
|
||||
call("ExportCertificate", "{\"CertificateArn\":\"" + exportable + "\",\"Passphrase\":\"c2VjcmV0\"}")
|
||||
.statusCode(200)
|
||||
.body("PrivateKey", notNullValue());
|
||||
}
|
||||
|
||||
@Test
|
||||
void requestCertificateWithoutExportIsNotExportable() {
|
||||
String plain = request("{\"DomainName\":\"plain.example.com\",\"Options\":{\"CertificateTransparencyLoggingPreference\":\"DISABLED\"}}");
|
||||
|
||||
call("DescribeCertificate", "{\"CertificateArn\":\"" + plain + "\"}")
|
||||
.statusCode(200)
|
||||
.body("Certificate.Options.Export", equalTo("DISABLED"))
|
||||
.body("Certificate.Options.CertificateTransparencyLoggingPreference", equalTo("DISABLED"));
|
||||
|
||||
call("ExportCertificate", "{\"CertificateArn\":\"" + plain + "\",\"Passphrase\":\"c2VjcmV0\"}")
|
||||
.statusCode(400)
|
||||
.body("__type", equalTo("ValidationException"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void requestCertificateRejectsAnUnknownExportValue() {
|
||||
call("RequestCertificate", "{\"DomainName\":\"bad.example.com\",\"Options\":{\"Export\":\"MAYBE\"}}")
|
||||
.statusCode(400)
|
||||
.body("__type", equalTo("ValidationException"));
|
||||
}
|
||||
|
||||
private String request(String body) {
|
||||
return call("RequestCertificate", body).statusCode(200).extract().jsonPath().getString("CertificateArn");
|
||||
}
|
||||
|
||||
private io.restassured.response.ValidatableResponse call(String action, String body) {
|
||||
return given()
|
||||
.header("X-Amz-Target", "CertificateManager." + action)
|
||||
.contentType(ACM_CONTENT_TYPE)
|
||||
.body(body)
|
||||
.when()
|
||||
.post("/")
|
||||
.then();
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
+109
@@ -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/"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user