com.netflix.genie.common.internal.aws.s3.S3ProtocolResolver Maven / Gradle / Ivy
/*
*
* Copyright 2018 Netflix, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package com.netflix.genie.common.internal.aws.s3;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3URI;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.springframework.core.io.ProtocolResolver;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.task.TaskExecutor;
import javax.annotation.Nullable;
import java.io.IOException;
import java.net.URI;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
/**
* This class implements the {@link ProtocolResolver} interface. When an instance of this class is added to a
* Spring Application context list of Protocol Resolvers via
* {@link org.springframework.context.ConfigurableApplicationContext#addProtocolResolver(ProtocolResolver)} allows
* valid S3 resources to be loaded using the Spring {@link ResourceLoader} abstraction.
*
* Leverages some work done by Spring Cloud AWS.
*
* @author tgianos
* @since 4.0.0
*/
@Slf4j
public class S3ProtocolResolver implements ProtocolResolver {
private static final String S3N_PROTOCOL = "s3n:";
private static final String S3A_PROTOCOL = "s3a:";
private static final String S3_REGEX = "s3.:";
private static final String S3_REPLACEMENT = "s3:";
private static final Pair NULL_RANGE = ImmutablePair.of(null, null);
private static final Pattern RANGE_HEADER_PATTERN = Pattern.compile("bytes=(\\d*)-(\\d*)");
private final S3ClientFactory s3ClientFactory;
private final TaskExecutor s3TaskExecutor;
/**
* Constructor.
*
* @param s3ClientFactory The S3 client factory to use to get S3 client instances
* @param s3TaskExecutor A task executor to use for uploading files to S3
*/
public S3ProtocolResolver(
final S3ClientFactory s3ClientFactory,
final TaskExecutor s3TaskExecutor
) {
this.s3ClientFactory = s3ClientFactory;
this.s3TaskExecutor = s3TaskExecutor;
}
/**
* TODO: It would be nice to use Spring's HttpRange for this parsing, but this module does not
* currently depend on spring-web. And this class cannot be moved to genie-web since it is used by
* {@link S3ProtocolResolver} which is shared with the genie-agent module.
*/
static Pair parseRangeHeader(@Nullable final String rangeHeader) {
if (StringUtils.isBlank(rangeHeader)) {
return NULL_RANGE;
}
final Matcher matcher = RANGE_HEADER_PATTERN.matcher(rangeHeader);
if (!matcher.matches()) {
return NULL_RANGE;
}
final String rangeStartString = matcher.group(1);
final String rangeEndString = matcher.group(2);
Integer rangeStart = null;
Integer rangeEnd = null;
if (!StringUtils.isBlank(rangeStartString)) {
rangeStart = Integer.parseInt(rangeStartString);
}
if (!StringUtils.isBlank(rangeEndString)) {
rangeEnd = Integer.parseInt(rangeEndString);
}
return ImmutablePair.of(rangeStart, rangeEnd);
}
/**
* {@inheritDoc}
*/
@Override
public Resource resolve(final String location, final ResourceLoader resourceLoader) {
log.debug("Attempting to resolve if {} is a S3 resource or not", location);
final String normalizedLocation;
// Rewrite s3n:// and s3a:// URIs as s3:// for backward compatibility
if (location.startsWith(S3N_PROTOCOL) || location.startsWith(S3A_PROTOCOL)) {
normalizedLocation = location.replaceFirst(S3_REGEX, S3_REPLACEMENT);
} else {
normalizedLocation = location;
}
final AmazonS3URI s3URI;
final URI uri;
try {
s3URI = new AmazonS3URI(normalizedLocation);
uri = URI.create(location);
} catch (final IllegalArgumentException iae) {
log.debug("{} is not a valid S3 resource (Error message: {}).", normalizedLocation, iae.getMessage());
return null;
}
// Remove the fragment portion of the URI path (which stores the range requested, if any)
final int fragmentIndex = s3URI.getKey().lastIndexOf("#");
final String normalizedKey;
if (fragmentIndex == -1) {
normalizedKey = s3URI.getKey();
} else {
normalizedKey = s3URI.getKey().substring(0, fragmentIndex);
}
final String rangeHeader = uri.getFragment();
final Pair range = parseRangeHeader(rangeHeader);
final AmazonS3 client = this.s3ClientFactory.getClient(s3URI);
log.debug("{} is a valid S3 resource.", location);
// TODO: This implementation from Spring Cloud AWS always wraps the passed in client with a proxy that follows
// redirects. I'm not sure if we want that or not. Probably ok for now but maybe revisit later?
try {
return new SimpleStorageRangeResource(
client,
s3URI.getBucket(),
normalizedKey,
s3URI.getVersionId(),
this.s3TaskExecutor,
range
);
} catch (IOException e) {
log.error("Failed to create S3 resource: " + location + ": " + e.getMessage());
return null;
}
}
}