All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.apache.kafka.streams.state.internals.StoreQueryUtils Maven / Gradle / Ivy

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.streams.state.internals;

import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.errors.ProcessorStateException;
import org.apache.kafka.streams.kstream.Windowed;
import org.apache.kafka.streams.processor.StateStore;
import org.apache.kafka.streams.processor.StateStoreContext;
import org.apache.kafka.streams.processor.api.RecordMetadata;
import org.apache.kafka.streams.query.FailureReason;
import org.apache.kafka.streams.query.KeyQuery;
import org.apache.kafka.streams.query.MultiVersionedKeyQuery;
import org.apache.kafka.streams.query.Position;
import org.apache.kafka.streams.query.PositionBound;
import org.apache.kafka.streams.query.Query;
import org.apache.kafka.streams.query.QueryConfig;
import org.apache.kafka.streams.query.QueryResult;
import org.apache.kafka.streams.query.RangeQuery;
import org.apache.kafka.streams.query.ResultOrder;
import org.apache.kafka.streams.query.VersionedKeyQuery;
import org.apache.kafka.streams.query.WindowKeyQuery;
import org.apache.kafka.streams.query.WindowRangeQuery;
import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.KeyValueStore;
import org.apache.kafka.streams.state.SessionStore;
import org.apache.kafka.streams.state.StateSerdes;
import org.apache.kafka.streams.state.VersionedKeyValueStore;
import org.apache.kafka.streams.state.VersionedRecord;
import org.apache.kafka.streams.state.VersionedRecordIterator;
import org.apache.kafka.streams.state.WindowStore;
import org.apache.kafka.streams.state.WindowStoreIterator;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;

public final class StoreQueryUtils {

    /**
     * a utility interface to facilitate stores' query dispatch logic,
     * allowing them to generically store query execution logic as the values
     * in a map.
     */
    @FunctionalInterface
    public interface QueryHandler {
        QueryResult apply(
            final Query query,
            final PositionBound positionBound,
            final QueryConfig config,
            final StateStore store
        );
    }

    @SuppressWarnings("rawtypes")
    private static final Map QUERY_HANDLER_MAP =
        mkMap(
            mkEntry(
                RangeQuery.class,
                StoreQueryUtils::runRangeQuery
            ),
            mkEntry(
                KeyQuery.class,
                StoreQueryUtils::runKeyQuery
            ),
            mkEntry(
                WindowKeyQuery.class,
                StoreQueryUtils::runWindowKeyQuery
            ),
            mkEntry(
                WindowRangeQuery.class,
                StoreQueryUtils::runWindowRangeQuery
            ),
            mkEntry(
                VersionedKeyQuery.class,
                StoreQueryUtils::runVersionedKeyQuery
            ),
            mkEntry(
                MultiVersionedKeyQuery.class,
                StoreQueryUtils::runMultiVersionedKeyQuery
            )
        );

    // make this class uninstantiable

    private StoreQueryUtils() {
    }
    @SuppressWarnings("unchecked")
    public static  QueryResult handleBasicQueries(
        final Query query,
        final PositionBound positionBound,
        final QueryConfig config,
        final StateStore store,
        final Position position,
        final StateStoreContext context
    ) {

        final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L;
        final QueryResult result;

        final QueryHandler handler = QUERY_HANDLER_MAP.get(query.getClass());
        if (handler == null) {
            result = QueryResult.forUnknownQueryType(query, store);
        } else if (context == null || !isPermitted(position, positionBound, context.taskId().partition())) {
            result = QueryResult.notUpToBound(
                position,
                positionBound,
                context == null ? null : context.taskId().partition()
            );
        } else {
            result = (QueryResult) handler.apply(
                query,
                positionBound,
                config,
                store
            );
        }
        if (config.isCollectExecutionInfo()) {
            result.addExecutionInfo(
                "Handled in " + store.getClass() + " in " + (System.nanoTime() - start) + "ns"
            );
        }
        result.setPosition(position);
        return result;
    }

    public static void updatePosition(
        final Position position,
        final StateStoreContext stateStoreContext) {

        if (stateStoreContext != null && stateStoreContext.recordMetadata().isPresent()) {
            final RecordMetadata meta = stateStoreContext.recordMetadata().get();
            if (meta.topic() != null) {
                position.withComponent(meta.topic(), meta.partition(), meta.offset());
            }
        }
    }

    public static boolean isPermitted(
        final Position position,
        final PositionBound positionBound,
        final int partition
    ) {
        final Position bound = positionBound.position();
        for (final String topic : bound.getTopics()) {
            final Map partitionBounds = bound.getPartitionPositions(topic);
            final Map seenPartitionPositions = position.getPartitionPositions(topic);
            if (!partitionBounds.containsKey(partition)) {
                // this topic isn't bounded for our partition, so just skip over it.
            } else {
                if (!seenPartitionPositions.containsKey(partition)) {
                    // we haven't seen a partition that we have a bound for
                    return false;
                } else if (seenPartitionPositions.get(partition) < partitionBounds.get(partition)) {
                    // our current position is behind the bound
                    return false;
                }
            }
        }
        return true;
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runRangeQuery(
        final Query query,
        final PositionBound positionBound,
        final QueryConfig config,
        final StateStore store
    ) {
        if (!(store instanceof KeyValueStore)) {
            return QueryResult.forUnknownQueryType(query, store);
        }
        final KeyValueStore kvStore = (KeyValueStore) store;
        final RangeQuery rangeQuery = (RangeQuery) query;
        final Optional lowerRange = rangeQuery.getLowerBound();
        final Optional upperRange = rangeQuery.getUpperBound();
        final ResultOrder order = rangeQuery.resultOrder();
        final KeyValueIterator iterator;
        try {
            if (!lowerRange.isPresent() && !upperRange.isPresent() && !order.equals(ResultOrder.DESCENDING)) {
                iterator = kvStore.all();
            } else if (!order.equals(ResultOrder.DESCENDING)) {
                iterator = kvStore.range(lowerRange.orElse(null), upperRange.orElse(null));
            } else if (!lowerRange.isPresent() && !upperRange.isPresent()) {
                iterator = kvStore.reverseAll();
            } else {
                iterator = kvStore.reverseRange(lowerRange.orElse(null), upperRange.orElse(null));
            }
            final R result = (R) iterator;
            return QueryResult.forResult(result);
        } catch (final Exception e) {
            final String message = parseStoreException(e, store, query);
            return QueryResult.forFailure(
                FailureReason.STORE_EXCEPTION,
                message
            );
        }
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runKeyQuery(final Query query,
                                                  final PositionBound positionBound,
                                                  final QueryConfig config,
                                                  final StateStore store) {

        if (store instanceof KeyValueStore) {
            final KeyQuery rawKeyQuery = (KeyQuery) query;
            final KeyValueStore keyValueStore =
                (KeyValueStore) store;
            try {
                final byte[] bytes = keyValueStore.get(rawKeyQuery.getKey());
                return (QueryResult) QueryResult.forResult(bytes);
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(
                    FailureReason.STORE_EXCEPTION,
                    message
                );
            }
        } else {
            return QueryResult.forUnknownQueryType(query, store);
        }
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runWindowKeyQuery(final Query query,
                                                        final PositionBound positionBound,
                                                        final QueryConfig config,
                                                        final StateStore store) {
        if (store instanceof WindowStore) {
            final WindowKeyQuery windowKeyQuery =
                (WindowKeyQuery) query;
            final WindowStore windowStore = (WindowStore) store;
            try {
                if (windowKeyQuery.getTimeFrom().isPresent() && windowKeyQuery.getTimeTo().isPresent()) {
                    final WindowStoreIterator iterator = windowStore.fetch(
                        windowKeyQuery.getKey(),
                        windowKeyQuery.getTimeFrom().get(),
                        windowKeyQuery.getTimeTo().get()
                    );
                    return (QueryResult) QueryResult.forResult(iterator);
                } else {
                    return QueryResult.forFailure(
                        FailureReason.UNKNOWN_QUERY_TYPE,
                        "This store (" + store.getClass() + ") doesn't know how to"
                            + " execute the given query (" + query + ") because it only supports"
                            + " closed-range queries."
                            + " Contact the store maintainer if you need support"
                            + " for a new query type."
                    );
                }
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(FailureReason.STORE_EXCEPTION, message);
            }
        } else {
            return QueryResult.forUnknownQueryType(query, store);
        }
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runWindowRangeQuery(final Query query,
                                                          final PositionBound positionBound,
                                                          final QueryConfig config,
                                                          final StateStore store) {
        if (store instanceof WindowStore) {
            final WindowRangeQuery windowRangeQuery =
                (WindowRangeQuery) query;
            final WindowStore windowStore = (WindowStore) store;
            try {
                // There's no store API for open time ranges
                if (windowRangeQuery.getTimeFrom().isPresent() && windowRangeQuery.getTimeTo().isPresent()) {
                    final KeyValueIterator, byte[]> iterator =
                        windowStore.fetchAll(
                            windowRangeQuery.getTimeFrom().get(),
                            windowRangeQuery.getTimeTo().get()
                        );
                    return (QueryResult) QueryResult.forResult(iterator);
                } else {
                    return QueryResult.forFailure(
                        FailureReason.UNKNOWN_QUERY_TYPE,
                        "This store (" + store.getClass() + ") doesn't know how to"
                            + " execute the given query (" + query + ") because"
                            + " WindowStores only supports WindowRangeQuery.withWindowStartRange."
                            + " Contact the store maintainer if you need support"
                            + " for a new query type."
                    );
                }
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(
                    FailureReason.STORE_EXCEPTION,
                    message
                );
            }
        } else if (store instanceof SessionStore) {
            final WindowRangeQuery windowRangeQuery =
                (WindowRangeQuery) query;
            final SessionStore sessionStore = (SessionStore) store;
            try {
                if (windowRangeQuery.getKey().isPresent()) {
                    final KeyValueIterator, byte[]> iterator = sessionStore.fetch(
                        windowRangeQuery.getKey().get());
                    return (QueryResult) QueryResult.forResult(iterator);
                } else {
                    return QueryResult.forFailure(
                        FailureReason.UNKNOWN_QUERY_TYPE,
                        "This store (" + store.getClass() + ") doesn't know how to"
                            + " execute the given query (" + query + ") because"
                            + " SessionStores only support WindowRangeQuery.withKey."
                            + " Contact the store maintainer if you need support"
                            + " for a new query type."
                    );
                }
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(
                    FailureReason.STORE_EXCEPTION,
                    message
                );
            }
        } else {
            return QueryResult.forUnknownQueryType(query, store);
        }
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runVersionedKeyQuery(final Query query,
                                                           final PositionBound positionBound,
                                                           final QueryConfig config,
                                                           final StateStore store) {
        if (store instanceof VersionedKeyValueStore) {
            final VersionedKeyValueStore versionedKeyValueStore =
                (VersionedKeyValueStore) store;
            final VersionedKeyQuery rawKeyQuery =
                (VersionedKeyQuery) query;
            try {
                final VersionedRecord bytes;
                if (((VersionedKeyQuery) query).asOfTimestamp().isPresent()) {
                    bytes = versionedKeyValueStore.get(rawKeyQuery.key(),
                        ((VersionedKeyQuery) query).asOfTimestamp().get().toEpochMilli());
                } else {
                    bytes = versionedKeyValueStore.get(rawKeyQuery.key());
                }
                return (QueryResult) QueryResult.forResult(bytes);
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(
                    FailureReason.STORE_EXCEPTION,
                    message
                );
            }

        } else {
            return QueryResult.forUnknownQueryType(query, store);
        }
    }

    @SuppressWarnings("unchecked")
    private static  QueryResult runMultiVersionedKeyQuery(final Query query,
                                                                final PositionBound positionBound,
                                                                final QueryConfig config,
                                                                final StateStore store) {

        if (store instanceof VersionedKeyValueStore) {
            final RocksDBVersionedStore rocksDBVersionedStore = (RocksDBVersionedStore) store;
            final MultiVersionedKeyQuery rawKeyQuery = (MultiVersionedKeyQuery) query;
            try {
                final VersionedRecordIterator segmentIterator =
                        rocksDBVersionedStore.get(rawKeyQuery.key(),
                                                  rawKeyQuery.fromTime().get().toEpochMilli(),
                                                  rawKeyQuery.toTime().get().toEpochMilli(),
                                                  rawKeyQuery.resultOrder());
                return (QueryResult) QueryResult.forResult(segmentIterator);
            } catch (final Exception e) {
                final String message = parseStoreException(e, store, query);
                return QueryResult.forFailure(FailureReason.STORE_EXCEPTION, message);
            }
        } else {
            return QueryResult.forUnknownQueryType(query, store);
        }
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    public static  Function getDeserializeValue(final StateSerdes serdes, final StateStore wrapped) {
        final Serde valueSerde = serdes.valueSerde();
        final boolean timestamped = WrappedStateStore.isTimestamped(wrapped) || isAdapter(wrapped);
        final Deserializer deserializer;
        if (!timestamped && valueSerde instanceof ValueAndTimestampSerde) {
            final ValueAndTimestampDeserializer valueAndTimestampDeserializer =
                (ValueAndTimestampDeserializer) ((ValueAndTimestampSerde) valueSerde).deserializer();
            deserializer = (Deserializer) valueAndTimestampDeserializer.valueDeserializer;
        } else {
            deserializer = valueSerde.deserializer();
        }
        return byteArray -> deserializer.deserialize(serdes.topic(), byteArray);
    }

    public static boolean isAdapter(final StateStore stateStore) {
        if (stateStore instanceof KeyValueToTimestampedKeyValueByteStoreAdapter) {
            return true;
        } else if (stateStore instanceof WrappedStateStore) {
            return isAdapter(((WrappedStateStore) stateStore).wrapped());
        } else {
            return false;
        }
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    public static  Function, VersionedRecord> getDeserializeValue(final StateSerdes serdes) {
        final Serde valueSerde = serdes.valueSerde();
        final Deserializer deserializer = valueSerde.deserializer();
        return rawVersionedRecord -> rawVersionedRecord.validTo().isPresent() ? new VersionedRecord<>(deserializer.deserialize(serdes.topic(), rawVersionedRecord.value()),
                                                                                                      rawVersionedRecord.timestamp(),
                                                                                                      rawVersionedRecord.validTo().get())
                                                                              : new VersionedRecord<>(deserializer.deserialize(serdes.topic(), rawVersionedRecord.value()),
                                                                                                      rawVersionedRecord.timestamp());
    }

    public static  VersionedRecord deserializeVersionedRecord(final StateSerdes serdes, final VersionedRecord rawVersionedRecord) {
        final Deserializer valueDeserializer = serdes.valueDeserializer();
        final V value = valueDeserializer.deserialize(serdes.topic(), rawVersionedRecord.value());
        return rawVersionedRecord.validTo().isPresent() ? new VersionedRecord<>(value, rawVersionedRecord.timestamp(), rawVersionedRecord.validTo().get())
                                                        : new VersionedRecord<>(value, rawVersionedRecord.timestamp());
    }

    public static void checkpointPosition(final OffsetCheckpoint checkpointFile, final Position position) {
        try {
            checkpointFile.write(positionToTopicPartitionMap(position));
        } catch (final IOException e) {
            throw new ProcessorStateException("Error writing checkpoint file", e);
        }
    }

    public static Position readPositionFromCheckpoint(final OffsetCheckpoint checkpointFile) {
        try {
            return topicPartitionMapToPosition(checkpointFile.read());
        } catch (final IOException e) {
            throw new ProcessorStateException("Error reading checkpoint file", e);
        }
    }

    private static Map positionToTopicPartitionMap(final Position position) {
        final Map topicPartitions = new HashMap<>();
        final Set topics = position.getTopics();
        for (final String t : topics) {
            final Map partitions = position.getPartitionPositions(t);
            for (final Entry e : partitions.entrySet()) {
                final TopicPartition tp = new TopicPartition(t, e.getKey());
                topicPartitions.put(tp, e.getValue());
            }
        }
        return topicPartitions;
    }

    private static Position topicPartitionMapToPosition(final Map topicPartitions) {
        final Map> pos = new HashMap<>();
        for (final Entry e : topicPartitions.entrySet()) {
            pos
                .computeIfAbsent(e.getKey().topic(), t -> new HashMap<>())
                .put(e.getKey().partition(), e.getValue());
        }
        return Position.fromMap(pos);
    }

    private static  String parseStoreException(final Exception e, final StateStore store, final Query query) {
        final StringWriter stringWriter = new StringWriter();
        final PrintWriter printWriter = new PrintWriter(stringWriter);
        printWriter.println(
            store.getClass() + " failed to handle query " + query + ":");
        e.printStackTrace(printWriter);
        printWriter.flush();
        return stringWriter.toString();
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy