Skip to content

Commit

Permalink
Add s3.storage-class property in S3FileSystemConfig
Browse files Browse the repository at this point in the history
Add s3.storage-class property integration in S3 operations

Modify configuration tests to use storage class property

Update Trino documentation to describe s3.storage-class configuration

Implement whitelist for supported storage classes and map to SDK options
  • Loading branch information
EdenKik committed Feb 13, 2025
1 parent 3d37aba commit 218a639
Show file tree
Hide file tree
Showing 6 changed files with 49 additions and 2 deletions.
2 changes: 2 additions & 0 deletions docs/src/main/sphinx/object-storage/file-system-s3.md
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ support:
- Required region name for S3.
* - `s3.path-style-access`
- Use path-style access for all requests to S3
* - `s3.storage-class`
- Set the S3 storage class to use when writing the data, defaults to STANDARD
* - `s3.exclusive-create`
- Whether conditional write is supported by the S3-compatible storage. Defaults to `true`.
* - `s3.canned-acl`
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import io.trino.filesystem.s3.S3FileSystemConfig.ObjectCannedAcl;
import io.trino.filesystem.s3.S3FileSystemConfig.S3SseType;
import io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType;
import io.trino.spi.security.ConnectorIdentity;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
Expand All @@ -37,6 +38,7 @@ record S3Context(
boolean requesterPays,
S3SseContext s3SseContext,
Optional<AwsCredentialsProvider> credentialsProviderOverride,
StorageClassType storageClass,
ObjectCannedAcl cannedAcl,
boolean exclusiveWriteSupported)
{
Expand All @@ -56,7 +58,7 @@ public RequestPayer requestPayer()

public S3Context withKmsKeyId(String kmsKeyId)
{
return new S3Context(partSize, requesterPays, S3SseContext.withKmsKeyId(kmsKeyId), credentialsProviderOverride, cannedAcl, exclusiveWriteSupported);
return new S3Context(partSize, requesterPays, S3SseContext.withKmsKeyId(kmsKeyId), credentialsProviderOverride, storageClass, cannedAcl, exclusiveWriteSupported);
}

public S3Context withCredentials(ConnectorIdentity identity)
Expand All @@ -73,7 +75,7 @@ public S3Context withCredentials(ConnectorIdentity identity)

public S3Context withSseCustomerKey(String key)
{
return new S3Context(partSize, requesterPays, S3SseContext.withSseCustomerKey(key), credentialsProviderOverride, cannedAcl, exclusiveWriteSupported);
return new S3Context(partSize, requesterPays, S3SseContext.withSseCustomerKey(key), credentialsProviderOverride, storageClass, cannedAcl, exclusiveWriteSupported);
}

public S3Context withCredentialsProviderOverride(AwsCredentialsProvider credentialsProviderOverride)
Expand All @@ -83,6 +85,7 @@ public S3Context withCredentialsProviderOverride(AwsCredentialsProvider credenti
requesterPays,
s3SseContext,
Optional.of(credentialsProviderOverride),
storageClass,
cannedAcl,
exclusiveWriteSupported);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import jakarta.validation.constraints.Size;
import software.amazon.awssdk.retries.api.RetryStrategy;
import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
import software.amazon.awssdk.services.s3.model.StorageClass;

import java.util.Optional;
import java.util.Set;
Expand All @@ -46,6 +47,22 @@ public enum S3SseType
NONE, S3, KMS, CUSTOMER
}

public enum StorageClassType
{
STANDARD,
STANDARD_IA,
INTELLIGENT_TIERING;

public static StorageClass getStorageClass(S3FileSystemConfig.StorageClassType storageClass)
{
return switch (storageClass) {
case STANDARD -> StorageClass.STANDARD;
case STANDARD_IA -> StorageClass.STANDARD_IA;
case INTELLIGENT_TIERING -> StorageClass.INTELLIGENT_TIERING;
};
}
}

public enum ObjectCannedAcl
{
NONE,
Expand Down Expand Up @@ -91,6 +108,7 @@ public static RetryStrategy getRetryStrategy(RetryMode retryMode)
private String endpoint;
private String region;
private boolean pathStyleAccess;
private StorageClassType storageClass = StorageClassType.STANDARD;
private String iamRole;
private String roleSessionName = "trino-filesystem";
private String externalId;
Expand Down Expand Up @@ -182,6 +200,19 @@ public S3FileSystemConfig setPathStyleAccess(boolean pathStyleAccess)
return this;
}

public StorageClassType getStorageClass()
{
return storageClass;
}

@Config("s3.storage-class")
@ConfigDescription("The S3 storage class to use when writing the data")
public S3FileSystemConfig setStorageClass(StorageClassType storageClass)
{
this.storageClass = storageClass;
return this;
}

public String getIamRole()
{
return iamRole;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ private S3FileSystemLoader(Optional<S3SecurityMappingProvider> mappingProvider,
config.getSseKmsKeyId(),
config.getSseCustomerKey()),
Optional.empty(),
config.getStorageClass(),
config.getCannedAcl(),
config.isSupportsExclusiveCreate());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.RequestPayer;
import software.amazon.awssdk.services.s3.model.S3Exception;
import software.amazon.awssdk.services.s3.model.StorageClass;
import software.amazon.awssdk.services.s3.model.UploadPartRequest;
import software.amazon.awssdk.services.s3.model.UploadPartResponse;

Expand All @@ -46,6 +47,7 @@
import static com.google.common.base.Verify.verify;
import static io.trino.filesystem.s3.S3FileSystemConfig.ObjectCannedAcl.getCannedAcl;
import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.NONE;
import static io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType.getStorageClass;
import static io.trino.filesystem.s3.S3SseCUtils.encoded;
import static io.trino.filesystem.s3.S3SseCUtils.md5Checksum;
import static io.trino.filesystem.s3.S3SseRequestConfigurator.setEncryptionSettings;
Expand All @@ -68,6 +70,7 @@ final class S3OutputStream
private final S3Context context;
private final int partSize;
private final RequestPayer requestPayer;
private final StorageClass storageClass;
private final ObjectCannedACL cannedAcl;
private final boolean exclusiveCreate;
private final Optional<EncryptionKey> key;
Expand Down Expand Up @@ -97,6 +100,7 @@ public S3OutputStream(AggregatedMemoryContext memoryContext, Executor uploadExec
this.context = requireNonNull(context, "context is null");
this.partSize = context.partSize();
this.requestPayer = context.requestPayer();
this.storageClass = getStorageClass(context.storageClass());
this.cannedAcl = getCannedAcl(context.cannedAcl());
this.key = requireNonNull(key, "key is null");

Expand Down Expand Up @@ -214,6 +218,7 @@ private void flushBuffer(boolean finished)
.requestPayer(requestPayer)
.bucket(location.bucket())
.key(location.key())
.storageClass(storageClass)
.contentLength((long) bufferSize)
.applyMutation(builder -> {
if (exclusiveCreate) {
Expand Down Expand Up @@ -304,6 +309,7 @@ private CompletedPart uploadPage(byte[] data, int length)
.requestPayer(requestPayer)
.bucket(location.bucket())
.key(location.key())
.storageClass(storageClass)
.applyMutation(builder ->
key.ifPresentOrElse(
encryption ->
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import io.airlift.units.Duration;
import io.trino.filesystem.s3.S3FileSystemConfig.ObjectCannedAcl;
import io.trino.filesystem.s3.S3FileSystemConfig.S3SseType;
import io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType;
import jakarta.validation.constraints.AssertTrue;
import org.junit.jupiter.api.Test;

Expand Down Expand Up @@ -49,6 +50,7 @@ public void testDefaults()
.setExternalId(null)
.setStsEndpoint(null)
.setStsRegion(null)
.setStorageClass(StorageClassType.STANDARD)
.setCannedAcl(ObjectCannedAcl.NONE)
.setSseType(S3SseType.NONE)
.setRetryMode(LEGACY)
Expand Down Expand Up @@ -88,6 +90,7 @@ public void testExplicitPropertyMappings()
.put("s3.external-id", "myid")
.put("s3.sts.endpoint", "sts.example.com")
.put("s3.sts.region", "us-west-2")
.put("s3.storage-class", "STANDARD_IA")
.put("s3.canned-acl", "BUCKET_OWNER_FULL_CONTROL")
.put("s3.retry-mode", "STANDARD")
.put("s3.max-error-retries", "12")
Expand Down Expand Up @@ -124,6 +127,7 @@ public void testExplicitPropertyMappings()
.setExternalId("myid")
.setStsEndpoint("sts.example.com")
.setStsRegion("us-west-2")
.setStorageClass(StorageClassType.STANDARD_IA)
.setCannedAcl(ObjectCannedAcl.BUCKET_OWNER_FULL_CONTROL)
.setStreamingPartSize(DataSize.of(42, MEGABYTE))
.setRetryMode(STANDARD)
Expand Down

0 comments on commit 218a639

Please sign in to comment.