@Test
public void testGZUncompressRetries()
throws ServiceException, IOException, SegmentLoadingException {
final String bucket = "bucket";
final String keyPrefix = "prefix/dir/0";
final RestS3Service s3Client = EasyMock.createStrictMock(RestS3Service.class);
final byte[] value = bucket.getBytes("utf8");
final File tmpFile = temporaryFolder.newFile("gzTest.gz");
try (OutputStream outputStream = new GZIPOutputStream(new FileOutputStream(tmpFile))) {
outputStream.write(value);
}
S3Object object0 = new S3Object();
object0.setBucketName(bucket);
object0.setKey(keyPrefix + "/renames-0.gz");
object0.setLastModifiedDate(new Date(0));
object0.setDataInputStream(new FileInputStream(tmpFile));
File tmpDir = temporaryFolder.newFolder("gzTestDir");
S3ServiceException exception = new S3ServiceException();
exception.setErrorCode("NoSuchKey");
exception.setResponseCode(404);
EasyMock.expect(
s3Client.getObjectDetails(
EasyMock.eq(object0.getBucketName()), EasyMock.eq(object0.getKey())))
.andReturn(null)
.once();
EasyMock.expect(
s3Client.getObjectDetails(
EasyMock.eq(object0.getBucketName()), EasyMock.eq(object0.getKey())))
.andReturn(object0)
.once();
EasyMock.expect(s3Client.getObject(EasyMock.eq(bucket), EasyMock.eq(object0.getKey())))
.andThrow(exception)
.once();
EasyMock.expect(
s3Client.getObjectDetails(
EasyMock.eq(object0.getBucketName()), EasyMock.eq(object0.getKey())))
.andReturn(object0)
.once();
EasyMock.expect(s3Client.getObject(EasyMock.eq(bucket), EasyMock.eq(object0.getKey())))
.andReturn(object0)
.once();
S3DataSegmentPuller puller = new S3DataSegmentPuller(s3Client);
EasyMock.replay(s3Client);
FileUtils.FileCopyResult result =
puller.getSegmentFiles(new S3DataSegmentPuller.S3Coords(bucket, object0.getKey()), tmpDir);
EasyMock.verify(s3Client);
Assert.assertEquals(value.length, result.size());
File expected = new File(tmpDir, "renames-0");
Assert.assertTrue(expected.exists());
Assert.assertEquals(value.length, expected.length());
}
@Override
public void getSegmentFiles(DataSegment segment, File outDir) throws SegmentLoadingException {
S3Coords s3Coords = new S3Coords(segment);
log.info("Pulling index at path[%s] to outDir[%s]", s3Coords, outDir);
if (!isObjectInBucket(s3Coords)) {
throw new SegmentLoadingException("IndexFile[%s] does not exist.", s3Coords);
}
if (!outDir.exists()) {
outDir.mkdirs();
}
if (!outDir.isDirectory()) {
throw new ISE("outDir[%s] must be a directory.", outDir);
}
long startTime = System.currentTimeMillis();
S3Object s3Obj = null;
try {
s3Obj = s3Client.getObject(new S3Bucket(s3Coords.bucket), s3Coords.path);
InputStream in = null;
try {
in = s3Obj.getDataInputStream();
final String key = s3Obj.getKey();
if (key.endsWith(".zip")) {
CompressionUtils.unzip(in, outDir);
} else if (key.endsWith(".gz")) {
final File outFile = new File(outDir, toFilename(key, ".gz"));
ByteStreams.copy(new GZIPInputStream(in), Files.newOutputStreamSupplier(outFile));
} else {
ByteStreams.copy(
in, Files.newOutputStreamSupplier(new File(outDir, toFilename(key, ""))));
}
log.info(
"Pull of file[%s] completed in %,d millis",
s3Obj, System.currentTimeMillis() - startTime);
} catch (IOException e) {
FileUtils.deleteDirectory(outDir);
throw new SegmentLoadingException(e, "Problem decompressing object[%s]", s3Obj);
} finally {
Closeables.closeQuietly(in);
}
} catch (Exception e) {
throw new SegmentLoadingException(e, e.getMessage());
} finally {
S3Utils.closeStreamsQuietly(s3Obj);
}
}
protected URLConnection openConnection(URL u)
throws IOException {
// This looking for accessKey id and accessKey secret code is based
// on code from hadoop S3.
String accessKey = null;
String secretAccessKey = null;
String userInfo = u.getUserInfo();
if (userInfo != null) {
int index = userInfo.indexOf(':');
if (index != -1) {
accessKey = userInfo.substring(0, index);
secretAccessKey = userInfo.substring(index + 1);
} else {
accessKey = userInfo;
}
}
if (accessKey == null) {
accessKey = System.getProperty("aws.access.key.id");
}
if (secretAccessKey == null) {
secretAccessKey = System.getProperty("aws.access.key.secret");
}
if (accessKey == null && secretAccessKey == null) {
throw new IllegalArgumentException("AWS " +
"Access Key ID and Secret Access Key " +
"must be specified as the username " +
"or password (respectively) of a s3 URL, " +
"or by setting the " +
"aws.access.key.id or " +
"aws.access.key.secret properties (respectively).");
} else if (accessKey == null) {
throw new IllegalArgumentException("AWS " +
"Access Key ID must be specified " +
"as the username of a s3 URL, or by setting the " +
"aws.access.key.id property.");
} else if (secretAccessKey == null) {
throw new IllegalArgumentException("AWS " +
"Secret Access Key must be specified " +
"as the password of a s3 URL, or by setting the " +
"aws.access.key.secret property.");
}
RestS3Service s3Service;
try {
s3Service = new RestS3Service(
new AWSCredentials(accessKey, secretAccessKey));
} catch (S3ServiceException e) {
e.printStackTrace();
throw new IOException(e.toString());
}
InputStream is = null;
try {
// This opens the stream to the bucket/key object.
S3Object s3obj = s3Service.getObject(new S3Bucket(u.getHost()),
u.getPath().substring(1) /* Skip starting '/' character */);
is = s3obj.getDataInputStream();
} catch (S3ServiceException e) {
e.printStackTrace();
throw new IOException(e.toString());
}
final InputStream inputStream = is;
return new URLConnection(u) {
private InputStream is = inputStream;
@Override
public InputStream getInputStream() throws IOException {
return this.is;
}
@Override
public void connect() throws IOException {
// Nothing to do. When we give back this object, we're
// connected.
}
};
}