Search in sources :

Example 1 with BatchGetItemOutcome

use of com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome in project nifi by apache.

the class GetDynamoDB method onTrigger.

@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;
    }
    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);
    }
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) AmazonClientException(com.amazonaws.AmazonClientException) DynamoDB(com.amazonaws.services.dynamodbv2.document.DynamoDB) KeysAndAttributes(com.amazonaws.services.dynamodbv2.model.KeysAndAttributes) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) BatchGetItemOutcome(com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome) AmazonServiceException(com.amazonaws.AmazonServiceException) AmazonClientException(com.amazonaws.AmazonClientException) Item(com.amazonaws.services.dynamodbv2.document.Item) ByteArrayInputStream(java.io.ByteArrayInputStream) AmazonServiceException(com.amazonaws.AmazonServiceException) List(java.util.List) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) HashMap(java.util.HashMap) Map(java.util.Map)

Example 2 with BatchGetItemOutcome

use of com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome in project nifi by apache.

the class GetDynamoDBTest method setUp.

@Before
public void setUp() {
    outcome = new BatchGetItemOutcome(result);
    KeysAndAttributes kaa = new KeysAndAttributes();
    Map<String, AttributeValue> map = new HashMap<>();
    map.put("hashS", new AttributeValue("h1"));
    map.put("rangeS", new AttributeValue("r1"));
    kaa.withKeys(map);
    unprocessed = new HashMap<>();
    unprocessed.put(stringHashStringRangeTableName, kaa);
    result.withUnprocessedKeys(unprocessed);
    Map<String, List<Map<String, AttributeValue>>> responses = new HashMap<>();
    List<Map<String, AttributeValue>> items = new ArrayList<>();
    responses.put("StringHashStringRangeTable", items);
    result.withResponses(responses);
    final DynamoDB mockDynamoDB = new DynamoDB(Regions.AP_NORTHEAST_1) {

        @Override
        public BatchGetItemOutcome batchGetItem(TableKeysAndAttributes... tableKeysAndAttributes) {
            return outcome;
        }
    };
    getDynamoDB = new GetDynamoDB() {

        @Override
        protected DynamoDB getDynamoDB() {
            return mockDynamoDB;
        }
    };
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) DynamoDB(com.amazonaws.services.dynamodbv2.document.DynamoDB) KeysAndAttributes(com.amazonaws.services.dynamodbv2.model.KeysAndAttributes) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) BatchGetItemOutcome(com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome) ArrayList(java.util.ArrayList) List(java.util.List) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) HashMap(java.util.HashMap) Map(java.util.Map) Before(org.junit.Before)

Example 3 with BatchGetItemOutcome

use of com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome in project nifi by apache.

the class GetDynamoDBTest method testStringHashStringRangeGetNotFound.

@Test
public void testStringHashStringRangeGetNotFound() {
    result.clearResponsesEntries();
    result.clearUnprocessedKeysEntries();
    final BatchGetItemOutcome notFoundOutcome = new BatchGetItemOutcome(result);
    Map<String, List<Map<String, AttributeValue>>> responses = new HashMap<>();
    List<Map<String, AttributeValue>> items = new ArrayList<>();
    responses.put(stringHashStringRangeTableName, items);
    result.withResponses(responses);
    final DynamoDB notFoundMockDynamoDB = new DynamoDB(Regions.AP_NORTHEAST_1) {

        @Override
        public BatchGetItemOutcome batchGetItem(TableKeysAndAttributes... tableKeysAndAttributes) {
            return notFoundOutcome;
        }
    };
    final GetDynamoDB getDynamoDB = new GetDynamoDB() {

        @Override
        protected DynamoDB getDynamoDB() {
            return notFoundMockDynamoDB;
        }
    };
    final TestRunner getRunner = TestRunners.newTestRunner(getDynamoDB);
    getRunner.setProperty(AbstractDynamoDBProcessor.ACCESS_KEY, "abcd");
    getRunner.setProperty(AbstractDynamoDBProcessor.SECRET_KEY, "cdef");
    getRunner.setProperty(AbstractDynamoDBProcessor.REGION, REGION);
    getRunner.setProperty(AbstractDynamoDBProcessor.TABLE, stringHashStringRangeTableName);
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_NAME, "rangeS");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_NAME, "hashS");
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_VALUE, "r1");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_VALUE, "h1");
    getRunner.setProperty(AbstractDynamoDBProcessor.JSON_DOCUMENT, "j1");
    getRunner.enqueue(new byte[] {});
    getRunner.run(1);
    getRunner.assertAllFlowFilesTransferred(GetDynamoDB.REL_NOT_FOUND, 1);
    List<MockFlowFile> flowFiles = getRunner.getFlowFilesForRelationship(AbstractDynamoDBProcessor.REL_UNPROCESSED);
    for (MockFlowFile flowFile : flowFiles) {
        assertNotNull(flowFile.getAttribute(AbstractDynamoDBProcessor.DYNAMODB_KEY_ERROR_NOT_FOUND));
    }
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) TestRunner(org.apache.nifi.util.TestRunner) ArrayList(java.util.ArrayList) DynamoDB(com.amazonaws.services.dynamodbv2.document.DynamoDB) BatchGetItemOutcome(com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome) MockFlowFile(org.apache.nifi.util.MockFlowFile) ArrayList(java.util.ArrayList) List(java.util.List) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) HashMap(java.util.HashMap) Map(java.util.Map) Test(org.junit.Test)

Example 4 with BatchGetItemOutcome

use of com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome in project nifi by apache.

the class GetDynamoDBTest method testStringHashStringRangeGetJsonObjectValid.

@Test
public void testStringHashStringRangeGetJsonObjectValid() throws IOException {
    outcome = new BatchGetItemOutcome(result);
    KeysAndAttributes kaa = new KeysAndAttributes();
    Map<String, AttributeValue> map = new HashMap<>();
    map.put("hashS", new AttributeValue("h1"));
    map.put("rangeS", new AttributeValue("r1"));
    kaa.withKeys(map);
    unprocessed = new HashMap<>();
    result.withUnprocessedKeys(unprocessed);
    Map<String, List<Map<String, AttributeValue>>> responses = new HashMap<>();
    List<Map<String, AttributeValue>> items = new ArrayList<>();
    Map<String, AttributeValue> item = new HashMap<String, AttributeValue>();
    String jsonDocument = "{\"name\": \"john\"}";
    item.put("j1", new AttributeValue(jsonDocument));
    item.put("hashS", new AttributeValue("h1"));
    item.put("rangeS", new AttributeValue("r1"));
    items.add(item);
    responses.put("StringHashStringRangeTable", items);
    result.withResponses(responses);
    final DynamoDB mockDynamoDB = new DynamoDB(Regions.AP_NORTHEAST_1) {

        @Override
        public BatchGetItemOutcome batchGetItem(TableKeysAndAttributes... tableKeysAndAttributes) {
            return outcome;
        }
    };
    getDynamoDB = new GetDynamoDB() {

        @Override
        protected DynamoDB getDynamoDB() {
            return mockDynamoDB;
        }
    };
    final TestRunner getRunner = TestRunners.newTestRunner(getDynamoDB);
    getRunner.setProperty(AbstractDynamoDBProcessor.ACCESS_KEY, "abcd");
    getRunner.setProperty(AbstractDynamoDBProcessor.SECRET_KEY, "cdef");
    getRunner.setProperty(AbstractDynamoDBProcessor.REGION, REGION);
    getRunner.setProperty(AbstractDynamoDBProcessor.TABLE, stringHashStringRangeTableName);
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_NAME, "rangeS");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_NAME, "hashS");
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_VALUE, "r1");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_VALUE, "h1");
    getRunner.setProperty(AbstractDynamoDBProcessor.JSON_DOCUMENT, "j1");
    getRunner.enqueue(new byte[] {});
    getRunner.run(1);
    getRunner.assertAllFlowFilesTransferred(AbstractDynamoDBProcessor.REL_SUCCESS, 1);
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) TestRunner(org.apache.nifi.util.TestRunner) ArrayList(java.util.ArrayList) DynamoDB(com.amazonaws.services.dynamodbv2.document.DynamoDB) KeysAndAttributes(com.amazonaws.services.dynamodbv2.model.KeysAndAttributes) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) BatchGetItemOutcome(com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome) ArrayList(java.util.ArrayList) List(java.util.List) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) HashMap(java.util.HashMap) Map(java.util.Map) Test(org.junit.Test)

Example 5 with BatchGetItemOutcome

use of com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome in project nifi by apache.

the class GetDynamoDBTest method testStringHashStringRangeGetJsonObjectNull.

@Test
public void testStringHashStringRangeGetJsonObjectNull() {
    outcome = new BatchGetItemOutcome(result);
    KeysAndAttributes kaa = new KeysAndAttributes();
    Map<String, AttributeValue> map = new HashMap<>();
    map.put("hashS", new AttributeValue("h1"));
    map.put("rangeS", new AttributeValue("r1"));
    kaa.withKeys(map);
    unprocessed = new HashMap<>();
    result.withUnprocessedKeys(unprocessed);
    Map<String, List<Map<String, AttributeValue>>> responses = new HashMap<>();
    List<Map<String, AttributeValue>> items = new ArrayList<>();
    Map<String, AttributeValue> item = new HashMap<String, AttributeValue>();
    item.put("j1", null);
    item.put("hashS", new AttributeValue("h1"));
    item.put("rangeS", new AttributeValue("r1"));
    items.add(item);
    responses.put("StringHashStringRangeTable", items);
    result.withResponses(responses);
    final DynamoDB mockDynamoDB = new DynamoDB(Regions.AP_NORTHEAST_1) {

        @Override
        public BatchGetItemOutcome batchGetItem(TableKeysAndAttributes... tableKeysAndAttributes) {
            return outcome;
        }
    };
    getDynamoDB = new GetDynamoDB() {

        @Override
        protected DynamoDB getDynamoDB() {
            return mockDynamoDB;
        }
    };
    final TestRunner getRunner = TestRunners.newTestRunner(getDynamoDB);
    getRunner.setProperty(AbstractDynamoDBProcessor.ACCESS_KEY, "abcd");
    getRunner.setProperty(AbstractDynamoDBProcessor.SECRET_KEY, "cdef");
    getRunner.setProperty(AbstractDynamoDBProcessor.REGION, REGION);
    getRunner.setProperty(AbstractDynamoDBProcessor.TABLE, stringHashStringRangeTableName);
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_NAME, "rangeS");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_NAME, "hashS");
    getRunner.setProperty(AbstractDynamoDBProcessor.RANGE_KEY_VALUE, "r1");
    getRunner.setProperty(AbstractDynamoDBProcessor.HASH_KEY_VALUE, "h1");
    getRunner.setProperty(AbstractDynamoDBProcessor.JSON_DOCUMENT, "j1");
    getRunner.enqueue(new byte[] {});
    getRunner.run(1);
    getRunner.assertAllFlowFilesTransferred(AbstractDynamoDBProcessor.REL_SUCCESS, 1);
    List<MockFlowFile> flowFiles = getRunner.getFlowFilesForRelationship(AbstractDynamoDBProcessor.REL_SUCCESS);
    for (MockFlowFile flowFile : flowFiles) {
        assertNull(flowFile.getContentClaim());
    }
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) TestRunner(org.apache.nifi.util.TestRunner) ArrayList(java.util.ArrayList) DynamoDB(com.amazonaws.services.dynamodbv2.document.DynamoDB) KeysAndAttributes(com.amazonaws.services.dynamodbv2.model.KeysAndAttributes) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) BatchGetItemOutcome(com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome) MockFlowFile(org.apache.nifi.util.MockFlowFile) ArrayList(java.util.ArrayList) List(java.util.List) TableKeysAndAttributes(com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes) HashMap(java.util.HashMap) Map(java.util.Map) Test(org.junit.Test)

Aggregations

TableKeysAndAttributes (com.amazonaws.services.dynamodbv2.document.TableKeysAndAttributes)9 DynamoDB (com.amazonaws.services.dynamodbv2.document.DynamoDB)8 BatchGetItemOutcome (com.amazonaws.services.dynamodbv2.document.BatchGetItemOutcome)6 TestRunner (org.apache.nifi.util.TestRunner)6 Test (org.junit.Test)6 AttributeValue (com.amazonaws.services.dynamodbv2.model.AttributeValue)5 KeysAndAttributes (com.amazonaws.services.dynamodbv2.model.KeysAndAttributes)5 HashMap (java.util.HashMap)5 List (java.util.List)5 Map (java.util.Map)5 MockFlowFile (org.apache.nifi.util.MockFlowFile)5 ArrayList (java.util.ArrayList)4 AmazonClientException (com.amazonaws.AmazonClientException)2 AmazonServiceException (com.amazonaws.AmazonServiceException)2 Item (com.amazonaws.services.dynamodbv2.document.Item)2 ByteArrayInputStream (java.io.ByteArrayInputStream)1 IOException (java.io.IOException)1 FlowFile (org.apache.nifi.flowfile.FlowFile)1 Before (org.junit.Before)1