Extract STS client getter

trinodb · Jun 11, 2024 · 48e6ae4 · 48e6ae4
1 parent 94daecd
commit 48e6ae4
Showing 1 changed file with 17 additions and 11 deletions.
diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java
@@ -66,30 +66,25 @@ public S3FileSystemFactory(OpenTelemetry openTelemetry, S3FileSystemConfig confi
                 .retryPolicy(retryPolicy)
                 .build());
 
-        Optional<StaticCredentialsProvider> staticCredentialsProvider = getStaticCredentialsProvider(config);
-        staticCredentialsProvider.ifPresent(s3::credentialsProvider);
-
         Optional.ofNullable(config.getRegion()).map(Region::of).ifPresent(s3::region);
         Optional.ofNullable(config.getEndpoint()).map(URI::create).ifPresent(s3::endpointOverride);
         s3.forcePathStyle(config.isPathStyleAccess());
 
-        if (config.getIamRole() != null) {
-            StsClientBuilder sts = StsClient.builder();
-            Optional.ofNullable(config.getStsEndpoint()).map(URI::create).ifPresent(sts::endpointOverride);
-            Optional.ofNullable(config.getStsRegion())
-                    .or(() -> Optional.ofNullable(config.getRegion()))
-                    .map(Region::of).ifPresent(sts::region);
-            staticCredentialsProvider.ifPresent(sts::credentialsProvider);
+        Optional<StaticCredentialsProvider> staticCredentialsProvider = getStaticCredentialsProvider(config);
 
+        if (config.getIamRole() != null) {
             s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder()
                     .refreshRequest(request -> request
                             .roleArn(config.getIamRole())
                             .roleSessionName(config.getRoleSessionName())
                             .externalId(config.getExternalId()))
-                    .stsClient(sts.build())
+                    .stsClient(getStsClient(config, staticCredentialsProvider))
                     .asyncCredentialUpdateEnabled(true)
                     .build());
         }
+        else {
+            staticCredentialsProvider.ifPresent(s3::credentialsProvider);
+        }
 
         ApacheHttpClient.Builder httpClient = ApacheHttpClient.builder()
                 .maxConnections(config.getMaxConnections())
@@ -150,4 +145,15 @@ private static Optional<StaticCredentialsProvider> getStaticCredentialsProvider(
         }
         return Optional.empty();
     }
+
+    private static StsClient getStsClient(S3FileSystemConfig config, Optional<StaticCredentialsProvider> staticCredentialsProvider)
+    {
+        StsClientBuilder sts = StsClient.builder();
+        Optional.ofNullable(config.getStsEndpoint()).map(URI::create).ifPresent(sts::endpointOverride);
+        Optional.ofNullable(config.getStsRegion())
+                .or(() -> Optional.ofNullable(config.getRegion()))
+                .map(Region::of).ifPresent(sts::region);
+        staticCredentialsProvider.ifPresent(sts::credentialsProvider);
+        return sts.build();
+    }
 }