Java 类com.amazonaws.services.dynamodbv2.model.KeysAndAttributes 实例源码

@Test
public void execute() {
    Map<String, AttributeValue> key = new HashMap<String, AttributeValue>();
    key.put("1", new AttributeValue("Key_1"));
    Map<String, AttributeValue> unprocessedKey = new HashMap<String, AttributeValue>();
    unprocessedKey.put("1", new AttributeValue("UNPROCESSED_KEY"));
    Map<String, KeysAndAttributes> keysAndAttributesMap = new HashMap<String, KeysAndAttributes>();
    KeysAndAttributes keysAndAttributes = new KeysAndAttributes().withKeys(key);
    keysAndAttributesMap.put("DOMAIN1", keysAndAttributes);
    exchange.getIn().setHeader(DdbConstants.BATCH_ITEMS, keysAndAttributesMap);

    command.execute();

    assertEquals(keysAndAttributesMap, ddbClient.batchGetItemRequest.getRequestItems());


    List<Map<String, AttributeValue>> batchResponse = (List<Map<String, AttributeValue>>)exchange.getIn().getHeader(DdbConstants.BATCH_RESPONSE, Map.class).get("DOMAIN1");
    AttributeValue value = batchResponse.get(0).get("attrName");

    KeysAndAttributes unProcessedAttributes = (KeysAndAttributes)exchange.getIn().getHeader(
            DdbConstants.UNPROCESSED_KEYS, Map.class).get("DOMAIN1");
    Map<String, AttributeValue> next = unProcessedAttributes.getKeys().iterator().next();

    assertEquals(new AttributeValue("attrValue"), value);
    assertEquals(unprocessedKey, next);
}

@SuppressWarnings("unchecked")
@Override
public BatchGetItemResult batchGetItem(BatchGetItemRequest batchGetItemRequest) {
    this.batchGetItemRequest = batchGetItemRequest;
    Map<String, List<Map<String, AttributeValue>>> responseMap = new HashMap<String, List<Map<String, AttributeValue>>>();
    List<Map<String, AttributeValue>> p = new ArrayList<Map<String, AttributeValue>>();
    p.add(getAttributes());
    responseMap.put("DOMAIN1", p);
    Map<String, AttributeValue> keysMap = new HashMap<String, AttributeValue>();
    keysMap.put("1", new AttributeValue("UNPROCESSED_KEY"));
    Map<String, KeysAndAttributes> unprocessedKeys = new HashMap<String, KeysAndAttributes>();
    unprocessedKeys.put("DOMAIN1", new KeysAndAttributes().withKeys(keysMap));

    return new BatchGetItemResult()
            .withResponses(responseMap)
            .withUnprocessedKeys(unprocessedKeys);
}

private KeysAndAttributes convertResultKeys(
        Map<String, Set<String>> resultKeys) {
    KeysAndAttributes keys = new KeysAndAttributes();

    for (final String s : resultKeys.keySet()) {
        for (final String value : resultKeys.get(s)) {
            keys.withKeys(new HashMap<String, AttributeValue>() {
                {
                    put(labelAttribute, new AttributeValue().withS(s));
                    put(dateAttribute, new AttributeValue().withS(value));
                }
            });
        }
    }

    return keys;
}

private List<Map<String, AttributeValue>> batchGetDataByKeys(
        final String tableName, final KeysAndAttributes keys) {
    Map<String, KeysAndAttributes> requestMap = new HashMap<>();
    keys.setConsistentRead(true);
    requestMap.put(tableName, keys);

    BatchGetItemResult result = null;
    try {
        result = dynamoClient.batchGetItem(new BatchGetItemRequest(
                requestMap));
    } catch (AmazonServiceException e) {
        LOG.error(e);
        throw e;
    }

    return result.getResponses().get(this.tableName);
}

public CompletableFuture<Map<String, List<Map<String, Object>>>> batchGetItem(final Map<String, KeysAndAttributes> requestItems) {
    return asyncExecutor.execute(new Callable<Map<String, List<Map<String, Object>>>>() {
        @Override
        public Map<String, List<Map<String, Object>>> call() throws Exception {
            return dbExecutor.batchGetItem(requestItems);
        }
    });
}

public CompletableFuture<Map<String, List<Map<String, Object>>>> batchGetItem(final Map<String, KeysAndAttributes> requestItems,
        final String returnConsumedCapacity) {
    return asyncExecutor.execute(new Callable<Map<String, List<Map<String, Object>>>>() {
        @Override
        public Map<String, List<Map<String, Object>>> call() throws Exception {
            return dbExecutor.batchGetItem(requestItems, returnConsumedCapacity);
        }
    });
}

public <T> CompletableFuture<Map<String, List<T>>> batchGetItem(final Class<T> targetClass, final Map<String, KeysAndAttributes> requestItems) {
    return asyncExecutor.execute(new Callable<Map<String, List<T>>>() {
        @Override
        public Map<String, List<T>> call() throws Exception {
            return dbExecutor.batchGetItem(targetClass, requestItems);
        }
    });
}

public <T> CompletableFuture<Map<String, List<T>>> batchGetItem(final Class<T> targetClass, final Map<String, KeysAndAttributes> requestItems,
        final String returnConsumedCapacity) {
    return asyncExecutor.execute(new Callable<Map<String, List<T>>>() {
        @Override
        public Map<String, List<T>> call() throws Exception {
            return dbExecutor.batchGetItem(targetClass, requestItems, returnConsumedCapacity);
        }
    });
}

@Test
public void test_batchGetItem_WithAllParameters() throws Exception {
  String TEST_ATTRIBUTE_2 = "Attribute2";
  String TEST_ATTRIBUTE_VALUE_2 = "AttributeValue2";

  createTable();
  putItem(TEST_ATTRIBUTE, TEST_ATTRIBUTE_VALUE);
  putItem(TEST_ATTRIBUTE_2, TEST_ATTRIBUTE_VALUE_2);

  List<Map<String, AttributeValue>> keys = new ArrayList<Map<String, AttributeValue>>();
  Map<String, AttributeValue> key1 = new HashMap<String, AttributeValue>();
  key1.put(TEST_ATTRIBUTE, new AttributeValue()
    .withS(TEST_ATTRIBUTE_VALUE));
  keys.add(key1);
  Map<String, AttributeValue> key2 = new HashMap<String, AttributeValue>();
  key2.put(TEST_ATTRIBUTE_2, new AttributeValue()
    .withS(TEST_ATTRIBUTE_VALUE_2));
  keys.add(key2);

  Map<String, KeysAndAttributes> requestItems = new HashMap<String, KeysAndAttributes>();
  requestItems.put(TEST_TABLE_NAME, new KeysAndAttributes()
    .withKeys(keys));
  String returnConsumedCapacity = "";

  BatchGetItemResult result = dynamoDb.batchGetItem(requestItems,returnConsumedCapacity);
  String tableName = result.getResponses().keySet().toArray(new String[1])[0];
  List<Map<String, AttributeValue>> items = result.getResponses().get(tableName);
  AttributeValue value1 = items.get(0).get(TEST_ATTRIBUTE);
  AttributeValue value2 = items.get(1).get(TEST_ATTRIBUTE_2);

  assertThat(tableName, equalTo(TEST_TABLE_NAME));
  assertThat(items.size(), equalTo(keys.size()));
  assertThat(value1.getS(), equalTo(TEST_ATTRIBUTE_VALUE));
  assertThat(value2.getS(), equalTo(TEST_ATTRIBUTE_VALUE_2));
}

/**
 * Reads multiple items from DynamoDB, in batch.
 * @param <P> type of object
 * @param kna a map of row key->data
 * @param results a map of ID->ParaObject
 */
protected static <P extends ParaObject> void batchGet(Map<String, KeysAndAttributes> kna, Map<String, P> results) {
    if (kna == null || kna.isEmpty() || results == null) {
        return;
    }
    try {
        BatchGetItemResult result = getClient().batchGetItem(new BatchGetItemRequest().
                withReturnConsumedCapacity(ReturnConsumedCapacity.TOTAL).withRequestItems(kna));
        if (result == null) {
            return;
        }

        List<Map<String, AttributeValue>> res = result.getResponses().get(kna.keySet().iterator().next());

        for (Map<String, AttributeValue> item : res) {
            P obj = fromRow(item);
            if (obj != null) {
                results.put(obj.getId(), obj);
            }
        }
        logger.debug("batchGet(): total {}, cc {}", res.size(), result.getConsumedCapacity());

        if (result.getUnprocessedKeys() != null && !result.getUnprocessedKeys().isEmpty()) {
            Thread.sleep(1000);
            logger.warn("{} UNPROCESSED read requests!", result.getUnprocessedKeys().size());
            batchGet(result.getUnprocessedKeys(), results);
        }
    } catch (Exception e) {
        logger.error(null, e);
    }
}

private Collection<OAuth2AccessToken> loadTokensByClientAndUserIndex(Map<String, Condition> keyCondition, boolean filterOutNullUsers) {
    List<OAuth2AccessToken> accessTokens = new ArrayList<OAuth2AccessToken>();

    List<String> accessTokenIds = null;
    try {
        accessTokenIds = dynamoDBTemplate.query(schema.getAccessTableName(), schema.getAccessIndexClientIdAndUserName(), keyCondition, //
                new ObjectExtractor<String>() {

                    public String extract(Map<String, AttributeValue> values) {
                        return values.get(schema.getAccessColumnTokenId()).getS();
                    }
                }, schema.getAccessColumnTokenId());

        List<Map<String, AttributeValue>> keys = new ArrayList<Map<String, AttributeValue>>(accessTokenIds.size());
        for (String accessTokenId : accessTokenIds) {
            keys.add(Collections.singletonMap(schema.getAccessColumnTokenId(), new AttributeValue(accessTokenId)));
        }
        if (filterOutNullUsers) {
            accessTokens = dynamoDBTemplate.batchGet(schema.getAccessTableName(), // 
                    new KeysAndAttributes().withKeys(keys).withConsistentRead(true).withAttributesToGet(schema.getAccessColumnTokenId(), schema.getAccessColumnToken(), schema.getAccessColumnIsNullUser()), // 
                    new NonNullUserSafeAccessTokenExtractor());
        } else {
            accessTokens = dynamoDBTemplate.batchGet(schema.getAccessTableName(), // 
                    new KeysAndAttributes().withKeys(keys).withConsistentRead(true).withAttributesToGet(schema.getAccessColumnTokenId(), schema.getAccessColumnToken()), // 
                    new SafeAccessTokenExtractor());
        }
    } catch (EmptyResultDataAccessException e) {
        if (LOG.isInfoEnabled()) {
            LOG.info("Failed to find access token for " + keyCondition.toString());
        }
    }
    accessTokens = removeNulls(accessTokens);

    return accessTokens;
}

public <T> List<T> batchGet(String tableName, KeysAndAttributes keysAndAttributes, final ObjectExtractor<T> extractor) throws EmptyResultDataAccessException {
    Assert.notNull(tableName, "Table must not be null");
    Assert.notNull(extractor, "ObjectExtractor must not be null");
    if (logger.isDebugEnabled()) {
        logger.debug("Executing batch get on " + tableName + " for " + keysAndAttributes.toString());
    }

    List<T> results = new ArrayList<T>(keysAndAttributes.getKeys().size());

    Map<String, KeysAndAttributes> unprocessedKeys = Collections.singletonMap(tableName, keysAndAttributes);
    while (unprocessedKeys.size() > 0) {
        BatchGetItemResult result = client.batchGetItem(unprocessedKeys);
        List<Map<String, AttributeValue>> items = result.getResponses().get(tableName);
        if (items != null) {
            for (Map<String, AttributeValue> item : items) {
                results.add(extractor.extract(item));
            }
        }

        unprocessedKeys = result.getUnprocessedKeys();
    }

    if (results.size() == 0) {
        throw new EmptyResultDataAccessException("No results found in " + tableName + "for " + keysAndAttributes.toString());
    }

    return results;
}

@SuppressWarnings("rawtypes")
public Map<String, List<Map<String, Object>>> batchGetItem(final Map<String, KeysAndAttributes> requestItems) {
    return (Map) batchGetItem(Map.class, requestItems);
}

@SuppressWarnings("rawtypes")
public Map<String, List<Map<String, Object>>> batchGetItem(final Map<String, KeysAndAttributes> requestItems, final String returnConsumedCapacity) {
    return (Map) batchGetItem(Map.class, requestItems, returnConsumedCapacity);
}

public <T> Map<String, List<T>> batchGetItem(final Class<T> targetClass, final Map<String, KeysAndAttributes> requestItems) {
    return toEntities(targetClass, dynamoDB.batchGetItem(requestItems).getResponses());
}

public <T> Map<String, List<T>> batchGetItem(final Class<T> targetClass, final Map<String, KeysAndAttributes> requestItems,
        final String returnConsumedCapacity) {
    return toEntities(targetClass, dynamoDB.batchGetItem(requestItems, returnConsumedCapacity).getResponses());
}

@SuppressWarnings("unchecked")
private Map<String, KeysAndAttributes> determineBatchItems() {
    return exchange.getIn().getHeader(DdbConstants.BATCH_ITEMS, Map.class);
}

@Override
public <P extends ParaObject> Map<String, P> readAll(String appid, List<String> keys, boolean getAllColumns) {
    if (keys == null || keys.isEmpty() || StringUtils.isBlank(appid)) {
        return new LinkedHashMap<>();
    }

    // DynamoDB doesn't allow duplicate keys in batch requests
    Set<String> keySet = new TreeSet<>(keys);
    if (keySet.size() < keys.size() && !keySet.isEmpty()) {
        logger.debug("Duplicate keys found - readAll({})", keys);
    }

    Map<String, P> results = new LinkedHashMap<>(keySet.size(), 0.75f, true);
    ArrayList<Map<String, AttributeValue>> keyz = new ArrayList<>(MAX_KEYS_PER_READ);

    try {
        int batchSteps = 1;
        if ((keySet.size() > MAX_KEYS_PER_READ)) {
            batchSteps = (keySet.size() / MAX_KEYS_PER_READ)
                    + ((keySet.size() % MAX_KEYS_PER_READ > 0) ? 1 : 0);
        }

        Iterator<String> it = keySet.iterator();
        String tableName = getTableNameForAppid(appid);
        int j = 0;

        for (int i = 0; i < batchSteps; i++) {
            while (it.hasNext() && j < MAX_KEYS_PER_READ) {
                String key = it.next();
                results.put(key, null);
                keyz.add(Collections.singletonMap(Config._KEY, new AttributeValue(getKeyForAppid(key, appid))));
                j++;
            }

            KeysAndAttributes kna = new KeysAndAttributes().withKeys(keyz);
            if (!getAllColumns) {
                kna.setAttributesToGet(Arrays.asList(Config._ID, Config._KEY, Config._TYPE));
            }

            batchGet(Collections.singletonMap(tableName, kna), results);
            keyz.clear();
            j = 0;
        }
        logger.debug("DAO.readAll({}) {}", keySet, results.size());
    } catch (Exception e) {
        logger.error("Failed to readAll({}): {}", keys, e);
    }
    return results;
}

private static void retrieveMultipleItemsBatchGet() {        

    try {

        TableKeysAndAttributes forumTableKeysAndAttributes = new TableKeysAndAttributes(forumTableName);
        forumTableKeysAndAttributes.addHashOnlyPrimaryKeys("Name", "Amazon S3", "Amazon DynamoDB");

        TableKeysAndAttributes threadTableKeysAndAttributes = new TableKeysAndAttributes(threadTableName);
        threadTableKeysAndAttributes.addHashAndRangePrimaryKeys("ForumName", "Subject", 
            "Amazon DynamoDB","DynamoDB Thread 1",
            "Amazon DynamoDB","DynamoDB Thread 2",
            "Amazon S3","S3 Thread 1");

         Map<String, TableKeysAndAttributes> requestItems = new HashMap<String, TableKeysAndAttributes>();
         requestItems.put(forumTableName, forumTableKeysAndAttributes);
         requestItems.put(threadTableName, threadTableKeysAndAttributes);

         System.out.println("Making the request.");

         BatchGetItemOutcome outcome = dynamoDB.batchGetItem(forumTableKeysAndAttributes,
                        threadTableKeysAndAttributes);

        do {

            for (String tableName : outcome.getTableItems().keySet()) {
                System.out.println("Items in table " + tableName);
                List<Item> items = outcome.getTableItems().get(tableName);
                for (Item item : items) {
                    System.out.println(item.toJSONPretty());
                }
            }

            // Check for unprocessed keys which could happen if you exceed provisioned
            // throughput or reach the limit on response size.

            Map<String, KeysAndAttributes> unprocessedKeys = outcome.getUnprocessedKeys();

            if (outcome.getUnprocessedKeys().size() == 0) {
                System.out.println("No unprocessed keys found");
            } else {
                System.out.println("Retrieving the unprocessed keys");
                outcome = dynamoDB.batchGetItemUnprocessed(unprocessedKeys);
            }

        } while (outcome.getUnprocessedKeys().size() > 0);


    }  catch (Exception e) {
        System.err.println("Failed to retrieve items.");
        System.err.println(e.getMessage());
    }  

}