001package com.nimbusds.infinispan.persistence.dynamodb.query; 002 003 004import java.util.Map; 005import java.util.function.Consumer; 006 007import com.amazonaws.services.dynamodbv2.document.Index; 008import com.amazonaws.services.dynamodbv2.document.ItemCollection; 009import com.amazonaws.services.dynamodbv2.document.QueryOutcome; 010import com.amazonaws.services.dynamodbv2.document.ScanOutcome; 011import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; 012import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec; 013import com.amazonaws.services.dynamodbv2.document.utils.NameMap; 014import com.amazonaws.services.dynamodbv2.document.utils.ValueMap; 015import com.nimbusds.infinispan.persistence.common.InfinispanEntry; 016import com.nimbusds.infinispan.persistence.common.query.MatchQuery; 017import com.nimbusds.infinispan.persistence.common.query.Query; 018import com.nimbusds.infinispan.persistence.common.query.SimpleMatchQuery; 019import com.nimbusds.infinispan.persistence.common.query.UnsupportedQueryException; 020import net.jcip.annotations.ThreadSafe; 021 022 023/** 024 * Simple match DynamoDB query executor. Accepts queries of type 025 * {@link SimpleMatchQuery} and {@link MatchQuery} (converts it to 026 * {@link SimpleMatchQuery by taking the first match pair}, where both key name 027 * and value must be of type {@link String}. 028 */ 029@ThreadSafe 030public class SimpleMatchQueryExecutor<K, V> implements DynamoDBQueryExecutor<K, V> { 031 032 033 /** 034 * The initialisation context. 035 */ 036 private DynamoDBQueryExecutorInitContext<K, V> initCtx; 037 038 039 @Override 040 public void init(final DynamoDBQueryExecutorInitContext<K, V> initCtx) { 041 if (initCtx == null) { 042 throw new IllegalArgumentException("The init context must not be null"); 043 } 044 this.initCtx = initCtx; 045 } 046 047 048 /** 049 * Returns the DynamoDB query spec for the specified simple match 050 * query. 051 * 052 * @param simpleMatchQuery The simple match query. Must not be 053 * {@code null}. 054 * 055 * @return The DynamoDB query spec. 056 */ 057 static QuerySpec toQuerySpec(final SimpleMatchQuery<String, String> simpleMatchQuery) { 058 059 return new QuerySpec() 060 .withKeyConditionExpression("#k = :value") 061 .withNameMap(new NameMap() 062 .with("#k", simpleMatchQuery.getKey())) 063 .withValueMap(new ValueMap() 064 .withString(":value", simpleMatchQuery.getValue())); 065 } 066 067 /** 068 * Returns the DynamoDB scan spec for the specified simple match 069 * query. 070 * 071 * @param simpleMatchQuery The simple match query. Must not be 072 * {@code null}. 073 * 074 * @return The DynamoDB scan spec. 075 */ 076 static ScanSpec toScanSpec(final SimpleMatchQuery<String, String> simpleMatchQuery) { 077 078 return new ScanSpec() 079 .withFilterExpression("#k = :value") 080 .withNameMap(new NameMap() 081 .with("#k", simpleMatchQuery.getKey())) 082 .withValueMap(new ValueMap() 083 .withString(":value", simpleMatchQuery.getValue())) 084 .withConsistentRead(false); 085 } 086 087 088 /** 089 * Converts the specified match query to a simple match query. 090 * 091 * @param matchQuery The match query. 092 * 093 * @return The simple match query. 094 */ 095 static SimpleMatchQuery<String, String> toSimpleMatchQuery(final MatchQuery<String, String> matchQuery) { 096 097 if (matchQuery instanceof SimpleMatchQuery) { 098 return (SimpleMatchQuery<String, String>)matchQuery; 099 } 100 101 Map.Entry<String, String> entry = matchQuery.getMatchMap().entrySet().iterator().next(); 102 103 return new SimpleMatchQuery<>(entry.getKey(), entry.getValue()); 104 } 105 106 107 @Override 108 @SuppressWarnings("unchecked") 109 public void executeQuery(final Query query, final Consumer<InfinispanEntry<K, V>> consumer) { 110 111 if (! (query instanceof MatchQuery)) { 112 throw new UnsupportedQueryException(query); 113 } 114 115 MatchQuery<String,String> matchQuery = (MatchQuery<String, String>)query; 116 117 if (matchQuery.getMatchMap().isEmpty()) { 118 throw new UnsupportedQueryException(query); 119 } 120 121 SimpleMatchQuery<String, String> simpleMatchQuery = toSimpleMatchQuery(matchQuery); 122 123 Index index = initCtx.getDynamoDBIndex(simpleMatchQuery.getKey()); 124 125 if (index != null) { 126 // Query via GSI 127 ItemCollection<QueryOutcome> items = index.query(toQuerySpec(simpleMatchQuery)); 128 items.forEach(item -> consumer.accept(initCtx.getDynamoDBItemTransformer().toInfinispanEntry(item))); 129 } else { 130 // Fall back to scan with filter expression 131 ItemCollection<ScanOutcome> scanOutcome = initCtx.getDynamoDBTable().scan(toScanSpec(simpleMatchQuery)); 132 scanOutcome.forEach(item -> consumer.accept(initCtx.getDynamoDBItemTransformer().toInfinispanEntry(item))); 133 } 134 } 135}