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}