embulk-output-s3v2 0.1.0

data/src/main/java/org/embulk/output/s3v2/s3/S3MultiPartStatus.java

@@ -0,0 +1,30 @@
+package org.embulk.output.s3v2.s3;
+
+public class S3MultiPartStatus
+{
+    private final int maxConcurrentRequests;
+    private final String multipartChunksize;
+    private final String multipartThreshold;
+
+    public S3MultiPartStatus(int maxConcurrentRequests, String multipartChunksize, String multipartThreshold)
+    {
+        this.maxConcurrentRequests = maxConcurrentRequests;
+        this.multipartChunksize = multipartChunksize;
+        this.multipartThreshold = multipartThreshold;
+    }
+
+    public int getMaxConcurrentRequests()
+    {
+        return maxConcurrentRequests;
+    }
+
+    public String getMultipartChunksize()
+    {
+        return multipartChunksize;
+    }
+
+    public String getMultipartThreshold()
+    {
+        return multipartThreshold;
+    }
+}

data/src/main/java/org/embulk/output/s3v2/strategy/AbstractStrategy.java

@@ -0,0 +1,41 @@
+package org.embulk.output.s3v2.strategy;
+
+import org.embulk.output.s3v2.PluginTask;
+import org.embulk.output.s3v2.s3.S3ClientManager;
+import org.embulk.output.s3v2.s3.S3MultiPartStatus;
+import org.embulk.spi.TransactionalFileOutput;
+
+abstract class AbstractStrategy implements TransactionalFileOutput
+{
+    protected S3ClientManager s3;
+    protected PluginTask task;
+    protected int taskIndex;
+
+    public AbstractStrategy(PluginTask task, int taskIndex)
+    {
+        s3 = new S3ClientManager(task.getRegion());
+        this.task = task;
+        this.taskIndex = taskIndex;
+
+        if (!validate()) {
+            throw new IllegalArgumentException("Unsupported parameters combination");
+        }
+    }
+
+    /**
+     * Validation for PluginTask parameters combination.
+     * @see PluginTask
+     */
+    protected abstract boolean validate();
+
+    protected final String getFileExtension()
+    {
+        return task.getExtension().startsWith(".") ? task.getExtension() : "." + task.getExtension();
+    }
+
+    protected final S3MultiPartStatus setUpS3MultiPartStatus()
+    {
+        return new S3MultiPartStatus(task.getMaxConcurrentRequests(), task.getMultipartChunksize(),
+                task.getMultipartThreshold());
+    }
+}

data/src/main/java/org/embulk/output/s3v2/strategy/BufferedStrategy.java

@@ -0,0 +1,76 @@
+package org.embulk.output.s3v2.strategy;
+
+import org.embulk.config.TaskReport;
+import org.embulk.output.s3v2.PluginTask;
+import org.embulk.spi.Buffer;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.nio.ByteBuffer;
+
+public class BufferedStrategy extends AbstractStrategy
+{
+    private final Logger logger = LoggerFactory.getLogger(BufferedStrategy.class);
+    private final String bucket;
+    private String s3ObjectKey;
+    private ByteBuffer byteBuffer;
+
+    public BufferedStrategy(PluginTask task, int taskIndex)
+    {
+        super(task, taskIndex);
+        this.bucket = task.getBucket();
+    }
+
+    @Override
+    protected boolean validate()
+    {
+        if (task.getEnableMultiPartUpload()) {
+            throw new UnsupportedOperationException("Buffering strategy does not support S3 multi-part upload.");
+        }
+        return true;
+    }
+
+    @Override
+    public void nextFile()
+    {
+        String ext = getFileExtension();
+        s3ObjectKey = task.getObjectKeyPrefix() + "-" + taskIndex + ext;
+        logger.info("[task:" + taskIndex + "] Temporary file is not created.");
+    }
+
+    @Override
+    public void add(Buffer buffer)
+    {
+        byteBuffer = ByteBuffer.wrap(buffer.array(), buffer.offset(), buffer.limit());
+    }
+
+    @Override
+    public void finish()
+    {
+        // Do nothing.
+    }
+
+    @Override
+    public void close()
+    {
+        // Do nothing.
+    }
+
+    @Override
+    public void abort()
+    {
+        if (s3.existsObject(bucket, s3ObjectKey)) {
+            s3.deleteObject(bucket, s3ObjectKey);
+            logger.info("Delete s3://" + bucket + "/" + s3ObjectKey);
+        }
+    }
+
+    @Override
+    public TaskReport commit()
+    {
+        s3.putObject(bucket, s3ObjectKey, byteBuffer);
+        logger.info("Put into s3://" + bucket + "/" + s3ObjectKey);
+        return Exec.newTaskReport();
+    }
+}

data/src/main/java/org/embulk/output/s3v2/strategy/FileOutputStrategy.java

@@ -0,0 +1,166 @@
+package org.embulk.output.s3v2.strategy;
+
+import org.embulk.config.TaskReport;
+import org.embulk.output.s3v2.PluginTask;
+import org.embulk.output.s3v2.s3.S3MultiPartStatus;
+import org.embulk.output.s3v2.util.ThresholdComputation;
+import org.embulk.spi.Buffer;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.BufferedOutputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Objects;
+
+public class FileOutputStrategy extends AbstractStrategy
+{
+    private final Logger logger = LoggerFactory.getLogger(FileOutputStrategy.class);
+    private final String bucket;
+    private String s3ObjectKey;
+    private BufferedOutputStream bufferStream;
+    private Path tempFilePath;
+
+    public FileOutputStrategy(PluginTask task, int taskIndex)
+    {
+        super(task, taskIndex);
+        this.bucket = task.getBucket();
+    }
+
+    @Override
+    protected boolean validate()
+    {
+        return true;
+    }
+
+    @Override
+    public void nextFile()
+    {
+        String ext = getFileExtension();
+        createTempFile(ext);
+        s3ObjectKey = task.getObjectKeyPrefix() + "-" + taskIndex + ext;
+    }
+
+    private void createTempFile(String ext)
+    {
+        try {
+            Path tempPath = Paths.get(task.getTempPath());
+            if (Files.notExists(tempPath)) {
+                Files.createDirectory(tempPath);
+            }
+
+            String dirSeparator = tempPath.endsWith("/") ? "" : "/";
+            tempFilePath = Paths.get(task.getTempPath() + dirSeparator
+                    + task.getTempFilePrefix() + "-" + taskIndex
+                    + ext);
+            if (Files.exists(tempFilePath)) {
+                Files.delete(tempFilePath);
+                logger.info("Deleted previous: " + tempFilePath);
+            }
+            Files.createFile(tempFilePath);
+            logger.info("Created: " + tempFilePath);
+
+            bufferStream = new BufferedOutputStream(new FileOutputStream(tempFilePath.toFile()));
+        }
+        catch (IOException ex) {
+            closeBuffer();
+            ex.printStackTrace();
+            throw new RuntimeException("Failed to create temp file: " + tempFilePath);
+        }
+    }
+
+    @Override
+    public void add(Buffer buffer)
+    {
+        try {
+            bufferStream.write(buffer.array(), buffer.offset(), buffer.limit());
+        }
+        catch (IOException ex) {
+            closeBuffer();
+            ex.printStackTrace();
+            throw new RuntimeException("Failed to buffer data.");
+        }
+    }
+
+    @Override
+    public void finish()
+    {
+        try {
+            bufferStream.flush();
+        }
+        catch (IOException ex) {
+            closeBuffer();
+            ex.printStackTrace();
+            throw new RuntimeException("Failed to write out data.");
+        }
+    }
+
+    @Override
+    public void close()
+    {
+        closeBuffer();
+    }
+
+    @Override
+    public void abort()
+    {
+        if (s3.existsObject(bucket, s3ObjectKey)) {
+            s3.deleteObject(bucket, s3ObjectKey);
+            logger.info("Deleted s3://" + bucket + "/" + s3ObjectKey);
+        }
+
+        try {
+            Files.delete(tempFilePath);
+            logger.info("Deleted " + tempFilePath);
+        }
+        catch (IOException ex) {
+            ex.printStackTrace();
+            throw new RuntimeException("Failed to delete: " + tempFilePath);
+        }
+        finally {
+            closeBuffer();
+        }
+    }
+
+    @Override
+    public TaskReport commit()
+    {
+        long start = System.currentTimeMillis();
+        try {
+            if (task.getEnableMultiPartUpload()
+                    && ThresholdComputation.largerThanThreshold(task.getMultipartThreshold(), Files.size(tempFilePath))) {
+                S3MultiPartStatus status = setUpS3MultiPartStatus();
+                s3.multiPartUpload(bucket, s3ObjectKey, tempFilePath, status);
+                logger.info("[Done] Multipart upload: s3://" + bucket + "/" + s3ObjectKey);
+            }
+            else {
+                s3.putObject(bucket, s3ObjectKey, tempFilePath);
+                logger.info("[Done] Put: s3://" + bucket + "/" + s3ObjectKey);
+            }
+        }
+        catch (IOException ex) {
+            ex.printStackTrace();
+            abort();
+        }
+        long end = System.currentTimeMillis();
+        logger.info("Time taken to upload s3://" + bucket + "/" + s3ObjectKey + ": " + (end - start) + "ms");
+        return Exec.newTaskReport();
+    }
+
+    private void closeBuffer()
+    {
+        if (!Objects.isNull(bufferStream)) {
+            try {
+                bufferStream.close();
+            }
+            catch (IOException ex) {
+                ex.printStackTrace();
+                throw new RuntimeException("Failed to close BufferedOutputStream.");
+            }
+        }
+    }
+}

data/src/main/java/org/embulk/output/s3v2/util/AbstractUnitComputation.java

@@ -0,0 +1,39 @@
+package org.embulk.output.s3v2.util;
+
+abstract class AbstractUnitComputation
+{
+    protected enum ComputeUnits
+    {
+        KB(Math.pow(10, 3)),
+        MB(Math.pow(10, 6)),
+        GB(Math.pow(10, 9)),
+        TB(Math.pow(10, 12));
+
+        private final double unit;
+
+        private ComputeUnits(double unit)
+        {
+            this.unit = unit;
+        }
+
+        public double getUnit()
+        {
+            return unit;
+        }
+    }
+
+    protected static String getSize(String value)
+    {
+        return value.replaceAll("[^0-9]", "");
+    }
+
+    protected static String getUnit(String value)
+    {
+        return value.replaceAll("[0-9]", "");
+    }
+
+    protected static boolean validateValue(String value)
+    {
+        return value.matches(getSize(value) + getUnit(value));
+    }
+}

data/src/main/java/org/embulk/output/s3v2/util/ChunksizeComputation.java

@@ -0,0 +1,26 @@
+package org.embulk.output.s3v2.util;
+
+public class ChunksizeComputation extends AbstractUnitComputation
+{
+    private ChunksizeComputation()
+    {
+        // Do nothing.
+    }
+
+    public static int getChunksizeBytes(String chunksize)
+    {
+        if (!validateValue(chunksize)) {
+            throw new IllegalArgumentException("Unrecognized value of multipart_chunksize: " + chunksize);
+        }
+
+        String sizePartOfChunksize = getSize(chunksize);
+        String unitPartOfChunksize = getUnit(chunksize);
+        ComputeUnits e = ComputeUnits.valueOf(unitPartOfChunksize.toUpperCase());
+        double chunksizeValue = (Double.valueOf(sizePartOfChunksize) * e.getUnit());
+        if (chunksizeValue < 5.0 * ComputeUnits.MB.getUnit() || 2.0 * ComputeUnits.GB.getUnit() < chunksizeValue) {
+            throw new IllegalArgumentException("Unrecognized range of value multipart_chunksize: " + chunksize);
+        }
+
+        return (int) chunksizeValue;
+    }
+}

data/src/main/java/org/embulk/output/s3v2/util/ThresholdComputation.java

@@ -0,0 +1,23 @@
+package org.embulk.output.s3v2.util;
+
+public class ThresholdComputation extends AbstractUnitComputation
+{
+    private ThresholdComputation()
+    {
+        // Do nothing.
+    }
+
+    public static boolean largerThanThreshold(String threshold, long computationTargetSize)
+    {
+        if (!validateValue(threshold)) {
+            throw new IllegalArgumentException("Unrecognized value of multipart_threshold: " + threshold);
+        }
+
+        String thresholdSize = getSize(threshold);
+        String thresholdUnit = getUnit(threshold);
+        ComputeUnits e = ComputeUnits.valueOf(thresholdUnit.toUpperCase());
+        double thresholdValue = (Double.valueOf(thresholdSize) * e.getUnit());
+
+        return (double) computationTargetSize > thresholdValue;
+    }
+}

data/src/test/java/org/embulk/output/s3v2/strategy/BufferedStrategyTests.java

@@ -0,0 +1,45 @@
+package org.embulk.output.s3v2.strategy;
+
+import org.embulk.output.s3v2.PluginTask;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mockito;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+/**
+ * @see BufferedStrategy
+ */
+@ExtendWith(MockitoExtension.class)
+public class BufferedStrategyTests
+{
+    private PluginTask task;
+
+    @BeforeEach
+    public void setUp()
+    {
+        task = Mockito.mock(PluginTask.class);
+        Mockito.doReturn("ap-northeast-1").when(task).getRegion();
+    }
+
+    @Test
+    @DisplayName("Test validate true")
+    public void testValidateTrue() throws Exception
+    {
+        Mockito.doReturn(false).when(task).getEnableMultiPartUpload();
+        BufferedStrategy output = new BufferedStrategy(task, 0);
+        Assertions.assertTrue(output.validate());
+    }
+
+    @Test
+    @DisplayName("Test validate false")
+    public void testValidateInvalidCase() throws Exception
+    {
+        Mockito.doReturn(true).when(task).getEnableMultiPartUpload();
+        UnsupportedOperationException ex = Assertions.assertThrows(UnsupportedOperationException.class,
+                () -> new BufferedStrategy(task, 0));
+        Assertions.assertEquals("Buffering strategy does not support S3 multi-part upload.", ex.getMessage());
+    }
+}