@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);
    }
  }
Ejemplo n.º 3
0
    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.
            }
        };
    }