This is an automated email from the ASF dual-hosted git repository.
JNSimba pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-spark-connector.git
The following commit(s) were added to refs/heads/master by this push:
new 9675be3 [feat] Support IAM roles and gzip for S3 TVF sink (#382)
9675be3 is described below
commit 9675be392bd4876de717dc2c2d0e2f7afb5387aa
Author: wudi <[email protected]>
AuthorDate: Fri Sep 18 10:03:18 2026 +0800
[feat] Support IAM roles and gzip for S3 TVF sink (#382)
---
.../spark-doris-connector-base/pom.xml | 15 +++
.../client/write/tvf/S3ClientObjectStore.java | 71 +++++++++-
.../spark/client/write/tvf/S3TvfCommitter.java | 4 +-
.../spark/client/write/tvf/S3TvfSqlBuilder.java | 23 +++-
.../doris/spark/client/write/tvf/S3TvfWriter.java | 20 ++-
.../apache/doris/spark/config/DorisOptions.java | 4 +
.../apache/doris/spark/config/S3TvfOptions.java | 74 +++++++++-
.../client/write/tvf/S3ClientObjectStoreTest.java | 14 ++
.../spark/client/write/tvf/S3TvfCommitterTest.java | 57 ++++++++
.../client/write/tvf/S3TvfSqlBuilderTest.java | 48 +++++++
.../doris/spark/config/S3TvfOptionsTest.java | 86 ++++++++++++
.../doris/spark/sql/S3TvfIamRoleITCase.scala | 150 +++++++++++++++++++++
.../apache/doris/spark/sql/S3TvfSinkITCase.scala | 38 +++++-
13 files changed, 584 insertions(+), 20 deletions(-)
diff --git a/spark-doris-connector/spark-doris-connector-base/pom.xml
b/spark-doris-connector/spark-doris-connector-base/pom.xml
index 4e24e7f..ea44406 100644
--- a/spark-doris-connector/spark-doris-connector-base/pom.xml
+++ b/spark-doris-connector/spark-doris-connector-base/pom.xml
@@ -223,6 +223,21 @@
</exclusion>
</exclusions>
</dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>sts</artifactId>
+ <version>${awssdk.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>apache-client</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>netty-nio-client</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>url-connection-client</artifactId>
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStore.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStore.java
index 882adf6..ccb3854 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStore.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStore.java
@@ -19,6 +19,8 @@ package org.apache.doris.spark.client.write.tvf;
import org.apache.doris.spark.config.S3TvfOptions;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
@@ -26,6 +28,9 @@ import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.S3Configuration;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.sts.StsClient;
+import
software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
import java.io.ByteArrayInputStream;
import java.io.IOException;
@@ -34,12 +39,17 @@ import java.net.URI;
/** AWS SDK based implementation for S3-compatible object storage. */
public final class S3ClientObjectStore implements S3ObjectStore {
private static final String JSON_LINES_CONTENT_TYPE =
"application/x-ndjson";
+ private static final String ROLE_SESSION_NAME = "doris-spark-connector";
private final S3Client client;
private final String bucket;
+ private DefaultCredentialsProvider defaultCredentialsProvider;
+ private StsClient stsClient;
+ private StsAssumeRoleCredentialsProvider assumeRoleCredentialsProvider;
public S3ClientObjectStore(S3TvfOptions options) {
- this(createClient(options), options.getBucket());
+ this.bucket = options.getBucket();
+ this.client = createClient(options);
}
S3ClientObjectStore(S3Client client, String bucket) {
@@ -47,14 +57,11 @@ public final class S3ClientObjectStore implements
S3ObjectStore {
this.bucket = bucket;
}
- private static S3Client createClient(S3TvfOptions options) {
+ private S3Client createClient(S3TvfOptions options) {
return S3Client.builder()
.endpointOverride(URI.create(options.getEndpoint()))
.region(Region.of(options.getRegion()))
- .credentialsProvider(
- StaticCredentialsProvider.create(
- AwsBasicCredentials.create(
- options.getAccessKey(),
options.getSecretKey())))
+ .credentialsProvider(createCredentialsProvider(options))
.httpClientBuilder(UrlConnectionHttpClient.builder())
.serviceConfiguration(
S3Configuration.builder()
@@ -63,6 +70,44 @@ public final class S3ClientObjectStore implements
S3ObjectStore {
.build();
}
+ private AwsCredentialsProvider createCredentialsProvider(S3TvfOptions
options) {
+ if (!options.hasRoleArn()) {
+ return staticCredentialsProvider(options);
+ }
+ AwsCredentialsProvider sourceCredentialsProvider;
+ if (options.hasStaticCredentials()) {
+ sourceCredentialsProvider = staticCredentialsProvider(options);
+ } else {
+ defaultCredentialsProvider =
DefaultCredentialsProvider.builder().build();
+ sourceCredentialsProvider = defaultCredentialsProvider;
+ }
+ stsClient = StsClient.builder()
+ .region(Region.of(options.getRegion()))
+ .credentialsProvider(sourceCredentialsProvider)
+ .httpClientBuilder(UrlConnectionHttpClient.builder())
+ .build();
+ assumeRoleCredentialsProvider =
StsAssumeRoleCredentialsProvider.builder()
+ .stsClient(stsClient)
+ .refreshRequest(buildAssumeRoleRequest(options))
+ .build();
+ return assumeRoleCredentialsProvider;
+ }
+
+ static AssumeRoleRequest buildAssumeRoleRequest(S3TvfOptions options) {
+ AssumeRoleRequest.Builder request = AssumeRoleRequest.builder()
+ .roleArn(options.getRoleArn())
+ .roleSessionName(ROLE_SESSION_NAME);
+ if (options.getExternalId() != null) {
+ request.externalId(options.getExternalId());
+ }
+ return request.build();
+ }
+
+ private static StaticCredentialsProvider
staticCredentialsProvider(S3TvfOptions options) {
+ return StaticCredentialsProvider.create(
+ AwsBasicCredentials.create(options.getAccessKey(),
options.getSecretKey()));
+ }
+
@Override
public void put(String objectKey, byte[] content) throws IOException {
PutObjectRequest request = PutObjectRequest.builder()
@@ -82,6 +127,18 @@ public final class S3ClientObjectStore implements
S3ObjectStore {
@Override
public void close() {
- client.close();
+ try {
+ client.close();
+ } finally {
+ if (assumeRoleCredentialsProvider != null) {
+ assumeRoleCredentialsProvider.close();
+ }
+ if (stsClient != null) {
+ stsClient.close();
+ }
+ if (defaultCredentialsProvider != null) {
+ defaultCredentialsProvider.close();
+ }
+ }
}
}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitter.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitter.java
index 540972b..6c1e330 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitter.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitter.java
@@ -36,6 +36,7 @@ public final class S3TvfCommitter implements AutoCloseable {
private static final String PARTIAL_COLUMNS = "partial_columns";
private static final String FORMAT = "format";
private static final String READ_JSON_BY_LINE = "read_json_by_line";
+ private static final String COMPRESS_TYPE = "compress_type";
private static final String ENABLE_UNIQUE_KEY_PARTIAL_UPDATE =
"enable_unique_key_partial_update";
@@ -92,7 +93,8 @@ public final class S3TvfCommitter implements AutoCloseable {
if (!COLUMNS.equals(name)
&& !PARTIAL_COLUMNS.equals(name)
&& !FORMAT.equals(name)
- && !READ_JSON_BY_LINE.equals(name)) {
+ && !READ_JSON_BY_LINE.equals(name)
+ && !COMPRESS_TYPE.equals(name)) {
values.put(name, entry.getValue());
}
}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilder.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilder.java
index 0fbd025..2f8ea64 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilder.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilder.java
@@ -41,6 +41,7 @@ public final class S3TvfSqlBuilder {
}
String columnSql = joinIdentifiers(columns);
String uri = buildUri(objectKeys);
+ String credentials = buildCredentials();
return "INSERT INTO "
+ TvfSqlUtils.quoteIdentifier(committable.getDatabase())
+ "."
@@ -58,18 +59,34 @@ public final class S3TvfSqlBuilder {
+ ","
+ property("read_json_by_line", "true")
+ ","
+ + (options.isGzipCompressionEnabled()
+ ? property("compress_type", "gz") + ","
+ : "")
+ property("s3.endpoint", options.getEndpoint())
+ ","
+ property("s3.region", options.getRegion())
+ ","
- + property("s3.access_key", options.getAccessKey())
- + ","
- + property("s3.secret_key", options.getSecretKey())
+ + credentials
+ ","
+ property("use_path_style",
Boolean.toString(options.isPathStyleAccess()))
+ ")";
}
+ private String buildCredentials() {
+ StringJoiner credentials = new StringJoiner(",");
+ if (options.hasStaticCredentials()) {
+ credentials.add(property("s3.access_key", options.getAccessKey()));
+ credentials.add(property("s3.secret_key", options.getSecretKey()));
+ }
+ if (options.hasRoleArn()) {
+ credentials.add(property("s3.role_arn", options.getRoleArn()));
+ if (options.getExternalId() != null) {
+ credentials.add(property("s3.external_id",
options.getExternalId()));
+ }
+ }
+ return credentials.toString();
+ }
+
private String buildUri(List<String> objectKeys) {
if (objectKeys.size() == 1) {
return "s3://" + options.getBucket() + "/" + objectKeys.get(0);
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfWriter.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfWriter.java
index 135f08e..3e4d754 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfWriter.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/tvf/S3TvfWriter.java
@@ -32,6 +32,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
+import java.util.zip.GZIPOutputStream;
/** Writes one logical Spark partition as deterministic JSON Lines objects. */
public final class S3TvfWriter implements AutoCloseable {
@@ -104,22 +105,25 @@ public final class S3TvfWriter implements AutoCloseable {
return;
}
byte[] content = buffer.toByteArray();
+ boolean gzipEnabled = options.isGzipCompressionEnabled();
int currentFileNumber = fileNumber++;
String fileName = String.format(
- "%s_%s_%s_%d_%d.json",
+ "%s_%s_%s_%d_%d.json%s",
labelPrefix,
normalizedTable,
batchUuid,
partitionId,
- currentFileNumber);
+ currentFileNumber,
+ gzipEnabled ? ".gz" : "");
String prefix = options.getPrefix();
String objectKey = prefix + (prefix.endsWith("/") ? "" : "/") +
fileName;
long uploadStartedAtNanos = System.nanoTime();
try {
- objectStore.put(objectKey, content);
+ byte[] uploadContent = gzipEnabled ? gzip(content) : content;
+ objectStore.put(objectKey, uploadContent);
LOG.info("S3 TVF object upload completed, objectKey={},
sizeBytes={}, uploadTimeMs={}.",
objectKey,
- content.length,
+ uploadContent.length,
TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
uploadStartedAtNanos));
} catch (IOException | RuntimeException e) {
LOG.warn("S3 TVF object upload failed, objectKey={}, sizeBytes={},
uploadTimeMs={}.",
@@ -134,6 +138,14 @@ public final class S3TvfWriter implements AutoCloseable {
recordCount = 0;
}
+ private static byte[] gzip(byte[] content) throws IOException {
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (GZIPOutputStream gzip = new GZIPOutputStream(output)) {
+ gzip.write(content);
+ }
+ return output.toByteArray();
+ }
+
private String label() {
return labelPrefix
+ "_"
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
index 3d1b3a2..5afcfa5 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
@@ -146,6 +146,10 @@ public class DorisOptions {
public static final ConfigOption<String> DORIS_SINK_S3_SECRET_KEY =
ConfigOptions.name("doris.sink.s3.secret-key").stringType().withoutDefaultValue().withDescription("Secret
key of the S3-compatible object storage.");
+ public static final ConfigOption<String> DORIS_SINK_S3_ROLE_ARN =
ConfigOptions.name("doris.sink.s3.role-arn").stringType().withoutDefaultValue().withDescription("AWS
IAM role ARN used to access the S3 object storage.");
+
+ public static final ConfigOption<String> DORIS_SINK_S3_EXTERNAL_ID =
ConfigOptions.name("doris.sink.s3.external-id").stringType().withoutDefaultValue().withDescription("External
ID used when assuming the AWS IAM role.");
+
public static final ConfigOption<Boolean> DORIS_SINK_S3_PATH_STYLE_ACCESS
=
ConfigOptions.name("doris.sink.s3.path-style-access").booleanType().defaultValue(false).withDescription("Whether
to use path-style access for object storage.");
public static final ConfigOption<Integer> DORIS_THRIFT_MAX_MESSAGE_SIZE =
ConfigOptions.name("doris.thrift.max.message.size").intType().defaultValue(Integer.MAX_VALUE).withDescription("")
;
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/S3TvfOptions.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/S3TvfOptions.java
index 399a3ba..ee2ce07 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/S3TvfOptions.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/S3TvfOptions.java
@@ -28,6 +28,7 @@ public final class S3TvfOptions implements Serializable {
private static final long serialVersionUID = 1L;
private static final String FORMAT = "format";
private static final String READ_JSON_BY_LINE = "read_json_by_line";
+ private static final String COMPRESS_TYPE = "compress_type";
private final String endpoint;
private final String region;
@@ -35,6 +36,9 @@ public final class S3TvfOptions implements Serializable {
private final String prefix;
private final String accessKey;
private final String secretKey;
+ private final String roleArn;
+ private final String externalId;
+ private final boolean gzipCompressionEnabled;
private final boolean pathStyleAccess;
private S3TvfOptions(
@@ -44,6 +48,9 @@ public final class S3TvfOptions implements Serializable {
String prefix,
String accessKey,
String secretKey,
+ String roleArn,
+ String externalId,
+ boolean gzipCompressionEnabled,
boolean pathStyleAccess) {
this.endpoint = endpoint;
this.region = region;
@@ -51,6 +58,9 @@ public final class S3TvfOptions implements Serializable {
this.prefix = prefix;
this.accessKey = accessKey;
this.secretKey = secretKey;
+ this.roleArn = roleArn;
+ this.externalId = externalId;
+ this.gzipCompressionEnabled = gzipCompressionEnabled;
this.pathStyleAccess = pathStyleAccess;
}
@@ -60,8 +70,11 @@ public final class S3TvfOptions implements Serializable {
String region = required(config, DorisOptions.DORIS_SINK_S3_REGION);
String bucket = required(config, DorisOptions.DORIS_SINK_S3_BUCKET);
String prefix = required(config, DorisOptions.DORIS_SINK_S3_PREFIX);
- String accessKey = required(config,
DorisOptions.DORIS_SINK_S3_ACCESS_KEY);
- String secretKey = required(config,
DorisOptions.DORIS_SINK_S3_SECRET_KEY);
+ String accessKey = optional(config,
DorisOptions.DORIS_SINK_S3_ACCESS_KEY);
+ String secretKey = optional(config,
DorisOptions.DORIS_SINK_S3_SECRET_KEY);
+ String roleArn = optional(config, DorisOptions.DORIS_SINK_S3_ROLE_ARN);
+ String externalId = optional(config,
DorisOptions.DORIS_SINK_S3_EXTERNAL_ID);
+ validateCredentials(accessKey, secretKey, roleArn, externalId);
validatePrefix(prefix);
return new S3TvfOptions(
@@ -71,9 +84,37 @@ public final class S3TvfOptions implements Serializable {
prefix,
accessKey,
secretKey,
+ roleArn,
+ externalId,
+ isGzipCompressionEnabled(config.getSinkProperties()),
config.getValue(DorisOptions.DORIS_SINK_S3_PATH_STYLE_ACCESS));
}
+ private static String optional(DorisConfig config, ConfigOption<String>
option)
+ throws OptionRequiredException {
+ if (!config.contains(option)) {
+ return null;
+ }
+ String value = config.getValue(option).trim();
+ return value.isEmpty() ? null : value;
+ }
+
+ private static void validateCredentials(
+ String accessKey, String secretKey, String roleArn, String
externalId) {
+ if ((accessKey == null) != (secretKey == null)) {
+ throw new IllegalArgumentException(
+ "doris.sink.s3.access-key and doris.sink.s3.secret-key
must be configured together");
+ }
+ if (accessKey == null && roleArn == null) {
+ throw new IllegalArgumentException(
+ "S3 TVF requires either access/secret keys or
doris.sink.s3.role-arn");
+ }
+ if (externalId != null && roleArn == null) {
+ throw new IllegalArgumentException(
+ "doris.sink.s3.external-id requires
doris.sink.s3.role-arn");
+ }
+ }
+
private static String required(DorisConfig config, ConfigOption<String>
option)
throws OptionRequiredException {
String value = config.getValue(option).trim();
@@ -93,6 +134,15 @@ public final class S3TvfOptions implements Serializable {
throw new IllegalArgumentException(
"TVF write mode requires
'doris.sink.properties.read_json_by_line' to be true");
}
+ String compressType = loadProperties.getOrDefault(COMPRESS_TYPE,
"gz").trim();
+ if (!compressType.isEmpty() && !"gz".equalsIgnoreCase(compressType)) {
+ throw new IllegalArgumentException(
+ "TVF write mode only supports 'gz' or an empty
compress_type");
+ }
+ }
+
+ private static boolean isGzipCompressionEnabled(Map<String, String>
loadProperties) {
+ return
"gz".equalsIgnoreCase(loadProperties.getOrDefault(COMPRESS_TYPE, "gz").trim());
}
private static void validatePrefix(String prefix) {
@@ -129,6 +179,26 @@ public final class S3TvfOptions implements Serializable {
return secretKey;
}
+ public String getRoleArn() {
+ return roleArn;
+ }
+
+ public String getExternalId() {
+ return externalId;
+ }
+
+ public boolean hasRoleArn() {
+ return roleArn != null;
+ }
+
+ public boolean hasStaticCredentials() {
+ return accessKey != null;
+ }
+
+ public boolean isGzipCompressionEnabled() {
+ return gzipCompressionEnabled;
+ }
+
public boolean isPathStyleAccess() {
return pathStyleAccess;
}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStoreTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStoreTest.java
index ae84933..24ef12b 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStoreTest.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3ClientObjectStoreTest.java
@@ -17,18 +17,32 @@
package org.apache.doris.spark.client.write.tvf;
+import org.apache.doris.spark.config.S3TvfOptions;
+import org.apache.doris.spark.config.S3TvfOptionsTest;
import org.junit.Assert;
import org.junit.Test;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
import java.io.DataInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
public class S3ClientObjectStoreTest {
+ @Test
+ public void buildsAssumeRoleRequest() throws Exception {
+ S3TvfOptions options =
S3TvfOptions.fromConfig(S3TvfOptionsTest.configWithRole());
+
+ AssumeRoleRequest request =
S3ClientObjectStore.buildAssumeRoleRequest(options);
+
+ Assert.assertEquals("arn:aws:iam::123456789012:role/doris",
request.roleArn());
+ Assert.assertEquals("external-id", request.externalId());
+ Assert.assertEquals("doris-spark-connector",
request.roleSessionName());
+ }
+
@Test
public void uploadBodyIsRepeatableWithoutCopying() throws Exception {
byte[] content = "{\"id\":1}\n".getBytes(StandardCharsets.UTF_8);
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitterTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitterTest.java
new file mode 100644
index 0000000..81d8b67
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfCommitterTest.java
@@ -0,0 +1,57 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.spark.client.write.tvf;
+
+import org.apache.doris.spark.config.S3TvfOptions;
+import org.apache.doris.spark.config.S3TvfOptionsTest;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.sql.SQLException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class S3TvfCommitterTest {
+
+ @Test
+ public void excludesCompressTypeFromSessionVariables() throws Exception {
+ Map<String, String> loadProperties = new HashMap<>();
+ loadProperties.put("compress_type", "");
+ AtomicReference<Map<String, String>> captured = new
AtomicReference<>();
+ S3TvfLoadClient loadClient = new S3TvfLoadClient() {
+ @Override
+ public void executeInsert(String sql, Map<String, String>
sessionVariables)
+ throws SQLException {
+ captured.set(sessionVariables);
+ }
+
+ @Override
+ public void close() throws IOException {}
+ };
+ S3TvfOptions options =
S3TvfOptions.fromConfig(S3TvfOptionsTest.configWithRole());
+ S3TvfCommitter committer = new S3TvfCommitter(options, loadProperties,
loadClient);
+
+ committer.commit(new S3TvfCommittable(
+ "db", "tbl", "label", Arrays.asList("file.json"),
Arrays.asList("id")));
+
+ Assert.assertFalse(captured.get().containsKey("compress_type"));
+ }
+}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilderTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilderTest.java
new file mode 100644
index 0000000..171b830
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/write/tvf/S3TvfSqlBuilderTest.java
@@ -0,0 +1,48 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.spark.client.write.tvf;
+
+import org.apache.doris.spark.config.S3TvfOptions;
+import org.apache.doris.spark.config.S3TvfOptionsTest;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+
+public class S3TvfSqlBuilderTest {
+
+ @Test
+ public void buildsIamRoleAndGzipProperties() throws Exception {
+ S3TvfOptions options =
S3TvfOptions.fromConfig(S3TvfOptionsTest.configWithRole());
+ S3TvfCommittable committable =
+ new S3TvfCommittable(
+ "db",
+ "tbl",
+ "label",
+ Arrays.asList("spark/file.json.gz"),
+ Arrays.asList("id"));
+
+ String sql = new S3TvfSqlBuilder(options).buildInsertSql(committable);
+
+ Assert.assertTrue(sql.contains("'s3.role_arn' =
'arn:aws:iam::123456789012:role/doris'"));
+ Assert.assertTrue(sql.contains("'s3.external_id' = 'external-id'"));
+ Assert.assertTrue(sql.contains("'compress_type' = 'gz'"));
+ Assert.assertFalse(sql.contains("s3.access_key"));
+ Assert.assertFalse(sql.contains("s3.secret_key"));
+ }
+}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/config/S3TvfOptionsTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/config/S3TvfOptionsTest.java
new file mode 100644
index 0000000..f77c2cd
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/config/S3TvfOptionsTest.java
@@ -0,0 +1,86 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.spark.config;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class S3TvfOptionsTest {
+
+ @Test
+ public void supportsIamRoleAndDefaultGzip() throws Exception {
+ S3TvfOptions options = S3TvfOptions.fromConfig(configWithRole());
+
+ Assert.assertEquals("arn:aws:iam::123456789012:role/doris",
options.getRoleArn());
+ Assert.assertEquals("external-id", options.getExternalId());
+ Assert.assertTrue(options.isGzipCompressionEnabled());
+ Assert.assertFalse(options.hasStaticCredentials());
+ }
+
+ @Test
+ public void allowsStaticSourceCredentialsAndDisablingGzip() throws
Exception {
+ Map<String, String> values = baseOptions();
+ values.put("doris.sink.s3.access-key", "access-key");
+ values.put("doris.sink.s3.secret-key", "secret-key");
+ values.put("doris.sink.s3.role-arn",
"arn:aws:iam::123456789012:role/doris");
+ values.put("doris.sink.properties.compress_type", "");
+
+ S3TvfOptions options =
S3TvfOptions.fromConfig(DorisConfig.fromMap(values, false));
+
+ Assert.assertTrue(options.hasStaticCredentials());
+ Assert.assertFalse(options.isGzipCompressionEnabled());
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void rejectsMissingCredentialsAndRole() throws Exception {
+ S3TvfOptions.fromConfig(DorisConfig.fromMap(baseOptions(), false));
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void rejectsUnsupportedCompression() throws Exception {
+ Map<String, String> values = baseOptions();
+ values.put("doris.sink.s3.role-arn",
"arn:aws:iam::123456789012:role/doris");
+ values.put("doris.sink.properties.compress_type", "zstd");
+ S3TvfOptions.fromConfig(DorisConfig.fromMap(values, false));
+ }
+
+ public static DorisConfig configWithRole() throws Exception {
+ Map<String, String> values = baseOptions();
+ values.put("doris.sink.s3.role-arn",
"arn:aws:iam::123456789012:role/doris");
+ values.put("doris.sink.s3.external-id", "external-id");
+ return DorisConfig.fromMap(values, false);
+ }
+
+ private static Map<String, String> baseOptions() {
+ Map<String, String> values = new HashMap<>();
+ values.put("doris.fenodes", "localhost:8030");
+ values.put("doris.query.port", "9030");
+ values.put("doris.table.identifier", "db.tbl");
+ values.put("doris.user", "root");
+ values.put("doris.password", "");
+ values.put("doris.sink.mode", "tvf");
+ values.put("doris.sink.s3.endpoint", "https://s3.example.com";);
+ values.put("doris.sink.s3.region", "us-east-1");
+ values.put("doris.sink.s3.bucket", "staging");
+ values.put("doris.sink.s3.prefix", "spark/orders");
+ return values;
+ }
+}
diff --git
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfIamRoleITCase.scala
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfIamRoleITCase.scala
new file mode 100644
index 0000000..c780018
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfIamRoleITCase.scala
@@ -0,0 +1,150 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.spark.sql
+
+import org.apache.doris.spark.container.ContainerUtils
+import org.apache.doris.spark.container.instance.DorisCustomerContainer
+import org.apache.spark.sql.{SaveMode, SparkSession}
+import org.junit.Assert.assertEquals
+import org.junit.{AfterClass, Assume, BeforeClass, Test}
+import org.slf4j.LoggerFactory
+
+import java.sql.Connection
+import java.util
+import java.util.UUID
+import scala.collection.JavaConverters._
+
+/** Opt-in integration test for Spark DataFrame S3 TVF writes with an AWS IAM
role. */
+class S3TvfIamRoleITCase {
+ import S3TvfIamRoleITCase._
+
+ @Test
+ def testWritesThroughIamRole(): Unit = {
+ val table = "iam_role_" + UUID.randomUUID().toString.replace("-", "")
+ try {
+ executeSql(
+ s"CREATE DATABASE IF NOT EXISTS `$database`",
+ s"CREATE TABLE `$database`.`$table` (`id` INT, `name` VARCHAR(64)) " +
+ "DUPLICATE KEY(`id`) DISTRIBUTED BY HASH(`id`) BUCKETS 1 " +
+ "PROPERTIES (\"replication_num\" = \"1\")")
+
+ val session = SparkSession.builder()
+ .appName("s3-tvf-iam-role-it")
+ .master("local[1]")
+ .config("spark.ui.enabled", "false")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.driver.bindAddress", "127.0.0.1")
+ .getOrCreate()
+ try {
+ import session.implicits._
+ Seq((1, "doris"), (2, "spark"))
+ .toDF("id", "name")
+ .write
+ .format("doris")
+ .options(sinkOptions(table).asScala)
+ .mode(SaveMode.Append)
+ .save()
+ } finally {
+ session.stop()
+ }
+
+ val connection = doris.getQueryConnection
+ val rows = try {
+ ContainerUtils.executeSQLStatement(
+ connection,
+ LOG,
+ s"SELECT id,name FROM `$database`.`$table` ORDER BY id",
+ 2)
+ } finally {
+ connection.close()
+ }
+ assertEquals(util.Arrays.asList("1,doris", "2,spark"), rows)
+ } finally {
+ executeSql(s"DROP TABLE IF EXISTS `$database`.`$table`")
+ }
+ }
+
+ private def sinkOptions(table: String): util.Map[String, String] = {
+ val options = new util.HashMap[String, String]()
+ options.put("doris.fenodes", doris.getFenodes)
+ options.put("doris.query.port", doris.getQueryPort.toString)
+ options.put("doris.table.identifier", s"$database.$table")
+ options.put("user", doris.getUsername)
+ options.put("password", doris.getPassword)
+ options.put("doris.sink.mode", "tvf")
+ options.put("doris.sink.label.prefix", "iam_role_" + UUID.randomUUID())
+ options.put("doris.sink.s3.endpoint", requiredProperty("s3_endpoint"))
+ options.put("doris.sink.s3.region", requiredProperty("s3_region"))
+ options.put("doris.sink.s3.bucket", requiredProperty("s3_bucket"))
+ options.put("doris.sink.s3.prefix", System.getProperty("s3_prefix",
"doris-spark-connector-it"))
+ options.put("doris.sink.s3.role-arn", requiredProperty("s3_role_arn"))
+ optionalProperty("s3_external_id")
+ .foreach(options.put("doris.sink.s3.external-id", _))
+ options
+ }
+
+ private def executeSql(sql: String*): Unit = {
+ val connection: Connection = doris.getQueryConnection
+ try {
+ ContainerUtils.executeSQLStatement(connection, LOG, sql: _*)
+ } finally {
+ connection.close()
+ }
+ }
+
+ private def requiredProperty(name: String): String = {
+ optionalProperty(name).getOrElse(
+ throw new IllegalArgumentException("Missing required system property: "
+ name))
+ }
+
+ private def optionalProperty(name: String): Option[String] = {
+ Option(System.getProperty(name)).map(_.trim).filter(_.nonEmpty)
+ }
+}
+
+object S3TvfIamRoleITCase {
+ private val LOG = LoggerFactory.getLogger(classOf[S3TvfIamRoleITCase])
+ private val database = "test_s3_tvf_iam_role"
+ private var doris: DorisCustomerContainer = _
+
+ @BeforeClass
+ def useExternalEnvironment(): Unit = {
+ Assume.assumeTrue(
+ "IAM role ITCase requires -Ds3_tvf_iam_role_it=true",
+ java.lang.Boolean.getBoolean("s3_tvf_iam_role_it"))
+ Assume.assumeTrue(
+ "IAM role ITCase requires -Dcustomer_env=true",
+ java.lang.Boolean.getBoolean("customer_env"))
+ Seq("s3_endpoint", "s3_region", "s3_bucket", "s3_role_arn")
+ .foreach(name => {
+ val value = System.getProperty(name)
+ if (value == null || value.trim.isEmpty) {
+ throw new IllegalArgumentException("Missing required system
property: " + name)
+ }
+ })
+ doris = new DorisCustomerContainer()
+ doris.startContainer()
+ }
+
+ @AfterClass
+ def closeExternalEnvironment(): Unit = {
+ if (doris != null) {
+ doris.close()
+ }
+ }
+}
diff --git
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfSinkITCase.scala
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfSinkITCase.scala
index 0f5176e..360df0c 100644
---
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfSinkITCase.scala
+++
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/S3TvfSinkITCase.scala
@@ -80,11 +80,43 @@ class S3TvfSinkITCase extends AbstractS3TvfTestBase {
assertEquals(2, keys.size)
val keyPattern =
Pattern.quote(prefix + labelPrefix + "_" + table + "_") +
- "[0-9a-f-]{36}_0_[0-9]+\\.json"
+ "[0-9a-f-]{36}_0_[0-9]+\\.json\\.gz"
assertTrue(keys.forall(_.matches(keyPattern)))
assertFalse(keys.exists(_.contains("//")))
}
+ @Test
+ def testDisablesCompressionWithEmptyCompressType(): Unit = {
+ val table = uniqueName("uncompressed")
+ val prefix = uniqueName("uncompressed_objects")
+ val labelPrefix = uniqueName("uncompressed_label")
+ createDuplicateTable(table, "`id` INT, `name` VARCHAR(128)")
+
+ withSpark("local[1]") { session =>
+ import session.implicits._
+ val options = s3TvfSinkOptions(
+ s"$database.$table", prefix, labelPrefix, 100)
+ options.put("doris.sink.properties.compress_type", "")
+ Seq((1, "doris"))
+ .toDF("id", "name")
+ .write
+ .format("doris")
+ .options(options.asScala)
+ .mode(SaveMode.Append)
+ .save()
+ }
+
+ assertResult(
+ table,
+ "id,name",
+ util.Collections.singletonList("1,doris"),
+ columnCount = 2)
+ val keys = listObjectKeys(prefix + "/").asScala
+ assertEquals(1, keys.size)
+ assertTrue(keys.head.endsWith(".json"))
+ assertFalse(keys.head.endsWith(".json.gz"))
+ }
+
@Test
def testParallelTaskWrite(): Unit = {
val table = uniqueName("parallel")
@@ -111,7 +143,7 @@ class S3TvfSinkITCase extends AbstractS3TvfTestBase {
assertEquals(2, keys.size)
val keyPattern =
(Pattern.quote(prefix + "/" + labelPrefix + "_" + table + "_") +
- "([0-9a-f-]{36})_([0-9]+)_0\\.json").r
+ "([0-9a-f-]{36})_([0-9]+)_0\\.json\\.gz").r
val taskFiles = keys.map {
case keyPattern(uuid, partition) => uuid -> partition
case key => throw new AssertionError("Unexpected S3 TVF object key: " +
key)
@@ -247,7 +279,7 @@ class S3TvfSinkITCase extends AbstractS3TvfTestBase {
val keys = listObjectKeys(prefix + "/").asScala
val keyPattern =
(Pattern.quote(prefix + "/" + labelPrefix + "_" + table + "_") +
- "([0-9a-f-]{36})_[0-9]+_[0-9]+\\.json").r
+ "([0-9a-f-]{36})_[0-9]+_[0-9]+\\.json\\.gz").r
val uuids = keys.map {
case keyPattern(uuid) => uuid
case key => throw new AssertionError("Unexpected S3 TVF object key: " +
key)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
