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}