embulk-output-dynamodb 0.1.0

data/src/main/java/org/embulk/output/dynamodb/DynamodbUtils.java

@@ -0,0 +1,361 @@
+package org.embulk.output.dynamodb;
+
+import com.amazonaws.AmazonClientException;
+import com.amazonaws.AmazonServiceException;
+import com.amazonaws.ClientConfiguration;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.regions.Regions;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
+import com.amazonaws.services.dynamodbv2.document.BatchWriteItemOutcome;
+import com.amazonaws.services.dynamodbv2.document.DynamoDB;
+import com.amazonaws.services.dynamodbv2.document.Item;
+import com.amazonaws.services.dynamodbv2.document.Table;
+import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
+import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
+import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
+import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
+import com.amazonaws.services.dynamodbv2.model.KeyType;
+import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
+import com.amazonaws.services.dynamodbv2.model.ResourceNotFoundException;
+import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
+import com.amazonaws.services.dynamodbv2.model.TableDescription;
+import com.amazonaws.services.dynamodbv2.model.WriteRequest;
+import com.google.common.base.Optional;
+import com.google.inject.Inject;
+import org.embulk.config.ConfigException;
+import org.embulk.config.UserDataException;
+import org.embulk.spi.Exec;
+import org.jruby.embed.ScriptingContainer;
+import org.slf4j.Logger;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+public class DynamodbUtils
+{
+    private final Logger log;
+
+    @Inject
+    public DynamodbUtils()
+    {
+        log = Exec.getLogger(getClass());
+    }
+
+    protected DynamoDB createDynamoDB(DynamodbOutputPlugin.PluginTask task)
+    {
+        DynamoDB dynamoDB;
+        try {
+            AmazonDynamoDBClient client = new AmazonDynamoDBClient(
+                    getCredentialsProvider(task),
+                    getClientConfiguration(task)
+            ).withRegion(Regions.fromName(task.getRegion()));
+
+            if (task.getEndpoint().isPresent()) {
+                client.setEndpoint(task.getEndpoint().get());
+            }
+
+            dynamoDB = new DynamoDB(client);
+            dynamoDB.getTable(task.getTable());
+        }
+        catch (AmazonServiceException ex) {
+            int statusCode = ex.getStatusCode();
+            if (statusCode == 400) {
+                throw new ConfigException(ex);
+            }
+            else {
+                throw new ConnectionException(ex);
+            }
+        }
+        catch (AmazonClientException ex) {
+            throw new ConnectionException(ex);
+        }
+        return dynamoDB;
+    }
+
+    protected ClientConfiguration getClientConfiguration(DynamodbOutputPlugin.PluginTask task)
+    {
+        ClientConfiguration clientConfig = new ClientConfiguration();
+
+        //clientConfig.setProtocol(Protocol.HTTP);
+        clientConfig.setMaxConnections(50); // SDK default: 50
+        clientConfig.setMaxErrorRetry(3); // SDK default: 3
+        clientConfig.setSocketTimeout(8 * 60 * 1000); // SDK default: 50*1000
+
+        return clientConfig;
+    }
+
+    private AWSCredentialsProvider getCredentialsProvider(DynamodbOutputPlugin.PluginTask task)
+    {
+        return AwsCredentials.getAWSCredentialsProvider(task);
+    }
+
+    protected void configCheck(DynamodbOutputPlugin.PluginTask task)
+    {
+        // @see http://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_BatchWriteItem.html
+        if (task.getMode().equals(DynamodbOutputPlugin.Mode.UPSERT)) {
+            if (task.getMaxPutItems() > 25) {
+                throw new ConfigException("'max_put_items' must less than or equal to 25");
+            }
+        }
+
+        if (task.getMode().equals(DynamodbOutputPlugin.Mode.UPSERT_WITH_EXPRESSION)) {
+            if (!task.getUpdateExpression().isPresent()) {
+                throw new ConfigException("'update_expression' is required when update mode");
+            }
+        }
+    }
+
+    protected void batchWriteItem(DynamoDB dynamoDB, TableWriteItems items)
+    {
+        BatchWriteItemOutcome outcome = dynamoDB.batchWriteItem(items);
+        int retryCount = 0;
+        try {
+            do {
+                Map<String, List<WriteRequest>> unprocessedItems = outcome.getUnprocessedItems();
+                // @see http://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_BatchWriteItem.html
+                // If DynamoDB returns any unprocessed items, you should retry the batch operation on those items.
+                // However, we strongly recommend that you use an exponential backoff algorithm
+                if (outcome.getUnprocessedItems().size() > 0) {
+                    retryCount++;
+                    if (retryCount >= 5) {
+                        throw new ConnectionException("Retry count expired while executing batchWriteItem");
+                    }
+                    Thread.sleep(500 * retryCount);
+                    log.warn("Retrieving the unprocessed items");
+                    outcome = dynamoDB.batchWriteItemUnprocessed(unprocessedItems);
+                }
+            } while (outcome.getUnprocessedItems().size() > 0);
+        }
+        catch (InterruptedException ex) {
+            throw new ConnectionException("Retry batchWriteItem was interrupted");
+        }
+    }
+
+    protected void updateItem(DynamoDB dynamoDB, String tableName, Item item, String primaryKey, Optional<String> expression)
+    {
+        Object primaryKeyValue = null;
+        Map<String, String> attributeNames = new HashMap<>();
+        Map<String, Object> attributeValues = new HashMap<>();
+
+        Map<String, Object> itemMap = item.asMap();
+        for (Map.Entry<String, Object> e : itemMap.entrySet()) {
+            String keyName = e.getKey();
+            if (keyName.equals(primaryKey)) {
+                primaryKeyValue = e.getValue();
+            }
+            else {
+                if (expression.get().indexOf(keyName) > 0) {
+                    attributeNames.put("#" + keyName, keyName);
+                    attributeValues.put(":" + keyName, e.getValue());
+                }
+            }
+        }
+        log.debug("attribute names: " + attributeNames.toString());
+        log.debug("attribute values: " + attributeValues.toString());
+        log.debug(String.format("primary key %s:%s", primaryKey, primaryKeyValue));
+        Table table = dynamoDB.getTable(tableName);
+        table.updateItem(primaryKey, primaryKeyValue, expression.get(), attributeNames, attributeValues);
+    }
+
+    protected String getPrimaryKeyName(DynamoDB dynamoDB, String tableName)
+    {
+        Table table = dynamoDB.getTable(tableName);
+
+        TableDescription description = table.describe();
+        Iterator<KeySchemaElement> schema = description.getKeySchema().iterator();
+        String primaryKey = null;
+        while (schema.hasNext()) {
+            KeySchemaElement element = schema.next();
+            primaryKey = element.getAttributeName();
+        }
+        return primaryKey;
+    }
+
+    protected void createTable(DynamoDB dynamoDB, DynamodbOutputPlugin.PluginTask task)
+            throws InterruptedException
+    {
+        ArrayList<KeySchemaElement> keySchema = getKeySchemaElements(task);
+        ArrayList<AttributeDefinition> attributeDefinitions = getAttributeDefinitions(task);
+        ProvisionedThroughput provisionedThroughput = new ProvisionedThroughput()
+                .withReadCapacityUnits(task.getReadCapacityUnits().get().getNormal().get())
+                .withWriteCapacityUnits(task.getWriteCapacityUnits().get().getNormal().get());
+
+        dynamoDB.createTable(new CreateTableRequest()
+                .withTableName(task.getTable())
+                .withKeySchema(keySchema)
+                .withAttributeDefinitions(attributeDefinitions)
+                .withProvisionedThroughput(provisionedThroughput)
+        );
+
+        Table table = dynamoDB.getTable(task.getTable());
+        table.waitForActive();
+        log.info(String.format("Created table '%s'", task.getTable()));
+    }
+
+    protected void deleteTable(DynamoDB dynamoDB, String tableName)
+            throws InterruptedException
+    {
+        Table table = dynamoDB.getTable(tableName);
+        table.delete();
+        table.waitForDelete();
+        log.info(String.format("Deleted table '%s'", tableName));
+    }
+
+    protected boolean isExistsTable(DynamoDB dynamoDB, String tableName)
+            throws InterruptedException
+    {
+        Table table = dynamoDB.getTable(tableName);
+        TableDescription description = null;
+        try {
+            switch (table.describe().getTableStatus()) {
+                case "CREATING":
+                case "UPDATING":
+                    table.waitForActive();
+                    return true;
+                case "DELETING":
+                    table.waitForDelete();
+                    return true;
+                default:
+                    return true;
+            }
+        }
+        catch (ResourceNotFoundException e) {
+            return false;
+        }
+        catch (AmazonClientException e) {
+            return false;
+        }
+    }
+
+    protected void updateTableProvision(DynamoDB dynamoDB, DynamodbOutputPlugin.PluginTask task, boolean isRaise)
+            throws InterruptedException
+    {
+        if (!task.getReadCapacityUnits().isPresent() && !task.getWriteCapacityUnits().isPresent()) {
+            return;
+        }
+
+        Boolean isNeedChange = false;
+
+        Table table = dynamoDB.getTable(task.getTable());
+        TableDescription description = table.describe();
+        long currentReadCapacityUnit = description.getProvisionedThroughput().getReadCapacityUnits();
+        long currentWriteCapacityUnit = description.getProvisionedThroughput().getWriteCapacityUnits();
+
+        ProvisionedThroughput throughput = new ProvisionedThroughput();
+        Optional<Long> readUnits = (isRaise) ? task.getReadCapacityUnits().get().getRaise() : task.getReadCapacityUnits().get().getNormal();
+        if (readUnits.isPresent()) {
+            Long readUnitsLong = readUnits.get();
+            if (currentReadCapacityUnit != readUnitsLong) {
+                throughput.withReadCapacityUnits(readUnitsLong);
+                isNeedChange = true;
+            }
+        }
+        Optional<Long> writeUnits = (isRaise) ? task.getWriteCapacityUnits().get().getRaise() : task.getWriteCapacityUnits().get().getNormal();
+        if (writeUnits.isPresent()) {
+            Long writeUnitsLong = writeUnits.get();
+            if (currentWriteCapacityUnit != writeUnitsLong) {
+                throughput.withWriteCapacityUnits(writeUnitsLong);
+                isNeedChange = true;
+            }
+        }
+
+        if (isNeedChange) {
+            table.updateTable(throughput);
+            log.info(String.format("Updated Provisioned Throughput of table[%s]. read_capacity_unit[%s], write_capacity_unit[%s]",
+                    task.getTable(), readUnits.orNull(), writeUnits.orNull())
+            );
+            table.waitForActive();
+        }
+        else {
+            log.info(String.format("No Provisioned Throughput update is needed for table[%s]. Current value is read_capacity_unit[%s], write_capacity_unit[%s]",
+                    task.getTable(), currentReadCapacityUnit, currentWriteCapacityUnit)
+            );
+        }
+    }
+
+    // Parse like "table_%Y_%m"(include pattern or not) format using Java is difficult. So use jRuby.
+    public String generateTableName(String tableName)
+    {
+        ScriptingContainer jruby = new ScriptingContainer();
+        return jruby.runScriptlet("Time.now.strftime('" + tableName + "')").toString();
+    }
+
+    private ArrayList<KeySchemaElement> getKeySchemaElements(DynamodbOutputPlugin.PluginTask task)
+    {
+        ArrayList<KeySchemaElement> keySchema = new ArrayList<>();
+        keySchema.add(new KeySchemaElement().withAttributeName(task.getPrimaryKey().get()).withKeyType(KeyType.HASH));
+        if (task.getSortKey().isPresent()) {
+            String sortKey = task.getSortKey().get();
+            keySchema.add(new KeySchemaElement().withAttributeName(sortKey).withKeyType(KeyType.RANGE));
+        }
+        return keySchema;
+    }
+
+    private ArrayList<AttributeDefinition> getAttributeDefinitions(DynamodbOutputPlugin.PluginTask task)
+    {
+        ArrayList<AttributeDefinition> attributeDefinitions = new ArrayList<>();
+        attributeDefinitions.add(
+                new AttributeDefinition()
+                        .withAttributeName(task.getPrimaryKey().get())
+                        .withAttributeType(getAttributeType(task.getPrimaryKeyType().get())));
+        if (task.getSortKey().isPresent()) {
+            String sortKey = task.getSortKey().get();
+            attributeDefinitions.add(
+                    new AttributeDefinition()
+                            .withAttributeName(sortKey)
+                            .withAttributeType(getAttributeType(task.getSortKeyType().get())));
+        }
+        return attributeDefinitions;
+    }
+
+    private ScalarAttributeType getAttributeType(String type)
+    {
+        switch (type.toLowerCase()) {
+            case "string":
+                return ScalarAttributeType.S;
+            case "number":
+                return ScalarAttributeType.N;
+            case "binary":
+                return ScalarAttributeType.B;
+            default:
+                throw new UnknownScalarAttributeTypeException(type + " is invalid key type");
+        }
+    }
+
+    public class ConnectionException extends RuntimeException implements UserDataException
+    {
+        protected ConnectionException()
+        {
+        }
+
+        public ConnectionException(String message)
+        {
+            super(message);
+        }
+
+        public ConnectionException(Throwable cause)
+        {
+            super(cause);
+        }
+    }
+
+    public class UnknownScalarAttributeTypeException extends RuntimeException implements UserDataException
+    {
+        protected UnknownScalarAttributeTypeException()
+        {
+        }
+
+        public UnknownScalarAttributeTypeException(String message)
+        {
+            super(message);
+        }
+
+        public UnknownScalarAttributeTypeException(Throwable cause)
+        {
+            super(cause);
+        }
+    }
+}

data/src/test/java/org/embulk/output/dynamodb/TestDynamodbOutputPlugin.java

@@ -0,0 +1,249 @@
+package org.embulk.output.dynamodb;
+
+import com.amazonaws.services.dynamodbv2.document.DynamoDB;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import org.embulk.EmbulkTestRuntime;
+import org.embulk.config.ConfigException;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.TaskReport;
+import org.embulk.config.TaskSource;
+import org.embulk.output.dynamodb.DynamodbOutputPlugin.PluginTask;
+import org.embulk.spi.Exec;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageTestUtils;
+import org.embulk.spi.Schema;
+import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
+import org.embulk.spi.TransactionalPageOutput;
+import org.embulk.standards.CsvParserPlugin;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+public class TestDynamodbOutputPlugin
+{
+    private static String PATH_PREFIX;
+
+    private MockPageOutput pageOutput;
+
+    @BeforeClass
+    public static void initializeConstant()
+    {
+        PATH_PREFIX = DynamodbOutputPlugin.class.getClassLoader().getResource("sample_01.csv").getPath();
+    }
+
+    @Rule
+    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
+    private DynamodbOutputPlugin plugin;
+
+    @Before
+    public void createResources() throws Exception
+    {
+        ConfigSource config = config();
+        plugin = new DynamodbOutputPlugin();
+        PluginTask task = config.loadConfig(PluginTask.class);
+        pageOutput = new MockPageOutput();
+
+        DynamodbUtils dynamoDbUtils = new DynamodbUtils();
+        DynamoDB dynamoDB = null;
+        try {
+            dynamoDB = dynamoDbUtils.createDynamoDB(task);
+            if (dynamoDbUtils.isExistsTable(dynamoDB, task.getTable())) {
+                dynamoDbUtils.deleteTable(dynamoDB, task.getTable());
+            }
+            dynamoDbUtils.createTable(dynamoDB, task);
+        }
+        finally {
+            if (dynamoDB != null) {
+                dynamoDB.shutdown();
+            }
+        }
+    }
+
+    @Test
+    public void testDefaultValues()
+    {
+        ConfigSource config = config();
+        DynamodbOutputPlugin.PluginTask task = config.loadConfig(PluginTask.class);
+        assertEquals("us-west-1", task.getRegion());
+    }
+
+    @Test
+    public void testTransaction()
+    {
+        ConfigSource config = config();
+        Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
+        plugin.transaction(config, schema, 0, new OutputPlugin.Control()
+        {
+            @Override
+            public List<TaskReport> run(TaskSource taskSource)
+            {
+                return Lists.newArrayList(Exec.newTaskReport());
+            }
+        });
+        // no error happens
+    }
+
+    @Test
+    public void testResume()
+    {
+        ConfigSource config = config();
+        Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
+        PluginTask task = config.loadConfig(PluginTask.class);
+        plugin.resume(task.dump(), schema, 0, new OutputPlugin.Control()
+        {
+            @Override
+            public List<TaskReport> run(TaskSource taskSource)
+            {
+                return Lists.newArrayList(Exec.newTaskReport());
+            }
+        });
+    }
+
+    @Test
+    public void testCleanup()
+    {
+        ConfigSource config = config();
+        Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
+        PluginTask task = config.loadConfig(PluginTask.class);
+        plugin.cleanup(task.dump(), schema, 0, Arrays.asList(Exec.newTaskReport()));
+        // no error happens
+    }
+
+    @Test
+    public void testOutputByOpen() throws Exception
+    {
+        ConfigSource config = config();
+        Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
+        PluginTask task = config.loadConfig(PluginTask.class);
+        plugin.transaction(config, schema, 0, new OutputPlugin.Control()
+        {
+            @Override
+            public List<TaskReport> run(TaskSource taskSource)
+            {
+                return Lists.newArrayList(Exec.newTaskReport());
+            }
+        });
+        TransactionalPageOutput output = plugin.open(task.dump(), schema, 0);
+
+        List<Page> pages = PageTestUtils.buildPage(runtime.getBufferAllocator(), schema, 1L, 32864L, "2015-01-27T19:23:49", "2015-01-27T00:00:00",  true, 123.45, "embulk");
+        assertEquals(1, pages.size());
+        for (Page page : pages) {
+            output.add(page);
+        }
+
+//        output.finish();
+//        output.commit();
+//
+//        DynamodbUtils dynamoDbUtils = new DynamodbUtils();
+//        DynamoDB dynamoDB = null;
+//        try {
+//            dynamoDB = dynamoDbUtils.createDynamoDB(task);
+//
+//            Table table = dynamoDB.getTable(task.getTable());
+//            ItemCollection<ScanOutcome> items = table.scan();
+//
+//            while (items.iterator().hasNext()) {
+//                Map<String, Object> item = items.iterator().next().asMap();
+//                assertEquals(1, item.get("id"));
+//                assertEquals(32864, item.get("account"));
+//                assertEquals("2015-01-27T19:23:49", item.get("time"));
+//                assertEquals("2015-01-27T00:00:00", item.get("purchase"));
+//                assertEquals(true, item.get("flg"));
+//                assertEquals(123.45, item.get("score"));
+//                assertEquals("embulk", item.get("comment"));
+//            }
+//        }
+//        finally {
+//            if (dynamoDB != null) {
+//                dynamoDB.shutdown();
+//            }
+//        }
+    }
+
+    @Test
+    public void testMode()
+    {
+        assertEquals(2, DynamodbOutputPlugin.Mode.values().length);
+        assertEquals(DynamodbOutputPlugin.Mode.UPSERT, DynamodbOutputPlugin.Mode.valueOf("UPSERT"));
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testModeThrowsConfigException()
+    {
+        DynamodbOutputPlugin.Mode.fromString("non-exists-mode");
+    }
+
+    private ConfigSource config()
+    {
+        return Exec.newConfigSource()
+                .set("in", inputConfig())
+                .set("parser", parserConfig(schemaConfig()))
+                .set("type", "dynamodb")
+                .set("mode", "upsert")
+                .set("region", "us-west-1")
+                .set("table", "dummy")
+                .set("primary_key", "id")
+                .set("primary_key_type", "string")
+                .set("read_capacity_units", capacityUnitConfig())
+                .set("write_capacity_units", capacityUnitConfig())
+                .set("auth_method", "basic")
+                .set("access_key_id", "dummy")
+                .set("secret_access_key", "dummy")
+                .set("endpoint", "http://localhost:8000");
+    }
+
+    private ImmutableMap<String, Object> capacityUnitConfig()
+    {
+        ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+        builder.put("normal", 5L);
+        builder.put("raise", 8L);
+        return builder.build();
+    }
+
+    private ImmutableMap<String, Object> inputConfig()
+    {
+        ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+        builder.put("type", "file");
+        builder.put("path_prefix", PATH_PREFIX);
+        builder.put("last_path", "");
+        return builder.build();
+    }
+
+    private ImmutableMap<String, Object> parserConfig(ImmutableList<Object> schemaConfig)
+    {
+        ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+        builder.put("type", "csv");
+        builder.put("newline", "CRLF");
+        builder.put("delimiter", ",");
+        builder.put("quote", "\"");
+        builder.put("escape", "\"");
+        builder.put("trim_if_not_quoted", false);
+        builder.put("skip_header_lines", 1);
+        builder.put("allow_extra_columns", false);
+        builder.put("allow_optional_columns", false);
+        builder.put("columns", schemaConfig);
+        return builder.build();
+    }
+
+    private ImmutableList<Object> schemaConfig()
+    {
+        ImmutableList.Builder<Object> builder = new ImmutableList.Builder<>();
+        builder.add(ImmutableMap.of("name", "id", "type", "long"));
+        builder.add(ImmutableMap.of("name", "account", "type", "long"));
+        builder.add(ImmutableMap.of("name", "time", "type", "timestamp", "format", "%Y-%m-%d %H:%M:%S"));
+        builder.add(ImmutableMap.of("name", "purchase", "type", "timestamp", "format", "%Y%m%d"));
+        builder.add(ImmutableMap.of("name", "flg", "type", "boolean"));
+        builder.add(ImmutableMap.of("name", "score", "type", "double"));
+        builder.add(ImmutableMap.of("name", "comment", "type", "string"));
+        return builder.build();
+    }
+}