Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
seratch committed Jan 14, 2025
1 parent 40feb5c commit a890288
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 22 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,12 @@
import software.amazon.awssdk.core.ResponseBytes;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.core.sync.ResponseTransformer;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.*;

import java.io.IOException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Optional;
Expand All @@ -26,23 +28,38 @@
public class AmazonS3InstallationService implements InstallationService {

private final String bucketName;
private AwsCredentialsProvider credentialsProvider;

private final AwsCredentialsProvider credentialsProvider;
private final Region region;
private final URI endpointOverride;

private boolean historicalDataEnabled;

public AmazonS3InstallationService(String bucketName) {
this(bucketName, null);
this(bucketName, DefaultCredentialsProvider.create(), null, null);
}

public AmazonS3InstallationService(String bucketName, AwsCredentialsProvider credentialsProvider) {
this(bucketName, credentialsProvider, null, null);
}

public AmazonS3InstallationService(
String bucketName,
AwsCredentialsProvider credentialsProvider,
Region region,
String endpointOverride
) {
this.bucketName = bucketName;
this.credentialsProvider = credentialsProvider;
this.region = region != null ? region : Region.of(System.getenv("AWS_REGION"));
this.endpointOverride = (endpointOverride != null && !endpointOverride.isEmpty()) ? URI.create(endpointOverride) : null;
}

@Override
public Initializer initializer() {
return (app) -> {
// The first access to S3 tends to be slow on AWS Lambda.
AwsCredentials credentials = createCredentials(getAwsCredentialsProvider());
AwsCredentials credentials = createCredentials(this.credentialsProvider);
if (credentials == null || credentials.accessKeyId() == null) {
throw new IllegalStateException("AWS credentials not found");
}
Expand Down Expand Up @@ -303,19 +320,16 @@ private Installer toInstaller(ResponseBytes<GetObjectResponse> s3Object) throws
return JsonOps.fromJson(json, DefaultInstaller.class);
}

protected AwsCredentialsProvider getAwsCredentialsProvider() {
if (this.credentialsProvider == null) {
this.credentialsProvider = DefaultCredentialsProvider.create();
}
return this.credentialsProvider;
}

protected AwsCredentials createCredentials(AwsCredentialsProvider provider) {
return provider.resolveCredentials();
}

protected S3Client createS3Client() {
return S3Client.builder().credentialsProvider(this.credentialsProvider).build();
return S3Client.builder()
.credentialsProvider(this.credentialsProvider)
.region(this.region)
.endpointOverride(this.endpointOverride)
.build();
}

private String getInstallerKey(Installer i) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,14 @@
import lombok.extern.slf4j.Slf4j;
import software.amazon.awssdk.auth.credentials.AwsCredentials;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.core.ResponseBytes;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.core.sync.ResponseTransformer;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.*;

import java.net.URI;
import java.nio.charset.StandardCharsets;

/**
Expand All @@ -23,22 +24,36 @@
public class AmazonS3OAuthStateService implements OAuthStateService {

private final String bucketName;
private AwsCredentialsProvider credentialsProvider;

private final AwsCredentialsProvider credentialsProvider;
private final Region region;
private final URI endpointOverride;

public AmazonS3OAuthStateService(String bucketName) {
this(bucketName, null);
}

public AmazonS3OAuthStateService(String bucketName, AwsCredentialsProvider credentialsProvider) {
this(bucketName, credentialsProvider, null, null);
}

public AmazonS3OAuthStateService(
String bucketName,
AwsCredentialsProvider credentialsProvider,
Region region,
String endpointOverride
) {
this.bucketName = bucketName;
this.credentialsProvider = credentialsProvider;
this.region = region != null ? region : Region.of(System.getenv("AWS_REGION"));
this.endpointOverride = (endpointOverride != null && !endpointOverride.isEmpty()) ? URI.create(endpointOverride) : null;
}

@Override
public Initializer initializer() {
return (app) -> {
// The first access to S3 tends to be slow on AWS Lambda.
AwsCredentials credentials = createCredentials(getAwsCredentialsProvider());
AwsCredentials credentials = createCredentials(this.credentialsProvider);
if (credentials == null || credentials.accessKeyId() == null) {
throw new IllegalStateException("AWS credentials not found");
}
Expand Down Expand Up @@ -102,19 +117,16 @@ public void deleteStateFromDatastore(String state) throws Exception {
}
}

protected AwsCredentialsProvider getAwsCredentialsProvider() {
if (this.credentialsProvider == null) {
this.credentialsProvider = DefaultCredentialsProvider.create();
}
return this.credentialsProvider;
}

protected AwsCredentials createCredentials(AwsCredentialsProvider provider) {
return provider.resolveCredentials();
}

protected S3Client createS3Client() {
return S3Client.builder().credentialsProvider(this.credentialsProvider).build();
return S3Client.builder()
.credentialsProvider(this.credentialsProvider)
.region(this.region)
.endpointOverride(this.endpointOverride)
.build();
}

private String getKey(String state) {
Expand Down

0 comments on commit a890288

Please sign in to comment.