diff --git a/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponse.java b/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponse.java index 05d64a235891..59db196244f5 100644 --- a/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponse.java +++ b/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponse.java @@ -24,19 +24,25 @@ import org.apache.iceberg.FileScanTask; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.rest.PlanStatus; +import org.apache.iceberg.rest.credentials.Credential; public class FetchPlanningResultResponse extends BaseScanTaskResponse { private final PlanStatus planStatus; + private final List credentials; private FetchPlanningResultResponse( PlanStatus planStatus, List planTasks, List fileScanTasks, List deleteFiles, - Map specsById) { + Map specsById, + List credentials) { super(planTasks, fileScanTasks, deleteFiles, specsById); this.planStatus = planStatus; + this.credentials = credentials; validate(); } @@ -44,6 +50,10 @@ public PlanStatus planStatus() { return planStatus; } + public List credentials() { + return credentials != null ? credentials : ImmutableList.of(); + } + public static Builder builder() { return new Builder(); } @@ -66,16 +76,22 @@ public static class Builder private Builder() {} private PlanStatus planStatus; + private final List credentials = Lists.newArrayList(); public Builder withPlanStatus(PlanStatus status) { this.planStatus = status; return this; } + public Builder withCredentials(List credentialsToAdd) { + credentials.addAll(credentialsToAdd); + return this; + } + @Override public FetchPlanningResultResponse build() { return new FetchPlanningResultResponse( - planStatus, planTasks(), fileScanTasks(), deleteFiles(), specsById()); + planStatus, planTasks(), fileScanTasks(), deleteFiles(), specsById(), credentials); } } } diff --git a/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponseParser.java b/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponseParser.java index 5400ef1dd13d..4a523d3c023b 100644 --- a/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponseParser.java +++ b/core/src/main/java/org/apache/iceberg/rest/responses/FetchPlanningResultResponseParser.java @@ -30,11 +30,14 @@ import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.rest.PlanStatus; import org.apache.iceberg.rest.TableScanResponseParser; +import org.apache.iceberg.rest.credentials.Credential; +import org.apache.iceberg.rest.credentials.CredentialParser; import org.apache.iceberg.util.JsonUtil; public class FetchPlanningResultResponseParser { private static final String STATUS = "status"; private static final String PLAN_TASKS = "plan-tasks"; + private static final String STORAGE_CREDENTIALS = "storage-credentials"; private FetchPlanningResultResponseParser() {} @@ -59,6 +62,15 @@ public static void toJson(FetchPlanningResultResponse response, JsonGenerator ge JsonUtil.writeStringArray(PLAN_TASKS, response.planTasks(), gen); } + if (!response.credentials().isEmpty()) { + gen.writeArrayFieldStart(STORAGE_CREDENTIALS); + for (Credential credential : response.credentials()) { + CredentialParser.toJson(credential, gen); + } + + gen.writeEndArray(); + } + TableScanResponseParser.serializeScanTasks( response.fileScanTasks(), response.deleteFiles(), response.specsById(), gen); gen.writeEndObject(); @@ -82,11 +94,18 @@ public static FetchPlanningResultResponse fromJson( List deleteFiles = TableScanResponseParser.parseDeleteFiles(json, specsById); List fileScanTasks = TableScanResponseParser.parseFileScanTasks(json, deleteFiles, specsById, caseSensitive); - return FetchPlanningResultResponse.builder() - .withPlanStatus(planStatus) - .withPlanTasks(planTasks) - .withFileScanTasks(fileScanTasks) - .withSpecsById(specsById) - .build(); + + FetchPlanningResultResponse.Builder builder = + FetchPlanningResultResponse.builder() + .withPlanStatus(planStatus) + .withPlanTasks(planTasks) + .withFileScanTasks(fileScanTasks) + .withSpecsById(specsById); + + if (json.hasNonNull(STORAGE_CREDENTIALS)) { + builder.withCredentials(LoadCredentialsResponseParser.fromJson(json).credentials()); + } + + return builder.build(); } } diff --git a/core/src/test/java/org/apache/iceberg/rest/responses/TestFetchPlanningResultResponseParser.java b/core/src/test/java/org/apache/iceberg/rest/responses/TestFetchPlanningResultResponseParser.java index 5d9d2cad828c..5fdfdc281f4f 100644 --- a/core/src/test/java/org/apache/iceberg/rest/responses/TestFetchPlanningResultResponseParser.java +++ b/core/src/test/java/org/apache/iceberg/rest/responses/TestFetchPlanningResultResponseParser.java @@ -40,8 +40,12 @@ import org.apache.iceberg.SchemaParser; import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.expressions.ResidualEvaluator; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.rest.PlanStatus; import org.apache.iceberg.rest.RESTSerializers; +import org.apache.iceberg.rest.credentials.Credential; +import org.apache.iceberg.rest.credentials.ImmutableCredential; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -140,7 +144,6 @@ public void roundTripSerdeWithInvalidPlanStatusSubmittedWithTasksPresent() { @Test public void roundTripSerdeWithInvalidPlanStatusSubmittedWithDeleteFilesNoFileScanTasksPresent() { - PlanStatus planStatus = PlanStatus.fromName("submitted"); assertThatThrownBy( () -> { @@ -225,4 +228,106 @@ public void roundTripSerdeWithValidStatusAndFileScanTasks() throws JsonProcessin assertThat(FetchPlanningResultResponseParser.toJson(copyResponse, false)) .isEqualTo(expectedToJson); } + + @Test + public void emptyOrInvalidCredentials() { + assertThat( + FetchPlanningResultResponseParser.fromJson( + "{\"status\": \"completed\",\"storage-credentials\": null}", + PARTITION_SPECS_BY_ID, + false) + .credentials()) + .isEmpty(); + + assertThat( + FetchPlanningResultResponseParser.fromJson( + "{\"status\": \"completed\",\"storage-credentials\": []}", + PARTITION_SPECS_BY_ID, + false) + .credentials()) + .isEmpty(); + + assertThatThrownBy( + () -> + FetchPlanningResultResponseParser.fromJson( + "{\"status\": \"completed\",\"storage-credentials\": \"invalid\"}", + PARTITION_SPECS_BY_ID, + false)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot parse credentials from non-array: \"invalid\""); + } + + @Test + public void roundTripSerdeWithCredentials() { + List credentials = + ImmutableList.of( + ImmutableCredential.builder() + .prefix("s3://custom-uri") + .config( + ImmutableMap.of( + "s3.access-key-id", + "keyId", + "s3.secret-access-key", + "accessKey", + "s3.session-token", + "sessionToken")) + .build(), + ImmutableCredential.builder() + .prefix("gs://custom-uri") + .config( + ImmutableMap.of( + "gcs.oauth2.token", "gcsToken1", "gcs.oauth2.token-expires-at", "1000")) + .build(), + ImmutableCredential.builder() + .prefix("gs") + .config( + ImmutableMap.of( + "gcs.oauth2.token", "gcsToken2", "gcs.oauth2.token-expires-at", "2000")) + .build()); + + FetchPlanningResultResponse response = + FetchPlanningResultResponse.builder() + .withPlanStatus(PlanStatus.COMPLETED) + .withCredentials(credentials) + .build(); + + String expectedJson = + "{\n" + + " \"status\" : \"completed\",\n" + + " \"storage-credentials\" : [ {\n" + + " \"prefix\" : \"s3://custom-uri\",\n" + + " \"config\" : {\n" + + " \"s3.access-key-id\" : \"keyId\",\n" + + " \"s3.secret-access-key\" : \"accessKey\",\n" + + " \"s3.session-token\" : \"sessionToken\"\n" + + " }\n" + + " }, {\n" + + " \"prefix\" : \"gs://custom-uri\",\n" + + " \"config\" : {\n" + + " \"gcs.oauth2.token\" : \"gcsToken1\",\n" + + " \"gcs.oauth2.token-expires-at\" : \"1000\"\n" + + " }\n" + + " }, {\n" + + " \"prefix\" : \"gs\",\n" + + " \"config\" : {\n" + + " \"gcs.oauth2.token\" : \"gcsToken2\",\n" + + " \"gcs.oauth2.token-expires-at\" : \"2000\"\n" + + " }\n" + + " } ]\n" + + "}"; + + String json = FetchPlanningResultResponseParser.toJson(response, true); + assertThat(json).isEqualTo(expectedJson); + + FetchPlanningResultResponse fromResponse = + FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false); + FetchPlanningResultResponse copyResponse = + FetchPlanningResultResponse.builder() + .withPlanStatus(fromResponse.planStatus()) + .withCredentials(credentials) + .build(); + + assertThat(FetchPlanningResultResponseParser.toJson(copyResponse, true)) + .isEqualTo(expectedJson); + } }