TryDax.java - Amazon DynamoDB

TryDax.java

import java.util.Map; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.BillingMode; import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest; import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest; import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement; import software.amazon.awssdk.services.dynamodb.model.KeyType; import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; import software.amazon.awssdk.services.dynamodb.model.QueryRequest; import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType; import software.amazon.dax.ClusterDaxAsyncClient; import software.amazon.dax.Configuration; public class TryDax { public static void main(String[] args) throws Exception { DynamoDbAsyncClient ddbClient = DynamoDbAsyncClient.builder() .build(); DynamoDbAsyncClient daxClient = null; if (args.length >= 1) { daxClient = ClusterDaxAsyncClient.builder() .overrideConfiguration(Configuration.builder() .url(args[0]) // e.g. dax://my-cluster.l6fzcv.dax-clusters.us-east-1.amazonaws.com .build()) .build(); } String tableName = "TryDaxTable"; System.out.println("Creating table..."); createTable(tableName, ddbClient); System.out.println("Populating table..."); writeData(tableName, ddbClient, 100, 10); DynamoDbAsyncClient testClient = null; if (daxClient != null) { testClient = daxClient; } else { testClient = ddbClient; } System.out.println("Running GetItem and Query tests..."); System.out.println("First iteration of each test will result in cache misses"); System.out.println("Next iterations are cache hits\n"); // GetItem getItemTest(tableName, testClient, 100, 10, 5); // Query queryTest(tableName, testClient, 100, 2, 9, 5); System.out.println("Deleting table..."); deleteTable(tableName, ddbClient); } private static void createTable(String tableName, DynamoDbAsyncClient client) { try { System.out.println("Attempting to create table; please wait..."); client.createTable(CreateTableRequest.builder() .tableName(tableName) .keySchema(KeySchemaElement.builder() .keyType(KeyType.HASH) .attributeName("pk") .build(), KeySchemaElement.builder() .keyType(KeyType.RANGE) .attributeName("sk") .build()) .attributeDefinitions(AttributeDefinition.builder() .attributeName("pk") .attributeType(ScalarAttributeType.N) .build(), AttributeDefinition.builder() .attributeName("sk") .attributeType(ScalarAttributeType.N) .build()) .billingMode(BillingMode.PAY_PER_REQUEST) .build()).get(); client.waiter().waitUntilTableExists(DescribeTableRequest.builder() .tableName(tableName) .build()).get(); System.out.println("Successfully created table."); } catch (Exception e) { System.err.println("Unable to create table: "); e.printStackTrace(); } } private static void deleteTable(String tableName, DynamoDbAsyncClient client) { try { System.out.println("\nAttempting to delete table; please wait..."); client.deleteTable(DeleteTableRequest.builder() .tableName(tableName) .build()).get(); client.waiter().waitUntilTableNotExists(DescribeTableRequest.builder() .tableName(tableName) .build()).get(); System.out.println("Successfully deleted table."); } catch (Exception e) { System.err.println("Unable to delete table: "); e.printStackTrace(); } } private static void writeData(String tableName, DynamoDbAsyncClient client, int pkmax, int skmax) { System.out.println("Writing data to the table..."); int stringSize = 1000; StringBuilder sb = new StringBuilder(stringSize); for (int i = 0; i < stringSize; i++) { sb.append('X'); } String someData = sb.toString(); try { for (int ipk = 1; ipk <= pkmax; ipk++) { System.out.println(("Writing " + skmax + " items for partition key: " + ipk)); for (int isk = 1; isk <= skmax; isk++) { client.putItem(PutItemRequest.builder() .tableName(tableName) .item(Map.of("pk", attr(ipk), "sk", attr(isk), "someData", attr(someData))) .build()).get(); } } } catch (Exception e) { System.err.println("Unable to write item:"); e.printStackTrace(); } } private static AttributeValue attr(int n) { return AttributeValue.builder().n(String.valueOf(n)).build(); } private static AttributeValue attr(String s) { return AttributeValue.builder().s(s).build(); } private static void getItemTest(String tableName, DynamoDbAsyncClient client, int pk, int sk, int iterations) { long startTime, endTime; System.out.println("GetItem test - partition key 1-" + pk + " and sort keys 1-" + sk); for (int i = 0; i < iterations; i++) { startTime = System.nanoTime(); try { for (int ipk = 1; ipk <= pk; ipk++) { for (int isk = 1; isk <= sk; isk++) { client.getItem(GetItemRequest.builder() .tableName(tableName) .key(Map.of("pk", attr(ipk), "sk", attr(isk))) .build()).get(); } } } catch (Exception e) { System.err.println("Unable to get item:"); e.printStackTrace(); } endTime = System.nanoTime(); printTime(startTime, endTime, pk * sk); } } private static void queryTest(String tableName, DynamoDbAsyncClient client, int pk, int sk1, int sk2, int iterations) { long startTime, endTime; System.out.println("Query test - partition key 1-" + pk + " and sort keys between " + sk1 + " and " + sk2); for (int i = 0; i < iterations; i++) { startTime = System.nanoTime(); for (int ipk = 1; ipk <= pk; ipk++) { try { // Pagination API for Query. client.queryPaginator(QueryRequest.builder() .tableName(tableName) .keyConditionExpression("pk = :pkval and sk between :skval1 and :skval2") .expressionAttributeValues(Map.of(":pkval", attr(ipk), ":skval1", attr(sk1), ":skval2", attr(sk2))) .build()).items().subscribe((item) -> { }).get(); } catch (Exception e) { System.err.println("Unable to query table:"); e.printStackTrace(); } } endTime = System.nanoTime(); printTime(startTime, endTime, pk); } } private static void printTime(long startTime, long endTime, int iterations) { System.out.format("\tTotal time: %.3f ms - ", (endTime - startTime) / (1000000.0)); System.out.format("Avg time: %.3f ms\n", (endTime - startTime) / (iterations * 1000000.0)); } }