Wait for DynamoDB table creation

orbit · Jan 13, 2016 · e70b510 · e70b510
1 parent d2646a4
commit e70b510
Showing 1 changed file with 39 additions and 16 deletions.
diff --git a/...amodb/src/main/java/com/ea/orbit/actors/extensions/dynamodb/DynamoDBStorageExtension.java b/...amodb/src/main/java/com/ea/orbit/actors/extensions/dynamodb/DynamoDBStorageExtension.java
@@ -81,6 +81,7 @@ public enum AmazonCredentialType
     private String secretKey;
     private String sessionToken;
     private boolean shouldCreateTables = true;
+    private String dbPrefix = "orbit";
 
 
     public DynamoDBStorageExtension()
@@ -189,8 +190,10 @@ public Task<Void> writeState(final RemoteReference<?> reference, final Object st
         }
     }
 
-    private Task<Table> getOrCreateTable(final String tableName)
+    private Task<Table> getOrCreateTable(final String inputName)
     {
+        final String tableName = getDbPrefix() + "_" + inputName.toLowerCase();
+
         final Table table = tableHashMap.get(tableName);
         if (table != null)
         {
@@ -203,21 +206,31 @@ private Task<Table> getOrCreateTable(final String tableName)
                     .thenApply(descriptor -> dynamoDB.getTable(descriptor.getTableName()))
                     .exceptionally(e ->
                     {
-                        if (shouldCreateTables && ExceptionUtils.isCauseInChain(ResourceNotFoundException.class, e))
-                        {
-                            final Table newTable = dynamoDB.createTable(tableName,
-                                    Collections.singletonList(
-                                            new KeySchemaElement("_id", KeyType.HASH)),
-                                    Collections.singletonList(
-                                            new AttributeDefinition("_id", ScalarAttributeType.S)),
-                                    new ProvisionedThroughput(10L, 10L));
-                            tableHashMap.putIfAbsent(tableName, newTable);
-                            return newTable;
-                        }
-                        else
-                        {
-                            throw new UncheckedException(e);
-                        }
+                            if (shouldCreateTables && ExceptionUtils.isCauseInChain(ResourceNotFoundException.class, e))
+                            {
+                                final Table newTable = dynamoDB.createTable(tableName,
+                                        Collections.singletonList(
+                                                new KeySchemaElement("_id", KeyType.HASH)),
+                                        Collections.singletonList(
+                                                new AttributeDefinition("_id", ScalarAttributeType.S)),
+                                        new ProvisionedThroughput(10L, 10L));
+
+                                try
+                                {
+                                    newTable.waitForActive();
+                                }
+                                catch(Exception ex)
+                                {
+                                    throw new UncheckedException(ex);
+                                }
+
+                                tableHashMap.putIfAbsent(tableName, newTable);
+                                return newTable;
+                            }
+                            else
+                            {
+                                throw new UncheckedException(e);
+                            }
                     });
         }
     }
@@ -287,4 +300,14 @@ public void setShouldCreateTables(boolean shouldCreateTables)
         this.shouldCreateTables = shouldCreateTables;
     }
 
+    public String getDbPrefix()
+    {
+        return dbPrefix;
+    }
+
+    public void setDbPrefix(String dbPrefix)
+    {
+        this.dbPrefix = dbPrefix;
+    }
+
 }