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 this.initCtx = initCtx; 042 } 043 044 045 /** 046 * Returns the DynamoDB query spec for the specified simple match 047 * query. 048 * 049 * @param simpleMatchQuery The simple match query. Must not be 050 * {@code null}. 051 * 052 * @return The DynamoDB query spec. 053 */ 054 static QuerySpec toQuerySpec(final SimpleMatchQuery<String, String> simpleMatchQuery) { 055 056 return new QuerySpec() 057 .withKeyConditionExpression("#k = :value") 058 .withNameMap(new NameMap() 059 .with("#k", simpleMatchQuery.getKey())) 060 .withValueMap(new ValueMap() 061 .withString(":value", simpleMatchQuery.getValue())); 062 } 063 064 /** 065 * Returns the DynamoDB scan spec for the specified simple match 066 * query. 067 * 068 * @param simpleMatchQuery The simple match query. Must not be 069 * {@code null}. 070 * 071 * @return The DynamoDB scan spec. 072 */ 073 static ScanSpec toScanSpec(final SimpleMatchQuery<String, String> simpleMatchQuery) { 074 075 return new ScanSpec() 076 .withFilterExpression("#k = :value") 077 .withNameMap(new NameMap() 078 .with("#k", simpleMatchQuery.getKey())) 079 .withValueMap(new ValueMap() 080 .withString(":value", simpleMatchQuery.getValue())) 081 .withConsistentRead(false); 082 } 083 084 085 /** 086 * Converts the specified match query to a simple match query. 087 * 088 * @param matchQuery The match query. 089 * 090 * @return The simple match query. 091 */ 092 static SimpleMatchQuery<String, String> toSimpleMatchQuery(final MatchQuery<String, String> matchQuery) { 093 094 if (matchQuery instanceof SimpleMatchQuery) { 095 return (SimpleMatchQuery<String, String>)matchQuery; 096 } 097 098 Map.Entry<String, String> entry = matchQuery.getMatchMap().entrySet().iterator().next(); 099 100 return new SimpleMatchQuery<>(entry.getKey(), entry.getValue()); 101 } 102 103 104 @Override 105 @SuppressWarnings("unchecked") 106 public void executeQuery(final Query query, final Consumer<InfinispanEntry<K, V>> consumer) { 107 108 if (! (query instanceof MatchQuery)) { 109 throw new UnsupportedQueryException(query); 110 } 111 112 MatchQuery<String,String> matchQuery = (MatchQuery<String, String>)query; 113 114 if (matchQuery.getMatchMap().isEmpty()) { 115 throw new UnsupportedQueryException(query); 116 } 117 118 SimpleMatchQuery<String, String> simpleMatchQuery = toSimpleMatchQuery(matchQuery); 119 120 Index index = initCtx.getDynamoDBIndex(simpleMatchQuery.getKey()); 121 122 if (index != null) { 123 // Query via GSI 124 ItemCollection<QueryOutcome> items = index.query(toQuerySpec(simpleMatchQuery)); 125 items.forEach(item -> consumer.accept(initCtx.getDynamoDBItemTransformer().toInfinispanEntry(item))); 126 } else { 127 // Fall back to scan with filter expression 128 ItemCollection<ScanOutcome> scanOutcome = initCtx.getDynamoDBTable().scan(toScanSpec(simpleMatchQuery)); 129 scanOutcome.forEach(item -> consumer.accept(initCtx.getDynamoDBItemTransformer().toInfinispanEntry(item))); 130 } 131 } 132}