Skip to content

Commit

Permalink
refactor(S3ClientProvider): Reduce complexity of generateClient (#248)
Browse files Browse the repository at this point in the history
* refactor(S3ClientProvider): Extract getClientForBucket method, call from generateSyncClient and generateAsyncClient

* refactor(S3ClientProvider): Extract getBucketRegionFromResponse

* refactor(S3ClientProvider): Move getClientForRegion call outside of the try

* refactor(S3ClientProvider): extract determineBucketLocation method

* logs(S3ClientProvider): update

* refactor(S3ClientProvider): extract methods to check status codes of S3Exceptions

* refactor(S3ClientProvider): extract getBucketLocation and getBucketLocationFromHead methods

* logger(S3ClientProvider): make it static private and final; use appropriate name

* cleanup(S3ClientProvider): Remove constructor used only in tests

* cleanup(S3ClientProvider): Reduce visibility of methods
  • Loading branch information
guicamest authored Nov 1, 2023
1 parent 9766afe commit f552fd6
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 102 deletions.
163 changes: 69 additions & 94 deletions src/main/java/software/amazon/nio/spi/s3/S3ClientProvider.java
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
import java.time.Duration;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.function.Function;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.auth.credentials.AwsCredentials;
Expand All @@ -25,6 +27,7 @@
import software.amazon.awssdk.core.retry.conditions.RetryOnStatusCodeCondition;
import software.amazon.awssdk.core.retry.conditions.RetryOnThrottlingCondition;
import software.amazon.awssdk.http.HttpStatusCode;
import software.amazon.awssdk.http.SdkHttpResponse;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3AsyncClient;
import software.amazon.awssdk.services.s3.S3Client;
Expand Down Expand Up @@ -96,16 +99,12 @@ public class S3ClientProvider {
);
}

Logger logger = LoggerFactory.getLogger("S3ClientStoreProvider");
static private final Logger logger = LoggerFactory.getLogger(S3ClientProvider.class);

public S3ClientProvider(S3NioSpiConfiguration c) {
this.configuration = (c == null) ? new S3NioSpiConfiguration() : c;
}

public S3ClientProvider() {
this(null);
}

public S3CrtAsyncClientBuilder asyncClientBuilder() {
return asyncClientBuilder;
}
Expand All @@ -122,7 +121,7 @@ public void asyncClientBuilder(final S3CrtAsyncClientBuilder builder) {
*
* @return a S3Client not bound to a region
*/
public S3Client universalClient() {
S3Client universalClient() {
return universalClient(false);
}

Expand All @@ -134,7 +133,7 @@ public S3Client universalClient() {
* @param <T> type of AwsClient
* @return a S3Client not bound to a region
*/
public <T extends AwsClient> T universalClient(boolean async) {
<T extends AwsClient> T universalClient(boolean async) {
return (T)((async) ? DEFAULT_ASYNC_CLIENT : DEFAULT_CLIENT);
}

Expand All @@ -161,54 +160,8 @@ protected S3AsyncClient generateAsyncClient(String bucket) {
*
* @return an S3 client appropriate for the region of the named bucket
*/
protected S3Client generateClient (String bucketName, S3Client locationClient) {
logger.debug("generating client for bucket: '{}'", bucketName);
S3Client bucketSpecificClient = null;

if ((configuration.getEndpoint() == null) || configuration.getEndpoint().isBlank()) {
//
// we try to locate a bucket only if no endpoint is provided, which
// means we are dealing with AWS S3 buckets
//
String bucketLocation = null;
try {
logger.debug("determining bucket location with getBucketLocation");
bucketLocation = locationClient.getBucketLocation(builder -> builder.bucket(bucketName)).locationConstraintAsString();
bucketSpecificClient = this.clientForRegion(bucketLocation);
} catch (S3Exception e) {
if(e.statusCode() == 403) {
logger.debug("Cannot determine location of '{}' bucket directly. Attempting to obtain bucket location with headBucket operation", bucketName);
try {
final HeadBucketResponse headBucketResponse = locationClient.headBucket(builder -> builder.bucket(bucketName));
bucketSpecificClient = this.clientForRegion(headBucketResponse.
sdkHttpResponse()
.firstMatchingHeader("x-amz-bucket-region")
.orElseThrow(() -> new NoSuchElementException("Head Bucket Response doesn't include the header 'x-amz-bucket-region'")));
} catch (S3Exception e2) {
if (e2.statusCode() == 301) {
bucketSpecificClient = this.clientForRegion(e2.awsErrorDetails().
sdkHttpResponse()
.firstMatchingHeader("x-amz-bucket-region")
.orElseThrow(() -> new NoSuchElementException("Head Bucket Response doesn't include the header 'x-amz-bucket-region'")));
} else {
throw e2;
}
}
} else {
throw e;
}
}

//
// if here, no S3 nor other client has been created yet and we do not
// have a location; we'll let it figure out from the profile region
//
logger.warn("Unable to determine the region of bucket: '{}'. Generating a client for the profile region.", bucketName);
}

return (bucketSpecificClient != null)
? bucketSpecificClient
: clientForRegion(configuration.getRegion());
S3Client generateSyncClient(String bucketName, S3Client locationClient) {
return getClientForBucket(bucketName, locationClient, this::clientForRegion);
}

/**
Expand All @@ -221,55 +174,77 @@ protected S3Client generateClient (String bucketName, S3Client locationClient) {
*
* @return an S3 client appropriate for the region of the named bucket
*/
protected S3AsyncClient generateAsyncClient (String bucketName, S3Client locationClient) {
logger.debug("generating asynchronous client for bucket: '{}'", bucketName);
S3AsyncClient bucketSpecificClient = null;
S3AsyncClient generateAsyncClient (String bucketName, S3Client locationClient) {
return getClientForBucket(bucketName, locationClient, this::asyncClientForRegion);
}

private <T extends AwsClient> T getClientForBucket(
String bucketName,
S3Client locationClient,
Function<String, T> getClientForRegion
) {
logger.debug("generating client for bucket: '{}'", bucketName);
T bucketSpecificClient = null;

if ((configuration.getEndpoint() == null) || configuration.getEndpoint().isBlank()) {
//
// we try to locate a bucket only if no endpoint is provided, which
// means we are dealing with AWS S3 buckets
//
String bucketLocation = null;
try {
logger.debug("determining bucket location with getBucketLocation");
bucketLocation = locationClient.getBucketLocation(builder -> builder.bucket(bucketName)).locationConstraintAsString();
bucketSpecificClient = this.asyncClientForRegion(bucketLocation);
} catch (S3Exception e) {
if(e.statusCode() == 403) {
logger.debug("Cannot determine location of '{}' bucket directly. Attempting to obtain bucket location with headBucket operation", bucketName);
try {
final HeadBucketResponse headBucketResponse = locationClient.headBucket(builder -> builder.bucket(bucketName));
bucketSpecificClient = this.asyncClientForRegion(headBucketResponse.sdkHttpResponse()
.firstMatchingHeader("x-amz-bucket-region")
.orElseThrow(() -> new NoSuchElementException("Head Bucket Response doesn't include the header 'x-amz-bucket-region'")));
} catch (S3Exception e2) {
if (e2.statusCode() == 301) {
bucketSpecificClient = this.asyncClientForRegion(e2
.awsErrorDetails()
.sdkHttpResponse()
.firstMatchingHeader("x-amz-bucket-region")
.orElseThrow(() -> new NoSuchElementException("Head Bucket Response doesn't include the header 'x-amz-bucket-region'"))
);
} else {
throw e2;
}
}
} else {
throw e;
}
String bucketLocation = determineBucketLocation(bucketName, locationClient);

if ( bucketLocation != null) {
bucketSpecificClient = getClientForRegion.apply(bucketLocation);
} else {
// if here, no S3 nor other client has been created yet, and we do not
// have a location; we'll let it figure out from the profile region
logger.warn("Unable to determine the region of bucket: '{}'. Generating a client for the profile region.", bucketName);
}

//
// if here, no S3 nor other client has been created yet and we do not
// have a location; we'll let it figure out from the profile region
//
logger.warn("Unable to determine the region of bucket: '{}'. Generating a client for the profile region.", bucketName);
}

return (bucketSpecificClient != null)
? bucketSpecificClient
: asyncClientForRegion(configuration.getRegion());
? bucketSpecificClient
: getClientForRegion.apply(configuration.getRegion());
}

private String determineBucketLocation(String bucketName, S3Client locationClient) {
try {
return getBucketLocation(bucketName, locationClient);
} catch (S3Exception e) {
if(isForbidden(e)) {
return getBucketLocationFromHead(bucketName, locationClient);
} else {
throw e;
}
}
}

private String getBucketLocation(String bucketName, S3Client locationClient) {
logger.debug("determining bucket location with getBucketLocation");
return locationClient.getBucketLocation(builder -> builder.bucket(bucketName)).locationConstraintAsString();
}

private String getBucketLocationFromHead(String bucketName, S3Client locationClient) {
try {
logger.debug("Cannot determine location of '{}' bucket directly. Attempting to obtain bucket location with headBucket operation", bucketName);
final HeadBucketResponse headBucketResponse = locationClient.headBucket(builder -> builder.bucket(bucketName));
return getBucketRegionFromResponse(headBucketResponse.sdkHttpResponse());
} catch (S3Exception e) {
if (isRedirect(e)) {
return getBucketRegionFromResponse(e.awsErrorDetails().sdkHttpResponse());
} else {
throw e;
}
}
}

private boolean isForbidden(S3Exception e) { return e.statusCode() == 403; }
private boolean isRedirect(S3Exception e) { return e.statusCode() == 301; }

private String getBucketRegionFromResponse(SdkHttpResponse response) {
return response.firstMatchingHeader("x-amz-bucket-region").orElseThrow(() ->
new NoSuchElementException("Head Bucket Response doesn't include the header 'x-amz-bucket-region'")
);
}

private S3Client clientForRegion(String regionName) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,12 @@ public class FixedS3ClientProvider extends S3ClientProvider {
final public AwsClient client;

public FixedS3ClientProvider(S3AsyncClient client) {
super(null);
this.client = client;
}

@Override
public S3Client universalClient() {
S3Client universalClient() {
return (S3Client)client;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,13 @@ public class S3ClientProviderTest {
S3ClientProvider provider;

@BeforeEach
public void before() throws Exception {
provider = new S3ClientProvider();
public void before() {
provider = new S3ClientProvider(null);
}

@Test
public void initialization() {
final S3ClientProvider P = new S3ClientProvider();
final S3ClientProvider P = new S3ClientProvider(null);

assertNotNull(P.configuration);

Expand Down Expand Up @@ -80,7 +80,7 @@ public void testGenerateClientWith403Response() {
.build());

// which should get you a client
final S3Client s3Client = provider.generateClient("test-bucket", mockClient);
final S3Client s3Client = provider.generateSyncClient("test-bucket", mockClient);
assertNotNull(s3Client);

final InOrder inOrder = inOrder(mockClient);
Expand Down Expand Up @@ -159,7 +159,7 @@ public void testGenerateClientWith403Then301ResponsesNoHeader(){
);

// then you should get a NoSuchElement exception when you try to get the header
assertThrows(NoSuchElementException.class, () -> provider.generateClient("test-bucket", mockClient));
assertThrows(NoSuchElementException.class, () -> provider.generateSyncClient("test-bucket", mockClient));

final InOrder inOrder = inOrder(mockClient);
inOrder.verify(mockClient).getBucketLocation(anyConsumer());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,8 @@ public void isReadOnly() {

@Test
public void getAndSetClientProvider() {
final S3ClientProvider P1 = new S3ClientProvider();
final S3ClientProvider P2 = new S3ClientProvider();
final S3ClientProvider P1 = new S3ClientProvider(null);
final S3ClientProvider P2 = new S3ClientProvider(null);
s3FileSystem.clientProvider(P1); then(s3FileSystem.clientProvider()).isSameAs(P1);
s3FileSystem.clientProvider(P2); then(s3FileSystem.clientProvider()).isSameAs(P2);
}
Expand Down

0 comments on commit f552fd6

Please sign in to comment.