diff --git a/.gitignore b/.gitignore index 37d101daaa..a6480b1401 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ package-lock.json cdk.out .env **/*.dylib +*/dynamodb-local-metadata.json diff --git a/athena-dynamodb/pom.xml b/athena-dynamodb/pom.xml index fd44faefe9..2ae4c5c4e9 100644 --- a/athena-dynamodb/pom.xml +++ b/athena-dynamodb/pom.xml @@ -8,6 +8,17 @@ 4.0.0 athena-dynamodb 2022.47.1 + + + + software.amazon.awssdk + bom + 2.25.4 + pom + import + + + com.amazonaws @@ -20,23 +31,30 @@ athena-federation-integ-test 2022.47.1 test - - - com.amazonaws - aws-java-sdk-dynamodb - ${aws-sdk.version} - - commons-logging - commons-logging + com.amazonaws + aws-java-sdk-sts + + software.amazon.awssdk + dynamodb + + + software.amazon.awssdk + dynamodb-enhanced + com.amazonaws DynamoDBLocal - 2.0.0 + LATEST + test + + + software.amazon.awssdk + url-connection-client test @@ -93,6 +111,14 @@ test-jar test + + software.amazon.awssdk + sdk-core + + + software.amazon.awssdk + sts + diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java index e2379cce9d..ef15f9683d 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java @@ -19,7 +19,6 @@ */ package com.amazonaws.athena.connectors.dynamodb; -import com.amazonaws.athena.connector.credentials.CrossAccountCredentialsProvider; import com.amazonaws.athena.connector.lambda.QueryStatusChecker; import com.amazonaws.athena.connector.lambda.ThrottlingInvoker; import com.amazonaws.athena.connector.lambda.data.Block; @@ -43,6 +42,7 @@ import com.amazonaws.athena.connector.lambda.metadata.glue.GlueFieldLexer; import com.amazonaws.athena.connector.lambda.security.EncryptionKeyFactory; import com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants; +import com.amazonaws.athena.connectors.dynamodb.credentials.CrossAccountCredentialsProviderV2; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBPaginatedTables; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable; @@ -53,10 +53,6 @@ import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; import com.amazonaws.athena.connectors.dynamodb.util.IncrementingValueNameProducer; import com.amazonaws.services.athena.AmazonAthena; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder; -import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.amazonaws.services.glue.AWSGlue; import com.amazonaws.services.glue.model.Database; import com.amazonaws.services.glue.model.Table; @@ -69,6 +65,9 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.util.ArrayList; import java.util.Collections; @@ -99,6 +98,7 @@ import static com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants.SEGMENT_ID_PROPERTY; import static com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants.TABLE_METADATA; import static com.amazonaws.athena.connectors.dynamodb.throttling.DynamoDBExceptionFilter.EXCEPTION_FILTER; +import static com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils.toAttributeValue; /** * Handles metadata requests for the Athena DynamoDB Connector. @@ -131,15 +131,15 @@ public class DynamoDBMetadataHandler private static final DatabaseFilter DB_FILTER = (Database database) -> (database.getLocationUri() != null && database.getLocationUri().contains(DYNAMO_DB_FLAG)); private final ThrottlingInvoker invoker; - private final AmazonDynamoDB ddbClient; + private final DynamoDbClient ddbClient; private final AWSGlue glueClient; private final DynamoDBTableResolver tableResolver; public DynamoDBMetadataHandler(java.util.Map configOptions) { super(SOURCE_TYPE, configOptions); - this.ddbClient = AmazonDynamoDBClientBuilder.standard() - .withCredentials(CrossAccountCredentialsProvider.getCrossAccountCredentialsIfPresent(configOptions, "DynamoDBMetadataHandler_CrossAccountRoleSession")) + this.ddbClient = DynamoDbClient.builder() + .credentialsProvider(CrossAccountCredentialsProviderV2.getCrossAccountCredentialsIfPresent(configOptions, "DynamoDBMetadataHandler_CrossAccountRoleSession")) .build(); this.glueClient = getAwsGlue(); this.invoker = ThrottlingInvoker.newDefaultBuilder(EXCEPTION_FILTER, configOptions).build(); @@ -153,7 +153,7 @@ public DynamoDBMetadataHandler(java.util.Map configOptions) AmazonAthena athena, String spillBucket, String spillPrefix, - AmazonDynamoDB ddbClient, + DynamoDbClient ddbClient, AWSGlue glueClient, java.util.Map configOptions) { @@ -400,8 +400,9 @@ private void precomputeAdditionalMetadata(Set columnsToIgnore, Map splitMetadata = new HashMap<>(partitionMetadata); Object hashKeyValue = DDBTypeUtils.convertArrowTypeIfNecessary(hashKeyName, hashKeyValueReader.readObject()); - String hashKeyValueJSON = Jackson.toJsonString(ItemUtils.toAttributeValue(hashKeyValue)); - splitMetadata.put(hashKeyName, hashKeyValueJSON); + splitMetadata.put(hashKeyName, DDBTypeUtils.attributeToJson(toAttributeValue(hashKeyValue), hashKeyName)); splits.add(new Split(spillLocation, makeEncryptionKey(), splitMetadata)); diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandler.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandler.java index f52f36c290..ce74d92cd1 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandler.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandler.java @@ -19,8 +19,6 @@ */ package com.amazonaws.athena.connectors.dynamodb; -import com.amazonaws.AmazonWebServiceRequest; -import com.amazonaws.athena.connector.credentials.CrossAccountCredentialsProvider; import com.amazonaws.athena.connector.lambda.QueryStatusChecker; import com.amazonaws.athena.connector.lambda.ThrottlingInvoker; import com.amazonaws.athena.connector.lambda.data.Block; @@ -31,30 +29,31 @@ import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; import com.amazonaws.athena.connector.lambda.handlers.RecordHandler; import com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest; +import com.amazonaws.athena.connectors.dynamodb.credentials.CrossAccountCredentialsProviderV2; import com.amazonaws.athena.connectors.dynamodb.resolver.DynamoDBFieldResolver; import com.amazonaws.athena.connectors.dynamodb.util.DDBPredicateUtils; import com.amazonaws.athena.connectors.dynamodb.util.DDBRecordMetadata; import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; import com.amazonaws.services.athena.AmazonAthena; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.QueryRequest; -import com.amazonaws.services.dynamodbv2.model.QueryResult; -import com.amazonaws.services.dynamodbv2.model.ScanRequest; -import com.amazonaws.services.dynamodbv2.model.ScanResult; import com.amazonaws.services.s3.AmazonS3; import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.util.json.Jackson; import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.annotations.VisibleForTesting; import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; -import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; +import software.amazon.awssdk.services.dynamodb.model.ScanRequest; +import software.amazon.awssdk.services.dynamodb.model.ScanResponse; import java.io.IOException; import java.util.ArrayList; @@ -103,14 +102,14 @@ public class DynamoDBRecordHandler private static final TypeReference> ATTRIBUTE_VALUE_MAP_TYPE_REFERENCE = new TypeReference>() {}; private final LoadingCache invokerCache; - private final AmazonDynamoDB ddbClient; + private final DynamoDbClient ddbClient; public DynamoDBRecordHandler(java.util.Map configOptions) { super(sourceType, configOptions); - this.ddbClient = AmazonDynamoDBClientBuilder.standard() - .withCredentials(CrossAccountCredentialsProvider.getCrossAccountCredentialsIfPresent(configOptions, "DynamoDBRecordHandler_CrossAccountRoleSession")) - .build(); + this.ddbClient = DynamoDbClient.builder() + .credentialsProvider(CrossAccountCredentialsProviderV2.getCrossAccountCredentialsIfPresent(configOptions, "DynamoDBMetadataHandler_CrossAccountRoleSession")) + .build(); this.invokerCache = CacheBuilder.newBuilder().build( new CacheLoader() { @Override @@ -124,7 +123,7 @@ public ThrottlingInvoker load(String tableName) } @VisibleForTesting - DynamoDBRecordHandler(AmazonDynamoDB ddbClient, AmazonS3 amazonS3, AWSSecretsManager secretsManager, AmazonAthena athena, String sourceType, java.util.Map configOptions) + DynamoDBRecordHandler(DynamoDbClient ddbClient, AmazonS3 amazonS3, AWSSecretsManager secretsManager, AmazonAthena athena, String sourceType, java.util.Map configOptions) { super(amazonS3, secretsManager, athena, sourceType, configOptions); this.ddbClient = ddbClient; @@ -259,10 +258,15 @@ private List getRangeValues(String rangeKeyFilter) return rangeValues; } + private boolean isQueryRequest(Split split) + { + return split.getProperty(SEGMENT_ID_PROPERTY) == null; + } + /* - Converts a split into a Query or Scan request + Converts a split into a Query */ - private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, Schema schema, Constraints constraints, boolean disableProjectionAndCasing) + private QueryRequest buildQueryRequest(Split split, String tableName, Schema schema, Constraints constraints, boolean disableProjectionAndCasing, Map exclusiveStartKey) { validateExpectedMetadata(split.getProperties()); // prepare filters @@ -273,7 +277,7 @@ private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, if (rangeKeyFilter != null || nonKeyFilter != null) { try { expressionAttributeNames.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_NAMES_METADATA), STRING_MAP_TYPE_REFERENCE)); - expressionAttributeValues.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_VALUES_METADATA), ATTRIBUTE_VALUE_MAP_TYPE_REFERENCE)); + expressionAttributeValues.putAll(EnhancedDocument.fromJson(split.getProperty(EXPRESSION_VALUES_METADATA)).toMap()); } catch (IOException e) { throw new RuntimeException(e); @@ -290,58 +294,89 @@ private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, }) .collect(Collectors.joining(",")); - boolean isQuery = split.getProperty(SEGMENT_ID_PROPERTY) == null; - - if (isQuery) { - // prepare key condition expression - String indexName = split.getProperty(INDEX_METADATA); - String hashKeyName = split.getProperty(HASH_KEY_NAME_METADATA); - String hashKeyAlias = DDBPredicateUtils.aliasColumn(hashKeyName); - String keyConditionExpression = hashKeyAlias + " = " + HASH_KEY_VALUE_ALIAS; - if (rangeKeyFilter != null) { - if (rangeFilterHasIn(rangeKeyFilter)) { - List rangeKeyValues = getRangeValues(rangeKeyFilter); - for (String value : rangeKeyValues) { - expressionAttributeValues.remove(value); - } - } - else { - keyConditionExpression += " AND " + rangeKeyFilter; + // prepare key condition expression + String indexName = split.getProperty(INDEX_METADATA); + String hashKeyName = split.getProperty(HASH_KEY_NAME_METADATA); + String hashKeyAlias = DDBPredicateUtils.aliasColumn(hashKeyName); + String keyConditionExpression = hashKeyAlias + " = " + HASH_KEY_VALUE_ALIAS; + if (rangeKeyFilter != null) { + if (rangeFilterHasIn(rangeKeyFilter)) { + List rangeKeyValues = getRangeValues(rangeKeyFilter); + for (String value : rangeKeyValues) { + expressionAttributeValues.remove(value); } } - expressionAttributeNames.put(hashKeyAlias, hashKeyName); - expressionAttributeValues.put(HASH_KEY_VALUE_ALIAS, Jackson.fromJsonString(split.getProperty(hashKeyName), AttributeValue.class)); - - QueryRequest queryRequest = new QueryRequest() - .withTableName(tableName) - .withIndexName(indexName) - .withKeyConditionExpression(keyConditionExpression) - .withFilterExpression(nonKeyFilter) - .withExpressionAttributeNames(expressionAttributeNames) - .withExpressionAttributeValues(expressionAttributeValues) - .withProjectionExpression(projectionExpression); - if (canApplyLimit(constraints)) { - queryRequest.setLimit((int) constraints.getLimit()); + else { + keyConditionExpression += " AND " + rangeKeyFilter; } - return queryRequest; } - else { - int segmentId = Integer.parseInt(split.getProperty(SEGMENT_ID_PROPERTY)); - int segmentCount = Integer.parseInt(split.getProperty(SEGMENT_COUNT_METADATA)); - - ScanRequest scanRequest = new ScanRequest() - .withTableName(tableName) - .withSegment(segmentId) - .withTotalSegments(segmentCount) - .withFilterExpression(nonKeyFilter) - .withExpressionAttributeNames(expressionAttributeNames.isEmpty() ? null : expressionAttributeNames) - .withExpressionAttributeValues(expressionAttributeValues.isEmpty() ? null : expressionAttributeValues) - .withProjectionExpression(projectionExpression); - if (canApplyLimit(constraints)) { - scanRequest.setLimit((int) constraints.getLimit()); + expressionAttributeNames.put(hashKeyAlias, hashKeyName); + + AttributeValue hashKeyAttribute = DDBTypeUtils.jsonToAttributeValue(split.getProperty(hashKeyName), hashKeyName); + expressionAttributeValues.put(HASH_KEY_VALUE_ALIAS, hashKeyAttribute); + + QueryRequest.Builder queryRequestBuilder = QueryRequest.builder() + .tableName(tableName) + .indexName(indexName) + .keyConditionExpression(keyConditionExpression) + .filterExpression(nonKeyFilter) + .expressionAttributeNames(expressionAttributeNames) + .expressionAttributeValues(expressionAttributeValues) + .projectionExpression(projectionExpression) + .exclusiveStartKey(exclusiveStartKey); + if (canApplyLimit(constraints)) { + queryRequestBuilder.limit((int) constraints.getLimit()); + } + return queryRequestBuilder.build(); + } + + /* + Converts a split into a Scan Request + */ + private ScanRequest buildScanRequest(Split split, String tableName, Schema schema, Constraints constraints, boolean disableProjectionAndCasing, Map exclusiveStartKey) + { + validateExpectedMetadata(split.getProperties()); + // prepare filters + String rangeKeyFilter = split.getProperty(RANGE_KEY_FILTER_METADATA); + String nonKeyFilter = split.getProperty(NON_KEY_FILTER_METADATA); + Map expressionAttributeNames = new HashMap<>(); + Map expressionAttributeValues = new HashMap<>(); + if (rangeKeyFilter != null || nonKeyFilter != null) { + try { + expressionAttributeNames.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_NAMES_METADATA), STRING_MAP_TYPE_REFERENCE)); + expressionAttributeValues.putAll(EnhancedDocument.fromJson(split.getProperty(EXPRESSION_VALUES_METADATA)).toMap()); + } + catch (IOException e) { + throw new RuntimeException(e); } - return scanRequest; } + + // Only read columns that are needed in the query + String projectionExpression = disableProjectionAndCasing ? null : schema.getFields() + .stream() + .map(field -> { + String aliasedName = DDBPredicateUtils.aliasColumn(field.getName()); + expressionAttributeNames.put(aliasedName, field.getName()); + return aliasedName; + }) + .collect(Collectors.joining(",")); + + int segmentId = Integer.parseInt(split.getProperty(SEGMENT_ID_PROPERTY)); + int segmentCount = Integer.parseInt(split.getProperty(SEGMENT_COUNT_METADATA)); + + ScanRequest.Builder scanRequestBuilder = ScanRequest.builder() + .tableName(tableName) + .segment(segmentId) + .totalSegments(segmentCount) + .filterExpression(nonKeyFilter) + .expressionAttributeNames(expressionAttributeNames.isEmpty() ? null : expressionAttributeNames) + .expressionAttributeValues(expressionAttributeValues.isEmpty() ? null : expressionAttributeValues) + .projectionExpression(projectionExpression) + .exclusiveStartKey(exclusiveStartKey); + if (canApplyLimit(constraints)) { + scanRequestBuilder.limit((int) constraints.getLimit()); + } + return scanRequestBuilder.build(); } /* @@ -349,7 +384,6 @@ private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, */ private Iterator> getIterator(Split split, String tableName, Schema schema, Constraints constraints, boolean disableProjectionAndCasing) { - AmazonWebServiceRequest request = buildReadRequest(split, tableName, schema, constraints, disableProjectionAndCasing); return new Iterator>() { AtomicReference> lastKeyEvaluated = new AtomicReference<>(); AtomicReference>> currentPageIterator = new AtomicReference<>(); @@ -359,7 +393,7 @@ public boolean hasNext() { return currentPageIterator.get() == null || currentPageIterator.get().hasNext() - || lastKeyEvaluated.get() != null; + || ((lastKeyEvaluated.get() != null && !lastKeyEvaluated.get().isEmpty())); } @Override @@ -370,19 +404,19 @@ public Map next() } Iterator> iterator; try { - if (request instanceof QueryRequest) { - QueryRequest paginatedRequest = ((QueryRequest) request).withExclusiveStartKey(lastKeyEvaluated.get()); + if (isQueryRequest(split)) { + QueryRequest request = buildQueryRequest(split, tableName, schema, constraints, disableProjectionAndCasing, lastKeyEvaluated.get()); logger.info("Invoking DDB with Query request: {}", request); - QueryResult queryResult = invokerCache.get(tableName).invoke(() -> ddbClient.query(paginatedRequest)); - lastKeyEvaluated.set(queryResult.getLastEvaluatedKey()); - iterator = queryResult.getItems().iterator(); + QueryResponse response = invokerCache.get(tableName).invoke(() -> ddbClient.query(request)); + lastKeyEvaluated.set(response.lastEvaluatedKey()); + iterator = response.items().iterator(); } else { - ScanRequest paginatedRequest = ((ScanRequest) request).withExclusiveStartKey(lastKeyEvaluated.get()); + ScanRequest request = buildScanRequest(split, tableName, schema, constraints, disableProjectionAndCasing, lastKeyEvaluated.get()); logger.info("Invoking DDB with Scan request: {}", request); - ScanResult scanResult = invokerCache.get(tableName).invoke(() -> ddbClient.scan(paginatedRequest)); - lastKeyEvaluated.set(scanResult.getLastEvaluatedKey()); - iterator = scanResult.getItems().iterator(); + ScanResponse response = invokerCache.get(tableName).invoke(() -> ddbClient.scan(request)); + lastKeyEvaluated.set(response.lastEvaluatedKey()); + iterator = response.items().iterator(); } } catch (TimeoutException | ExecutionException e) { diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/credentials/CrossAccountCredentialsProviderV2.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/credentials/CrossAccountCredentialsProviderV2.java new file mode 100644 index 0000000000..98ee948edb --- /dev/null +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/credentials/CrossAccountCredentialsProviderV2.java @@ -0,0 +1,58 @@ +/*- + * #%L + * athena-dynamodb + * %% + * Copyright (C) 2019 - 2024 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connectors.dynamodb.credentials; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; +import software.amazon.awssdk.services.sts.model.Credentials; + +import java.util.Map; + +public class CrossAccountCredentialsProviderV2 +{ + private static final String CROSS_ACCOUNT_ROLE_ARN_CONFIG = "cross_account_role_arn"; + private static final Logger logger = LoggerFactory.getLogger(CrossAccountCredentialsProviderV2.class); + + private CrossAccountCredentialsProviderV2() {} + + public static AwsCredentialsProvider getCrossAccountCredentialsIfPresent(Map configOptions, String roleSessionName) + { + if (configOptions.containsKey(CROSS_ACCOUNT_ROLE_ARN_CONFIG)) { + logger.debug("Found cross-account role arn to assume."); + StsClient stsClient = StsClient.builder().build(); + AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest.builder() + .roleArn(configOptions.get(CROSS_ACCOUNT_ROLE_ARN_CONFIG)) + .roleSessionName(roleSessionName) + .build(); + AssumeRoleResponse assumeRoleResponse = stsClient.assumeRole(assumeRoleRequest); + Credentials credentials = assumeRoleResponse.credentials(); + AwsSessionCredentials sessionCredentials = AwsSessionCredentials.create(credentials.accessKeyId(), credentials.secretAccessKey(), credentials.sessionToken()); + return StaticCredentialsProvider.create(sessionCredentials); + } + return DefaultCredentialsProvider.create(); + } +} diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBIndex.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBIndex.java index b1129f3101..4ae6bff18b 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBIndex.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBIndex.java @@ -19,7 +19,7 @@ */ package com.amazonaws.athena.connectors.dynamodb.model; -import com.amazonaws.services.dynamodbv2.model.ProjectionType; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; import java.util.List; import java.util.Objects; diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBTable.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBTable.java index bfee037723..8e2ded34f9 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBTable.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/model/DynamoDBTable.java @@ -19,8 +19,8 @@ */ package com.amazonaws.athena.connectors.dynamodb.model; -import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; import com.google.common.collect.ImmutableList; +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; import java.util.List; import java.util.Objects; diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBTableResolver.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBTableResolver.java index 6a1acef51d..a3d0b73e3e 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBTableResolver.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBTableResolver.java @@ -23,15 +23,15 @@ import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBPaginatedTables; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable; import com.amazonaws.athena.connectors.dynamodb.util.DDBTableUtils; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.model.ListTablesRequest; -import com.amazonaws.services.dynamodbv2.model.ListTablesResult; -import com.amazonaws.services.dynamodbv2.model.ResourceNotFoundException; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.Multimap; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.ListTablesRequest; +import software.amazon.awssdk.services.dynamodb.model.ListTablesResponse; +import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException; import java.util.ArrayList; import java.util.Collection; @@ -55,11 +55,11 @@ public class DynamoDBTableResolver { private static final Logger logger = LoggerFactory.getLogger(DynamoDBTableResolver.class); - private AmazonDynamoDB ddbClient; + private DynamoDbClient ddbClient; // used to handle Throttling events using an AIMD strategy for congestion control. private ThrottlingInvoker invoker; - public DynamoDBTableResolver(ThrottlingInvoker invoker, AmazonDynamoDB ddbClient) + public DynamoDBTableResolver(ThrottlingInvoker invoker, DynamoDbClient ddbClient) { this.invoker = invoker; this.ddbClient = ddbClient; @@ -85,11 +85,13 @@ private DynamoDBPaginatedTables listPaginatedTables(String token, int pageSize) limit = 100; } do { - ListTablesRequest ddbRequest = new ListTablesRequest() - .withExclusiveStartTableName(nextToken).withLimit(limit); - ListTablesResult result = invoker.invoke(() -> ddbClient.listTables(ddbRequest)); - tables.addAll(result.getTableNames()); - nextToken = result.getLastEvaluatedTableName(); + ListTablesRequest ddbRequest = ListTablesRequest.builder() + .exclusiveStartTableName(nextToken) + .limit(limit) + .build(); + ListTablesResponse response = invoker.invoke(() -> ddbClient.listTables(ddbRequest)); + tables.addAll(response.tableNames()); + nextToken = response.lastEvaluatedTableName(); } while (nextToken != null && pageSize == UNLIMITED_PAGE_SIZE_VALUE); logger.info("{} tables returned with pagination", tables.size()); diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/throttling/DynamoDBExceptionFilter.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/throttling/DynamoDBExceptionFilter.java index 0e030b060e..a6d95d4f6c 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/throttling/DynamoDBExceptionFilter.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/throttling/DynamoDBExceptionFilter.java @@ -20,9 +20,9 @@ package com.amazonaws.athena.connectors.dynamodb.throttling; import com.amazonaws.athena.connector.lambda.ThrottlingInvoker; -import com.amazonaws.services.dynamodbv2.model.LimitExceededException; -import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughputExceededException; -import com.amazonaws.services.dynamodbv2.model.RequestLimitExceededException; +import software.amazon.awssdk.services.dynamodb.model.LimitExceededException; +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputExceededException; +import software.amazon.awssdk.services.dynamodb.model.RequestLimitExceededException; /** * Used by {@link ThrottlingInvoker} to determine which DynamoDB exceptions are thrown for throttling. diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBPredicateUtils.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBPredicateUtils.java index c4f5f891c5..c8601ad26f 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBPredicateUtils.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBPredicateUtils.java @@ -26,11 +26,10 @@ import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable; -import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.ProjectionType; import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; import java.util.ArrayList; import java.util.HashSet; @@ -168,7 +167,7 @@ public static String aliasColumn(String columnName) */ private static void bindValue(String columnName, Object value, List accumulator, DDBRecordMetadata recordMetadata) { - accumulator.add(ItemUtils.toAttributeValue(DDBTypeUtils.convertArrowTypeIfNecessary(columnName, value, recordMetadata))); + accumulator.add(DDBTypeUtils.toAttributeValue(DDBTypeUtils.convertArrowTypeIfNecessary(columnName, value, recordMetadata))); } /* diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTableUtils.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTableUtils.java index dfed880bfc..9c0c97ead1 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTableUtils.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTableUtils.java @@ -23,26 +23,25 @@ import com.amazonaws.athena.connector.lambda.data.SchemaBuilder; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.model.AmazonDynamoDBException; -import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.DescribeTableRequest; -import com.amazonaws.services.dynamodbv2.model.GlobalSecondaryIndexDescription; -import com.amazonaws.services.dynamodbv2.model.IndexStatus; -import com.amazonaws.services.dynamodbv2.model.KeySchemaElement; -import com.amazonaws.services.dynamodbv2.model.KeyType; -import com.amazonaws.services.dynamodbv2.model.LocalSecondaryIndexDescription; -import com.amazonaws.services.dynamodbv2.model.ProjectionType; -import com.amazonaws.services.dynamodbv2.model.ScanRequest; -import com.amazonaws.services.dynamodbv2.model.ScanResult; -import com.amazonaws.services.dynamodbv2.model.TableDescription; import com.google.common.collect.ImmutableList; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; +import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndexDescription; +import software.amazon.awssdk.services.dynamodb.model.IndexStatus; +import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement; +import software.amazon.awssdk.services.dynamodb.model.KeyType; +import software.amazon.awssdk.services.dynamodb.model.LocalSecondaryIndexDescription; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputDescription; +import software.amazon.awssdk.services.dynamodb.model.ScanRequest; +import software.amazon.awssdk.services.dynamodb.model.ScanResponse; +import software.amazon.awssdk.services.dynamodb.model.TableDescription; import java.util.HashSet; import java.util.List; @@ -76,37 +75,42 @@ private DDBTableUtils() {} * @param ddbClient the DDB client to use * @return the table metadata */ - public static DynamoDBTable getTable(String tableName, ThrottlingInvoker invoker, AmazonDynamoDB ddbClient) + public static DynamoDBTable getTable(String tableName, ThrottlingInvoker invoker, DynamoDbClient ddbClient) throws TimeoutException { - DescribeTableRequest request = new DescribeTableRequest().withTableName(tableName); - TableDescription table = invoker.invoke(() -> ddbClient.describeTable(request).getTable()); + DescribeTableRequest request = DescribeTableRequest.builder().tableName(tableName).build(); - KeyNames keys = getKeys(table.getKeySchema()); + TableDescription table = invoker.invoke(() -> ddbClient.describeTable(request).table()); + + KeyNames keys = getKeys(table.keySchema()); // get data statistics - long approxTableSizeInBytes = table.getTableSizeBytes(); - long approxItemCount = table.getItemCount(); - final long provisionedReadCapacity = table.getProvisionedThroughput() != null ? table.getProvisionedThroughput().getReadCapacityUnits() : PSUEDO_CAPACITY_FOR_ON_DEMAND; + long approxTableSizeInBytes = table.tableSizeBytes(); + long approxItemCount = table.itemCount(); + + ProvisionedThroughputDescription provisionedThroughputDescription = table.provisionedThroughput(); + // #todo: from the Documentation; this doesn't seem to be returning null from the looks of it; but test; + final long provisionedReadCapacity = provisionedThroughputDescription != null ? provisionedThroughputDescription.readCapacityUnits() : PSUEDO_CAPACITY_FOR_ON_DEMAND; // get secondary indexes - List localSecondaryIndexes = table.getLocalSecondaryIndexes() != null ? table.getLocalSecondaryIndexes() : ImmutableList.of(); - List globalSecondaryIndexes = table.getGlobalSecondaryIndexes() != null ? table.getGlobalSecondaryIndexes() : ImmutableList.of(); + List localSecondaryIndexes = table.hasLocalSecondaryIndexes() ? table.localSecondaryIndexes() : ImmutableList.of(); + List globalSecondaryIndexes = table.hasGlobalSecondaryIndexes() ? table.globalSecondaryIndexes() : ImmutableList.of(); ImmutableList.Builder indices = ImmutableList.builder(); + localSecondaryIndexes.forEach(i -> { - KeyNames indexKeys = getKeys(i.getKeySchema()); + KeyNames indexKeys = getKeys(i.keySchema()); // DynamoDB automatically fetches all attributes from the table for local secondary index, so ignore projected attributes - indices.add(new DynamoDBIndex(i.getIndexName(), indexKeys.getHashKey(), indexKeys.getRangeKey(), ProjectionType.ALL, ImmutableList.of())); + indices.add(new DynamoDBIndex(i.indexName(), indexKeys.getHashKey(), indexKeys.getRangeKey(), ProjectionType.ALL, ImmutableList.of())); }); globalSecondaryIndexes.stream() - .filter(i -> IndexStatus.fromValue(i.getIndexStatus()).equals(IndexStatus.ACTIVE)) + .filter(i -> i.indexStatus().equals(IndexStatus.ACTIVE)) .forEach(i -> { - KeyNames indexKeys = getKeys(i.getKeySchema()); - indices.add(new DynamoDBIndex(i.getIndexName(), indexKeys.getHashKey(), indexKeys.getRangeKey(), ProjectionType.fromValue(i.getProjection().getProjectionType()), - i.getProjection().getNonKeyAttributes() == null ? ImmutableList.of() : i.getProjection().getNonKeyAttributes())); + KeyNames indexKeys = getKeys(i.keySchema()); + indices.add(new DynamoDBIndex(i.indexName(), indexKeys.getHashKey(), indexKeys.getRangeKey(), i.projection().projectionType(), + i.projection().nonKeyAttributes() == null ? ImmutableList.of() : i.projection().nonKeyAttributes())); }); - return new DynamoDBTable(tableName, keys.getHashKey(), keys.getRangeKey(), table.getAttributeDefinitions(), indices.build(), approxTableSizeInBytes, approxItemCount, provisionedReadCapacity); + return new DynamoDBTable(tableName, keys.getHashKey(), keys.getRangeKey(), table.attributeDefinitions(), indices.build(), approxTableSizeInBytes, approxItemCount, provisionedReadCapacity); } /* @@ -117,11 +121,11 @@ private static KeyNames getKeys(List keys) String hashKey = null; String rangeKey = null; for (KeySchemaElement key : keys) { - if (key.getKeyType().equals(KeyType.HASH.toString())) { - hashKey = key.getAttributeName(); + if (key.keyType().equals(KeyType.HASH)) { + hashKey = key.attributeName(); } - else if (key.getKeyType().equals(KeyType.RANGE.toString())) { - rangeKey = key.getAttributeName(); + else if (key.keyType().equals(KeyType.RANGE)) { + rangeKey = key.attributeName(); } } return new KeyNames(hashKey, rangeKey); @@ -137,20 +141,25 @@ else if (key.getKeyType().equals(KeyType.RANGE.toString())) { * @param ddbClient the DDB client to use * @return the table's derived schema */ - public static Schema peekTableForSchema(String tableName, ThrottlingInvoker invoker, AmazonDynamoDB ddbClient) + public static Schema peekTableForSchema(String tableName, ThrottlingInvoker invoker, DynamoDbClient ddbClient) throws TimeoutException { - ScanRequest scanRequest = new ScanRequest().withTableName(tableName).withLimit(SCHEMA_INFERENCE_NUM_RECORDS); + ScanRequest scanRequest = ScanRequest.builder() + .tableName(tableName) + .limit(SCHEMA_INFERENCE_NUM_RECORDS) + .build(); SchemaBuilder schemaBuilder = new SchemaBuilder(); + try { - ScanResult scanResult = invoker.invoke(() -> ddbClient.scan(scanRequest)); - List> items = scanResult.getItems(); - Set discoveredColumns = new HashSet<>(); - if (!items.isEmpty()) { + ScanResponse scanResponse = invoker.invoke(() -> ddbClient.scan(scanRequest)); + if (!scanResponse.items().isEmpty()) { + List> items = scanResponse.items(); + Set discoveredColumns = new HashSet<>(); + for (Map item : items) { for (Map.Entry column : item.entrySet()) { if (!discoveredColumns.contains(column.getKey())) { - Field field = DDBTypeUtils.inferArrowField(column.getKey(), ItemUtils.toSimpleValue(column.getValue())); + Field field = DDBTypeUtils.inferArrowField(column.getKey(), column.getValue()); if (field != null) { schemaBuilder.addField(field); discoveredColumns.add(column.getKey()); @@ -163,16 +172,16 @@ public static Schema peekTableForSchema(String tableName, ThrottlingInvoker invo // there's no items, so use any attributes defined in the table metadata DynamoDBTable table = getTable(tableName, invoker, ddbClient); for (AttributeDefinition attributeDefinition : table.getKnownAttributeDefinitions()) { - schemaBuilder.addField(DDBTypeUtils.getArrowFieldFromDDBType(attributeDefinition.getAttributeName(), attributeDefinition.getAttributeType())); + schemaBuilder.addField(DDBTypeUtils.getArrowFieldFromDDBType(attributeDefinition.attributeName(), attributeDefinition.attributeType().toString())); } } } - catch (AmazonDynamoDBException amazonDynamoDBException) { - if (amazonDynamoDBException.getMessage().contains("AWSKMSException")) { - logger.warn("Failed to retrieve table schema due to KMS issue, empty schema for table: {}. Error Message: {}", tableName, amazonDynamoDBException.getMessage()); + catch (RuntimeException runtimeException) { + if (runtimeException.getMessage().contains("AWSKMSException")) { + logger.warn("Failed to retrieve table schema due to KMS issue, empty schema for table: {}. Error Message: {}", tableName, runtimeException.getMessage()); } else { - throw amazonDynamoDBException; + throw runtimeException; } } return schemaBuilder.build(); diff --git a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTypeUtils.java b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTypeUtils.java index e81e1e8ccc..3408be4a8e 100644 --- a/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTypeUtils.java +++ b/athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTypeUtils.java @@ -31,8 +31,6 @@ import com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarBinaryHolder; import com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector; import com.amazonaws.athena.connectors.dynamodb.resolver.DynamoDBFieldResolver; -import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.types.Types; @@ -43,20 +41,32 @@ import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.enhanced.dynamodb.internal.converter.attribute.BigDecimalAttributeConverter; +import software.amazon.awssdk.enhanced.dynamodb.internal.converter.attribute.BooleanAttributeConverter; +import software.amazon.awssdk.enhanced.dynamodb.internal.converter.attribute.ByteArrayAttributeConverter; +import software.amazon.awssdk.enhanced.dynamodb.internal.converter.attribute.EnhancedAttributeValue; +import software.amazon.awssdk.enhanced.dynamodb.internal.converter.attribute.StringAttributeConverter; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.utils.ImmutableMap; import java.math.BigDecimal; +import java.nio.ByteBuffer; import java.time.LocalDateTime; import java.time.ZoneId; import java.time.format.DateTimeFormatter; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.TreeMap; +import java.util.stream.Collectors; /** * Provides utility methods relating to type handling. @@ -86,53 +96,72 @@ private DDBTypeUtils() {} * @param value the value of the field * @return the inferred Arrow field */ - public static Field inferArrowField(String key, Object value) + public static Field inferArrowField(String key, AttributeValue value) { logger.debug("inferArrowField invoked for key {} of class {}", key, - value != null ? value.getClass() : null); - if (value == null) { + value != null ? value.toString() : null); + + EnhancedAttributeValue enhancedAttributeValue = EnhancedAttributeValue.fromAttributeValue(value); + if (enhancedAttributeValue == null || enhancedAttributeValue.isNull()) { return null; } - if (value instanceof String) { + if (enhancedAttributeValue.isString()) { return new Field(key, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); } - else if (value instanceof byte[]) { + else if (enhancedAttributeValue.isBytes()) { return new Field(key, FieldType.nullable(Types.MinorType.VARBINARY.getType()), null); } - else if (value instanceof Boolean) { + else if (enhancedAttributeValue.isBoolean()) { return new Field(key, FieldType.nullable(Types.MinorType.BIT.getType()), null); } - else if (value instanceof BigDecimal) { + else if (enhancedAttributeValue.isNumber()) { return new Field(key, FieldType.nullable(new ArrowType.Decimal(38, 9)), null); } - else if (value instanceof List || value instanceof Set) { + else if (enhancedAttributeValue.isSetOfBytes() || enhancedAttributeValue.isSetOfNumbers() || enhancedAttributeValue.isSetOfStrings()) { + Field child = null; + + if (enhancedAttributeValue.isSetOfBytes()) { + child = new Field(key, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); + } + else { + child = new Field(key, FieldType.nullable(Types.MinorType.VARBINARY.getType()), null); + } + + return child == null + ? null + : new Field(key, FieldType.nullable(Types.MinorType.LIST.getType()), + Collections.singletonList(child)); + } + else if (enhancedAttributeValue.isListOfAttributeValues()) { Field child = null; - if (((Collection) value).isEmpty()) { + List listOfAttributes = enhancedAttributeValue.asListOfAttributeValues(); + + if (listOfAttributes.isEmpty()) { logger.warn("Automatic schema inference encountered empty List or Set {}. Unable to determine element types. Falling back to VARCHAR representation", key); - child = inferArrowField("", ""); + child = inferArrowField("", EnhancedAttributeValue.nullValue().toAttributeValue()); } else { - Iterator iterator = ((Collection) value).iterator(); - Object previousValue = iterator.next(); + Iterator iterator = listOfAttributes.iterator(); + EnhancedAttributeValue previousValue = EnhancedAttributeValue.fromAttributeValue(iterator.next()); boolean allElementsAreSameType = true; while (iterator.hasNext()) { - Object currentValue = iterator.next(); + EnhancedAttributeValue currentValue = EnhancedAttributeValue.fromAttributeValue(iterator.next()); // null is considered the same as any prior type - if (previousValue != null && currentValue != null && !previousValue.getClass().equals(currentValue.getClass())) { + if (!previousValue.isNull() && !currentValue.isNull() && !previousValue.type().equals(currentValue.type())) { allElementsAreSameType = false; break; } // only update the previousValue if the currentValue is not null otherwise we will end // up inferring the type as null if the currentValue is null and is the last value. - previousValue = (currentValue == null) ? previousValue : currentValue; + previousValue = (currentValue.isNull()) ? previousValue : currentValue; } if (allElementsAreSameType) { - child = inferArrowField(key + ".element", previousValue); + child = inferArrowField(key + ".element", previousValue.toAttributeValue()); } else { logger.warn("Automatic schema inference encountered List or Set {} containing multiple element types. Falling back to VARCHAR representation of elements", key); - child = inferArrowField("", ""); + child = inferArrowField("", AttributeValue.builder().s("").build()); } } return child == null @@ -140,12 +169,12 @@ else if (value instanceof List || value instanceof Set) { : new Field(key, FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(child)); } - else if (value instanceof Map) { + else if (enhancedAttributeValue.isMap()) { List children = new ArrayList<>(); // keys are always Strings in DDB's case - Map doc = (Map) value; + Map doc = value.m(); for (String childKey : doc.keySet()) { - Object childVal = doc.get(childKey); + AttributeValue childVal = doc.get(childKey); Field child = inferArrowField(childKey, childVal); if (child != null) { children.add(child); @@ -161,6 +190,7 @@ else if (value instanceof Map) { return new Field(key, FieldType.nullable(Types.MinorType.STRUCT.getType()), children); } + // #todo; fix this String className = (value == null || value.getClass() == null) ? "null" : value.getClass().getName(); throw new RuntimeException("Unknown type[" + className + "] for field[" + key + "]"); } @@ -401,7 +431,7 @@ public static Optional makeExtractor(Field field, DDBRecordMetadata r case DECIMAL: return Optional.of((DecimalExtractor) (Object context, NullableDecimalHolder dst) -> { - Object value = ItemUtils.toSimpleValue(contextAsMap(context, caseInsensitive).get(field.getName())); + Object value = toSimpleValue(contextAsMap(context, caseInsensitive).get(field.getName())); if (value != null) { dst.isSet = 1; dst.value = (BigDecimal) value; @@ -414,7 +444,7 @@ public static Optional makeExtractor(Field field, DDBRecordMetadata r return Optional.of((VarBinaryExtractor) (Object context, NullableVarBinaryHolder dst) -> { Map item = contextAsMap(context, caseInsensitive); - Object value = ItemUtils.toSimpleValue(item.get(field.getName())); + Object value = toSimpleValue(item.get(field.getName())); value = DDBTypeUtils.coerceValueToExpectedType(value, field, fieldType, recordMetadata); if (value != null) { @@ -431,7 +461,7 @@ public static Optional makeExtractor(Field field, DDBRecordMetadata r AttributeValue attributeValue = (contextAsMap(context, caseInsensitive)).get(field.getName()); if (attributeValue != null) { dst.isSet = 1; - dst.value = attributeValue.getBOOL() ? 1 : 0; + dst.value = attributeValue.bool() ? 1 : 0; } else { dst.isSet = 0; @@ -459,28 +489,19 @@ public static FieldWriterFactory makeFactory(Field field, DDBRecordMetadata reco (FieldWriter) (Object context, int rowNum) -> { Map item = contextAsMap(context, caseInsensitive); - Object value = ItemUtils.toSimpleValue(item.get(field.getName())); + Object value = toSimpleValue(item.get(field.getName())); List valueAsList = value != null ? DDBTypeUtils.coerceListToExpectedType(value, field, recordMetadata) : null; BlockUtils.setComplexValue(vector, rowNum, resolver, valueAsList); return true; }; case STRUCT: - return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> - (FieldWriter) (Object context, int rowNum) -> - { - Map item = contextAsMap(context, caseInsensitive); - Object value = ItemUtils.toSimpleValue(item.get(field.getName())); - value = DDBTypeUtils.coerceValueToExpectedType(value, field, fieldType, recordMetadata); - BlockUtils.setComplexValue(vector, rowNum, resolver, value); - return true; - }; case MAP: return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> { Map item = contextAsMap(context, caseInsensitive); - Object value = ItemUtils.toSimpleValue(item.get(field.getName())); + Object value = toSimpleValue(item.get(field.getName())); value = DDBTypeUtils.coerceValueToExpectedType(value, field, fieldType, recordMetadata); BlockUtils.setComplexValue(vector, rowNum, resolver, value); return true; @@ -491,11 +512,170 @@ public static FieldWriterFactory makeFactory(Field field, DDBRecordMetadata reco (FieldWriter) (Object context, int rowNum) -> { Map item = contextAsMap(context, caseInsensitive); - Object value = ItemUtils.toSimpleValue(item.get(field.getName())); + Object value = toSimpleValue(item.get(field.getName())); value = DDBTypeUtils.coerceValueToExpectedType(value, field, fieldType, recordMetadata); BlockUtils.setValue(vector, rowNum, value); return true; }; } } + + public static T toSimpleValue(AttributeValue value) + { + if (value == null) { + return null; + } + EnhancedAttributeValue enhancedAttributeValue = EnhancedAttributeValue.fromAttributeValue(value); + T result = null; + Object child = null; + + switch (enhancedAttributeValue.type()) { + case NULL: + break; + case BOOL: + result = (T) Boolean.valueOf(enhancedAttributeValue.asBoolean()); + break; + case S: + result = (T) StringAttributeConverter.create().transformTo(value); + break; + case N: + result = (T) BigDecimalAttributeConverter.create().transformTo(value); + break; + case B: + result = (T) enhancedAttributeValue.asBytes().asByteArray(); + break; + case SS: + result = (T) enhancedAttributeValue.asSetOfStrings(); + break; + case NS: + result = (T) value.ns().stream().map(BigDecimal::new).collect(Collectors.toList()); + break; + case BS: + result = (T) value.bs().stream().map(sdkBytes -> sdkBytes.asByteArray()).collect(Collectors.toList()); + break; + case L: + result = handleListAttribute(enhancedAttributeValue); + break; + case M: + result = handleMapAttribute(enhancedAttributeValue); + break; + } + return result; + } + + private static T handleMapAttribute(EnhancedAttributeValue enhancedAttributeValue) + { + Map valueMap = enhancedAttributeValue.asMap(); + if (valueMap.isEmpty()) { + return (T) Collections.emptyMap(); + } + Map result = new HashMap<>(valueMap.size()); + for (Map.Entry entry : valueMap.entrySet()) { + String key = entry.getKey(); + AttributeValue attributeValue = entry.getValue(); + result.put(key, toSimpleValue(attributeValue)); + } + return (T) result; + } + + private static T handleListAttribute(EnhancedAttributeValue enhancedAttributeValue) + { + List result = + enhancedAttributeValue.asListOfAttributeValues().stream().map(attributeValue -> toSimpleValue(attributeValue)).collect(Collectors.toList()); + return (T) result; + } + + public static AttributeValue toAttributeValue(Object value) + { + if (value == null) { + return AttributeValue.builder().nul(true).build(); + } + else if (value instanceof String) { + return StringAttributeConverter.create().transformFrom((String) value); + } + else if (value instanceof Boolean) { + return BooleanAttributeConverter.create().transformFrom((Boolean) value); + } + else if (value instanceof BigDecimal) { + return BigDecimalAttributeConverter.create().transformFrom((BigDecimal) value); + } + else if (value instanceof Number) { + // For other numbers, we can convert them to BigDecimal first + return BigDecimalAttributeConverter.create().transformFrom(BigDecimal.valueOf(((Number) value).doubleValue())); + } + else if (value instanceof byte[]) { + return ByteArrayAttributeConverter.create().transformFrom((byte[]) value); + } + else if (value instanceof ByteBuffer) { + SdkBytes b = SdkBytes.fromByteBuffer((ByteBuffer) value); + return AttributeValue.builder().b(b).build(); + } + else if (value instanceof Set) { + return handleSetType((Set) value); + } + else if (value instanceof List) { + return handleListType((List) value); + } + else if (value instanceof Map) { + return handleMapType((Map) value); + } + else { + throw new UnsupportedOperationException("Unsupported value type: " + value.getClass()); + } + } + + public static String attributeToJson(AttributeValue attributeValue, String key) + { + EnhancedDocument enhancedDocument = EnhancedDocument.fromAttributeValueMap(ImmutableMap.of(key, attributeValue)); + return enhancedDocument.toJson(); + } + + public static AttributeValue jsonToAttributeValue(String jsonString, String key) + { + EnhancedDocument enhancedDocument = EnhancedDocument.fromJson(jsonString); + if (!enhancedDocument.isPresent(key)) { + throw new RuntimeException("Unknown attribute Key"); + } + return enhancedDocument.toMap().get(key); + } + + private static AttributeValue handleSetType(Set value) + { + // Check the type of the first element as a sample + Object firstElement = value.iterator().next(); + if (firstElement instanceof String) { + // Handle String Set + Set stringSet = value.stream().map(e -> (String) e).collect(Collectors.toSet()); + return AttributeValue.builder().ss(stringSet).build(); + } + else if (firstElement instanceof Number) { + // Handle Number Set + Set numberSet = value.stream() + .map(e -> String.valueOf(((Number) e).doubleValue())) // Convert numbers to strings + .collect(Collectors.toSet()); + return AttributeValue.builder().ns(numberSet).build(); + } // Add other types if needed + + // Fallback for unsupported set types + throw new UnsupportedOperationException("Unsupported Set element type: " + firstElement.getClass()); + } + + private static AttributeValue handleListType(List value) + { + List attributeList = new ArrayList<>(); + for (Object element : value) { + attributeList.add(toAttributeValue(element)); + } + return AttributeValue.builder().l(attributeList).build(); + } + + private static AttributeValue handleMapType(Map value) + { + Map attributeMap = new HashMap<>(); + for (Map.Entry entry : value.entrySet()) { + // Convert each value in the map to an AttributeValue + attributeMap.put(entry.getKey(), toAttributeValue(entry.getValue())); + } + return AttributeValue.builder().m(attributeMap).build(); + } } diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBPredicateUtilsTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBPredicateUtilsTest.java index 8280354126..5b131fabc9 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBPredicateUtilsTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBPredicateUtilsTest.java @@ -26,8 +26,7 @@ import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex; import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable; import com.amazonaws.athena.connectors.dynamodb.util.DDBPredicateUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; -import com.amazonaws.services.dynamodbv2.model.ProjectionType; + import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.junit.Test; @@ -35,6 +34,9 @@ import org.mockito.junit.MockitoJUnitRunner; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; +import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType; import java.util.Optional; @@ -73,9 +75,9 @@ public void testGetBestIndexForPredicatesWithGSIProjectionTypeInclude() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-gsi", "col0", Optional.empty(), ProjectionType.INCLUDE, ImmutableList.of("col1")) ), 1000, 10, 5); @@ -93,9 +95,9 @@ public void testGetBestIndexForPredicatesWithGSIProjectionTypeKeysOnly() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-gsi", "col0", Optional.empty(), ProjectionType.KEYS_ONLY, ImmutableList.of()) ), 1000, 10, 5); @@ -112,9 +114,9 @@ public void testGetBestIndexForPredicatesWithNonEqualityPredicate() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-gsi", "col0", Optional.empty(), ProjectionType.KEYS_ONLY, ImmutableList.of()) ), 1000, 10, 5); @@ -129,9 +131,9 @@ public void testGetBestIndexForPredicatesWithGSIProjectionTypeAll() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-gsi", "col0", Optional.of("col1"), ProjectionType.ALL, ImmutableList.of()) ), 1000, 10, 5); @@ -154,9 +156,9 @@ public void testGetBestIndexForPredicatesWithLSI() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-lsi", "hashKey", Optional.of("col0"), ProjectionType.ALL, ImmutableList.of()) ), 1000, 10, 5); @@ -174,10 +176,10 @@ public void testGetBestIndexForPredicatesWithMultipleIndices() ValueSet singleValueSet = SortedRangeSet.of(Range.equal(new BlockAllocatorImpl(), VARCHAR.getType(), "value")); DynamoDBTable table = new DynamoDBTable("tableName", "hashKey", Optional.of("sortKey"), ImmutableList.of( - new AttributeDefinition("hashKey", "S"), - new AttributeDefinition("sortKey", "S"), - new AttributeDefinition("col0", "S"), - new AttributeDefinition("col1", "S")), + AttributeDefinition.builder().attributeName("hashKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("sortKey").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col0").attributeType(ScalarAttributeType.S).build(), + AttributeDefinition.builder().attributeName("col1").attributeType(ScalarAttributeType.S).build()), ImmutableList.of( new DynamoDBIndex("col0-gsi", "col0", Optional.empty(), ProjectionType.INCLUDE, ImmutableList.of("col1")), new DynamoDBIndex("col1-gsi", "col1", Optional.empty(), ProjectionType.INCLUDE, ImmutableList.of("col2")), diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBTypeUtilsTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBTypeUtilsTest.java index 3fb505f0e0..5fc61fdf05 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBTypeUtilsTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DDBTypeUtilsTest.java @@ -30,7 +30,7 @@ import com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder; import com.amazonaws.athena.connectors.dynamodb.util.DDBRecordMetadata; import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; + import com.google.common.collect.ImmutableMap; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.types.Types; @@ -44,6 +44,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.apache.arrow.vector.types.pojo.Field; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.io.IOException; import java.math.BigDecimal; @@ -89,10 +91,8 @@ public void makeDecimalExtractorTest() String literalValue = "12345"; String literalValue2 = "789.1234"; - AttributeValue myValue = new AttributeValue(); - myValue.setN(literalValue); - AttributeValue myValue2 = new AttributeValue(); - myValue2.setN(literalValue2); + AttributeValue myValue = AttributeValue.builder().n(literalValue).build(); + AttributeValue myValue2 = AttributeValue.builder().n(literalValue2).build(); Map testValue = ImmutableMap.of(col1, myValue, col2, myValue2); @@ -121,10 +121,13 @@ public void makeVarBinaryExtractorTest() byte[] byteValue2 = "World!".getBytes(); ByteBuffer byteBuffer2 = ByteBuffer.wrap(byteValue2); - AttributeValue myValue = new AttributeValue(); - myValue.setB(byteBuffer1); - AttributeValue myValue2 = new AttributeValue(); - myValue2.setB(byteBuffer2); + // Creating AttributeValue with binary data in SDK v2 + AttributeValue myValue = AttributeValue.builder() + .b(SdkBytes.fromByteBuffer(byteBuffer1)) + .build(); + AttributeValue myValue2 = AttributeValue.builder() + .b(SdkBytes.fromByteBuffer(byteBuffer2)) + .build(); Map testValue = ImmutableMap.of(col1, myValue, col2, myValue2); @@ -149,10 +152,12 @@ public void makeBitExtractorTest() .addField(col2, Types.MinorType.BIT.getType()) .build(); - AttributeValue myValue = new AttributeValue(); - myValue.setBOOL(true); - AttributeValue myValue2 = new AttributeValue(); - myValue2.setBOOL(false); + AttributeValue myValue = AttributeValue.builder() + .bool(true) + .build(); + AttributeValue myValue2 = AttributeValue.builder() + .bool(false) + .build(); Map testValue = ImmutableMap.of(col1, myValue, col2, myValue2); @@ -173,7 +178,7 @@ public void inferArrowFieldListWithNullTest() throws Exception inputArray.add(null); inputArray.add("value3"); - Field testField = DDBTypeUtils.inferArrowField("asdf", inputArray); + Field testField = DDBTypeUtils.inferArrowField("asdf", DDBTypeUtils.toAttributeValue(inputArray)); assertEquals("Type does not match!", ArrowType.List.INSTANCE, testField.getType()); assertEquals("Children Length Off!", 1, testField.getChildren().size()); diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DdbTableUtils.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DdbTableUtils.java index ccbf812044..ac489514ea 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DdbTableUtils.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DdbTableUtils.java @@ -20,12 +20,9 @@ package com.amazonaws.athena.connectors.dynamodb; import com.amazonaws.AmazonServiceException; -import com.amazonaws.services.dynamodbv2.model.ResourceNotFoundException; + import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awscdk.core.RemovalPolicy; @@ -34,9 +31,11 @@ import software.amazon.awscdk.services.dynamodb.AttributeType; import software.amazon.awscdk.services.dynamodb.BillingMode; import software.amazon.awscdk.services.dynamodb.Table; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; +import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException; -import java.util.ArrayList; -import java.util.List; import java.util.Map; /** @@ -52,10 +51,10 @@ public class DdbTableUtils { private static final long READ_CAPACITY_UNITS = 10L; private static final long WRITE_CAPACITY_UNITS = 10L; - private final AmazonDynamoDB client; + private final DynamoDbClient client; public DdbTableUtils() { - client = AmazonDynamoDBClientBuilder.defaultClient(); + client = DynamoDbClient.create(); } /** @@ -87,11 +86,15 @@ protected void putItem(String tableName, Map item) try { // Add record to table in DynamoDB service. logger.info("Add item attempt: {}", attempt); - client.putItem(tableName, item); + PutItemRequest putItemRequest = PutItemRequest.builder() + .tableName(tableName) + .item(item) + .build(); + client.putItem(putItemRequest); logger.info("Added item in {} attempt(s).", attempt); break; } catch (ResourceNotFoundException e) { - logger.info(e.getErrorMessage()); + logger.info(e.getMessage()); if (attempt < MAX_TRIES) { // Sleep for 10 seconds and try again. logger.info("Sleeping for 10 seconds..."); diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java index 5228664d75..2ed5bd668e 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java @@ -43,9 +43,10 @@ import com.amazonaws.athena.connector.lambda.metadata.MetadataRequestType; import com.amazonaws.athena.connector.lambda.metadata.MetadataResponse; import com.amazonaws.athena.connector.lambda.security.LocalKeyFactory; +import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; import com.amazonaws.services.athena.AmazonAthena; + import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.amazonaws.services.glue.AWSGlue; import com.amazonaws.services.glue.model.Column; import com.amazonaws.services.glue.model.Database; @@ -63,6 +64,7 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.After; import org.junit.Before; +import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TestName; @@ -71,6 +73,8 @@ import org.mockito.junit.MockitoJUnitRunner; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.time.Instant; import java.time.LocalDateTime; @@ -302,7 +306,6 @@ public void doGetTableLayoutScan() constraintsMap.put("col_3", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true) .add(true).build()); - GetTableLayoutRequest req = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, @@ -310,7 +313,6 @@ public void doGetTableLayoutScan() new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET); - GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req); logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema()); @@ -326,8 +328,8 @@ public void doGetTableLayoutScan() ImmutableMap expressionNames = ImmutableMap.of("#col_3", "col_3"); assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames))); - ImmutableMap expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null)); - assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues))); + ImmutableMap expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(true), ":v1", DDBTypeUtils.toAttributeValue(null)); + assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(EnhancedDocument.fromAttributeValueMap(expressionValues).toJson())); } @Test @@ -371,8 +373,8 @@ public void doGetTableLayoutQueryIndex() ImmutableMap expressionNames = ImmutableMap.of("#col_4", "col_4", "#col_5", "col_5"); assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames))); - ImmutableMap expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(startTime), ":v1", ItemUtils.toAttributeValue(endTime)); - assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues))); + ImmutableMap expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(startTime), ":v1", DDBTypeUtils.toAttributeValue(endTime)); + assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(EnhancedDocument.fromAttributeValueMap(expressionValues).toJson())); // Tests to validate that we correctly generate predicates that avoid this error: // "KeyConditionExpressions must only contain one condition per key" @@ -688,7 +690,7 @@ public void doGetTableLayoutScanWithTypeOverride() ImmutableMap expressionNames = ImmutableMap.of("#col3", "col3", "#col2", "col2"); assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames))); - ImmutableMap expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null)); - assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues))); + ImmutableMap expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(true), ":v1", DDBTypeUtils.toAttributeValue(null)); + assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(EnhancedDocument.fromAttributeValueMap(expressionValues).toJson())); } } diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java index a601c7728a..d9f3f421b1 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java @@ -37,8 +37,8 @@ import com.amazonaws.athena.connector.lambda.records.RecordResponse; import com.amazonaws.athena.connector.lambda.security.EncryptionKeyFactory; import com.amazonaws.athena.connector.lambda.security.LocalKeyFactory; +import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; import com.amazonaws.services.athena.AmazonAthena; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.amazonaws.services.glue.AWSGlue; import com.amazonaws.services.glue.model.Column; import com.amazonaws.services.glue.model.EntityNotFoundException; @@ -67,6 +67,8 @@ import org.mockito.junit.MockitoJUnitRunner; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.time.LocalDate; import java.time.LocalDateTime; @@ -248,14 +250,14 @@ public void testReadScanSplitFiltered() throws Exception { Map expressionNames = ImmutableMap.of("#col_6", "col_6"); - Map expressionValues = ImmutableMap.of(":v0", toAttributeValue(0), ":v1", toAttributeValue(1)); + Map expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(0), ":v1", DDBTypeUtils.toAttributeValue(1)); Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()) .add(TABLE_METADATA, TEST_TABLE) .add(SEGMENT_ID_PROPERTY, "0") .add(SEGMENT_COUNT_METADATA, "1") .add(NON_KEY_FILTER_METADATA, "NOT #col_6 IN (:v0,:v1)") .add(EXPRESSION_NAMES_METADATA, toJsonString(expressionNames)) - .add(EXPRESSION_VALUES_METADATA, toJsonString(expressionValues)) + .add(EXPRESSION_VALUES_METADATA, EnhancedDocument.fromAttributeValueMap(expressionValues).toJson()) .build(); ReadRecordsRequest request = new ReadRecordsRequest( @@ -285,14 +287,14 @@ public void testReadQuerySplit() throws Exception { Map expressionNames = ImmutableMap.of("#col_1", "col_1"); - Map expressionValues = ImmutableMap.of(":v0", toAttributeValue(1)); + Map expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(1)); Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()) .add(TABLE_METADATA, TEST_TABLE) .add(HASH_KEY_NAME_METADATA, "col_0") - .add("col_0", toJsonString(toAttributeValue("test_str_0"))) + .add("col_0", DDBTypeUtils.attributeToJson(DDBTypeUtils.toAttributeValue("test_str_0"), "col_0")) .add(RANGE_KEY_FILTER_METADATA, "#col_1 >= :v0") .add(EXPRESSION_NAMES_METADATA, toJsonString(expressionNames)) - .add(EXPRESSION_VALUES_METADATA, toJsonString(expressionValues)) + .add(EXPRESSION_VALUES_METADATA, EnhancedDocument.fromAttributeValueMap(expressionValues).toJson()) .build(); ReadRecordsRequest request = new ReadRecordsRequest( @@ -322,14 +324,14 @@ public void testReadQuerySplitWithLimit() throws Exception { Map expressionNames = ImmutableMap.of("#col_1", "col_1"); - Map expressionValues = ImmutableMap.of(":v0", toAttributeValue(1)); + Map expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(1)); Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()) .add(TABLE_METADATA, TEST_TABLE) .add(HASH_KEY_NAME_METADATA, "col_0") - .add("col_0", toJsonString(toAttributeValue("test_str_0"))) + .add("col_0", DDBTypeUtils.attributeToJson(DDBTypeUtils.toAttributeValue("test_str_0"), "col_0")) .add(RANGE_KEY_FILTER_METADATA, "#col_1 >= :v0") .add(EXPRESSION_NAMES_METADATA, toJsonString(expressionNames)) - .add(EXPRESSION_VALUES_METADATA, toJsonString(expressionValues)) + .add(EXPRESSION_VALUES_METADATA, EnhancedDocument.fromAttributeValueMap(expressionValues).toJson()) .build(); ReadRecordsRequest request = new ReadRecordsRequest( @@ -359,14 +361,14 @@ public void testZeroRowQuery() throws Exception { Map expressionNames = ImmutableMap.of("#col_1", "col_1"); - Map expressionValues = ImmutableMap.of(":v0", toAttributeValue(1)); + Map expressionValues = ImmutableMap.of(":v0", DDBTypeUtils.toAttributeValue(1)); Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()) .add(TABLE_METADATA, TEST_TABLE) .add(HASH_KEY_NAME_METADATA, "col_0") - .add("col_0", toJsonString(toAttributeValue("test_str_999999"))) + .add("col_0", DDBTypeUtils.attributeToJson(DDBTypeUtils.toAttributeValue("test_str_999999"), "col_0")) .add(RANGE_KEY_FILTER_METADATA, "#col_1 >= :v0") .add(EXPRESSION_NAMES_METADATA, toJsonString(expressionNames)) - .add(EXPRESSION_VALUES_METADATA, toJsonString(expressionValues)) + .add(EXPRESSION_VALUES_METADATA, EnhancedDocument.fromAttributeValueMap(expressionValues).toJson()) .build(); ReadRecordsRequest request = new ReadRecordsRequest( diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDbIntegTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDbIntegTest.java index a514058df6..4ab03112ea 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDbIntegTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDbIntegTest.java @@ -20,9 +20,7 @@ package com.amazonaws.athena.connectors.dynamodb; import com.amazonaws.athena.connector.integ.IntegrationTestBase; -import com.amazonaws.services.athena.model.Datum; import com.amazonaws.services.athena.model.Row; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.slf4j.Logger; @@ -33,6 +31,8 @@ import software.amazon.awscdk.services.iam.Effect; import software.amazon.awscdk.services.iam.PolicyDocument; import software.amazon.awscdk.services.iam.PolicyStatement; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -43,7 +43,6 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; -import static org.testng.AssertJUnit.assertNull; /** * Integration-tests for the DynamoDB connector using the Integration-test module. @@ -88,17 +87,18 @@ protected void addMovieItems() private void addDatatypeItems() { Map item = new ImmutableMap.Builder() - .put("int_type", new AttributeValue().withN(String.valueOf(TEST_DATATYPES_INT_VALUE))) - .put("smallint_type", new AttributeValue().withN(String.valueOf(TEST_DATATYPES_SHORT_VALUE))) - .put("bigint_type", new AttributeValue().withN(String.valueOf(TEST_DATATYPES_LONG_VALUE))) - .put("varchar_type", new AttributeValue(TEST_DATATYPES_VARCHAR_VALUE)) - .put("boolean_type", new AttributeValue().withBOOL(TEST_DATATYPES_BOOLEAN_VALUE)) - .put("float4_type", new AttributeValue(String.valueOf(TEST_DATATYPES_SINGLE_PRECISION_VALUE))) - .put("float8_type", new AttributeValue(String.valueOf(1E-130))) // smallest number dynamo can handle - .put("date_type", new AttributeValue(TEST_DATATYPES_DATE_VALUE)) - .put("timestamp_type", new AttributeValue(TEST_DATATYPES_TIMESTAMP_VALUE)) - .put("byte_type", new AttributeValue().withB(ByteBuffer.wrap(TEST_DATATYPES_BYTE_ARRAY_VALUE))) - .put("textarray_type", new AttributeValue(TEST_DATATYPES_VARCHAR_ARRAY_VALUE)).build(); + .put("int_type", AttributeValue.builder().n(String.valueOf(TEST_DATATYPES_INT_VALUE)).build()) + .put("smallint_type", AttributeValue.builder().n(String.valueOf(TEST_DATATYPES_SHORT_VALUE)).build()) + .put("bigint_type", AttributeValue.builder().n(String.valueOf(TEST_DATATYPES_LONG_VALUE)).build()) + .put("varchar_type", AttributeValue.builder().s(TEST_DATATYPES_VARCHAR_VALUE).build()) + .put("boolean_type", AttributeValue.builder().bool(TEST_DATATYPES_BOOLEAN_VALUE).build()) + .put("float4_type", AttributeValue.builder().s(String.valueOf(TEST_DATATYPES_SINGLE_PRECISION_VALUE)).build()) + .put("float8_type", AttributeValue.builder().s(String.valueOf(1E-130)).build()) // smallest number dynamo can handle + .put("date_type", AttributeValue.builder().s(TEST_DATATYPES_DATE_VALUE).build()) + .put("timestamp_type", AttributeValue.builder().s(TEST_DATATYPES_TIMESTAMP_VALUE).build()) + .put("byte_type", AttributeValue.builder().b(SdkBytes.fromByteBuffer(ByteBuffer.wrap(TEST_DATATYPES_BYTE_ARRAY_VALUE))).build()) + .put("textarray_type", AttributeValue.builder().s(TEST_DATATYPES_VARCHAR_ARRAY_VALUE).build()) + .build(); tableUtils.putItem(TEST_DATATYPES_TABLE_NAME, item); } @@ -112,25 +112,24 @@ private void addDatatypeItems() * @param genre List of movie generes. */ private Map createMovieItem(String year, String title, String director, - List cast, List genre) + List cast, List genre) { logger.info("Add item: year=[{}], title=[{}], director=[{}], cast={}, genre={}", year, title, director, cast, genre); List castAttrib = new ArrayList<>(); - cast.forEach(value -> castAttrib.add(new AttributeValue(value))); + cast.forEach(value -> castAttrib.add(AttributeValue.builder().s(value).build())); List genreAttrib = new ArrayList<>(); - genre.forEach(value -> genreAttrib.add(new AttributeValue(value))); + genre.forEach(value -> genreAttrib.add(AttributeValue.builder().s(value).build()));; Map item = ImmutableMap.of( - "year", new AttributeValue().withN(year), - "title", new AttributeValue(title), - "info", new AttributeValue().withM(ImmutableMap.of( - "director", new AttributeValue(director), - "cast", new AttributeValue().withL(castAttrib), - "genre", new AttributeValue().withL(genreAttrib)))); - + "year", AttributeValue.builder().n(year).build(), + "title", AttributeValue.builder().s(title).build(), + "info", AttributeValue.builder().m(ImmutableMap.of( + "director", AttributeValue.builder().s(director).build(), + "cast", AttributeValue.builder().l(castAttrib).build(), + "genre", AttributeValue.builder().l(genreAttrib).build())).build()); return item; } diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/TestBase.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/TestBase.java index f1801c2d1e..e36c43912b 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/TestBase.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/TestBase.java @@ -21,43 +21,64 @@ import com.amazonaws.athena.connector.lambda.ThrottlingInvoker; import com.amazonaws.athena.connector.lambda.domain.TableName; - import com.amazonaws.athena.connector.lambda.security.FederatedIdentity; import com.amazonaws.athena.connectors.dynamodb.util.DDBTableUtils; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; -import com.amazonaws.services.dynamodbv2.document.Index; -import com.amazonaws.services.dynamodbv2.document.Table; -import com.amazonaws.services.dynamodbv2.document.TableWriteItems; -import com.amazonaws.services.dynamodbv2.local.embedded.DynamoDBEmbedded; -import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.CreateGlobalSecondaryIndexAction; -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.Projection; -import com.amazonaws.services.dynamodbv2.model.ProjectionType; -import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput; -import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType; +import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils; + +import com.amazonaws.services.dynamodbv2.local.main.ServerRunner; +import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer; + import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.logging.log4j.Level; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.core.config.Configurator; import org.junit.AfterClass; import org.junit.BeforeClass; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.waiters.WaiterResponse; +import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemRequest; +import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemResponse; +import software.amazon.awssdk.services.dynamodb.model.CreateGlobalSecondaryIndexAction; +import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.CreateTableResponse; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse; +import software.amazon.awssdk.services.dynamodb.model.DynamoDbException; +import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndexUpdate; +import software.amazon.awssdk.services.dynamodb.model.IndexStatus; +import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement; +import software.amazon.awssdk.services.dynamodb.model.KeyType; +import software.amazon.awssdk.services.dynamodb.model.Projection; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput; +import software.amazon.awssdk.services.dynamodb.model.PutRequest; +import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType; +import software.amazon.awssdk.services.dynamodb.model.UpdateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.WriteRequest; +import software.amazon.awssdk.services.dynamodb.waiters.DynamoDbWaiter; + +import java.net.URI; import java.sql.Timestamp; import java.time.LocalDateTime; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants.DEFAULT_SCHEMA; import static com.amazonaws.athena.connectors.dynamodb.throttling.DynamoDBExceptionFilter.EXCEPTION_FILTER; -import static com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue; -import static com.amazonaws.services.dynamodbv2.document.ItemUtils.toItem; public class TestBase { @@ -81,277 +102,532 @@ public class TestBase protected static final TableName TEST_TABLE_7_NAME = new TableName(DEFAULT_SCHEMA, TEST_TABLE7); protected static final TableName TEST_TABLE_8_NAME = new TableName(DEFAULT_SCHEMA, TEST_TABLE8); - protected static AmazonDynamoDB ddbClient; + protected static DynamoDbClient ddbClient; protected static Schema schema; - protected static Table tableDdbNoGlue; + protected static DynamoDBProxyServer server; @BeforeClass public static void setupOnce() throws Exception { - ddbClient = setupDatabase(); + ddbClient = setupLocalDDB(); ThrottlingInvoker invoker = ThrottlingInvoker.newDefaultBuilder(EXCEPTION_FILTER, com.google.common.collect.ImmutableMap.of()).build(); schema = DDBTableUtils.peekTableForSchema(TEST_TABLE, invoker, ddbClient); } @AfterClass - public static void tearDownOnce() - { - ddbClient.shutdown(); + public static void tearDownOnce() throws Exception { + server.stop(); } - private static AmazonDynamoDB setupDatabase() throws InterruptedException + private static void waitForActive(DynamoDbClient ddbClient, String tableName, String indexName) { - System.setProperty("sqlite4java.library.path", "native-libs"); - AmazonDynamoDB client = DynamoDBEmbedded.create().amazonDynamoDB(); - DynamoDB ddb = new DynamoDB(client); + boolean isIndexActive = false; + while (!isIndexActive) { + DescribeTableResponse tableDescription = ddbClient.describeTable(DescribeTableRequest.builder() + .tableName(tableName) + .build()); + isIndexActive = tableDescription.table().globalSecondaryIndexes().stream() + .anyMatch(index -> index.indexName().equals(indexName) && index.indexStatus().equals(IndexStatus.ACTIVE)); + + if (!isIndexActive) { + try { + Thread.sleep(2000); // Wait for 20 seconds before the next check + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + // Handle interruption + } + } + } + } + + private static String waitForTableToBecomeActive(DynamoDbClient ddb, DynamoDbWaiter dbWaiter, CreateTableResponse createTableResponse, String tableName) { + System.out.println("------------------------------------------------------------------------"); + System.out.println("Entered: waitForTableToBecomeActive: for: " + tableName); + + try { + DescribeTableRequest tableRequest = DescribeTableRequest.builder() + .tableName(tableName) + .build(); + + // Wait until the Amazon DynamoDB table is created + dbWaiter.waitUntilTableExists(tableRequest); + String name = createTableResponse.tableDescription().tableName(); + System.out.println("Table :[ " + name + "] is active"); + System.out.println("------------------------------------------------------------------------"); + return name; + + } catch (DynamoDbException e) { + System.err.println(e.getMessage()); + System.exit(1); + } + return tableName; + } - ArrayList attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("col_0").withAttributeType("S")); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("col_1").withAttributeType("N")); + public static void waitForGsiToBecomeActive(DynamoDbClient ddb, String tableName, String indexName) { + System.out.println("------------------------------------------------------------------------"); + System.out.println("Waiting for GSI to become active: " + indexName + " in table: " + tableName); - ArrayList keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("col_0").withKeyType(KeyType.HASH)); - keySchema.add(new KeySchemaElement().withAttributeName("col_1").withKeyType(KeyType.RANGE)); + boolean isIndexActive = false; - ProvisionedThroughput provisionedThroughput = new ProvisionedThroughput() - .withReadCapacityUnits(5L) - .withWriteCapacityUnits(6L); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); + while (!isIndexActive) { + try { + //This is a local Dynamo DB, so it should be reasonably fast. + Thread.sleep(5_000); // 5-second wait between checks - Table table = ddb.createTable(createTableRequest); + DescribeTableResponse response = ddb.describeTable(DescribeTableRequest.builder().tableName(tableName).build()); + isIndexActive = response.table().globalSecondaryIndexes().stream() + .anyMatch(gsi -> gsi.indexName().equals(indexName) && gsi.indexStatus().equals(IndexStatus.ACTIVE)); - table.waitForActive(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while waiting for the GSI to become active.", e); + } + } + + System.out.println("GSI " + indexName + " is now ACTIVE in table " + tableName); + System.out.println("------------------------------------------------------------------------"); + } + + //Adapted from: https://github.com/awslabs/amazon-dynamodb-local-samples/blob/main/DynamoDBLocal + private static DynamoDbClient setupLocalDDB() throws Exception + { + DynamoDbWaiter dbWaiter; + DynamoDbClient ddbClient; + Configurator.setAllLevels(LogManager.getRootLogger().getName(), Level.ERROR); + + try { + String port = "8000"; + String uri = "http://localhost:" + port; + // Create an in-memory and in-process instance of DynamoDB Local that runs over HTTP + final String[] localArgs = {"-inMemory", "-port", port}; + System.out.println("Starting DynamoDB Local..."); + server = ServerRunner.createServerFromCommandLineArgs(localArgs); + server.start(); + System.out.println("Started DynamoDB Local..."); + + + // Create a client and connect to DynamoDB Local + // Note: This is a dummy key and secret and AWS_ACCESS_KEY_ID can contain only letters (A–Z, a–z) and numbers (0–9). + ddbClient = DynamoDbClient.builder() + .endpointOverride(URI.create(uri)) + .httpClient(UrlConnectionHttpClient.builder().build()) + .region(Region.US_WEST_2) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("dummyKey", "dummySecret"))) + .build(); + + dbWaiter = ddbClient.waiter(); + + // Create ProvisionedThroughput using the builder pattern + ProvisionedThroughput provisionedThroughput = ProvisionedThroughput.builder() + .readCapacityUnits(5L) + .writeCapacityUnits(6L) + .build(); + + setUpTable1And2(provisionedThroughput, ddbClient, dbWaiter); + setUpTable3(provisionedThroughput, ddbClient, dbWaiter); + setUpTable4(provisionedThroughput, ddbClient, dbWaiter); + setUpTable5(provisionedThroughput, ddbClient, dbWaiter); + setUpTable6(provisionedThroughput, ddbClient, dbWaiter); + setUpTable7(provisionedThroughput, ddbClient, dbWaiter); + setUpTable8(provisionedThroughput, ddbClient, dbWaiter); + + return ddbClient; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static void setUpTable1And2( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) throws InterruptedException { + + // For AttributeDefinitions + List attributeDefinitions = new ArrayList<>(); + attributeDefinitions.add(AttributeDefinition.builder() + .attributeName("col_0") + .attributeType("S") + .build()); + attributeDefinitions.add(AttributeDefinition.builder() + .attributeName("col_1") + .attributeType("N") + .build()); + + // For KeySchema + List keySchema = new ArrayList<>(); + keySchema.add(KeySchemaElement.builder() + .attributeName("col_0") + .keyType(KeyType.HASH) + .build()); + keySchema.add(KeySchemaElement.builder() + .attributeName("col_1") + .keyType(KeyType.RANGE) + .build()); + + // Create CreateTableRequest using the builder pattern + CreateTableRequest createTableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE) + .keySchema(keySchema) + .attributeDefinitions(attributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); + + CreateTableResponse table1CreateTableResponse = ddb.createTable(createTableRequest); + String newTableName = waitForTableToBecomeActive(ddb, dbWaiter, table1CreateTableResponse, TEST_TABLE); + System.out.println("New Table Name: " + newTableName); - TableWriteItems tableWriteItems = new TableWriteItems(TEST_TABLE); int len = 1000; LocalDateTime dateTime = LocalDateTime.of(2019, 9, 23, 11, 18, 37); + List writeRequests = new ArrayList<>(); for (int i = 0; i < len; i++) { Map item = new HashMap<>(); - item.put("col_0", toAttributeValue("test_str_" + (i - i % 3))); - item.put("col_1", toAttributeValue(i)); + // Populating the item map + item.put("col_0", AttributeValue.builder().s("test_str_" + (i - i % 3)).build()); + item.put("col_1", AttributeValue.builder().n(String.valueOf(i)).build()); double doubleVal = 200000.0 + i / 2.0; if (Math.floor(doubleVal) != doubleVal) { - item.put("col_2", toAttributeValue(200000.0 + i / 2.0)); + item.put("col_2", AttributeValue.builder().n(String.valueOf(doubleVal)).build()); } - item.put("col_3", toAttributeValue(ImmutableMap.of("modulo", i % 2, "nextModulos", ImmutableList.of((i + 1) % 2, ((i + 2) % 2))))); - item.put("col_4", toAttributeValue(dateTime.toLocalDate().toEpochDay())); - item.put("col_5", toAttributeValue(Timestamp.valueOf(dateTime).toInstant().toEpochMilli())); - item.put("col_6", toAttributeValue(i % 128 == 0 ? null : i % 128)); - item.put("col_7", toAttributeValue(ImmutableList.of(-i, String.valueOf(i)))); - item.put("col_8", toAttributeValue(ImmutableList.of(ImmutableMap.of("mostlyEmptyMap", - i % 128 == 0 ? ImmutableMap.of("subtractions", ImmutableSet.of(i - 100, i - 200)) : ImmutableMap.of())))); - item.put("col_9", toAttributeValue(100.0f + i)); - item.put("col_10", toAttributeValue(ImmutableList.of(ImmutableList.of(1 * i, 2 * i, 3 * i), + // Assuming you have appropriate methods to create Map and List AttributeValues + item.put("col_3", DDBTypeUtils.toAttributeValue(ImmutableMap.of("modulo", i % 2, "nextModulos", ImmutableList.of((i + 1) % 2, ((i + 2) % 2))))); + item.put("col_4", AttributeValue.builder().n(String.valueOf(dateTime.toLocalDate().toEpochDay())).build()); + item.put("col_5", AttributeValue.builder().n(String.valueOf(Timestamp.valueOf(dateTime).toInstant().toEpochMilli())).build()); + item.put("col_6", i % 128 == 0 ? AttributeValue.builder().nul(true).build() : AttributeValue.builder().n(String.valueOf(i % 128)).build()); + item.put("col_7", DDBTypeUtils.toAttributeValue(ImmutableList.of(-i,String.valueOf(i)))); + item.put("col_8", DDBTypeUtils.toAttributeValue((ImmutableList.of(ImmutableMap.of("mostlyEmptyMap", + i % 128 == 0 ? ImmutableMap.of("subtractions", ImmutableSet.of(i - 100, i - 200)) : ImmutableMap.of()))))); + item.put("col_9", AttributeValue.builder().n(String.valueOf(100.0f + i)).build()); + item.put("col_10", DDBTypeUtils.toAttributeValue(ImmutableList.of(ImmutableList.of(1 * i, 2 * i, 3 * i), ImmutableList.of(4 * i, 5 * i), ImmutableList.of(6 * i, 7 * i, 8 * i)))); - tableWriteItems.addItemToPut(toItem(item)); - if (tableWriteItems.getItemsToPut().size() == 25) { - ddb.batchWriteItem(tableWriteItems); - tableWriteItems = new TableWriteItems(TEST_TABLE); - } + writeRequests.add(WriteRequest.builder() + .putRequest(PutRequest.builder().item(item).build()) + .build()); + + if (writeRequests.size() == 25) { + BatchWriteItemRequest batchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE, writeRequests)) + .build(); + ddb.batchWriteItem(batchWriteItemRequest); + // Handle response, check for unprocessed items, etc. + writeRequests.clear(); // Clear the list for the next batch + } dateTime = dateTime.plusHours(26); } - CreateGlobalSecondaryIndexAction createIndexRequest = new CreateGlobalSecondaryIndexAction() - .withIndexName("test_index") - .withKeySchema( - new KeySchemaElement().withKeyType(KeyType.HASH).withAttributeName("col_4"), - new KeySchemaElement().withKeyType(KeyType.RANGE).withAttributeName("col_5")) - .withProjection(new Projection().withProjectionType(ProjectionType.ALL)) - .withProvisionedThroughput(provisionedThroughput); - Index gsi = table.createGSI(createIndexRequest, - new AttributeDefinition().withAttributeName("col_4").withAttributeType(ScalarAttributeType.N), - new AttributeDefinition().withAttributeName("col_5").withAttributeType(ScalarAttributeType.N)); - gsi.waitForActive(); + + CreateGlobalSecondaryIndexAction createIndexRequest = CreateGlobalSecondaryIndexAction.builder() + .indexName("test_index") + .keySchema(KeySchemaElement.builder() + .keyType(KeyType.HASH) + .attributeName("col_4") + .build(), + KeySchemaElement.builder() + .keyType(KeyType.RANGE) + .attributeName("col_5") + .build()) + .projection(Projection.builder() + .projectionType(ProjectionType.ALL) + .build()) + .provisionedThroughput(provisionedThroughput) + .build(); + + UpdateTableRequest updateTableRequest = UpdateTableRequest.builder() + .tableName(TEST_TABLE) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("col_4") + .attributeType(ScalarAttributeType.N) + .build(), + AttributeDefinition.builder() + .attributeName("col_5") + .attributeType(ScalarAttributeType.N) + .build()) + .globalSecondaryIndexUpdates(GlobalSecondaryIndexUpdate.builder() + .create(createIndexRequest) + .build()) + .build(); + + ddb.updateTable(updateTableRequest); + dbWaiter.waitUntilTableExists(DescribeTableRequest.builder().tableName(TEST_TABLE).build()); + waitForGsiToBecomeActive(ddb, TEST_TABLE, "test_index"); + // for case sensitivity testing - createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE2) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); - table = ddb.createTable(createTableRequest); - table.waitForActive(); - - ArrayList attributeDefinitionsGlue = new ArrayList<>(); - attributeDefinitionsGlue.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); - attributeDefinitionsGlue.add(new AttributeDefinition().withAttributeName("Col1").withAttributeType("S")); - - ArrayList keySchemaGlue = new ArrayList<>(); - keySchemaGlue.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); - keySchemaGlue.add(new KeySchemaElement().withAttributeName("Col1").withKeyType(KeyType.RANGE)); - - CreateTableRequest createTableRequestGlue = new CreateTableRequest() - .withTableName(TEST_TABLE3) - .withKeySchema(keySchemaGlue) - .withAttributeDefinitions(attributeDefinitionsGlue) - .withProvisionedThroughput(provisionedThroughput); - - Table tableGlue = ddb.createTable(createTableRequestGlue); - tableGlue.waitForActive(); - - tableWriteItems = new TableWriteItems(TEST_TABLE3); - Map item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("Col1", toAttributeValue("20200227S091227")); - item.put("Col2", toAttributeValue("2020-02-27T09:12:27Z")); - item.put("Col3", toAttributeValue("27/02/2020")); - item.put("Col4", toAttributeValue("2020-02-27")); - // below three columns are testing timestamp with timezone - // col5 with non-utc timezone, col6 with utc timezone, and c7 without timezone that will fall back to default - item.put("Col5", toAttributeValue("2015-12-21T17:42:34-05:00")); - item.put("Col6", toAttributeValue("2015-12-21T17:42:34Z")); - item.put("Col7", toAttributeValue("2015-12-21T17:42:34")); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); - - // Table 4 - attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); - - keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); - - createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE4) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); - - tableDdbNoGlue = ddb.createTable(createTableRequest); - tableDdbNoGlue.waitForActive(); - - Map col1 = new HashMap<>(); - col1.put("field1", "someField1"); - col1.put("field2", null); - - tableWriteItems = new TableWriteItems(TEST_TABLE4); - item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("Col1", toAttributeValue(col1)); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); - - // Table5 - attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); - - keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); - - createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE5) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); - - tableDdbNoGlue = ddb.createTable(createTableRequest); - tableDdbNoGlue.waitForActive(); + CreateTableRequest table2Request = CreateTableRequest.builder() + .tableName(TEST_TABLE2) + .keySchema(keySchema) + .attributeDefinitions(attributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); - Map> nestedCol = new HashMap<>(); - ArrayList value = new ArrayList(); - value.add("list1"); - value.add("list2"); - nestedCol.put("list", value); - Map> listStructCol = new HashMap<>(); - Map structVal = new HashMap<>(); - structVal.put("key1", "str1"); - structVal.put("key2", "str2"); - listStructCol.put("structKey", structVal); + waitForTableToBecomeActive(ddb, dbWaiter, ddb.createTable(table2Request), TEST_TABLE2); - tableWriteItems = new TableWriteItems(TEST_TABLE5); - item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("outermap", toAttributeValue(nestedCol)); - item.put("structcol", toAttributeValue(listStructCol)); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); - - setUpTable6(provisionedThroughput, ddb); - setUpTable7(provisionedThroughput, ddb); - setUpTable8(provisionedThroughput, ddb); - return client; } + private static void setUpTable3( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) { + ///Glue Table; + List attributeDefinitionsGlue = new ArrayList<>(); + attributeDefinitionsGlue.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType("S") + .build()); + attributeDefinitionsGlue.add(AttributeDefinition.builder() + .attributeName("Col1") + .attributeType("S") + .build()); + + List keySchemaGlue = new ArrayList<>(); + keySchemaGlue.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + keySchemaGlue.add(KeySchemaElement.builder() + .attributeName("Col1") + .keyType(KeyType.RANGE) + .build()); + CreateTableRequest table3Request = CreateTableRequest.builder() + .tableName(TEST_TABLE3) + .keySchema(keySchemaGlue) + .attributeDefinitions(attributeDefinitionsGlue) + .provisionedThroughput(provisionedThroughput) + .build(); + + waitForTableToBecomeActive(ddb, dbWaiter, ddb.createTable(table3Request), TEST_TABLE3); - private static void setUpTable6( + Map item = new HashMap<>(); + item.put("Col0", DDBTypeUtils.toAttributeValue("hashVal")); + item.put("Col1", DDBTypeUtils.toAttributeValue("20200227S091227")); + item.put("Col2", DDBTypeUtils.toAttributeValue("2020-02-27T09:12:27Z")); + item.put("Col3", DDBTypeUtils.toAttributeValue("27/02/2020")); + item.put("Col4", DDBTypeUtils.toAttributeValue("2020-02-27")); + item.put("Col5", DDBTypeUtils.toAttributeValue("2015-12-21T17:42:34-05:00")); + item.put("Col6", DDBTypeUtils.toAttributeValue("2015-12-21T17:42:34Z")); + item.put("Col7", DDBTypeUtils.toAttributeValue("2015-12-21T17:42:34")); + + List table3WriteRequest = new ArrayList<>(); + table3WriteRequest.add(WriteRequest.builder() + .putRequest(PutRequest.builder().item(item).build()) + .build()); + + BatchWriteItemRequest table3BatchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE3, table3WriteRequest)) + .build(); + ddb.batchWriteItem(table3BatchWriteItemRequest); + + } + private static void setUpTable4( ProvisionedThroughput provisionedThroughput, - DynamoDB ddb) - throws InterruptedException - { + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) { + + List table4AttributeDefinitions = new ArrayList<>(); + table4AttributeDefinitions.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType(ScalarAttributeType.S) + .build()); + + // Define key schema for Table 4 + List table4KeySchema = new ArrayList<>(); + table4KeySchema.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + + // Create Table Request for Table 4 + CreateTableRequest table4TableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE4) + .keySchema(table4KeySchema) + .attributeDefinitions(table4AttributeDefinitions) + .provisionedThroughput(ProvisionedThroughput.builder() // Assuming 'provisionedThroughput' is defined + .readCapacityUnits(5L) // Example capacity units + .writeCapacityUnits(5L) + .build()) + .build(); + + waitForTableToBecomeActive(ddb, dbWaiter, ddb.createTable(table4TableRequest), TEST_TABLE4); + + Map col1 = new HashMap<>(); + col1.put("field1", AttributeValue.builder().s("someField1").build()); + col1.put("field2", AttributeValue.builder().nul(true).build()); // Representing null - ArrayList attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); + Map item = new HashMap<>(); + item.put("Col0", AttributeValue.builder().s("hashVal").build()); + item.put("Col1", AttributeValue.builder().m(col1).build()); - ArrayList keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); + WriteRequest writeRequest = WriteRequest.builder() + .putRequest(PutRequest.builder().item(item).build()) + .build(); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE6) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); + BatchWriteItemRequest batchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE4, ImmutableList.of(writeRequest))) + .build(); - tableDdbNoGlue = ddb.createTable(createTableRequest); - tableDdbNoGlue.waitForActive(); + ddb.batchWriteItem(batchWriteItemRequest); + } - Map> nestedCol = new HashMap<>(); - ArrayList value = new ArrayList(); + private static void setUpTable5( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) { + + List table5AttributeDefinitions = new ArrayList<>(); + table5AttributeDefinitions.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType(ScalarAttributeType.S) + .build()); + + List table5KeySchema = new ArrayList<>(); + table5KeySchema.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + + CreateTableRequest table5CreateTableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE5) + .keySchema(table5KeySchema) + .attributeDefinitions(table5AttributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); + + waitForTableToBecomeActive(ddb, dbWaiter, ddb.createTable(table5CreateTableRequest), TEST_TABLE5); + + // Nested collection + Map nestedCol = new HashMap<>(); + ArrayList value = new ArrayList<>(); + value.add("list1"); + value.add("list2"); + nestedCol.put("list", AttributeValue.builder() + .l(value.stream().map(s -> AttributeValue.builder().s(s).build()).collect(Collectors.toList())) + .build()); + + // Structured collection + Map listStructCol = new HashMap<>(); + Map structVal = new HashMap<>(); + structVal.put("key1", AttributeValue.builder().s("str1").build()); + structVal.put("key2", AttributeValue.builder().s("str2").build()); + listStructCol.put("structKey", AttributeValue.builder().m(structVal).build()); + + // Item to write + Map item = new HashMap<>(); + item.put("Col0", AttributeValue.builder().s("hashVal").build()); + item.put("outermap", AttributeValue.builder().m(nestedCol).build()); + item.put("structcol", AttributeValue.builder().m(listStructCol).build()); + + // Batch write request + WriteRequest writeRequest = WriteRequest.builder() + .putRequest(PutRequest.builder().item(item).build()) + .build(); + BatchWriteItemRequest batchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE5, ImmutableList.of(writeRequest))) + .build(); + + // Perform the batch write operation + ddb.batchWriteItem(batchWriteItemRequest); + } + + private static void setUpTable6( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) + throws InterruptedException { + + // Define attribute definitions for Table 6 + List table6AttributeDefinitions = new ArrayList<>(); + table6AttributeDefinitions.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType(ScalarAttributeType.S) + .build()); + + // Define key schema for Table 6 + List table6KeySchema = new ArrayList<>(); + table6KeySchema.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + + // Create Table Request for Table 6 + CreateTableRequest table6CreateTableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE6) + .keySchema(table6KeySchema) + .attributeDefinitions(table6AttributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); + + + waitForTableToBecomeActive(ddb, dbWaiter, ddb.createTable(table6CreateTableRequest), TEST_TABLE6); // Use the same waitForTableToBecomeActive method + + // Prepare nested collections for batch write + ArrayList value = new ArrayList<>(); value.add("list1"); value.add("list2"); + Map> nestedCol = new HashMap<>(); nestedCol.put("list", value); - Map> listStructCol = new HashMap<>(); Map structVal = new HashMap<>(); structVal.put("key1", "str1"); structVal.put("key2", "str2"); + Map> listStructCol = new HashMap<>(); listStructCol.put("structKey", structVal); - TableWriteItems tableWriteItems = new TableWriteItems(TEST_TABLE6); - Map item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("outermap", toAttributeValue(nestedCol)); - item.put("structcol", toAttributeValue(listStructCol)); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); - } + // Batch writes items to Table 6 + Map table6Item = new HashMap<>(); + // Assuming DDBTypeUtils.toAttributeValue is properly defined and imported + table6Item.put("Col0", DDBTypeUtils.toAttributeValue("hashVal")); + table6Item.put("outermap", DDBTypeUtils.toAttributeValue(nestedCol)); + table6Item.put("structcol", DDBTypeUtils.toAttributeValue(listStructCol)); - private static void setUpTable7( - ProvisionedThroughput provisionedThroughput, - DynamoDB ddb) - throws InterruptedException - { - - ArrayList attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); + WriteRequest table6WriteRequest = WriteRequest.builder() + .putRequest(PutRequest.builder().item(table6Item).build()) + .build(); - ArrayList keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); + List table6WriteRequests = new ArrayList<>(); + table6WriteRequests.add(table6WriteRequest); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE7) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); + BatchWriteItemRequest table6BatchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE6, table6WriteRequests)) + .build(); - tableDdbNoGlue = ddb.createTable(createTableRequest); - tableDdbNoGlue.waitForActive(); + ddb.batchWriteItem(table6BatchWriteItemRequest); + } - ArrayList stringList = new ArrayList(); + private static void setUpTable7( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) + throws InterruptedException { + + List table7AttributeDefinitions = new ArrayList<>(); + table7AttributeDefinitions.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType(ScalarAttributeType.S) + .build()); + + List table7KeySchema = new ArrayList<>(); + table7KeySchema.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + + CreateTableRequest table7CreateTableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE7) + .keySchema(table7KeySchema) + .attributeDefinitions(table7AttributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); + + CreateTableResponse table7CreateTableResponse = ddb.createTable(table7CreateTableRequest); + waitForTableToBecomeActive(ddb, dbWaiter, table7CreateTableResponse, TEST_TABLE7); + + // Prepare data for batch write + ArrayList stringList = new ArrayList<>(); stringList.add("list1"); stringList.add("list2"); - ArrayList intList = new ArrayList(); + ArrayList intList = new ArrayList<>(); intList.add(0); intList.add(1); intList.add(2); - ArrayList> listStructCol = new ArrayList(); + ArrayList> listStructCol = new ArrayList<>(); Map structVal = new HashMap<>(); structVal.put("key1", "str1"); structVal.put("key2", "str2"); @@ -361,46 +637,75 @@ private static void setUpTable7( listStructCol.add(structVal); listStructCol.add(structVal2); - TableWriteItems tableWriteItems = new TableWriteItems(TEST_TABLE7); - Map item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("stringList", toAttributeValue(stringList)); - item.put("intList", toAttributeValue(intList)); - item.put("listStructCol", toAttributeValue(listStructCol)); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); - } + // Batch writes items to Table 7 + Map table7Item = new HashMap<>(); + table7Item.put("Col0", DDBTypeUtils.toAttributeValue("hashVal")); + table7Item.put("stringList", DDBTypeUtils.toAttributeValue(stringList)); + table7Item.put("intList", DDBTypeUtils.toAttributeValue(intList)); + table7Item.put("listStructCol", DDBTypeUtils.toAttributeValue(listStructCol)); - private static void setUpTable8( - ProvisionedThroughput provisionedThroughput, - DynamoDB ddb) - throws InterruptedException - { + WriteRequest table7WriteRequest = WriteRequest.builder() + .putRequest(PutRequest.builder().item(table7Item).build()) + .build(); - ArrayList attributeDefinitions = new ArrayList<>(); - attributeDefinitions.add(new AttributeDefinition().withAttributeName("Col0").withAttributeType("S")); + List table7WriteRequests = new ArrayList<>(); + table7WriteRequests.add(table7WriteRequest); - ArrayList keySchema = new ArrayList<>(); - keySchema.add(new KeySchemaElement().withAttributeName("Col0").withKeyType(KeyType.HASH)); + BatchWriteItemRequest table7BatchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE7, table7WriteRequests)) + .build(); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withTableName(TEST_TABLE8) - .withKeySchema(keySchema) - .withAttributeDefinitions(attributeDefinitions) - .withProvisionedThroughput(provisionedThroughput); - - tableDdbNoGlue = ddb.createTable(createTableRequest); - tableDdbNoGlue.waitForActive(); + ddb.batchWriteItem(table7BatchWriteItemRequest); + } + private static void setUpTable8( + ProvisionedThroughput provisionedThroughput, + DynamoDbClient ddb, DynamoDbWaiter dbWaiter) + throws InterruptedException { + + List table8AttributeDefinitions = new ArrayList<>(); + table8AttributeDefinitions.add(AttributeDefinition.builder() + .attributeName("Col0") + .attributeType(ScalarAttributeType.S) + .build()); + + List table8KeySchema = new ArrayList<>(); + table8KeySchema.add(KeySchemaElement.builder() + .attributeName("Col0") + .keyType(KeyType.HASH) + .build()); + + CreateTableRequest table8CreateTableRequest = CreateTableRequest.builder() + .tableName(TEST_TABLE8) + .keySchema(table8KeySchema) + .attributeDefinitions(table8AttributeDefinitions) + .provisionedThroughput(provisionedThroughput) + .build(); + + CreateTableResponse table8CreateTableResponse = ddb.createTable(table8CreateTableRequest); + waitForTableToBecomeActive(ddb, dbWaiter, table8CreateTableResponse, TEST_TABLE8); + + // Prepare data for batch write Map numMap = new HashMap<>(); numMap.put("key1", 1); numMap.put("key2", 2); - TableWriteItems tableWriteItems = new TableWriteItems(TEST_TABLE8); - Map item = new HashMap<>(); - item.put("Col0", toAttributeValue("hashVal")); - item.put("nummap", toAttributeValue(numMap)); - tableWriteItems.addItemToPut(toItem(item)); - ddb.batchWriteItem(tableWriteItems); + // Batch writes items to Table 8 + Map table8Item = new HashMap<>(); + table8Item.put("Col0", DDBTypeUtils.toAttributeValue("hashVal")); + table8Item.put("nummap", DDBTypeUtils.toAttributeValue(numMap)); + + WriteRequest table8WriteRequest = WriteRequest.builder() + .putRequest(PutRequest.builder().item(table8Item).build()) + .build(); + + List table8WriteRequests = new ArrayList<>(); + table8WriteRequests.add(table8WriteRequest); + + BatchWriteItemRequest table8BatchWriteItemRequest = BatchWriteItemRequest.builder() + .requestItems(ImmutableMap.of(TEST_TABLE8, table8WriteRequests)) + .build(); + + ddb.batchWriteItem(table8BatchWriteItemRequest); } }