generated from kestra-io/plugin-template
-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix(batch-script-runner): post-review fixes
- Loading branch information
1 parent
df2f557
commit 8e92d74
Showing
14 changed files
with
212 additions
and
223 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
135 changes: 135 additions & 0 deletions
135
src/main/java/io/kestra/plugin/aws/ConnectionUtils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,135 @@ | ||
| package io.kestra.plugin.aws; | ||
|
|
||
| import io.kestra.core.exceptions.IllegalVariableEvaluationException; | ||
| import org.apache.commons.lang3.StringUtils; | ||
| import software.amazon.awssdk.auth.credentials.*; | ||
| import software.amazon.awssdk.awscore.AwsClient; | ||
| import software.amazon.awssdk.awscore.client.builder.AwsAsyncClientBuilder; | ||
| import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder; | ||
| import software.amazon.awssdk.awscore.client.builder.AwsSyncClientBuilder; | ||
| import software.amazon.awssdk.http.apache.ApacheHttpClient; | ||
| import software.amazon.awssdk.regions.Region; | ||
| import software.amazon.awssdk.services.sts.StsClient; | ||
| import software.amazon.awssdk.services.sts.StsClientBuilder; | ||
| import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; | ||
| import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; | ||
|
|
||
| import java.net.URI; | ||
|
|
||
| public class ConnectionUtils { | ||
| /** | ||
| * Factory method for constructing a new {@link AwsCredentialsProvider} for the given AWS Client config. | ||
| * | ||
| * @param awsClientConfig The AwsClientConfig. | ||
| * @return a new {@link AwsCredentialsProvider} instance. | ||
| */ | ||
| public static AwsCredentialsProvider credentialsProvider(final AbstractConnection.AwsClientConfig awsClientConfig) { | ||
|
|
||
| // StsAssumeRoleCredentialsProvider | ||
| if (StringUtils.isNotEmpty(awsClientConfig.stsRoleArn())) { | ||
| return stsAssumeRoleCredentialsProvider(awsClientConfig); | ||
| } | ||
|
|
||
| // StaticCredentialsProvider | ||
| if (StringUtils.isNotEmpty(awsClientConfig.accessKeyId()) && | ||
| StringUtils.isNotEmpty(awsClientConfig.secretKeyId())) { | ||
| return staticCredentialsProvider(awsClientConfig); | ||
| } | ||
|
|
||
| // Otherwise, use DefaultCredentialsProvider | ||
| return DefaultCredentialsProvider.builder().build(); | ||
| } | ||
|
|
||
| public static StaticCredentialsProvider staticCredentialsProvider(final AbstractConnection.AwsClientConfig awsClientConfig) { | ||
| final AwsCredentials credentials; | ||
| if (StringUtils.isNotEmpty(awsClientConfig.sessionToken())) { | ||
| credentials = AwsSessionCredentials.create( | ||
| awsClientConfig.accessKeyId(), | ||
| awsClientConfig.secretKeyId(), | ||
| awsClientConfig.sessionToken() | ||
| ); | ||
| } else { | ||
| credentials = AwsBasicCredentials.create( | ||
| awsClientConfig.accessKeyId(), | ||
| awsClientConfig.secretKeyId() | ||
| ); | ||
| } | ||
| return StaticCredentialsProvider.create(credentials); | ||
| } | ||
|
|
||
| public static StsAssumeRoleCredentialsProvider stsAssumeRoleCredentialsProvider(final AbstractConnection.AwsClientConfig awsClientConfig) { | ||
|
|
||
| String roleSessionName = awsClientConfig.stsRoleSessionName(); | ||
| roleSessionName = roleSessionName != null ? roleSessionName : "kestra-plugin-s3-" + System.currentTimeMillis(); | ||
|
|
||
| final AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest.builder() | ||
| .roleArn(awsClientConfig.stsRoleArn()) | ||
| .roleSessionName(roleSessionName) | ||
| .durationSeconds((int) awsClientConfig.stsRoleSessionDuration().toSeconds()) | ||
| .externalId(awsClientConfig.stsRoleExternalId()) | ||
| .build(); | ||
|
|
||
| return StsAssumeRoleCredentialsProvider.builder() | ||
| .stsClient(stsClient(awsClientConfig)) | ||
| .refreshRequest(assumeRoleRequest) | ||
| .build(); | ||
| } | ||
|
|
||
| public static StsClient stsClient(final AbstractConnection.AwsClientConfig awsClientConfig) { | ||
| StsClientBuilder builder = StsClient.builder(); | ||
|
|
||
| final String stsEndpointOverride = awsClientConfig.stsEndpointOverride(); | ||
| if (stsEndpointOverride != null) { | ||
| builder.applyMutation(stsClientBuilder -> | ||
| stsClientBuilder.endpointOverride(URI.create(stsEndpointOverride))); | ||
| } | ||
|
|
||
| final String regionString = awsClientConfig.region(); | ||
| if (regionString != null) { | ||
| builder.applyMutation(stsClientBuilder -> | ||
| stsClientBuilder.region(Region.of(regionString))); | ||
| } | ||
| return builder.build(); | ||
| } | ||
|
|
||
| /** | ||
| * Configures and returns the given {@link AwsSyncClientBuilder}. | ||
| */ | ||
| public static <C extends AwsClient, B extends AwsClientBuilder<B, C> & AwsSyncClientBuilder<B, C>> B configureSyncClient( | ||
| final AbstractConnection.AwsClientConfig clientConfig, final B builder) throws IllegalVariableEvaluationException { | ||
|
|
||
| builder | ||
| // Use the httpClientBuilder to delegate the lifecycle management of the HTTP client to the AWS SDK | ||
| .httpClientBuilder(serviceDefaults -> ApacheHttpClient.builder().build()) | ||
| .credentialsProvider(ConnectionUtils.credentialsProvider(clientConfig)); | ||
|
|
||
| return configureClient(clientConfig, builder); | ||
| } | ||
|
|
||
| /** | ||
| * Configures and returns the given {@link AwsAsyncClientBuilder}. | ||
| */ | ||
| public static <C extends AwsClient, B extends AwsClientBuilder<B, C> & AwsAsyncClientBuilder<B, C>> B configureAsyncClient( | ||
| final AbstractConnection.AwsClientConfig clientConfig, final B builder) { | ||
|
|
||
| builder.credentialsProvider(ConnectionUtils.credentialsProvider(clientConfig)); | ||
| return configureClient(clientConfig, builder); | ||
| } | ||
|
|
||
| /** | ||
| * Configures and returns the given {@link AwsClientBuilder}. | ||
| */ | ||
| public static <C extends AwsClient, B extends AwsClientBuilder<B, C>> B configureClient( | ||
| final AbstractConnection.AwsClientConfig clientConfig, final B builder) { | ||
|
|
||
| builder.credentialsProvider(ConnectionUtils.credentialsProvider(clientConfig)); | ||
|
|
||
| if (clientConfig.region() != null) { | ||
| builder.region(Region.of(clientConfig.region())); | ||
| } | ||
| if (clientConfig.endpointOverride() != null) { | ||
| builder.endpointOverride(URI.create(clientConfig.endpointOverride())); | ||
| } | ||
| return builder; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.