Please wait. This can take some minutes ...
Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
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();
}
}