Example usage for com.amazonaws.services.dynamodbv2.document TableKeysAndAttributes TableKeysAndAttributes

List of usage examples for com.amazonaws.services.dynamodbv2.document TableKeysAndAttributes TableKeysAndAttributes

Introduction

In this page you can find the example usage for com.amazonaws.services.dynamodbv2.document TableKeysAndAttributes TableKeysAndAttributes.

Prototype

public TableKeysAndAttributes(String tableName) 

Source Link

Usage

From source file:com.clicktravel.infrastructure.persistence.aws.dynamodb.DynamoDocumentStoreTemplate.java

License:Apache License

private <T extends Item> List<T> executeQuery(final KeySetQuery query, final Class<T> itemClass) {
    final ItemConfiguration itemConfiguration = getItemConfiguration(itemClass);
    final String tableName = databaseSchemaHolder.schemaName() + "." + itemConfiguration.tableName();
    // max 100 keys per fetch
    final List<List<ItemId>> split_ids = split(new ArrayList<ItemId>(query.itemIds()), 100);
    final List<T> fetchedItems = new ArrayList<T>();
    for (final List<ItemId> ids : split_ids) {
        final TableKeysAndAttributes keys = new TableKeysAndAttributes(tableName);
        for (final ItemId id : ids) {
            keys.addPrimaryKey(getPrimaryKey(id, itemConfiguration));
        }/*www  . j a  va  2  s.  c o  m*/
        processBatchRead(dynamoDBClient.batchGetItem(keys), fetchedItems, tableName, itemClass);
    }
    return fetchedItems;
}

From source file:org.apache.nifi.processors.aws.dynamodb.GetDynamoDB.java

License:Apache License

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    List<FlowFile> flowFiles = session
            .get(context.getProperty(BATCH_SIZE).evaluateAttributeExpressions().asInteger());
    if (flowFiles == null || flowFiles.size() == 0) {
        return;/*from   w ww.  j  ava 2  s.  c o m*/
    }

    Map<ItemKeys, FlowFile> keysToFlowFileMap = new HashMap<>();

    final String table = context.getProperty(TABLE).evaluateAttributeExpressions().getValue();
    TableKeysAndAttributes tableKeysAndAttributes = new TableKeysAndAttributes(table);

    final String hashKeyName = context.getProperty(HASH_KEY_NAME).evaluateAttributeExpressions().getValue();
    final String rangeKeyName = context.getProperty(RANGE_KEY_NAME).evaluateAttributeExpressions().getValue();
    final String jsonDocument = context.getProperty(JSON_DOCUMENT).evaluateAttributeExpressions().getValue();

    for (FlowFile flowFile : flowFiles) {
        final Object hashKeyValue = getValue(context, HASH_KEY_VALUE_TYPE, HASH_KEY_VALUE, flowFile);
        final Object rangeKeyValue = getValue(context, RANGE_KEY_VALUE_TYPE, RANGE_KEY_VALUE, flowFile);

        if (!isHashKeyValueConsistent(hashKeyName, hashKeyValue, session, flowFile)) {
            continue;
        }

        if (!isRangeKeyValueConsistent(rangeKeyName, rangeKeyValue, session, flowFile)) {
            continue;
        }

        keysToFlowFileMap.put(new ItemKeys(hashKeyValue, rangeKeyValue), flowFile);

        if (rangeKeyValue == null || StringUtils.isBlank(rangeKeyValue.toString())) {
            tableKeysAndAttributes.addHashOnlyPrimaryKey(hashKeyName, hashKeyValue);
        } else {
            tableKeysAndAttributes.addHashAndRangePrimaryKey(hashKeyName, hashKeyValue, rangeKeyName,
                    rangeKeyValue);
        }
    }

    if (keysToFlowFileMap.isEmpty()) {
        return;
    }

    final DynamoDB dynamoDB = getDynamoDB();

    try {
        BatchGetItemOutcome result = dynamoDB.batchGetItem(tableKeysAndAttributes);

        // Handle processed items and get the json document
        List<Item> items = result.getTableItems().get(table);
        for (Item item : items) {
            ItemKeys itemKeys = new ItemKeys(item.get(hashKeyName), item.get(rangeKeyName));
            FlowFile flowFile = keysToFlowFileMap.get(itemKeys);

            if (item.get(jsonDocument) != null) {
                ByteArrayInputStream bais = new ByteArrayInputStream(item.getJSON(jsonDocument).getBytes());
                flowFile = session.importFrom(bais, flowFile);
            }

            session.transfer(flowFile, REL_SUCCESS);
            keysToFlowFileMap.remove(itemKeys);
        }

        // Handle unprocessed keys
        Map<String, KeysAndAttributes> unprocessedKeys = result.getUnprocessedKeys();
        if (unprocessedKeys != null && unprocessedKeys.size() > 0) {
            KeysAndAttributes keysAndAttributes = unprocessedKeys.get(table);
            List<Map<String, AttributeValue>> keys = keysAndAttributes.getKeys();

            for (Map<String, AttributeValue> unprocessedKey : keys) {
                Object hashKeyValue = getAttributeValue(context, HASH_KEY_VALUE_TYPE,
                        unprocessedKey.get(hashKeyName));
                Object rangeKeyValue = getAttributeValue(context, RANGE_KEY_VALUE_TYPE,
                        unprocessedKey.get(rangeKeyName));
                sendUnprocessedToUnprocessedRelationship(session, keysToFlowFileMap, hashKeyValue,
                        rangeKeyValue);
            }
        }

        // Handle any remaining items
        for (ItemKeys key : keysToFlowFileMap.keySet()) {
            FlowFile flowFile = keysToFlowFileMap.get(key);
            flowFile = session.putAttribute(flowFile, DYNAMODB_KEY_ERROR_NOT_FOUND,
                    DYNAMODB_KEY_ERROR_NOT_FOUND_MESSAGE + key.toString());
            session.transfer(flowFile, REL_NOT_FOUND);
            keysToFlowFileMap.remove(key);
        }

    } catch (AmazonServiceException exception) {
        getLogger().error("Could not process flowFiles due to service exception : " + exception.getMessage());
        List<FlowFile> failedFlowFiles = processServiceException(session, flowFiles, exception);
        session.transfer(failedFlowFiles, REL_FAILURE);
    } catch (AmazonClientException exception) {
        getLogger().error("Could not process flowFiles due to client exception : " + exception.getMessage());
        List<FlowFile> failedFlowFiles = processClientException(session, flowFiles, exception);
        session.transfer(failedFlowFiles, REL_FAILURE);
    } catch (Exception exception) {
        getLogger().error("Could not process flowFiles due to exception : " + exception.getMessage());
        List<FlowFile> failedFlowFiles = processException(session, flowFiles, exception);
        session.transfer(failedFlowFiles, REL_FAILURE);
    }
}