From 962bbd5b25f4595bd175b368dc61e797acb472a5 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Wed, 1 Jul 2026 19:53:33 -0400 Subject: [PATCH 01/17] Add ReadOnlyRecordIterator IQv2 result type (KIP-1356) --- .../streams/state/ReadOnlyRecordIterator.java | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java new file mode 100644 index 0000000000000..a4d5ce7964885 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java @@ -0,0 +1,45 @@ +/* + * 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; + +import org.apache.kafka.common.annotation.InterfaceStability.Evolving; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; + +import java.io.Closeable; +import java.util.Iterator; + + +/** + * Iterator interface of {@link ReadOnlyRecord ReadOnlyRecord}. + *

+ * This is the result type of the headers-aware range/window/session IQv2 query types: it yields + * {@link ReadOnlyRecord} elements directly, so each element carries its own key, value, timestamp, + * and headers. + *

+ * Users must call its {@code close} method explicitly upon completeness to release resources, + * or use try-with-resources statement (available since JDK7) for this {@link Closeable} class. + * Note that {@code remove()} is not supported. + * + * @param Type of keys + * @param Type of values + */ +@Evolving +public interface ReadOnlyRecordIterator extends Iterator>, Closeable { + + @Override + void close(); +} From 506aa6acce994742bb641a57274e9bfe30b441cd Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Fri, 3 Jul 2026 08:32:19 -0400 Subject: [PATCH 02/17] Add TimestampedRangeWithHeadersQuery IQv2 query type (KIP-1356) - Add the headers-aware range query TimestampedRangeWithHeadersQuery, returning a ReadOnlyRecordIterator of ReadOnlyRecord (key/value/timestamp/headers). - MeteredTimestampedKeyValueStoreWithHeaders handles it by forwarding a raw RangeQuery to the wrapped store and adapting each entry into a frozen-header ReadOnlyRecord; an entry whose stored timestamp is negative cannot be represented as a ReadOnlyRecord, so advancing the iterator throws StreamsException. - Also drop the query() override on RocksDBTimestampedStoreWithHeaders so the native header store serves the basic IQv2 queries via the inherited RocksDBStore handling. --- .../TimestampedRangeWithHeadersQuery.java | 154 ++++++++++++++++++ ...edTimestampedKeyValueStoreWithHeaders.java | 134 +++++++++++++++ .../RocksDBTimestampedStoreWithHeaders.java | 34 ---- 3 files changed, 288 insertions(+), 34 deletions(-) create mode 100644 streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java diff --git a/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java new file mode 100644 index 0000000000000..db49063b3fbc7 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java @@ -0,0 +1,154 @@ +/* + * 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.query; + +import org.apache.kafka.common.annotation.InterfaceStability.Evolving; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; +import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; + +import java.util.Optional; + +/** + * Interactive query for issuing range queries and scans over a + * {@link TimestampedKeyValueStoreWithHeaders}, returning each record together with its headers. + * + *

This is the headers-aware parallel of {@link TimestampedRangeQuery}: it returns a + * {@link ReadOnlyRecordIterator} of {@link ReadOnlyRecord} elements, each carrying the key, value, + * timestamp, and headers, whereas {@link TimestampedRangeQuery} returns a + * {@link org.apache.kafka.streams.state.KeyValueIterator} of + * {@link org.apache.kafka.streams.state.ValueAndTimestamp} (value and timestamp only, no headers). + * + *

A range query retrieves a set of records, specified using an upper and/or lower bound on the + * keys. A scan query (no bounds) retrieves all records contained in the store. Keys' order is based + * on the serialized {@code byte[]} of the keys, not the 'logical' key order. + * + *

Headers are persisted and returned only when the store is backed by a native headers store, + * i.e. built with a KIP-1271 {@code WithHeaders} byte-store supplier (e.g. + * {@code Stores.persistentTimestampedKeyValueStoreWithHeaders}). A {@code WithHeaders} store built + * over a legacy (non-headers) supplier cannot persist headers, so the store-served reads come back + * with empty {@code headers()}. + * + *

Against a plain store not built with the {@code WithHeaders} builder at all, this query type is + * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}. + * + *

Each element is a {@link ReadOnlyRecord}, whose timestamp is non-negative. If the backing store + * does not persist timestamps -- for example a {@code WithHeaders} store built over a plain + * {@link org.apache.kafka.streams.state.KeyValueStore} supplier, which surfaces every entry with + * {@code NO_TIMESTAMP} ({@code -1}) -- that entry cannot be represented, so advancing the returned + * {@link ReadOnlyRecordIterator} throws {@link org.apache.kafka.streams.errors.StreamsException} at + * that entry. Back the store with {@code Stores.persistentTimestampedKeyValueStoreWithHeaders(...)} to + * persist timestamps and headers, or use {@link TimestampedRangeQuery} if headers are not needed. + * + * @param Type of keys + * @param Type of values + */ +@Evolving +public final class TimestampedRangeWithHeadersQuery implements Query> { + + private final Optional lower; + private final Optional upper; + private final ResultOrder order; + + private TimestampedRangeWithHeadersQuery(final Optional lower, final Optional upper, final ResultOrder order) { + this.lower = lower; + this.upper = upper; + this.order = order; + } + + /** + * Interactive range query using a lower and upper bound to filter the keys returned. + * @param lower The key that specifies the lower bound of the range + * @param upper The key that specifies the upper bound of the range + * @param The key type + * @param The value type + */ + public static TimestampedRangeWithHeadersQuery withRange(final K lower, final K upper) { + return new TimestampedRangeWithHeadersQuery<>(Optional.ofNullable(lower), Optional.ofNullable(upper), ResultOrder.ANY); + } + + /** + * Interactive range query using an upper bound to filter the keys returned. + * @param upper The key that specifies the upper bound of the range + * @param The key type + * @param The value type + */ + public static TimestampedRangeWithHeadersQuery withUpperBound(final K upper) { + return new TimestampedRangeWithHeadersQuery<>(Optional.empty(), Optional.of(upper), ResultOrder.ANY); + } + + /** + * Interactive range query using a lower bound to filter the keys returned. + * @param lower The key that specifies the lower bound of the range + * @param The key type + * @param The value type + */ + public static TimestampedRangeWithHeadersQuery withLowerBound(final K lower) { + return new TimestampedRangeWithHeadersQuery<>(Optional.of(lower), Optional.empty(), ResultOrder.ANY); + } + + /** + * Interactive scan query that returns all records in the store. + * @param The key type + * @param The value type + */ + public static TimestampedRangeWithHeadersQuery withNoBounds() { + return new TimestampedRangeWithHeadersQuery<>(Optional.empty(), Optional.empty(), ResultOrder.ANY); + } + + /** + * Set the query to return the serialized {@code byte[]} of the keys in descending order. + * Order is based on the serialized {@code byte[]} of the keys, not the 'logical' key order. + * @return a new query instance with the descending flag set. + */ + public TimestampedRangeWithHeadersQuery withDescendingKeys() { + return new TimestampedRangeWithHeadersQuery<>(this.lower, this.upper, ResultOrder.DESCENDING); + } + + /** + * Set the query to return the serialized {@code byte[]} of the keys in ascending order. + * Order is based on the serialized {@code byte[]} of the keys, not the 'logical' key order. + * @return a new query instance with the ascending flag set. + */ + public TimestampedRangeWithHeadersQuery withAscendingKeys() { + return new TimestampedRangeWithHeadersQuery<>(this.lower, this.upper, ResultOrder.ASCENDING); + } + + /** + * The lower bound of the query, if specified. + */ + public Optional lowerBound() { + return lower; + } + + /** + * The upper bound of the query, if specified. + */ + public Optional upperBound() { + return upper; + } + + /** + * Determines if the serialized {@code byte[]} of the keys in ascending or descending or unordered order. + * Order is based on the serialized {@code byte[]} of the keys, not the 'logical' key order. + * @return the order of the returned records based on the serialized {@code byte[]} of the keys + * (can be unordered, ascending, or descending). + */ + public ResultOrder resultOrder() { + return order; + } +} diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java index 98b6a024b221c..f1f224f9f58ff 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java @@ -26,6 +26,7 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.errors.ProcessorStateException; +import org.apache.kafka.streams.errors.StreamsException; import org.apache.kafka.streams.processor.api.ReadOnlyRecord; import org.apache.kafka.streams.processor.api.Record; import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; @@ -41,10 +42,12 @@ import org.apache.kafka.streams.query.TimestampedKeyQuery; import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; import org.apache.kafka.streams.query.TimestampedRangeQuery; +import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; import org.apache.kafka.streams.query.internals.InternalQueryResultUtil; import org.apache.kafka.streams.state.KeyValueIterator; import org.apache.kafka.streams.state.KeyValueStore; import org.apache.kafka.streams.state.ReadOnlyKeyValueStore; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.ValueTimestampHeaders; @@ -106,6 +109,10 @@ public class MeteredTimestampedKeyValueStoreWithHeaders mkEntry( TimestampedRangeQuery.class, (query, positionBound, config, store) -> runTimestampedRangeQuery(query, positionBound, config) + ), + mkEntry( + TimestampedRangeWithHeadersQuery.class, + (query, positionBound, config, store) -> runTimestampedRangeWithHeadersQuery(query, positionBound, config) ) ); @@ -542,6 +549,50 @@ private QueryResult runTimestampedRangeQuery( return result; } + @SuppressWarnings("unchecked") + private QueryResult runTimestampedRangeWithHeadersQuery( + final Query query, + final PositionBound positionBound, + final QueryConfig config + ) { + final QueryResult result; + final TimestampedRangeWithHeadersQuery typedQuery = (TimestampedRangeWithHeadersQuery) query; + + RangeQuery rawRangeQuery; + final ResultOrder order = typedQuery.resultOrder(); + rawRangeQuery = RangeQuery.withRange( + serializeKey(typedQuery.lowerBound().orElse(null), internalContext.headers()), + serializeKey(typedQuery.upperBound().orElse(null), internalContext.headers()) + ); + if (order.equals(ResultOrder.DESCENDING)) { + rawRangeQuery = rawRangeQuery.withDescendingKeys(); + } + if (order.equals(ResultOrder.ASCENDING)) { + rawRangeQuery = rawRangeQuery.withAscendingKeys(); + } + + final QueryResult> rawResult = wrapped().query(rawRangeQuery, positionBound, config); + if (rawResult.isSuccess()) { + final KeyValueIterator iterator = rawResult.getResult(); + final ReadOnlyRecordIterator resultIterator = + new MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator( + iterator, + getSensor, + StoreQueryUtils.deserializeValue(serdes, wrapped()) + ); + final QueryResult> typedQueryResult = + InternalQueryResultUtil.copyAndSubstituteDeserializedResult( + rawResult, + resultIterator + ); + result = (QueryResult) typedQueryResult; + } else { + // the generic type doesn't matter, since failed queries have no result set. + result = (QueryResult) rawResult; + } + return result; + } + @Override public , P> KeyValueIterator> prefixScan( final P prefix, final PS prefixKeySerializer @@ -696,6 +747,89 @@ public K peekNextKey() { } } + /** + * Iterator backing {@link TimestampedRangeWithHeadersQuery}: yields each entry as a + * {@link ReadOnlyRecord} (implemented by {@link Record}) carrying key, value, timestamp, and the + * stored headers, with the headers frozen so a caller cannot mutate the read-only result. + * + *

A {@link ReadOnlyRecord} timestamp is contractually non-negative, so an entry with a negative + * stored timestamp cannot be represented. The dominant, deterministic cause is a store that does + * not persist timestamps: a {@code WithHeaders} store built over a plain {@link KeyValueStore} + * supplier surfaces every entry with {@code NO_TIMESTAMP} (-1). A genuine negative write is + * otherwise blocked (the source {@code RecordQueue} drops negative-timestamp records at ingestion), + * though {@link ValueAndTimestamp#make}/{@link ValueTimestampHeaders#make} do not themselves reject + * one written directly. + * + *

This mirrors the rule the point query {@link TimestampedKeyWithHeadersQuery} applies. But + * because a lazily-evaluated range has already returned a successful {@link QueryResult} before any + * entry is read, such an entry cannot be surfaced as a query-level failure; it is instead reported + * by throwing a {@link StreamsException} while advancing the iterator. + */ + private class MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator + implements ReadOnlyRecordIterator, MeteredIterator { + + private final KeyValueIterator iter; + private final Sensor sensor; + private final long startNs; + private final long startTimestampMs; + private final Function> valueTimestampHeadersDeserializer; + + private MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator( + final KeyValueIterator iter, + final Sensor sensor, + final Function> valueTimestampHeadersDeserializer + ) { + this.iter = iter; + this.sensor = sensor; + this.valueTimestampHeadersDeserializer = valueTimestampHeadersDeserializer; + this.startNs = time.nanoseconds(); + this.startTimestampMs = time.milliseconds(); + openIterators.add(this); + } + + @Override + public long startTimestamp() { + return startTimestampMs; + } + + @Override + public boolean hasNext() { + return iter.hasNext(); + } + + @Override + public ReadOnlyRecord next() { + final KeyValue keyValue = iter.next(); + final ValueTimestampHeaders valueTimestampHeaders = valueTimestampHeadersDeserializer.apply(keyValue.value); + final Headers headers = valueTimestampHeaders.headers(); + final K key = deserializeKey(keyValue.key.get(), headers); + if (valueTimestampHeaders.timestamp() < 0) { + throw new StreamsException( + "Cannot represent the stored record for key [" + key + "] as a ReadOnlyRecord: its " + + "timestamp (" + valueTimestampHeaders.timestamp() + ") is negative."); + } + final Record record = new Record<>( + key, + valueTimestampHeaders.value(), + valueTimestampHeaders.timestamp(), + headers); + ((RecordHeaders) record.headers()).setReadOnly(); + return record; + } + + @Override + public void close() { + try { + iter.close(); + } finally { + final long duration = time.nanoseconds() - startNs; + sensor.record(duration); + iteratorDurationSensor.record(duration); + openIterators.remove(this); + } + } + } + private class MeteredTimestampedKeyValueStoreWithHeadersIterator implements KeyValueIterator>, MeteredIterator { private final KeyValueIterator iter; private final Sensor sensor; diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeaders.java index d0ad49acfcd17..931889c58d205 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeaders.java @@ -18,11 +18,6 @@ package org.apache.kafka.streams.state.internals; import org.apache.kafka.streams.errors.ProcessorStateException; -import org.apache.kafka.streams.query.KeyQuery; -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.state.HeadersBytesStore; import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder; @@ -195,33 +190,4 @@ private void openFromTimestampedStore(final DBOptions dbOptions, } } - @SuppressWarnings("SynchronizeOnNonFinalField") - @Override - public QueryResult query(final Query query, - final PositionBound positionBound, - final QueryConfig config) { - // KIP-1356 (point query only): the headers-aware TimestampedKeyWithHeadersQuery forwards a raw - // KeyQuery to this native store, so enable KeyQuery via the inherited RocksDBStore handling. - // Every other query type (e.g. RangeQuery) stays UNKNOWN_QUERY_TYPE here, unchanged from before; - // those are enabled by their own KIP-1356 follow-ups. - if (query instanceof KeyQuery) { - return super.query(query, positionBound, config); - } - - final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L; - final QueryResult result; - - synchronized (position) { - result = QueryResult.forUnknownQueryType(query, this); - - if (config.isCollectExecutionInfo()) { - result.addExecutionInfo( - "Handled in " + this.getClass() + " in " + (System.nanoTime() - start) + "ns" - ); - } - result.setPosition(position.copy()); - } - return result; - } - } From 020ea21c8bce35f08ae7c382af8edde286948f4f Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Fri, 3 Jul 2026 08:32:19 -0400 Subject: [PATCH 03/17] Add range-query tests for TimestampedRangeWithHeadersQuery (KIP-1356) - TimestampedKeyValueStoreBuilderWithHeadersTest: end-to-end build-path coverage (native/in-memory return headers; adapter returns empty headers; plain-legacy and negative timestamps throw while iterating), execution-info collection, and native/adapter parity. - MeteredTimestampedKeyValueStoreWithHeadersTest: unit coverage of the metered handler forwarding the bounds and result order onto the raw RangeQuery. - RocksDBTimestampedStoreWithHeadersTest: the native store now serves RangeQuery. - TimestampedRangeWithHeadersQueryIntegrationTest: IQv2 integration coverage. --- ...dRangeWithHeadersQueryIntegrationTest.java | 372 ++++++++++++++++++ ...mestampedKeyValueStoreWithHeadersTest.java | 22 ++ ...ocksDBTimestampedStoreWithHeadersTest.java | 27 +- ...edKeyValueStoreBuilderWithHeadersTest.java | 327 +++++++++++---- 4 files changed, 670 insertions(+), 78 deletions(-) create mode 100644 streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java new file mode 100644 index 0000000000000..487774a87408b --- /dev/null +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java @@ -0,0 +1,372 @@ +/* + * 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.integration; + +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.serialization.IntegerDeserializer; +import org.apache.kafka.common.serialization.IntegerSerializer; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.common.serialization.StringSerializer; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.StreamsBuilder; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.errors.StreamsException; +import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; +import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; +import org.apache.kafka.streams.kstream.Consumed; +import org.apache.kafka.streams.kstream.Produced; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.query.FailureReason; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.StateQueryRequest; +import org.apache.kafka.streams.query.StateQueryResult; +import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; +import org.apache.kafka.streams.state.Stores; +import org.apache.kafka.streams.state.TimestampedKeyValueStore; +import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.ValueAndTimestamp; +import org.apache.kafka.streams.state.ValueTimestampHeaders; +import org.apache.kafka.test.TestUtils; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.stream.Collectors; + +import static org.apache.kafka.streams.query.StateQueryRequest.inStore; +import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@Tag("integration") +public class TimestampedRangeWithHeadersQueryIntegrationTest { + + private static final String STORE_NAME = "headers-store"; + + private String inputStream; + private String outputStream; + private long baseTimestamp; + + private KafkaStreams kafkaStreams; + + private static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1); + + private static final Headers HEADERS1 = new RecordHeaders() + .add("source", "test".getBytes()) + .add("version", "1.0".getBytes()); + + private static final Headers HEADERS2 = new RecordHeaders() + .add("source", "test".getBytes()) + .add("version", "2.0".getBytes()); + + private static final Headers EMPTY_HEADERS = new RecordHeaders(); + + public TestInfo testInfo; + + @BeforeAll + public static void before() throws IOException { + CLUSTER.start(); + } + + @AfterAll + public static void after() { + CLUSTER.stop(); + } + + @BeforeEach + public void beforeTest(final TestInfo testInfo) throws InterruptedException { + this.testInfo = testInfo; + final String uniqueTestName = safeUniqueTestName(testInfo); + inputStream = "input-stream-" + uniqueTestName; + outputStream = "output-stream-" + uniqueTestName; + // single partition so the query has exactly one partition result + CLUSTER.createTopic(inputStream, 1, 1); + CLUSTER.createTopic(outputStream, 1, 1); + baseTimestamp = CLUSTER.time.milliseconds(); + } + + @AfterEach + public void afterTest() { + if (kafkaStreams != null) { + kafkaStreams.close(Duration.ofSeconds(30L)); + kafkaStreams.cleanUp(); + } + } + + @Test + public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { + final StreamsBuilder streamsBuilder = new StreamsBuilder(); + streamsBuilder + .addStateStore( + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), + Serdes.Integer(), + Serdes.String() + ) + ) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(HeadersStoreWriterProcessor::new, STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); + + kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); + kafkaStreams.start(); + + // Write: keys 1,2 (HEADERS1), key 3 (empty headers), key 4 (then tombstone it). + int produced = 0; + produced += produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + produced += produceWithHeaders(baseTimestamp + 1, EMPTY_HEADERS, KeyValue.pair(3, "three")); + produced += produceWithHeaders(baseTimestamp + 2, HEADERS2, KeyValue.pair(4, "four")); + produced += produceWithHeaders(baseTimestamp + 3, EMPTY_HEADERS, KeyValue.pair(4, null)); // tombstone + + awaitProcessed(produced); + + // Full scan, ascending: keys 1, 2, 3 (key 4 tombstoned and omitted). + final List> ascending = runQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + assertEquals(List.of(1, 2, 3), keys(ascending)); + assertEquals(List.of("one", "two", "three"), values(ascending)); + + // Headers, timestamps, and key are carried on each record. + assertEquals(HEADERS1, ascending.get(0).headers()); + assertEquals(baseTimestamp, ascending.get(0).timestamp()); + assertEquals(HEADERS1, ascending.get(1).headers()); + // Record written without headers round-trips as empty (never null) headers. + final Headers emptyHeaders = ascending.get(2).headers(); + assertNotNull(emptyHeaders); + assertEquals(0, emptyHeaders.toArray().length); + assertEquals(baseTimestamp + 1, ascending.get(2).timestamp()); + + // Returned headers are a read-only snapshot. + assertTrue(headersAreReadOnly(ascending.get(0).headers())); + + // Full scan, descending. + final List> descending = + runQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withDescendingKeys()); + assertEquals(List.of(3, 2, 1), keys(descending)); + + // Bounded ranges. + assertEquals(List.of(2, 3), keys(runQuery(TimestampedRangeWithHeadersQuery.withRange(2, 3)))); + assertEquals(List.of(2, 3), keys(runQuery(TimestampedRangeWithHeadersQuery.withLowerBound(2)))); + assertEquals(List.of(1, 2), keys(runQuery(TimestampedRangeWithHeadersQuery.withUpperBound(2)))); + } + + @Test + public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Exception { + final StreamsBuilder streamsBuilder = new StreamsBuilder(); + streamsBuilder + .addStateStore( + Stores.timestampedKeyValueStoreBuilder( + Stores.persistentTimestampedKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String() + ) + ) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(NonHeadersStoreWriterProcessor::new, STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); + + kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); + kafkaStreams.start(); + + final int produced = produceWithHeaders(baseTimestamp, EMPTY_HEADERS, KeyValue.pair(1, "one")); + awaitProcessed(produced); + + final StateQueryRequest> request = + inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + final Map>> partitionResults = + result.getPartitionResults(); + assertFalse(partitionResults.isEmpty()); + for (final QueryResult> partitionResult : partitionResults.values()) { + assertTrue(partitionResult.isFailure()); + assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, partitionResult.getFailureReason()); + } + } + + @Test + public void shouldThrowForPlainSupplierWithNoRepresentableTimestamp() throws Exception { + // WithHeaders builder over a plain (non-timestamped) supplier: the store cannot persist a + // timestamp, so entries come back with timestamp = -1, which cannot be a ReadOnlyRecord. The + // query itself succeeds, but advancing the iterator throws. + final StreamsBuilder streamsBuilder = new StreamsBuilder(); + streamsBuilder + .addStateStore( + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String() + ) + ) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(HeadersStoreWriterProcessor::new, STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); + + kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); + kafkaStreams.start(); + + final int produced = produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + awaitProcessed(produced); + + final StateQueryRequest> request = + inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + final QueryResult> partitionResult = result.getOnlyPartitionResult(); + assertTrue(partitionResult.isSuccess()); + try (ReadOnlyRecordIterator iterator = partitionResult.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } + } + + private List> runQuery(final TimestampedRangeWithHeadersQuery query) { + final StateQueryRequest> request = inStore(STORE_NAME).withQuery(query); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + final List> records = new ArrayList<>(); + try (ReadOnlyRecordIterator iterator = result.getOnlyPartitionResult().getResult()) { + while (iterator.hasNext()) { + records.add(iterator.next()); + } + } + return records; + } + + private static List keys(final List> records) { + return records.stream().map(ReadOnlyRecord::key).collect(Collectors.toList()); + } + + private static List values(final List> records) { + return records.stream().map(ReadOnlyRecord::value).collect(Collectors.toList()); + } + + private static boolean headersAreReadOnly(final Headers headers) { + try { + headers.add("x", new byte[0]); + return false; + } catch (final IllegalStateException expected) { + return true; + } + } + + private void awaitProcessed(final int count) throws Exception { + IntegrationTestUtils.waitUntilMinRecordsReceived( + TestUtils.consumerConfig(CLUSTER.bootstrapServers(), IntegerDeserializer.class, IntegerDeserializer.class), + outputStream, + count); + } + + private Properties props(final String prefix) { + final String safeTestName = safeUniqueTestName(testInfo); + final Properties streamsConfiguration = new Properties(); + streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, prefix + safeTestName); + streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); + streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath()); + streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000L); + streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); + return streamsConfiguration; + } + + @SuppressWarnings("varargs") + @SafeVarargs + private final int produceWithHeaders(final long timestamp, + final Headers headers, + final KeyValue... keyValues) { + IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp( + inputStream, + Arrays.asList(keyValues), + TestUtils.producerConfig(CLUSTER.bootstrapServers(), IntegerSerializer.class, StringSerializer.class), + headers, + timestamp, + false); + return keyValues.length; + } + + /** + * Writes each incoming record into the headers-aware store (deleting on a null value), then + * forwards a marker downstream so the test can wait for processing to complete. + */ + private static class HeadersStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedKeyValueStoreWithHeaders store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + if (record.value() == null) { + store.delete(record.key()); + } else { + store.put(record.key(), + ValueTimestampHeaders.make(record.value(), record.timestamp(), record.headers())); + } + context.forward(record.withValue(1)); + } + } + + /** + * Writes each incoming record into a plain (non-headers) timestamped store, then forwards a + * marker downstream so the test can wait for processing to complete. + */ + private static class NonHeadersStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedKeyValueStore store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + if (record.value() == null) { + store.delete(record.key()); + } else { + store.put(record.key(), ValueAndTimestamp.make(record.value(), record.timestamp())); + } + context.forward(record.withValue(1)); + } + } +} diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java index 568a11a7226c7..d2f77345124fe 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java @@ -41,7 +41,10 @@ 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.TimestampedKeyWithHeadersQuery; +import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; import org.apache.kafka.streams.state.KeyValueIterator; import org.apache.kafka.streams.state.KeyValueStore; import org.apache.kafka.streams.state.ValueTimestampHeaders; @@ -302,6 +305,25 @@ public void shouldGetRangeFromInnerStoreAndRecordRangeMetric() { assertTrue((Double) metric.metricValue() > 0); } + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldPropagateBoundsAndDescendingOrderForTimestampedRangeWithHeadersQuery() { + setUp(); + final ArgumentCaptor captor = ArgumentCaptor.forClass(RangeQuery.class); + when(inner.query(captor.capture(), any(), any())).thenReturn( + QueryResult.forResult(new KeyValueIteratorStub<>(List.>of().iterator()))); + init(); + + final TimestampedRangeWithHeadersQuery query = + TimestampedRangeWithHeadersQuery.withLowerBound("a").withDescendingKeys(); + metered.query(query, PositionBound.unbounded(), new QueryConfig(false)); + + final RangeQuery rawQuery = captor.getValue(); + assertEquals(ResultOrder.DESCENDING, rawQuery.resultOrder()); + assertTrue(rawQuery.getLowerBound().isPresent()); + assertFalse(rawQuery.getUpperBound().isPresent()); + } + @Test public void shouldGetAllFromInnerStoreAndRecordAllMetric() { setUp(); diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeadersTest.java index be048ece6106a..c1b1259d7f8f8 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeadersTest.java @@ -1030,23 +1030,30 @@ public void shouldHandleKeyQuery() { } @Test - public void shouldReturnUnknownQueryTypeForRangeQuery() { + public void shouldHandleRangeQuery() { // Initialize the store rocksDBStore.init(context, rocksDBStore); - // KIP-1356 (point query only): this store enables KeyQuery but leaves RangeQuery (and every - // other non-KeyQuery type) as UNKNOWN_QUERY_TYPE, unchanged from before. RangeQuery support on - // the native header store is deferred to the range KIP-1356 follow-up. + // Store raw (serialized ValueTimestampHeaders) bytes for a key. + final Bytes key = new Bytes("test-key".getBytes()); + final byte[] storedBytes = "headers+timestamp+value".getBytes(); + rocksDBStore.put(key, storedBytes); + + // KIP-1356: the range follow-up enables RangeQuery on the native header store via the inherited + // RocksDBStore handling (returning the raw stored header-format bytes; the metered store does + // the header-aware deserialization), matching the adapter build path. final RangeQuery query = RangeQuery.withNoBounds(); final QueryResult> result = rocksDBStore.query(query, PositionBound.unbounded(), new QueryConfig(false)); - assertFalse(result.isSuccess(), "Expected RangeQuery to be unsupported"); - assertEquals( - FailureReason.UNKNOWN_QUERY_TYPE, - result.getFailureReason(), - "Expected UNKNOWN_QUERY_TYPE failure reason" - ); + assertTrue(result.isSuccess(), "Expected RangeQuery to succeed"); + try (KeyValueIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext(), "Expected the stored key in the range result"); + final KeyValue keyValue = iterator.next(); + assertEquals(key, keyValue.key); + assertArrayEquals(storedBytes, keyValue.value, "Expected the raw stored bytes to be returned"); + assertFalse(iterator.hasNext()); + } assertNotNull(result.getPosition(), "Expected position to be set"); } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java index cccc1174e53b9..81cf34ed081a0 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java @@ -24,6 +24,8 @@ import org.apache.kafka.common.utils.Bytes; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.common.utils.internals.LogContext; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.errors.StreamsException; import org.apache.kafka.streams.processor.StateStore; import org.apache.kafka.streams.processor.api.ReadOnlyRecord; import org.apache.kafka.streams.processor.internals.MockStreamsMetrics; @@ -37,8 +39,12 @@ import org.apache.kafka.streams.query.RangeQuery; import org.apache.kafka.streams.query.TimestampedKeyQuery; import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; +import org.apache.kafka.streams.query.TimestampedRangeQuery; +import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; import org.apache.kafka.streams.state.KeyValueBytesStoreSupplier; +import org.apache.kafka.streams.state.KeyValueIterator; import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.ValueTimestampHeaders; @@ -56,7 +62,9 @@ import org.mockito.quality.Strictness; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -481,58 +489,6 @@ public void shouldFailTimestampedKeyWithHeadersQueryForNegativeStoredTimestamp(f } } - @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void shouldReturnUnknownQueryTypeForRangeQueryOnHeadersStore(final boolean cachingEnabled) { - // KIP-1356 (point query only): the native header store enables KeyQuery but not RangeQuery, so a - // RangeQuery reports UNKNOWN_QUERY_TYPE (caching on or off). RangeQuery support on the native - // build is deferred to the range follow-up. - final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE, cachingEnabled); - try { - final QueryResult result = - store.query(RangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); - assertFalse(result.isSuccess(), "Expected RangeQuery to be unsupported on the native header store"); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getFailureReason()); - assertNotNull(result.getPosition(), "Expected position to be set"); - } finally { - store.close(); - } - } - - @ParameterizedTest - @CsvSource({"NATIVE", "ADAPTER", "IN_MEMORY"}) - public void shouldReturnEmptyPositionInitially(final StoreType storeType) { - final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); - try { - final Position position = ((WrappedStateStore) store).wrapped().getPosition(); - assertNotNull(position, "Expected non-null position"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); - } finally { - store.close(); - } - } - - @ParameterizedTest - @CsvSource({"NATIVE", "ADAPTER"}) - public void shouldCollectExecutionInfoWhenRequested(final StoreType storeType) { - final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); - try { - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - final QueryResult result = wrapped.query( - KeyQuery.withKey(new Bytes("k".getBytes())), PositionBound.unbounded(), new QueryConfig(true)); - - final String executionInfo = String.join("\n", result.getExecutionInfo()); - assertFalse(executionInfo.isEmpty(), "Expected execution info to be collected"); - assertTrue(executionInfo.contains("Handled in"), "Expected execution info to contain handling information"); - final String expectedClass = storeType == StoreType.NATIVE - ? RocksDBTimestampedStoreWithHeaders.class.getName() - : TimestampedToHeadersStoreAdapter.class.getName(); - assertTrue(executionInfo.contains(expectedClass), "Expected execution info to mention " + expectedClass); - } finally { - store.close(); - } - } - @Test public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenRequested() { // The typed query, run through the metered handler with execution info enabled, must carry both @@ -544,9 +500,9 @@ public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenReque store.put("bad", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x"))); final QueryResult> ok = - store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(true)); + store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(true)); final QueryResult> bad = - store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(true)); + store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(true)); assertTrue(ok.isSuccess()); assertFalse(bad.isSuccess()); @@ -555,13 +511,13 @@ public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenReque final String okInfo = String.join("\n", ok.getExecutionInfo()); final String badInfo = String.join("\n", bad.getExecutionInfo()); assertTrue( - okInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) - && okInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), - "success execution info missing an entry: " + okInfo); + okInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) + && okInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), + "success execution info missing an entry: " + okInfo); assertTrue( - badInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) - && badInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), - "failure execution info missing an entry (wrapped-store entry must be preserved): " + badInfo); + badInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) + && badInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), + "failure execution info missing an entry (wrapped-store entry must be preserved): " + badInfo); } finally { store.close(); } @@ -577,9 +533,9 @@ public void shouldNotCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenNo store.put("bad", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x"))); final QueryResult> ok = - store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(false)); + store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(false)); final QueryResult> bad = - store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(false)); + store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(false)); assertTrue(ok.isSuccess()); assertFalse(bad.isSuccess()); @@ -603,16 +559,251 @@ public void shouldReturnIdenticalResultsForNativeAndAdapterBuiltStores(final boo adapterStore.put("k", value); assertEquals( - nativeStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - adapterStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - "KeyQuery results should be identical across native and adapter build paths"); + nativeStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + adapterStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + "KeyQuery results should be identical across native and adapter build paths"); assertEquals( - nativeStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - adapterStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - "TimestampedKeyQuery results should be identical across native and adapter build paths"); + nativeStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + adapterStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + "TimestampedKeyQuery results should be identical across native and adapter build paths"); } finally { nativeStore.close(); adapterStore.close(); } } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void shouldHandleRangeQuery(final boolean cachingEnabled) { + // KIP-1356: the native header store now serves RangeQuery (header-stripped results), matching + // the adapter build path (caching on or off). + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE, cachingEnabled); + try { + final QueryResult> result = + store.query(RangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + assertTrue(result.isSuccess(), "Expected RangeQuery to be handled on the native header store"); + try (KeyValueIterator iterator = result.getResult()) { + assertFalse(iterator.hasNext(), "Expected empty result from an empty store"); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @ParameterizedTest + @CsvSource({"NATIVE", "IN_MEMORY"}) + public void shouldReturnHeadersForTimestampedRangeWithHeadersQueryOnHeaderPersistingStore(final StoreType storeType) { + // Range counterpart of shouldReturnHeadersForTimestampedKeyWithHeadersQueryOnHeaderPersistingStore. + // A range query reads the underlying store (it never consults the cache), so use a store-served + // (caching-disabled) build; the native and in-memory builds persist headers, so every element carries them. + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); + try { + final Headers headers = headersWith("h", "x"); + store.put("k", ValueTimestampHeaders.make("v", 123L, headers)); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess(), "Expected TimestampedRangeWithHeadersQuery to succeed"); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord record = iterator.next(); + assertEquals("k", record.key()); + assertEquals("v", record.value()); + assertEquals(123L, record.timestamp()); + assertEquals(headers, record.headers()); + // The IQ result is a read-only snapshot: its headers are immutable. + assertThrows(IllegalStateException.class, () -> record.headers().add("new", new byte[0]), + "IQ result headers should be read-only"); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldReturnEmptyHeadersForTimestampedRangeWithHeadersQueryOnAdapterStore() { + // The timestamped adapter keeps the timestamp but drops headers on write. A range query reads + // the underlying store (never the cache), so the headers always come back empty (never null) -- + // unlike the point query, whose warm-cache read can still return them. + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.ADAPTER, false); + try { + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord record = iterator.next(); + assertEquals("k", record.key()); + assertEquals("v", record.value()); + assertEquals(123L, record.timestamp()); + assertEquals(new RecordHeaders(), record.headers()); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainLegacyStore() { + // A plain (non-timestamped) legacy supplier surfaces every entry with timestamp = -1, which + // cannot be represented as a ReadOnlyRecord. A range query reads the underlying store, so the + // failure surfaces while iterating (there is no cache-served success path like the point query). + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.PLAIN_ADAPTER, false); + try { + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess(), "The range query itself succeeds; the failure surfaces while iterating"); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + assertThrows(StreamsException.class, iterator::next, + "An entry with ts=-1 cannot be represented as a ReadOnlyRecord"); + } + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForNegativeStoredTimestampForTimestampedRangeWithHeadersQuery() { + // A caller can store a negative timestamp directly. Unlike the point query (which fails the whole + // query with STORE_EXCEPTION), a lazily-evaluated range iterator has already been returned, so the + // failure surfaces by throwing while advancing. + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE, false); + try { + store.put("k", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x"))); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } + } finally { + store.close(); + } + } + + @Test + public void shouldCollectExecutionInfoForTimestampedRangeWithHeadersQueryWhenRequested() { + // The range query, run through the metered handler with execution info enabled, must carry both + // the wrapped store's entry and the metered handler's entry. (Unlike the point query there is no + // query-level failure path -- a negative timestamp surfaces while iterating, not as a failed result.) + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE, false); + try { + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(true)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + final String info = String.join("\n", result.getExecutionInfo()); + assertTrue( + info.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) + && info.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), + "execution info missing an entry: " + info); + } + } finally { + store.close(); + } + } + + @Test + public void shouldNotCollectExecutionInfoForTimestampedRangeWithHeadersQueryWhenNotRequested() { + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE, false); + try { + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); + + final QueryResult> result = + store.query(TimestampedRangeWithHeadersQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator iterator = result.getResult()) { + assertTrue(result.getExecutionInfo().isEmpty(), "Expected no execution info: " + result.getExecutionInfo()); + } + } finally { + store.close(); + } + } + + @Test + public void shouldReturnIdenticalRangeResultsForNativeAndAdapterBuiltStores() { + // Build-path parity for the existing (header-stripped) range query types. Range queries read the + // underlying store (never the cache), so use a store-served (caching-disabled) build. + final ValueTimestampHeaders value = ValueTimestampHeaders.make("v", 123L, headersWith("h", "x")); + final TimestampedKeyValueStoreWithHeaders nativeStore = buildAndInitStore(StoreType.NATIVE, false); + final TimestampedKeyValueStoreWithHeaders adapterStore = buildAndInitStore(StoreType.ADAPTER, false); + try { + nativeStore.put("k", value); + adapterStore.put("k", value); + + assertEquals( + drain(nativeStore.query(RangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)).getResult()), + drain(adapterStore.query(RangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)).getResult()), + "RangeQuery results should be identical across native and adapter build paths"); + assertEquals( + drain(nativeStore.query(TimestampedRangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)).getResult()), + drain(adapterStore.query(TimestampedRangeQuery.withNoBounds(), PositionBound.unbounded(), new QueryConfig(false)).getResult()), + "TimestampedRangeQuery results should be identical across native and adapter build paths"); + } finally { + nativeStore.close(); + adapterStore.close(); + } + } + + private static List> drain(final KeyValueIterator iterator) { + final List> out = new ArrayList<>(); + try (iterator) { + while (iterator.hasNext()) { + out.add(iterator.next()); + } + } + return out; + } + + @ParameterizedTest + @CsvSource({"NATIVE", "ADAPTER", "IN_MEMORY"}) + public void shouldReturnEmptyPositionInitially(final StoreType storeType) { + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); + try { + final Position position = ((WrappedStateStore) store).wrapped().getPosition(); + assertNotNull(position, "Expected non-null position"); + assertTrue(position.isEmpty(), "Expected position to be empty initially"); + } finally { + store.close(); + } + } + + @ParameterizedTest + @CsvSource({"NATIVE", "ADAPTER"}) + public void shouldCollectExecutionInfoWhenRequested(final StoreType storeType) { + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); + try { + final StateStore wrapped = ((WrappedStateStore) store).wrapped(); + final QueryResult result = wrapped.query( + KeyQuery.withKey(new Bytes("k".getBytes())), PositionBound.unbounded(), new QueryConfig(true)); + + final String executionInfo = String.join("\n", result.getExecutionInfo()); + assertFalse(executionInfo.isEmpty(), "Expected execution info to be collected"); + assertTrue(executionInfo.contains("Handled in"), "Expected execution info to contain handling information"); + final String expectedClass = storeType == StoreType.NATIVE + ? RocksDBTimestampedStoreWithHeaders.class.getName() + : TimestampedToHeadersStoreAdapter.class.getName(); + assertTrue(executionInfo.contains(expectedClass), "Expected execution info to mention " + expectedClass); + } finally { + store.close(); + } + } } From 79d3ce6864a2bc049328c665bc56ea891f08c695 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Fri, 3 Jul 2026 09:05:41 -0400 Subject: [PATCH 04/17] Add ascending-order range test and restore test indentation (KIP-1356) --- ...mestampedKeyValueStoreWithHeadersTest.java | 19 +++++++++++ ...edKeyValueStoreBuilderWithHeadersTest.java | 32 +++++++++---------- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java index d2f77345124fe..386d3f1a0b260 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java @@ -324,6 +324,25 @@ public void shouldPropagateBoundsAndDescendingOrderForTimestampedRangeWithHeader assertFalse(rawQuery.getUpperBound().isPresent()); } + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldPropagateBoundsAndAscendingOrderForTimestampedRangeWithHeadersQuery() { + setUp(); + final ArgumentCaptor captor = ArgumentCaptor.forClass(RangeQuery.class); + when(inner.query(captor.capture(), any(), any())).thenReturn( + QueryResult.forResult(new KeyValueIteratorStub<>(List.>of().iterator()))); + init(); + + final TimestampedRangeWithHeadersQuery query = + TimestampedRangeWithHeadersQuery.withUpperBound("z").withAscendingKeys(); + metered.query(query, PositionBound.unbounded(), new QueryConfig(false)); + + final RangeQuery rawQuery = captor.getValue(); + assertEquals(ResultOrder.ASCENDING, rawQuery.resultOrder()); + assertFalse(rawQuery.getLowerBound().isPresent()); + assertTrue(rawQuery.getUpperBound().isPresent()); + } + @Test public void shouldGetAllFromInnerStoreAndRecordAllMetric() { setUp(); diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java index 81cf34ed081a0..d8a500404b9c2 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java @@ -500,9 +500,9 @@ public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenReque store.put("bad", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x"))); final QueryResult> ok = - store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(true)); + store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(true)); final QueryResult> bad = - store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(true)); + store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(true)); assertTrue(ok.isSuccess()); assertFalse(bad.isSuccess()); @@ -511,13 +511,13 @@ public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenReque final String okInfo = String.join("\n", ok.getExecutionInfo()); final String badInfo = String.join("\n", bad.getExecutionInfo()); assertTrue( - okInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) - && okInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), - "success execution info missing an entry: " + okInfo); + okInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) + && okInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), + "success execution info missing an entry: " + okInfo); assertTrue( - badInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) - && badInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), - "failure execution info missing an entry (wrapped-store entry must be preserved): " + badInfo); + badInfo.contains(RocksDBTimestampedStoreWithHeaders.class.getName()) + && badInfo.contains(MeteredTimestampedKeyValueStoreWithHeaders.class.getName()), + "failure execution info missing an entry (wrapped-store entry must be preserved): " + badInfo); } finally { store.close(); } @@ -533,9 +533,9 @@ public void shouldNotCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenNo store.put("bad", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x"))); final QueryResult> ok = - store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(false)); + store.query(TimestampedKeyWithHeadersQuery.withKey("ok"), PositionBound.unbounded(), new QueryConfig(false)); final QueryResult> bad = - store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(false)); + store.query(TimestampedKeyWithHeadersQuery.withKey("bad"), PositionBound.unbounded(), new QueryConfig(false)); assertTrue(ok.isSuccess()); assertFalse(bad.isSuccess()); @@ -559,13 +559,13 @@ public void shouldReturnIdenticalResultsForNativeAndAdapterBuiltStores(final boo adapterStore.put("k", value); assertEquals( - nativeStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - adapterStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - "KeyQuery results should be identical across native and adapter build paths"); + nativeStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + adapterStore.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + "KeyQuery results should be identical across native and adapter build paths"); assertEquals( - nativeStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - adapterStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), - "TimestampedKeyQuery results should be identical across native and adapter build paths"); + nativeStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + adapterStore.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)).getResult(), + "TimestampedKeyQuery results should be identical across native and adapter build paths"); } finally { nativeStore.close(); adapterStore.close(); From f73cd7e204672f90ea3d802d78462e25156093ce Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Fri, 3 Jul 2026 09:24:58 -0400 Subject: [PATCH 05/17] Align range-query integration test with the IQv2 position-bound convention - Query with PositionBound.at(INPUT_POSITION) + iqv2WaitForResult --- ...dRangeWithHeadersQueryIntegrationTest.java | 60 ++++++++++--------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java index 487774a87408b..aa72f37da0de0 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java @@ -19,7 +19,6 @@ import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.common.header.Headers; import org.apache.kafka.common.header.internals.RecordHeaders; -import org.apache.kafka.common.serialization.IntegerDeserializer; import org.apache.kafka.common.serialization.IntegerSerializer; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.common.serialization.StringSerializer; @@ -37,6 +36,8 @@ import org.apache.kafka.streams.processor.api.ReadOnlyRecord; import org.apache.kafka.streams.processor.api.Record; import org.apache.kafka.streams.query.FailureReason; +import org.apache.kafka.streams.query.Position; +import org.apache.kafka.streams.query.PositionBound; import org.apache.kafka.streams.query.QueryResult; import org.apache.kafka.streams.query.StateQueryRequest; import org.apache.kafka.streams.query.StateQueryResult; @@ -82,6 +83,11 @@ public class TimestampedRangeWithHeadersQueryIntegrationTest { private String inputStream; private String outputStream; private long baseTimestamp; + // Accumulated position of the produced input records, so queries can require freshness via + // PositionBound.at(inputPosition) -- matching the IQv2StoreIntegrationTest convention rather than + // relying on output-topic consumption as the readiness signal. + private Position inputPosition; + private long nextInputOffset; private KafkaStreams kafkaStreams; @@ -119,6 +125,8 @@ public void beforeTest(final TestInfo testInfo) throws InterruptedException { CLUSTER.createTopic(inputStream, 1, 1); CLUSTER.createTopic(outputStream, 1, 1); baseTimestamp = CLUSTER.time.milliseconds(); + inputPosition = Position.emptyPosition(); + nextInputOffset = 0L; } @AfterEach @@ -138,7 +146,7 @@ public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), Serdes.Integer(), Serdes.String() - ) + ).withCachingDisabled() ) .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) .process(HeadersStoreWriterProcessor::new, STORE_NAME) @@ -148,13 +156,10 @@ public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { kafkaStreams.start(); // Write: keys 1,2 (HEADERS1), key 3 (empty headers), key 4 (then tombstone it). - int produced = 0; - produced += produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); - produced += produceWithHeaders(baseTimestamp + 1, EMPTY_HEADERS, KeyValue.pair(3, "three")); - produced += produceWithHeaders(baseTimestamp + 2, HEADERS2, KeyValue.pair(4, "four")); - produced += produceWithHeaders(baseTimestamp + 3, EMPTY_HEADERS, KeyValue.pair(4, null)); // tombstone - - awaitProcessed(produced); + produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + produceWithHeaders(baseTimestamp + 1, EMPTY_HEADERS, KeyValue.pair(3, "three")); + produceWithHeaders(baseTimestamp + 2, HEADERS2, KeyValue.pair(4, "four")); + produceWithHeaders(baseTimestamp + 3, EMPTY_HEADERS, KeyValue.pair(4, null)); // tombstone // Full scan, ascending: keys 1, 2, 3 (key 4 tombstoned and omitted). final List> ascending = runQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); @@ -194,7 +199,7 @@ public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Except Stores.persistentTimestampedKeyValueStore(STORE_NAME), Serdes.Integer(), Serdes.String() - ) + ).withCachingDisabled() ) .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) .process(NonHeadersStoreWriterProcessor::new, STORE_NAME) @@ -203,11 +208,11 @@ public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Except kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); kafkaStreams.start(); - final int produced = produceWithHeaders(baseTimestamp, EMPTY_HEADERS, KeyValue.pair(1, "one")); - awaitProcessed(produced); + produceWithHeaders(baseTimestamp, EMPTY_HEADERS, KeyValue.pair(1, "one")); final StateQueryRequest> request = - inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) + .withPositionBound(PositionBound.at(inputPosition)); final StateQueryResult> result = IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); @@ -232,7 +237,7 @@ public void shouldThrowForPlainSupplierWithNoRepresentableTimestamp() throws Exc Stores.persistentKeyValueStore(STORE_NAME), Serdes.Integer(), Serdes.String() - ) + ).withCachingDisabled() ) .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) .process(HeadersStoreWriterProcessor::new, STORE_NAME) @@ -241,11 +246,11 @@ public void shouldThrowForPlainSupplierWithNoRepresentableTimestamp() throws Exc kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); kafkaStreams.start(); - final int produced = produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); - awaitProcessed(produced); + produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); final StateQueryRequest> request = - inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) + .withPositionBound(PositionBound.at(inputPosition)); final StateQueryResult> result = IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); @@ -257,7 +262,8 @@ public void shouldThrowForPlainSupplierWithNoRepresentableTimestamp() throws Exc } private List> runQuery(final TimestampedRangeWithHeadersQuery query) { - final StateQueryRequest> request = inStore(STORE_NAME).withQuery(query); + final StateQueryRequest> request = + inStore(STORE_NAME).withQuery(query).withPositionBound(PositionBound.at(inputPosition)); final StateQueryResult> result = IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); final List> records = new ArrayList<>(); @@ -286,13 +292,6 @@ private static boolean headersAreReadOnly(final Headers headers) { } } - private void awaitProcessed(final int count) throws Exception { - IntegrationTestUtils.waitUntilMinRecordsReceived( - TestUtils.consumerConfig(CLUSTER.bootstrapServers(), IntegerDeserializer.class, IntegerDeserializer.class), - outputStream, - count); - } - private Properties props(final String prefix) { final String safeTestName = safeUniqueTestName(testInfo); final Properties streamsConfiguration = new Properties(); @@ -306,9 +305,9 @@ private Properties props(final String prefix) { @SuppressWarnings("varargs") @SafeVarargs - private final int produceWithHeaders(final long timestamp, - final Headers headers, - final KeyValue... keyValues) { + private final void produceWithHeaders(final long timestamp, + final Headers headers, + final KeyValue... keyValues) { IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp( inputStream, Arrays.asList(keyValues), @@ -316,7 +315,10 @@ private final int produceWithHeaders(final long timestamp, headers, timestamp, false); - return keyValues.length; + // Single-partition input topic with contiguous offsets starting at 0: advance to the latest + // produced offset so PositionBound.at(inputPosition) requires the store to have read this far. + nextInputOffset += keyValues.length; + inputPosition.withComponent(inputStream, 0, nextInputOffset - 1); } /** From 1f4bbd1db316620129bedd38c52a8bf9e2701035 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 11:48:13 -0400 Subject: [PATCH 06/17] Move range-query integration tests into IQv2HeadersStoreIntegrationTest --- .../IQv2HeadersStoreIntegrationTest.java | 147 ++++++- ...dRangeWithHeadersQueryIntegrationTest.java | 374 ------------------ 2 files changed, 142 insertions(+), 379 deletions(-) delete mode 100644 streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java index f8eb62da2e10f..f39808d6966d1 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java @@ -30,6 +30,7 @@ import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.StreamsBuilder; import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.errors.StreamsException; import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; import org.apache.kafka.streams.kstream.Consumed; @@ -45,6 +46,8 @@ import org.apache.kafka.streams.query.StateQueryRequest; import org.apache.kafka.streams.query.StateQueryResult; import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; +import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.StoreBuilder; import org.apache.kafka.streams.state.Stores; import org.apache.kafka.streams.state.TimestampedKeyValueStore; @@ -63,24 +66,28 @@ import java.io.IOException; import java.time.Duration; +import java.util.ArrayList; import java.util.LinkedList; import java.util.List; import java.util.Properties; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; import static org.apache.kafka.streams.query.StateQueryRequest.inStore; import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; /** * IQv2 integration tests for KIP-1271/KIP-1356 headers-aware state stores. * *

It builds a KIP-1271 {@code WithHeaders} store, writes records (with headers) into it through a processor, - * and queries it through IQv2. Currently covers {@link TimestampedKeyWithHeadersQuery}; the remaining KIP-1356 headers - * query types (range/window/session) are expected to extend this class as they land. + * and queries it through IQv2. Covers {@link TimestampedKeyWithHeadersQuery} and + * {@link TimestampedRangeWithHeadersQuery}; the remaining KIP-1356 headers query types (window/session) are + * expected to extend this class as they land. */ @Tag("integration") public class IQv2HeadersStoreIntegrationTest { @@ -217,6 +224,110 @@ public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Except assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); } + @Test + public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { + // Caching disabled: a range query reads the underlying store directly (it never consults the + // cache), so the writes must be store-served. + startStreams(); + + // keys 1,2 (headers), key 3 (empty headers), key 4 written then tombstoned + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, + KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 1, new RecordHeaders(), + KeyValue.pair(3, "three")); + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 2, HEADERS, + KeyValue.pair(4, "four"), KeyValue.pair(4, null)); + + // Full scan, ascending: keys 1, 2, 3 (key 4 tombstoned and omitted). + final List> ascending = rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + assertEquals(List.of(1, 2, 3), keys(ascending)); + assertEquals(List.of("one", "two", "three"), values(ascending)); + + // key/value/timestamp/headers carried on each element + assertEquals(HEADERS, ascending.get(0).headers()); + assertEquals(baseTimestamp, ascending.get(0).timestamp()); + // key 3 written without headers -> empty (never null) headers + assertEquals(new RecordHeaders(), ascending.get(2).headers()); + assertEquals(baseTimestamp + 1, ascending.get(2).timestamp()); + // returned headers are a read-only snapshot + assertThrows(IllegalStateException.class, () -> ascending.get(0).headers().add("x", new byte[0])); + + // Full scan, descending. + assertEquals(List.of(3, 2, 1), + keys(rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withDescendingKeys()))); + + // Bounded ranges. + assertEquals(List.of(2, 3), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withRange(2, 3)))); + assertEquals(List.of(2, 3), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withLowerBound(2)))); + assertEquals(List.of(1, 2), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withUpperBound(2)))); + } + + @Test + public void shouldFailWithUnknownQueryTypeForRangeQueryAgainstNonHeadersStore() throws Exception { + // store built WITHOUT a WithHeaders supplier -> the range-with-headers query type is unsupported + final StreamsBuilder builder = new StreamsBuilder(); + builder + .addStateStore( + Stores.timestampedKeyValueStoreBuilder( + Stores.persistentTimestampedKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String())) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(() -> new PlainStoreWriterProcessor(), STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); + + kafkaStreams = new KafkaStreams(builder.build(), props()); + IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); + + final StateQueryRequest> request = + inStore(STORE_NAME) + .withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + assertTrue(result.getOnlyPartitionResult().isFailure()); + assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); + } + + @Test + public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() throws Exception { + // WithHeaders builder over a plain (non-timestamped) supplier: entries come back with + // timestamp = -1, which cannot be a ReadOnlyRecord. The query succeeds, but iterating throws. + final StreamsBuilder builder = new StreamsBuilder(); + final StoreBuilder> storeBuilder = + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String()); + builder + .addStateStore(storeBuilder.withCachingDisabled()) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(() -> new HeadersStoreWriterProcessor(), STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); + + kafkaStreams = new KafkaStreams(builder.build(), props()); + IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, + KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + + final StateQueryRequest> request = + inStore(STORE_NAME) + .withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + final QueryResult> onlyResult = result.getOnlyPartitionResult(); + assertTrue(onlyResult.isSuccess()); + try (ReadOnlyRecordIterator iterator = onlyResult.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } + } + private void startStreams() throws Exception { // Caching disabled: every IQv2 query is forced down to the persistent // RocksDBTimestampedStoreWithHeaders layer, exercising its KeyQuery handling @@ -268,6 +379,28 @@ private ReadOnlyRecord querySkipCache(final int key) { return onlyResult == null ? null : onlyResult.getResult(); } + private List> rangeQuery(final TimestampedRangeWithHeadersQuery query) { + final StateQueryRequest> request = + inStore(STORE_NAME).withQuery(query).withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + final List> records = new ArrayList<>(); + try (ReadOnlyRecordIterator iterator = result.getOnlyPartitionResult().getResult()) { + while (iterator.hasNext()) { + records.add(iterator.next()); + } + } + return records; + } + + private static List keys(final List> records) { + return records.stream().map(ReadOnlyRecord::key).collect(Collectors.toList()); + } + + private static List values(final List> records) { + return records.stream().map(ReadOnlyRecord::value).collect(Collectors.toList()); + } + private Properties props() { final String safeTestName = safeUniqueTestName(testInfo); final Properties streamsConfiguration = new Properties(); @@ -319,9 +452,13 @@ public void init(final ProcessorContext context) { @Override public void process(final Record record) { - store.put( - record.key(), - ValueTimestampHeaders.make(record.value(), record.timestamp(), record.headers())); + if (record.value() == null) { + store.delete(record.key()); + } else { + store.put( + record.key(), + ValueTimestampHeaders.make(record.value(), record.timestamp(), record.headers())); + } context.forward(record); } } diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java deleted file mode 100644 index aa72f37da0de0..0000000000000 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedRangeWithHeadersQueryIntegrationTest.java +++ /dev/null @@ -1,374 +0,0 @@ -/* - * 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.integration; - -import org.apache.kafka.clients.consumer.ConsumerConfig; -import org.apache.kafka.common.header.Headers; -import org.apache.kafka.common.header.internals.RecordHeaders; -import org.apache.kafka.common.serialization.IntegerSerializer; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.common.serialization.StringSerializer; -import org.apache.kafka.streams.KafkaStreams; -import org.apache.kafka.streams.KeyValue; -import org.apache.kafka.streams.StreamsBuilder; -import org.apache.kafka.streams.StreamsConfig; -import org.apache.kafka.streams.errors.StreamsException; -import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; -import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; -import org.apache.kafka.streams.kstream.Consumed; -import org.apache.kafka.streams.kstream.Produced; -import org.apache.kafka.streams.processor.api.Processor; -import org.apache.kafka.streams.processor.api.ProcessorContext; -import org.apache.kafka.streams.processor.api.ReadOnlyRecord; -import org.apache.kafka.streams.processor.api.Record; -import org.apache.kafka.streams.query.FailureReason; -import org.apache.kafka.streams.query.Position; -import org.apache.kafka.streams.query.PositionBound; -import org.apache.kafka.streams.query.QueryResult; -import org.apache.kafka.streams.query.StateQueryRequest; -import org.apache.kafka.streams.query.StateQueryResult; -import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; -import org.apache.kafka.streams.state.ReadOnlyRecordIterator; -import org.apache.kafka.streams.state.Stores; -import org.apache.kafka.streams.state.TimestampedKeyValueStore; -import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; -import org.apache.kafka.streams.state.ValueAndTimestamp; -import org.apache.kafka.streams.state.ValueTimestampHeaders; -import org.apache.kafka.test.TestUtils; - -import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInfo; - -import java.io.IOException; -import java.time.Duration; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.stream.Collectors; - -import static org.apache.kafka.streams.query.StateQueryRequest.inStore; -import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -@Tag("integration") -public class TimestampedRangeWithHeadersQueryIntegrationTest { - - private static final String STORE_NAME = "headers-store"; - - private String inputStream; - private String outputStream; - private long baseTimestamp; - // Accumulated position of the produced input records, so queries can require freshness via - // PositionBound.at(inputPosition) -- matching the IQv2StoreIntegrationTest convention rather than - // relying on output-topic consumption as the readiness signal. - private Position inputPosition; - private long nextInputOffset; - - private KafkaStreams kafkaStreams; - - private static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1); - - private static final Headers HEADERS1 = new RecordHeaders() - .add("source", "test".getBytes()) - .add("version", "1.0".getBytes()); - - private static final Headers HEADERS2 = new RecordHeaders() - .add("source", "test".getBytes()) - .add("version", "2.0".getBytes()); - - private static final Headers EMPTY_HEADERS = new RecordHeaders(); - - public TestInfo testInfo; - - @BeforeAll - public static void before() throws IOException { - CLUSTER.start(); - } - - @AfterAll - public static void after() { - CLUSTER.stop(); - } - - @BeforeEach - public void beforeTest(final TestInfo testInfo) throws InterruptedException { - this.testInfo = testInfo; - final String uniqueTestName = safeUniqueTestName(testInfo); - inputStream = "input-stream-" + uniqueTestName; - outputStream = "output-stream-" + uniqueTestName; - // single partition so the query has exactly one partition result - CLUSTER.createTopic(inputStream, 1, 1); - CLUSTER.createTopic(outputStream, 1, 1); - baseTimestamp = CLUSTER.time.milliseconds(); - inputPosition = Position.emptyPosition(); - nextInputOffset = 0L; - } - - @AfterEach - public void afterTest() { - if (kafkaStreams != null) { - kafkaStreams.close(Duration.ofSeconds(30L)); - kafkaStreams.cleanUp(); - } - } - - @Test - public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { - final StreamsBuilder streamsBuilder = new StreamsBuilder(); - streamsBuilder - .addStateStore( - Stores.timestampedKeyValueStoreWithHeadersBuilder( - Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), - Serdes.Integer(), - Serdes.String() - ).withCachingDisabled() - ) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(HeadersStoreWriterProcessor::new, STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); - - kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); - kafkaStreams.start(); - - // Write: keys 1,2 (HEADERS1), key 3 (empty headers), key 4 (then tombstone it). - produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); - produceWithHeaders(baseTimestamp + 1, EMPTY_HEADERS, KeyValue.pair(3, "three")); - produceWithHeaders(baseTimestamp + 2, HEADERS2, KeyValue.pair(4, "four")); - produceWithHeaders(baseTimestamp + 3, EMPTY_HEADERS, KeyValue.pair(4, null)); // tombstone - - // Full scan, ascending: keys 1, 2, 3 (key 4 tombstoned and omitted). - final List> ascending = runQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); - assertEquals(List.of(1, 2, 3), keys(ascending)); - assertEquals(List.of("one", "two", "three"), values(ascending)); - - // Headers, timestamps, and key are carried on each record. - assertEquals(HEADERS1, ascending.get(0).headers()); - assertEquals(baseTimestamp, ascending.get(0).timestamp()); - assertEquals(HEADERS1, ascending.get(1).headers()); - // Record written without headers round-trips as empty (never null) headers. - final Headers emptyHeaders = ascending.get(2).headers(); - assertNotNull(emptyHeaders); - assertEquals(0, emptyHeaders.toArray().length); - assertEquals(baseTimestamp + 1, ascending.get(2).timestamp()); - - // Returned headers are a read-only snapshot. - assertTrue(headersAreReadOnly(ascending.get(0).headers())); - - // Full scan, descending. - final List> descending = - runQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withDescendingKeys()); - assertEquals(List.of(3, 2, 1), keys(descending)); - - // Bounded ranges. - assertEquals(List.of(2, 3), keys(runQuery(TimestampedRangeWithHeadersQuery.withRange(2, 3)))); - assertEquals(List.of(2, 3), keys(runQuery(TimestampedRangeWithHeadersQuery.withLowerBound(2)))); - assertEquals(List.of(1, 2), keys(runQuery(TimestampedRangeWithHeadersQuery.withUpperBound(2)))); - } - - @Test - public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Exception { - final StreamsBuilder streamsBuilder = new StreamsBuilder(); - streamsBuilder - .addStateStore( - Stores.timestampedKeyValueStoreBuilder( - Stores.persistentTimestampedKeyValueStore(STORE_NAME), - Serdes.Integer(), - Serdes.String() - ).withCachingDisabled() - ) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(NonHeadersStoreWriterProcessor::new, STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); - - kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); - kafkaStreams.start(); - - produceWithHeaders(baseTimestamp, EMPTY_HEADERS, KeyValue.pair(1, "one")); - - final StateQueryRequest> request = - inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) - .withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = - IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - - final Map>> partitionResults = - result.getPartitionResults(); - assertFalse(partitionResults.isEmpty()); - for (final QueryResult> partitionResult : partitionResults.values()) { - assertTrue(partitionResult.isFailure()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, partitionResult.getFailureReason()); - } - } - - @Test - public void shouldThrowForPlainSupplierWithNoRepresentableTimestamp() throws Exception { - // WithHeaders builder over a plain (non-timestamped) supplier: the store cannot persist a - // timestamp, so entries come back with timestamp = -1, which cannot be a ReadOnlyRecord. The - // query itself succeeds, but advancing the iterator throws. - final StreamsBuilder streamsBuilder = new StreamsBuilder(); - streamsBuilder - .addStateStore( - Stores.timestampedKeyValueStoreWithHeadersBuilder( - Stores.persistentKeyValueStore(STORE_NAME), - Serdes.Integer(), - Serdes.String() - ).withCachingDisabled() - ) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(HeadersStoreWriterProcessor::new, STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.Integer())); - - kafkaStreams = new KafkaStreams(streamsBuilder.build(), props("app-")); - kafkaStreams.start(); - - produceWithHeaders(baseTimestamp, HEADERS1, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); - - final StateQueryRequest> request = - inStore(STORE_NAME).withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) - .withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = - IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - - final QueryResult> partitionResult = result.getOnlyPartitionResult(); - assertTrue(partitionResult.isSuccess()); - try (ReadOnlyRecordIterator iterator = partitionResult.getResult()) { - assertThrows(StreamsException.class, iterator::next); - } - } - - private List> runQuery(final TimestampedRangeWithHeadersQuery query) { - final StateQueryRequest> request = - inStore(STORE_NAME).withQuery(query).withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = - IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - final List> records = new ArrayList<>(); - try (ReadOnlyRecordIterator iterator = result.getOnlyPartitionResult().getResult()) { - while (iterator.hasNext()) { - records.add(iterator.next()); - } - } - return records; - } - - private static List keys(final List> records) { - return records.stream().map(ReadOnlyRecord::key).collect(Collectors.toList()); - } - - private static List values(final List> records) { - return records.stream().map(ReadOnlyRecord::value).collect(Collectors.toList()); - } - - private static boolean headersAreReadOnly(final Headers headers) { - try { - headers.add("x", new byte[0]); - return false; - } catch (final IllegalStateException expected) { - return true; - } - } - - private Properties props(final String prefix) { - final String safeTestName = safeUniqueTestName(testInfo); - final Properties streamsConfiguration = new Properties(); - streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, prefix + safeTestName); - streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); - streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath()); - streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000L); - streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); - return streamsConfiguration; - } - - @SuppressWarnings("varargs") - @SafeVarargs - private final void produceWithHeaders(final long timestamp, - final Headers headers, - final KeyValue... keyValues) { - IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp( - inputStream, - Arrays.asList(keyValues), - TestUtils.producerConfig(CLUSTER.bootstrapServers(), IntegerSerializer.class, StringSerializer.class), - headers, - timestamp, - false); - // Single-partition input topic with contiguous offsets starting at 0: advance to the latest - // produced offset so PositionBound.at(inputPosition) requires the store to have read this far. - nextInputOffset += keyValues.length; - inputPosition.withComponent(inputStream, 0, nextInputOffset - 1); - } - - /** - * Writes each incoming record into the headers-aware store (deleting on a null value), then - * forwards a marker downstream so the test can wait for processing to complete. - */ - private static class HeadersStoreWriterProcessor implements Processor { - private ProcessorContext context; - private TimestampedKeyValueStoreWithHeaders store; - - @Override - public void init(final ProcessorContext context) { - this.context = context; - store = context.getStateStore(STORE_NAME); - } - - @Override - public void process(final Record record) { - if (record.value() == null) { - store.delete(record.key()); - } else { - store.put(record.key(), - ValueTimestampHeaders.make(record.value(), record.timestamp(), record.headers())); - } - context.forward(record.withValue(1)); - } - } - - /** - * Writes each incoming record into a plain (non-headers) timestamped store, then forwards a - * marker downstream so the test can wait for processing to complete. - */ - private static class NonHeadersStoreWriterProcessor implements Processor { - private ProcessorContext context; - private TimestampedKeyValueStore store; - - @Override - public void init(final ProcessorContext context) { - this.context = context; - store = context.getStateStore(STORE_NAME); - } - - @Override - public void process(final Record record) { - if (record.value() == null) { - store.delete(record.key()); - } else { - store.put(record.key(), ValueAndTimestamp.make(record.value(), record.timestamp())); - } - context.forward(record.withValue(1)); - } - } -} From 7590078678d3eab96cdae5f52a2a4cd9ee5c3560 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 12:12:31 -0400 Subject: [PATCH 07/17] Keep point-query BuilderTest methods in place to reduce diff noise --- ...edKeyValueStoreBuilderWithHeadersTest.java | 68 +++++++++---------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java index d8a500404b9c2..d64f2614b06c4 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java @@ -489,6 +489,40 @@ public void shouldFailTimestampedKeyWithHeadersQueryForNegativeStoredTimestamp(f } } + @ParameterizedTest + @CsvSource({"NATIVE", "ADAPTER", "IN_MEMORY"}) + public void shouldReturnEmptyPositionInitially(final StoreType storeType) { + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); + try { + final Position position = ((WrappedStateStore) store).wrapped().getPosition(); + assertNotNull(position, "Expected non-null position"); + assertTrue(position.isEmpty(), "Expected position to be empty initially"); + } finally { + store.close(); + } + } + + @ParameterizedTest + @CsvSource({"NATIVE", "ADAPTER"}) + public void shouldCollectExecutionInfoWhenRequested(final StoreType storeType) { + final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); + try { + final StateStore wrapped = ((WrappedStateStore) store).wrapped(); + final QueryResult result = wrapped.query( + KeyQuery.withKey(new Bytes("k".getBytes())), PositionBound.unbounded(), new QueryConfig(true)); + + final String executionInfo = String.join("\n", result.getExecutionInfo()); + assertFalse(executionInfo.isEmpty(), "Expected execution info to be collected"); + assertTrue(executionInfo.contains("Handled in"), "Expected execution info to contain handling information"); + final String expectedClass = storeType == StoreType.NATIVE + ? RocksDBTimestampedStoreWithHeaders.class.getName() + : TimestampedToHeadersStoreAdapter.class.getName(); + assertTrue(executionInfo.contains(expectedClass), "Expected execution info to mention " + expectedClass); + } finally { + store.close(); + } + } + @Test public void shouldCollectExecutionInfoForTimestampedKeyWithHeadersQueryWhenRequested() { // The typed query, run through the metered handler with execution info enabled, must carry both @@ -772,38 +806,4 @@ private static List> drain(final KeyValueIterator store = buildAndInitStore(storeType, false); - try { - final Position position = ((WrappedStateStore) store).wrapped().getPosition(); - assertNotNull(position, "Expected non-null position"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); - } finally { - store.close(); - } - } - - @ParameterizedTest - @CsvSource({"NATIVE", "ADAPTER"}) - public void shouldCollectExecutionInfoWhenRequested(final StoreType storeType) { - final TimestampedKeyValueStoreWithHeaders store = buildAndInitStore(storeType, false); - try { - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - final QueryResult result = wrapped.query( - KeyQuery.withKey(new Bytes("k".getBytes())), PositionBound.unbounded(), new QueryConfig(true)); - - final String executionInfo = String.join("\n", result.getExecutionInfo()); - assertFalse(executionInfo.isEmpty(), "Expected execution info to be collected"); - assertTrue(executionInfo.contains("Handled in"), "Expected execution info to contain handling information"); - final String expectedClass = storeType == StoreType.NATIVE - ? RocksDBTimestampedStoreWithHeaders.class.getName() - : TimestampedToHeadersStoreAdapter.class.getName(); - assertTrue(executionInfo.contains(expectedClass), "Expected execution info to mention " + expectedClass); - } finally { - store.close(); - } - } } From dafacbd95a9b343fadeaafd3806c7ac8a854914b Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 11:11:02 -0400 Subject: [PATCH 08/17] MINOR: Fix ReadOnlyRecordIterator javadoc link and import spacing --- .../org/apache/kafka/streams/state/ReadOnlyRecordIterator.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java index a4d5ce7964885..050b9b2315873 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java @@ -22,9 +22,8 @@ import java.io.Closeable; import java.util.Iterator; - /** - * Iterator interface of {@link ReadOnlyRecord ReadOnlyRecord}. + * Iterator interface of {@link ReadOnlyRecord}. *

* This is the result type of the headers-aware range/window/session IQv2 query types: it yields * {@link ReadOnlyRecord} elements directly, so each element carries its own key, value, timestamp, From 122cf5f7b8b4677bd8e6743688b749b2a76b44cc Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 12:30:59 -0400 Subject: [PATCH 09/17] MINOR: Add @InterfaceAudience.Public to ReadOnlyRecordIterator --- .../org/apache/kafka/streams/state/ReadOnlyRecordIterator.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java index 050b9b2315873..e186095b1db31 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.state; +import org.apache.kafka.common.annotation.InterfaceAudience; import org.apache.kafka.common.annotation.InterfaceStability.Evolving; import org.apache.kafka.streams.processor.api.ReadOnlyRecord; @@ -36,6 +37,7 @@ * @param Type of keys * @param Type of values */ +@InterfaceAudience.Public @Evolving public interface ReadOnlyRecordIterator extends Iterator>, Closeable { From 7c16ffd90eb751f570c406053c939475fa27ead7 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 13:27:14 -0400 Subject: [PATCH 10/17] Extract shared raw RangeQuery construction into helper --- ...edTimestampedKeyValueStoreWithHeaders.java | 59 ++++++++----------- 1 file changed, 23 insertions(+), 36 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java index f1f224f9f58ff..98d077310990c 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java @@ -57,6 +57,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.function.Function; import static org.apache.kafka.common.utils.Utils.mkEntry; @@ -458,6 +459,22 @@ record = headerRecord; return result; } + private RangeQuery rawRangeQuery(final Optional lowerBound, + final Optional upperBound, + final ResultOrder order) { + RangeQuery rawRangeQuery = RangeQuery.withRange( + serializeKey(lowerBound.orElse(null), internalContext.headers()), + serializeKey(upperBound.orElse(null), internalContext.headers()) + ); + if (order.equals(ResultOrder.DESCENDING)) { + rawRangeQuery = rawRangeQuery.withDescendingKeys(); + } + if (order.equals(ResultOrder.ASCENDING)) { + rawRangeQuery = rawRangeQuery.withAscendingKeys(); + } + return rawRangeQuery; + } + @SuppressWarnings("unchecked") private QueryResult runRangeQuery( final Query query, @@ -467,18 +484,8 @@ private QueryResult runRangeQuery( final QueryResult result; final RangeQuery typedQuery = (RangeQuery) query; - RangeQuery rawRangeQuery; - final ResultOrder order = typedQuery.resultOrder(); - rawRangeQuery = RangeQuery.withRange( - serializeKey(typedQuery.getLowerBound().orElse(null), internalContext.headers()), - serializeKey(typedQuery.getUpperBound().orElse(null), internalContext.headers()) - ); - if (order.equals(ResultOrder.DESCENDING)) { - rawRangeQuery = rawRangeQuery.withDescendingKeys(); - } - if (order.equals(ResultOrder.ASCENDING)) { - rawRangeQuery = rawRangeQuery.withAscendingKeys(); - } + final RangeQuery rawRangeQuery = + rawRangeQuery(typedQuery.getLowerBound(), typedQuery.getUpperBound(), typedQuery.resultOrder()); final QueryResult> rawResult = wrapped().query(rawRangeQuery, positionBound, config); if (rawResult.isSuccess()) { @@ -512,18 +519,8 @@ private QueryResult runTimestampedRangeQuery( final QueryResult result; final TimestampedRangeQuery typedQuery = (TimestampedRangeQuery) query; - RangeQuery rawRangeQuery; - final ResultOrder order = typedQuery.resultOrder(); - rawRangeQuery = RangeQuery.withRange( - serializeKey(typedQuery.lowerBound().orElse(null), internalContext.headers()), - serializeKey(typedQuery.upperBound().orElse(null), internalContext.headers()) - ); - if (order.equals(ResultOrder.DESCENDING)) { - rawRangeQuery = rawRangeQuery.withDescendingKeys(); - } - if (order.equals(ResultOrder.ASCENDING)) { - rawRangeQuery = rawRangeQuery.withAscendingKeys(); - } + final RangeQuery rawRangeQuery = + rawRangeQuery(typedQuery.lowerBound(), typedQuery.upperBound(), typedQuery.resultOrder()); final QueryResult> rawResult = wrapped().query(rawRangeQuery, positionBound, config); if (rawResult.isSuccess()) { @@ -558,18 +555,8 @@ private QueryResult runTimestampedRangeWithHeadersQuery( final QueryResult result; final TimestampedRangeWithHeadersQuery typedQuery = (TimestampedRangeWithHeadersQuery) query; - RangeQuery rawRangeQuery; - final ResultOrder order = typedQuery.resultOrder(); - rawRangeQuery = RangeQuery.withRange( - serializeKey(typedQuery.lowerBound().orElse(null), internalContext.headers()), - serializeKey(typedQuery.upperBound().orElse(null), internalContext.headers()) - ); - if (order.equals(ResultOrder.DESCENDING)) { - rawRangeQuery = rawRangeQuery.withDescendingKeys(); - } - if (order.equals(ResultOrder.ASCENDING)) { - rawRangeQuery = rawRangeQuery.withAscendingKeys(); - } + final RangeQuery rawRangeQuery = + rawRangeQuery(typedQuery.lowerBound(), typedQuery.upperBound(), typedQuery.resultOrder()); final QueryResult> rawResult = wrapped().query(rawRangeQuery, positionBound, config); if (rawResult.isSuccess()) { From 731641cb148bb3b947d9d80920c74ae4e5fdf14a Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 14:06:11 -0400 Subject: [PATCH 11/17] Track ReadOnlyRecordIterator in num-open-iterators metric --- .../internals/MeteredTimestampedKeyValueStoreWithHeaders.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java index 98d077310990c..8457f026a49f9 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java @@ -771,6 +771,7 @@ private MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator( this.valueTimestampHeadersDeserializer = valueTimestampHeadersDeserializer; this.startNs = time.nanoseconds(); this.startTimestampMs = time.milliseconds(); + numOpenIterators.increment(); openIterators.add(this); } @@ -812,6 +813,7 @@ public void close() { final long duration = time.nanoseconds() - startNs; sensor.record(duration); iteratorDurationSensor.record(duration); + numOpenIterators.decrement(); openIterators.remove(this); } } From 65ad6fddad6154f6b5bfb7599679e3b2aea342bf Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Mon, 6 Jul 2026 14:32:11 -0400 Subject: [PATCH 12/17] Track QueryIterator in num-open-iterators metric --- .../internals/MeteredTimestampedKeyValueStoreWithHeaders.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java index 8457f026a49f9..86b6957e3fa4f 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java @@ -677,6 +677,7 @@ private MeteredTimestampedKeyValueStoreWithHeadersQueryIterator( this.startNs = time.nanoseconds(); this.startTimestampMs = time.milliseconds(); this.returnPlainValue = returnPlainValue; + numOpenIterators.increment(); openIterators.add(this); } @@ -721,6 +722,7 @@ public void close() { final long duration = time.nanoseconds() - startNs; sensor.record(duration); iteratorDurationSensor.record(duration); + numOpenIterators.decrement(); openIterators.remove(this); } } From 578ae64062a799c7ee890d3d5a1d75774a9a4320 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Tue, 7 Jul 2026 09:17:23 -0400 Subject: [PATCH 13/17] add @InterfaceAudience.Public --- .../kafka/streams/query/TimestampedRangeWithHeadersQuery.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java index db49063b3fbc7..65762c72f57ac 100644 --- a/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java +++ b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.query; +import org.apache.kafka.common.annotation.InterfaceAudience; import org.apache.kafka.common.annotation.InterfaceStability.Evolving; import org.apache.kafka.streams.processor.api.ReadOnlyRecord; import org.apache.kafka.streams.state.ReadOnlyRecordIterator; @@ -58,6 +59,7 @@ * @param Type of values */ @Evolving +@InterfaceAudience.Public public final class TimestampedRangeWithHeadersQuery implements Query> { private final Optional lower; From 8bd56d97f62d5ccbb75836302141ce45edf72086 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Wed, 8 Jul 2026 09:36:15 -0400 Subject: [PATCH 14/17] Assert serialized range bounds and extract forwardedRawRangeQuery helper --- ...mestampedKeyValueStoreWithHeadersTest.java | 39 ++++++++----------- 1 file changed, 17 insertions(+), 22 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java index 386d3f1a0b260..c548d9fe76e84 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java @@ -305,42 +305,37 @@ public void shouldGetRangeFromInnerStoreAndRecordRangeMetric() { assertTrue((Double) metric.metricValue() > 0); } - @SuppressWarnings({"unchecked", "rawtypes"}) @Test public void shouldPropagateBoundsAndDescendingOrderForTimestampedRangeWithHeadersQuery() { setUp(); - final ArgumentCaptor captor = ArgumentCaptor.forClass(RangeQuery.class); - when(inner.query(captor.capture(), any(), any())).thenReturn( - QueryResult.forResult(new KeyValueIteratorStub<>(List.>of().iterator()))); init(); - - final TimestampedRangeWithHeadersQuery query = - TimestampedRangeWithHeadersQuery.withLowerBound("a").withDescendingKeys(); - metered.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - final RangeQuery rawQuery = captor.getValue(); + final RangeQuery rawQuery = forwardedRawRangeQuery( + TimestampedRangeWithHeadersQuery.withLowerBound("a").withDescendingKeys()); assertEquals(ResultOrder.DESCENDING, rawQuery.resultOrder()); - assertTrue(rawQuery.getLowerBound().isPresent()); + assertEquals(Bytes.wrap("a".getBytes()), rawQuery.getLowerBound().get()); assertFalse(rawQuery.getUpperBound().isPresent()); } - @SuppressWarnings({"unchecked", "rawtypes"}) @Test public void shouldPropagateBoundsAndAscendingOrderForTimestampedRangeWithHeadersQuery() { setUp(); - final ArgumentCaptor captor = ArgumentCaptor.forClass(RangeQuery.class); - when(inner.query(captor.capture(), any(), any())).thenReturn( - QueryResult.forResult(new KeyValueIteratorStub<>(List.>of().iterator()))); init(); - - final TimestampedRangeWithHeadersQuery query = - TimestampedRangeWithHeadersQuery.withUpperBound("z").withAscendingKeys(); - metered.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - final RangeQuery rawQuery = captor.getValue(); + final RangeQuery rawQuery = forwardedRawRangeQuery( + TimestampedRangeWithHeadersQuery.withUpperBound("z").withAscendingKeys()); assertEquals(ResultOrder.ASCENDING, rawQuery.resultOrder()); assertFalse(rawQuery.getLowerBound().isPresent()); - assertTrue(rawQuery.getUpperBound().isPresent()); + assertEquals(Bytes.wrap("z".getBytes()), rawQuery.getUpperBound().get()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private RangeQuery forwardedRawRangeQuery(final Query query) { + when(inner.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult( + new KeyValueIteratorStub<>(List.>of().iterator()))); + metered.query(query, PositionBound.unbounded(), new QueryConfig(false)); + final ArgumentCaptor captor = ArgumentCaptor.forClass(RangeQuery.class); + verify(inner).query(captor.capture(), any(PositionBound.class), any(QueryConfig.class)); + return captor.getValue(); } @Test From 1cef60674c70142b40bfd983e3d96de44bcbedb9 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Wed, 8 Jul 2026 14:06:01 -0400 Subject: [PATCH 15/17] Strengthen and refactor IQv2 headers range integration test --- .../IQv2HeadersStoreIntegrationTest.java | 221 ++++++++++-------- 1 file changed, 124 insertions(+), 97 deletions(-) diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java index f39808d6966d1..88483fe540132 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java @@ -37,11 +37,13 @@ import org.apache.kafka.streams.kstream.Produced; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; import org.apache.kafka.streams.processor.api.ReadOnlyRecord; import org.apache.kafka.streams.processor.api.Record; import org.apache.kafka.streams.query.FailureReason; 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.QueryResult; import org.apache.kafka.streams.query.StateQueryRequest; import org.apache.kafka.streams.query.StateQueryResult; @@ -140,7 +142,7 @@ public void afterTest() { @Test public void shouldHandleTimestampedKeyWithHeadersQuery() throws Exception { - startStreams(); + startStreamsWithHeadersStore(); // key 1 has headers, key 2 has empty headers, key 3 is tombstoned (null value) produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, @@ -178,7 +180,7 @@ public void shouldServeCacheHitWhenCachingEnabledAndRecordNotYetFlushed() throws // successful query therefore proves the result was served from the cache via // CachingKeyValueStoreWithHeaders -> the metered store's cache-hit path, end-to-end. commitIntervalMs = Duration.ofMinutes(10).toMillis(); - startStreams(true); + startStreamsWithHeadersStore(true); produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); @@ -195,42 +197,17 @@ public void shouldServeCacheHitWhenCachingEnabledAndRecordNotYetFlushed() throws } @Test - public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Exception { - // store built WITHOUT a WithHeaders supplier -> the query type is unsupported - final StreamsBuilder builder = new StreamsBuilder(); - builder - .addStateStore( - Stores.timestampedKeyValueStoreBuilder( - Stores.persistentTimestampedKeyValueStore(STORE_NAME), - Serdes.Integer(), - Serdes.String())) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(() -> new PlainStoreWriterProcessor(), STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); - - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); - - produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); - - final StateQueryRequest> request = - inStore(STORE_NAME) - .withQuery(TimestampedKeyWithHeadersQuery.withKey(1)) - .withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = - IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - - assertTrue(result.getOnlyPartitionResult().isFailure()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); + public void shouldFailWithUnknownQueryTypeForKeyQueryAgainstNonHeadersStore() throws Exception { + assertUnknownQueryTypeAgainstNonHeadersStore(TimestampedKeyWithHeadersQuery.withKey(1)); } @Test public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { // Caching disabled: a range query reads the underlying store directly (it never consults the // cache), so the writes must be store-served. - startStreams(); + startStreamsWithHeadersStore(); - // keys 1,2 (headers), key 3 (empty headers), key 4 written then tombstoned + // keys 1,2 (headers), key 3 (empty headers), key 4 written then tombstone produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); produceDataToTopicWithHeaders(inputStream, baseTimestamp + 1, new RecordHeaders(), @@ -238,78 +215,88 @@ public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { produceDataToTopicWithHeaders(inputStream, baseTimestamp + 2, HEADERS, KeyValue.pair(4, "four"), KeyValue.pair(4, null)); - // Full scan, ascending: keys 1, 2, 3 (key 4 tombstoned and omitted). - final List> ascending = rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds()); + // Full scan, ascending: keys 1, 2, 3 (key 4 tombstone and omitted). + final List> ascending = + rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withAscendingKeys()); assertEquals(List.of(1, 2, 3), keys(ascending)); assertEquals(List.of("one", "two", "three"), values(ascending)); - // key/value/timestamp/headers carried on each element - assertEquals(HEADERS, ascending.get(0).headers()); + assertEquals(HEADERS, ascending.get(0).headers()); // key 1 assertEquals(baseTimestamp, ascending.get(0).timestamp()); + assertEquals(HEADERS, ascending.get(1).headers()); // key 2 + assertEquals(baseTimestamp, ascending.get(1).timestamp()); // key 3 written without headers -> empty (never null) headers assertEquals(new RecordHeaders(), ascending.get(2).headers()); assertEquals(baseTimestamp + 1, ascending.get(2).timestamp()); - // returned headers are a read-only snapshot + // returned headers are a read-only snapshot: no mutation (add or remove) is allowed assertThrows(IllegalStateException.class, () -> ascending.get(0).headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> ascending.get(0).headers().remove("source")); - // Full scan, descending. - assertEquals(List.of(3, 2, 1), - keys(rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withDescendingKeys()))); - - // Bounded ranges. - assertEquals(List.of(2, 3), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withRange(2, 3)))); - assertEquals(List.of(2, 3), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withLowerBound(2)))); - assertEquals(List.of(1, 2), keys(rangeQuery(TimestampedRangeWithHeadersQuery.withUpperBound(2)))); + // Full scan, descending: keys reversed, payload still carried. + final List> descending = + rangeQuery(TimestampedRangeWithHeadersQuery.withNoBounds().withDescendingKeys()); + assertEquals(List.of(3, 2, 1), keys(descending)); + assertEquals(List.of("three", "two", "one"), values(descending)); + // key/value/timestamp/headers carried on each element + // key 3 written without headers -> empty (never null) headers + assertEquals(new RecordHeaders(), descending.get(0).headers()); // key 3 + assertEquals(baseTimestamp + 1, descending.get(0).timestamp()); + assertEquals(HEADERS, descending.get(1).headers()); // key 2 + assertEquals(baseTimestamp, descending.get(1).timestamp()); + assertEquals(HEADERS, descending.get(2).headers()); // key 1 + assertEquals(baseTimestamp, descending.get(2).timestamp()); + // returned headers are a read-only snapshot: no mutation (add or remove) is allowed + assertThrows(IllegalStateException.class, () -> descending.get(2).headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> descending.get(2).headers().remove("source")); + + // Bounded ranges (inclusive on both ends); same per-element checks as the full scans. + // withRange(2, 3) -> keys 2, 3. + final List> range = rangeQuery(TimestampedRangeWithHeadersQuery.withRange(2, 3)); + assertEquals(List.of(2, 3), keys(range)); + assertEquals(List.of("two", "three"), values(range)); + assertEquals(HEADERS, range.get(0).headers()); // key 2 + assertEquals(baseTimestamp, range.get(0).timestamp()); + // key 3 written without headers -> empty (never null) headers + assertEquals(new RecordHeaders(), range.get(1).headers()); // key 3 + assertEquals(baseTimestamp + 1, range.get(1).timestamp()); + // returned headers are a read-only snapshot: no mutation (add or remove) is allowed + assertThrows(IllegalStateException.class, () -> range.get(0).headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> range.get(0).headers().remove("source")); + + // withLowerBound(2) -> keys 2, 3. + final List> lowerBounded = + rangeQuery(TimestampedRangeWithHeadersQuery.withLowerBound(2)); + assertEquals(List.of(2, 3), keys(lowerBounded)); + assertEquals(List.of("two", "three"), values(lowerBounded)); + assertEquals(HEADERS, lowerBounded.get(0).headers()); // key 2 + assertEquals(baseTimestamp, lowerBounded.get(0).timestamp()); + assertEquals(new RecordHeaders(), lowerBounded.get(1).headers()); // key 3 + assertEquals(baseTimestamp + 1, lowerBounded.get(1).timestamp()); + assertThrows(IllegalStateException.class, () -> lowerBounded.get(0).headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> lowerBounded.get(0).headers().remove("source")); + + // withUpperBound(2) -> keys 1, 2. + final List> upperBounded = + rangeQuery(TimestampedRangeWithHeadersQuery.withUpperBound(2)); + assertEquals(List.of(1, 2), keys(upperBounded)); + assertEquals(List.of("one", "two"), values(upperBounded)); + assertEquals(HEADERS, upperBounded.get(0).headers()); // key 1 + assertEquals(baseTimestamp, upperBounded.get(0).timestamp()); + assertEquals(HEADERS, upperBounded.get(1).headers()); // key 2 + assertEquals(baseTimestamp, upperBounded.get(1).timestamp()); + assertThrows(IllegalStateException.class, () -> upperBounded.get(0).headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> upperBounded.get(0).headers().remove("source")); } @Test public void shouldFailWithUnknownQueryTypeForRangeQueryAgainstNonHeadersStore() throws Exception { - // store built WITHOUT a WithHeaders supplier -> the range-with-headers query type is unsupported - final StreamsBuilder builder = new StreamsBuilder(); - builder - .addStateStore( - Stores.timestampedKeyValueStoreBuilder( - Stores.persistentTimestampedKeyValueStore(STORE_NAME), - Serdes.Integer(), - Serdes.String())) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(() -> new PlainStoreWriterProcessor(), STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); - - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); - - produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); - - final StateQueryRequest> request = - inStore(STORE_NAME) - .withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) - .withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = - IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - - assertTrue(result.getOnlyPartitionResult().isFailure()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); + assertUnknownQueryTypeAgainstNonHeadersStore(TimestampedRangeWithHeadersQuery.withNoBounds()); } @Test public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() throws Exception { - // WithHeaders builder over a plain (non-timestamped) supplier: entries come back with - // timestamp = -1, which cannot be a ReadOnlyRecord. The query succeeds, but iterating throws. - final StreamsBuilder builder = new StreamsBuilder(); - final StoreBuilder> storeBuilder = - Stores.timestampedKeyValueStoreWithHeadersBuilder( - Stores.persistentKeyValueStore(STORE_NAME), - Serdes.Integer(), - Serdes.String()); - builder - .addStateStore(storeBuilder.withCachingDisabled()) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(() -> new HeadersStoreWriterProcessor(), STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); - - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + // The query succeeds, but iterating throws because timestamp = -1 cannot be a ReadOnlyRecord. + startStreamsWithPlainSupplierStore(); produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); @@ -328,28 +315,68 @@ public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() thro } } - private void startStreams() throws Exception { + private void startStreams(final StoreBuilder storeBuilder, + final ProcessorSupplier processorSupplier) throws Exception { + final StreamsBuilder builder = new StreamsBuilder(); + builder + .addStateStore(storeBuilder) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(processorSupplier, STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); + + kafkaStreams = new KafkaStreams(builder.build(), props()); + IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + } + + private void startStreamsWithHeadersStore() throws Exception { // Caching disabled: every IQv2 query is forced down to the persistent // RocksDBTimestampedStoreWithHeaders layer, exercising its KeyQuery handling // (rather than being short-circuited by a cache hit). - startStreams(false); + startStreamsWithHeadersStore(false); } - private void startStreams(final boolean cachingEnabled) throws Exception { - final StreamsBuilder builder = new StreamsBuilder(); + private void startStreamsWithHeadersStore(final boolean cachingEnabled) throws Exception { final StoreBuilder> storeBuilder = Stores.timestampedKeyValueStoreWithHeadersBuilder( Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), Serdes.Integer(), Serdes.String()); - builder - .addStateStore(cachingEnabled ? storeBuilder.withCachingEnabled() : storeBuilder.withCachingDisabled()) - .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) - .process(() -> new HeadersStoreWriterProcessor(), STORE_NAME) - .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); + startStreams( + cachingEnabled ? storeBuilder.withCachingEnabled() : storeBuilder.withCachingDisabled(), + HeadersStoreWriterProcessor::new); + } - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + private void startStreamsWithNonHeadersStore() throws Exception { + // A plain (non-WithHeaders) timestamped store: the headers-aware query types are unsupported here. + startStreams( + Stores.timestampedKeyValueStoreBuilder( + Stores.persistentTimestampedKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String()), + PlainStoreWriterProcessor::new); + } + + private void startStreamsWithPlainSupplierStore() throws Exception { + // A WithHeaders builder over a plain (non-timestamped) supplier: entries come back with + // timestamp = -1, which cannot be represented as a ReadOnlyRecord. + startStreams( + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + HeadersStoreWriterProcessor::new); + } + + private void assertUnknownQueryTypeAgainstNonHeadersStore(final Query query) throws Exception { + startStreamsWithNonHeadersStore(); + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); + + final StateQueryRequest request = + inStore(STORE_NAME).withQuery(query).withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult result = IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + assertTrue(result.getOnlyPartitionResult().isFailure()); + assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); } private ReadOnlyRecord query(final int key) { From 763418e1df2ddbca5334ad9dd06f60501f5f6dc6 Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Thu, 9 Jul 2026 14:32:44 -0400 Subject: [PATCH 16/17] Add TimestampedWindowKeyWithHeadersQuery IQv2 query type (KIP-1356) - Introduce the headers-aware windowed point query, returning a ReadOnlyRecordIterator, V> whose elements carry the stored record headers. - MeteredTimestampedWindowStoreWithHeaders handles it by forwarding a raw WindowKeyQuery to the wrapped store and wrapping each entry as a ReadOnlyRecord (Windowed key + value + timestamp + frozen headers). - The native header window stores (RocksDBTimestampedWindowStoreWithHeaders, RocksDBTimeOrderedWindowStoreWithHeaders) special-case WindowKeyQuery to super.query() so it is served via StoreQueryUtils; every other query type (incl. WindowRangeQuery) stays UNKNOWN_QUERY_TYPE, deferred to a follow-up, mirroring the KV KeyQuery PR. - Also fixes a latent WindowKeyQuery constructor bug that swapped timeFrom/timeTo. --- .../TimestampedWindowKeyWithHeadersQuery.java | 126 +++++++++++++++++ .../kafka/streams/query/WindowKeyQuery.java | 4 +- ...eredTimestampedWindowStoreWithHeaders.java | 133 ++++++++++++++++++ ...ksDBTimeOrderedWindowStoreWithHeaders.java | 14 +- ...ksDBTimestampedWindowStoreWithHeaders.java | 10 ++ 5 files changed, 281 insertions(+), 6 deletions(-) create mode 100644 streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java diff --git a/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java new file mode 100644 index 0000000000000..a8ed2bdfa6c2c --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java @@ -0,0 +1,126 @@ +/* + * 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.query; + +import org.apache.kafka.common.annotation.InterfaceAudience; +import org.apache.kafka.common.annotation.InterfaceStability.Evolving; +import org.apache.kafka.streams.kstream.Windowed; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; +import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; + +import java.time.Instant; +import java.util.Objects; +import java.util.Optional; + +/** + * Interactive query for retrieving records of a single key across a window-start range, including + * their record headers, from a {@link TimestampedWindowStoreWithHeaders}. + * + *

This is the headers-aware parallel of {@link WindowKeyQuery}: it returns a + * {@link ReadOnlyRecordIterator} of {@link ReadOnlyRecord} elements, each carrying the windowed key, + * value, timestamp, and headers, whereas {@link WindowKeyQuery} returns a + * {@link org.apache.kafka.streams.state.WindowStoreIterator} of values keyed by window-start + * timestamp (value only, no headers). Each element's key is a {@link Windowed} whose window + * describes the record's window; the record's stored event-time is exposed via + * {@link ReadOnlyRecord#timestamp()}. + * + *

As with {@link WindowKeyQuery}, a closed window-start range must be supplied (both + * {@code timeFrom} and {@code timeTo}); open-ended ranges are not supported. + * + *

Headers are persisted and returned only when the store is backed by a native headers store, + * i.e. built with a KIP-1271 {@code WithHeaders} byte-store supplier (e.g. + * {@code Stores.persistentTimestampedWindowStoreWithHeaders}). A {@code WithHeaders} store built + * over a legacy (non-headers) supplier cannot persist headers, so the store-served reads come back + * with empty {@code headers()}. + * + *

Against a plain store not built with the {@code WithHeaders} builder at all, this query type is + * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}. + * + *

Each element is a {@link ReadOnlyRecord}, whose timestamp is non-negative. If the backing store + * does not persist timestamps -- for example a {@code WithHeaders} store built over a plain + * window-store supplier that surfaces entries with {@code NO_TIMESTAMP} ({@code -1}) -- that entry + * cannot be represented, so advancing the returned {@link ReadOnlyRecordIterator} throws + * {@link org.apache.kafka.streams.errors.StreamsException} at that entry. Back the store with + * {@code Stores.persistentTimestampedWindowStoreWithHeaders(...)} to persist timestamps and headers, + * or use {@link WindowKeyQuery} if headers are not needed. + * + * @param Type of keys + * @param Type of values + */ +@Evolving +@InterfaceAudience.Public +public final class TimestampedWindowKeyWithHeadersQuery implements Query, V>> { + + private final K key; + private final Optional timeFrom; + private final Optional timeTo; + + private TimestampedWindowKeyWithHeadersQuery(final K key, + final Optional timeFrom, + final Optional timeTo) { + this.key = key; + this.timeFrom = timeFrom; + this.timeTo = timeTo; + } + + /** + * Creates a query that will retrieve the records (value, timestamp, and headers) identified by + * {@code key} whose window start falls within the closed range {@code [timeFrom, timeTo]}. + * @param key The key to retrieve + * @param timeFrom The inclusive lower bound of the window-start range + * @param timeTo The inclusive upper bound of the window-start range + * @param The type of the key + * @param The type of the value that will be retrieved + */ + public static TimestampedWindowKeyWithHeadersQuery withKeyAndWindowStartRange(final K key, + final Instant timeFrom, + final Instant timeTo) { + Objects.requireNonNull(key, "the key should not be null"); + return new TimestampedWindowKeyWithHeadersQuery<>(key, Optional.of(timeFrom), Optional.of(timeTo)); + } + + /** + * The key that was specified for this query. + */ + public K key() { + return key; + } + + /** + * The inclusive lower bound of the window-start range, if specified. + */ + public Optional timeFrom() { + return timeFrom; + } + + /** + * The inclusive upper bound of the window-start range, if specified. + */ + public Optional timeTo() { + return timeTo; + } + + @Override + public String toString() { + return "TimestampedWindowKeyWithHeadersQuery{" + + "key=" + key + + ", timeFrom=" + timeFrom + + ", timeTo=" + timeTo + + '}'; + } +} diff --git a/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java index ab14cfc728529..9d37ebf38d0ae 100644 --- a/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java +++ b/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java @@ -33,8 +33,8 @@ public class WindowKeyQuery implements Query> { private final Optional timeTo; private WindowKeyQuery(final K key, - final Optional timeTo, - final Optional timeFrom) { + final Optional timeFrom, + final Optional timeTo) { this.key = key; this.timeFrom = timeFrom; this.timeTo = timeTo; diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java index c1ab307001baa..0525296d2d677 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java @@ -23,8 +23,12 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.errors.ProcessorStateException; +import org.apache.kafka.streams.errors.StreamsException; import org.apache.kafka.streams.kstream.Windowed; +import org.apache.kafka.streams.kstream.internals.TimeWindow; import org.apache.kafka.streams.processor.StateStore; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.processor.api.Record; import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; import org.apache.kafka.streams.processor.internals.SerdeGetter; import org.apache.kafka.streams.query.FailureReason; @@ -32,10 +36,12 @@ 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.TimestampedWindowKeyWithHeadersQuery; import org.apache.kafka.streams.query.WindowKeyQuery; import org.apache.kafka.streams.query.WindowRangeQuery; import org.apache.kafka.streams.query.internals.InternalQueryResultUtil; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.TimestampedBytesStore; import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; @@ -61,6 +67,8 @@ public class MeteredTimestampedWindowStoreWithHeaders extends MeteredWindowStore> implements TimestampedWindowStoreWithHeaders { + private final long windowSizeMs; + MeteredTimestampedWindowStoreWithHeaders( final WindowStore inner, final long windowSizeMs, @@ -70,6 +78,7 @@ public class MeteredTimestampedWindowStoreWithHeaders final Serde> valueSerde ) { super(inner, windowSizeMs, metricScope, time, keySerde, valueSerde); + this.windowSizeMs = windowSizeMs; } @SuppressWarnings("unchecked") @@ -139,6 +148,8 @@ public QueryResult query( if (query instanceof WindowKeyQuery) { result = runWindowKeyQuery((WindowKeyQuery>) query, positionBound, config); + } else if (query instanceof TimestampedWindowKeyWithHeadersQuery) { + result = runTimestampedWindowKeyWithHeadersQuery((TimestampedWindowKeyWithHeadersQuery) query, positionBound, config); } else if (query instanceof WindowRangeQuery) { result = runWindowRangeQuery((WindowRangeQuery>) query, positionBound, config); } else { @@ -213,6 +224,47 @@ private QueryResult runWindowKeyQuery( return queryResult; } + /** + * Handles {@link TimestampedWindowKeyWithHeadersQuery} by forwarding a raw byte-level + * {@link WindowKeyQuery} to the wrapped store and surfacing each entry as a {@link ReadOnlyRecord} + * (via {@link Record}) whose key is a {@link Windowed} of the queried key and the entry's window, + * carrying the stored value, event-time timestamp, and headers. + */ + @SuppressWarnings("unchecked") + private QueryResult runTimestampedWindowKeyWithHeadersQuery( + final TimestampedWindowKeyWithHeadersQuery query, + final PositionBound positionBound, + final QueryConfig config + ) { + final QueryResult queryResult; + + if (query.timeFrom().isPresent() && query.timeTo().isPresent()) { + final WindowKeyQuery rawKeyQuery = + WindowKeyQuery.withKeyAndWindowStartRange( + serializeKey(query.key(), internalContext.headers()), + query.timeFrom().get(), + query.timeTo().get() + ); + final QueryResult> rawResult = wrapped().query(rawKeyQuery, positionBound, config); + if (rawResult.isSuccess()) { + final ReadOnlyRecordIterator, V> resultIterator = + new MeteredWindowStoreWithHeadersReadOnlyRecordIterator(rawResult.getResult(), query.key()); + queryResult = (QueryResult) InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, resultIterator); + } else { + queryResult = (QueryResult) rawResult; + } + } else { + queryResult = QueryResult.forFailure( + FailureReason.UNKNOWN_QUERY_TYPE, + "This 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." + ); + } + return queryResult; + } + private MeteredWindowStoreIterator meteredIterator( final QueryResult> rawResult, final Function valueDeserializer @@ -412,6 +464,87 @@ public Windowed peekNextKey() { } } + /** + * Iterator backing {@link TimestampedWindowKeyWithHeadersQuery}: yields each entry as a + * {@link ReadOnlyRecord} (implemented by {@link Record}) whose key is a {@link Windowed} of the + * queried key and the entry's window, carrying the value, stored event-time timestamp, and the + * stored headers, with the headers frozen so a caller cannot mutate the read-only result. + * + *

A {@link ReadOnlyRecord} timestamp is contractually non-negative, so an entry with a + * negative stored timestamp cannot be represented -- for example a {@code WithHeaders} store + * built over a plain, non-timestamped window supplier surfaces entries with + * {@code NO_TIMESTAMP} (-1). This mirrors the rule the point/range headers queries apply. But + * because a lazily-evaluated iterator has already returned a successful {@link QueryResult} + * before any entry is read, such an entry cannot be surfaced as a query-level failure; it is + * instead reported by throwing a {@link StreamsException} while advancing the iterator. + */ + private class MeteredWindowStoreWithHeadersReadOnlyRecordIterator + implements ReadOnlyRecordIterator, V>, MeteredIterator { + + private final WindowStoreIterator iter; + private final K key; + private final long startNs; + private final long startTimestampMs; + + private MeteredWindowStoreWithHeadersReadOnlyRecordIterator( + final WindowStoreIterator iter, + final K key + ) { + this.iter = iter; + this.key = key; + this.startNs = time.nanoseconds(); + this.startTimestampMs = time.milliseconds(); + numOpenIterators.increment(); + openIterators.add(this); + } + + @Override + public long startTimestamp() { + return startTimestampMs; + } + + @Override + public boolean hasNext() { + return iter.hasNext(); + } + + @Override + public ReadOnlyRecord, V> next() { + final KeyValue next = iter.next(); + final long windowStart = next.key; + final ValueTimestampHeaders valueTimestampHeaders = deserializeValue(next.value); + final Headers headers = valueTimestampHeaders.headers(); + if (valueTimestampHeaders.timestamp() < 0) { + throw new StreamsException( + "Cannot represent the stored record for key [" + key + "] and window start [" + + windowStart + "] as a ReadOnlyRecord: its timestamp (" + + valueTimestampHeaders.timestamp() + ") is negative."); + } + final Windowed windowedKey = + new Windowed<>(key, new TimeWindow(windowStart, windowStart + windowSizeMs)); + final Record, V> record = new Record<>( + windowedKey, + valueTimestampHeaders.value(), + valueTimestampHeaders.timestamp(), + headers); + ((RecordHeaders) record.headers()).setReadOnly(); + return record; + } + + @Override + public void close() { + try { + iter.close(); + } finally { + final long duration = time.nanoseconds() - startNs; + fetchSensor.record(duration); + iteratorDurationSensor.record(duration); + numOpenIterators.decrement(); + openIterators.remove(this); + } + } + } + private MeteredWindowedKeyValueIterator meteredWindowedIterator( final QueryResult, byte[]>> rawResult, final Function, ValueType> valueConverter diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java index 8d4d5a23b4419..08a8f1771b774 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java @@ -16,12 +16,12 @@ */ package org.apache.kafka.streams.state.internals; -import org.apache.kafka.streams.processor.StateStore; 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.WindowKeyQuery; import org.apache.kafka.streams.state.HeadersBytesStore; import org.apache.kafka.streams.state.TimestampedBytesStore; @@ -32,9 +32,6 @@ * {@link TimestampedBytesStore} (for timestamp support) and {@link HeadersBytesStore} * (for header support) marker interfaces. *

- * This store returns {@link QueryResult#forUnknownQueryType(Query, StateStore)} for all queries, - * as IQv2 query handling is done at the metered layer. - *

* The storage format for values is: [headersSize(varint)][headersBytes][timestamp(8)][value] * * @see RocksDBTimeOrderedWindowStore @@ -53,6 +50,15 @@ class RocksDBTimeOrderedWindowStoreWithHeaders extends RocksDBTimeOrderedWindowS public QueryResult query(final Query query, final PositionBound positionBound, final QueryConfig config) { + // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery + // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited + // RocksDBTimeOrderedWindowStore handling (StoreQueryUtils). Every other query type (e.g. + // WindowRangeQuery) stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by + // their own KIP-1356 follow-ups. + if (query instanceof WindowKeyQuery) { + return super.query(query, positionBound, config); + } + final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L; final QueryResult result; final Position position = getPosition(); diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java index 37b1994b02f73..f588aef3866ee 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java @@ -21,6 +21,7 @@ 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.WindowKeyQuery; import org.apache.kafka.streams.state.HeadersBytesStore; import org.apache.kafka.streams.state.TimestampedBytesStore; @@ -57,6 +58,15 @@ class RocksDBTimestampedWindowStoreWithHeaders extends RocksDBWindowStore implem public QueryResult query(final Query query, final PositionBound positionBound, final QueryConfig config) { + // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery + // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited + // RocksDBWindowStore handling (StoreQueryUtils). Every other query type (e.g. WindowRangeQuery) + // stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by their own KIP-1356 + // follow-ups. + if (query instanceof WindowKeyQuery) { + return super.query(query, positionBound, config); + } + final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L; final QueryResult result; final Position position = getPosition(); From c125059d4bcebda6130a1e5b250e44a9c56c9acd Mon Sep 17 00:00:00 2001 From: Jess Jin Date: Thu, 9 Jul 2026 16:38:52 -0400 Subject: [PATCH 17/17] Add tests for TimestampedWindowKeyWithHeadersQuery (KIP-1356) --- .../IQv2HeadersStoreIntegrationTest.java | 233 ++++++++++++-- ...TimestampedWindowStoreWithHeadersTest.java | 61 ++++ ...TimeOrderedWindowStoreWithHeadersTest.java | 26 +- ...TimestampedWindowStoreWithHeadersTest.java | 39 ++- ...mpedWindowStoreWithHeadersBuilderTest.java | 286 ++++++++++++++++++ 5 files changed, 602 insertions(+), 43 deletions(-) diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java index 88483fe540132..253e614c71223 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java @@ -35,6 +35,7 @@ import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; import org.apache.kafka.streams.kstream.Consumed; import org.apache.kafka.streams.kstream.Produced; +import org.apache.kafka.streams.kstream.Windowed; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.ProcessorSupplier; @@ -49,11 +50,14 @@ import org.apache.kafka.streams.query.StateQueryResult; import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.StoreBuilder; import org.apache.kafka.streams.state.Stores; import org.apache.kafka.streams.state.TimestampedKeyValueStore; import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.TimestampedWindowStore; +import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.test.TestUtils; @@ -68,6 +72,7 @@ import java.io.IOException; import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.LinkedList; import java.util.List; @@ -102,6 +107,11 @@ public class IQv2HeadersStoreIntegrationTest { .add("source", "test".getBytes()) .add("version", "1.0".getBytes()); + // Window-store test parameters (used by the window headers query tests). + private static final long WINDOW_SIZE_MS = 10L; + private static final Duration WINDOW_SIZE = Duration.ofMillis(WINDOW_SIZE_MS); + private static final Duration RETENTION = Duration.ofMinutes(1); + private String inputStream; private String outputStream; private long baseTimestamp; @@ -142,7 +152,7 @@ public void afterTest() { @Test public void shouldHandleTimestampedKeyWithHeadersQuery() throws Exception { - startStreamsWithHeadersStore(); + startStreamsWithKeyValueHeadersStore(); // key 1 has headers, key 2 has empty headers, key 3 is tombstoned (null value) produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, @@ -153,24 +163,24 @@ public void shouldHandleTimestampedKeyWithHeadersQuery() throws Exception { KeyValue.pair(3, "c0"), KeyValue.pair(3, null)); // key 1: key + value + timestamp + headers round-trip - final ReadOnlyRecord result1 = query(1); + final ReadOnlyRecord result1 = keyQuery(1); assertEquals(Integer.valueOf(1), result1.key()); assertEquals("a0", result1.value()); assertEquals(baseTimestamp, result1.timestamp()); assertEquals(HEADERS, result1.headers()); // key 2: written with no headers -> empty (never null) headers - final ReadOnlyRecord result2 = query(2); + final ReadOnlyRecord result2 = keyQuery(2); assertEquals(Integer.valueOf(2), result2.key()); assertEquals("b0", result2.value()); assertEquals(baseTimestamp + 1, result2.timestamp()); assertEquals(new RecordHeaders(), result2.headers()); // key 3: tombstoned -> null result, never a partially-populated wrapper - assertNull(query(3)); + assertNull(keyQuery(3)); // never-written key -> null result - assertNull(query(999)); + assertNull(keyQuery(999)); } @Test @@ -180,16 +190,16 @@ public void shouldServeCacheHitWhenCachingEnabledAndRecordNotYetFlushed() throws // successful query therefore proves the result was served from the cache via // CachingKeyValueStoreWithHeaders -> the metered store's cache-hit path, end-to-end. commitIntervalMs = Duration.ofMinutes(10).toMillis(); - startStreamsWithHeadersStore(true); + startStreamsWithKeyValueHeadersStore(true); produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); // Read-your-writes: the not-yet-flushed record is visible (served from the cache), with headers. - final ReadOnlyRecord result = query(1); + final ReadOnlyRecord result = keyQuery(1); // skipCache bypasses the cache and reads the persistent store directly. A null result positively // proves nothing has been flushed, so the read above was genuinely cache-served (not an accidental // store read) -- and it covers skipCache end-to-end. - assertNull(querySkipCache(1)); + assertNull(keyQuerySkipCache(1)); assertEquals(Integer.valueOf(1), result.key()); assertEquals("a0", result.value()); assertEquals(baseTimestamp, result.timestamp()); @@ -205,7 +215,7 @@ public void shouldFailWithUnknownQueryTypeForKeyQueryAgainstNonHeadersStore() th public void shouldHandleTimestampedRangeWithHeadersQuery() throws Exception { // Caching disabled: a range query reads the underlying store directly (it never consults the // cache), so the writes must be store-served. - startStreamsWithHeadersStore(); + startStreamsWithKeyValueHeadersStore(); // keys 1,2 (headers), key 3 (empty headers), key 4 written then tombstone produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, @@ -296,7 +306,7 @@ public void shouldFailWithUnknownQueryTypeForRangeQueryAgainstNonHeadersStore() @Test public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() throws Exception { // The query succeeds, but iterating throws because timestamp = -1 cannot be a ReadOnlyRecord. - startStreamsWithPlainSupplierStore(); + startStreamsWithKeyValuePlainSupplierStore(); produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); @@ -315,6 +325,82 @@ public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() thro } } + @Test + public void shouldHandleTimestampedWindowKeyWithHeadersQuery() throws Exception { + // Caching disabled: the window key query reads the underlying store directly. + startStreamsWithWindowHeadersStore(); + + final long window0 = baseTimestamp; + final long window1 = baseTimestamp + 100; + final long window2 = baseTimestamp + 200; + + // A single key across three window starts: window0 (headers), window1 (empty headers), + // window2 written then tombstone. + produceDataToTopicWithHeaders(inputStream, window0, HEADERS, KeyValue.pair(1, "v0")); + produceDataToTopicWithHeaders(inputStream, window1, new RecordHeaders(), KeyValue.pair(1, "v1")); + produceDataToTopicWithHeaders(inputStream, window2, HEADERS, KeyValue.pair(1, "v2"), KeyValue.pair(1, null)); + + final List, String>> records = + windowKeyQuery(1, Instant.ofEpochMilli(window0), Instant.ofEpochMilli(window2)); + + // window2 tombstone -> omitted; window0 and window1 returned in window-start order. + assertEquals(2, records.size()); + + // window0: windowed key + value + timestamp + headers round-trip, with the window in the key. + final ReadOnlyRecord, String> r0 = records.get(0); + assertEquals(Integer.valueOf(1), r0.key().key()); + assertEquals(window0, r0.key().window().start()); + assertEquals(window0 + WINDOW_SIZE_MS, r0.key().window().end()); + assertEquals("v0", r0.value()); + assertEquals(window0, r0.timestamp()); + assertEquals(HEADERS, r0.headers()); + // returned headers are a read-only snapshot: no mutation (add or remove) is allowed + assertThrows(IllegalStateException.class, () -> r0.headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> r0.headers().remove("source")); + + // window1: written without headers -> empty (never null) headers. + final ReadOnlyRecord, String> r1 = records.get(1); + assertEquals(Integer.valueOf(1), r1.key().key()); + assertEquals(window1, r1.key().window().start()); + assertEquals(window1 + WINDOW_SIZE_MS, r1.key().window().end()); + assertEquals("v1", r1.value()); + assertEquals(window1, r1.timestamp()); + assertEquals(new RecordHeaders(), r1.headers()); + // even the empty headers are a read-only snapshot: no mutation (add or remove) is allowed + assertThrows(IllegalStateException.class, () -> r1.headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> r1.headers().remove("source")); + } + + @Test + public void shouldFailWithUnknownQueryTypeForWindowKeyQueryAgainstNonHeadersStore() throws Exception { + startStreamsWithWindowNonHeadersStore(); + assertUnknownQueryType(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + 1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp + 1))); + } + + @Test + public void shouldThrowForTimestampedWindowKeyWithHeadersQueryOnPlainSupplier() throws Exception { + // A WithHeaders window builder over a plain (non-timestamped) window supplier: entries come back + // with timestamp = -1. The query succeeds, but iterating throws because -1 cannot be a ReadOnlyRecord. + startStreamsWithWindowPlainSupplierStore(); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one")); + + final StateQueryRequest, String>> request = + inStore(STORE_NAME) + .withQuery(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + 1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp + 1))) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult, String>> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + final QueryResult, String>> onlyResult = result.getOnlyPartitionResult(); + assertTrue(onlyResult.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = onlyResult.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } + } + private void startStreams(final StoreBuilder storeBuilder, final ProcessorSupplier processorSupplier) throws Exception { final StreamsBuilder builder = new StreamsBuilder(); @@ -328,14 +414,14 @@ private void startStreams(final StoreBuilder storeBuilder, IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); } - private void startStreamsWithHeadersStore() throws Exception { + private void startStreamsWithKeyValueHeadersStore() throws Exception { // Caching disabled: every IQv2 query is forced down to the persistent // RocksDBTimestampedStoreWithHeaders layer, exercising its KeyQuery handling // (rather than being short-circuited by a cache hit). - startStreamsWithHeadersStore(false); + startStreamsWithKeyValueHeadersStore(false); } - private void startStreamsWithHeadersStore(final boolean cachingEnabled) throws Exception { + private void startStreamsWithKeyValueHeadersStore(final boolean cachingEnabled) throws Exception { final StoreBuilder> storeBuilder = Stores.timestampedKeyValueStoreWithHeadersBuilder( Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), @@ -343,20 +429,20 @@ private void startStreamsWithHeadersStore(final boolean cachingEnabled) throws E Serdes.String()); startStreams( cachingEnabled ? storeBuilder.withCachingEnabled() : storeBuilder.withCachingDisabled(), - HeadersStoreWriterProcessor::new); + KeyValueHeadersStoreWriterProcessor::new); } - private void startStreamsWithNonHeadersStore() throws Exception { + private void startStreamsWithKeyValueNonHeadersStore() throws Exception { // A plain (non-WithHeaders) timestamped store: the headers-aware query types are unsupported here. startStreams( Stores.timestampedKeyValueStoreBuilder( Stores.persistentTimestampedKeyValueStore(STORE_NAME), Serdes.Integer(), Serdes.String()), - PlainStoreWriterProcessor::new); + KeyValuePlainStoreWriterProcessor::new); } - private void startStreamsWithPlainSupplierStore() throws Exception { + private void startStreamsWithKeyValuePlainSupplierStore() throws Exception { // A WithHeaders builder over a plain (non-timestamped) supplier: entries come back with // timestamp = -1, which cannot be represented as a ReadOnlyRecord. startStreams( @@ -364,11 +450,48 @@ private void startStreamsWithPlainSupplierStore() throws Exception { Stores.persistentKeyValueStore(STORE_NAME), Serdes.Integer(), Serdes.String()).withCachingDisabled(), - HeadersStoreWriterProcessor::new); + KeyValueHeadersStoreWriterProcessor::new); + } + + private void startStreamsWithWindowHeadersStore() throws Exception { + // Caching disabled: every IQv2 query is forced down to the persistent + // RocksDBTimestampedWindowStoreWithHeaders layer, exercising its inherited WindowKeyQuery + // handling (rather than being short-circuited by a cache hit). + startStreams( + Stores.timestampedWindowStoreWithHeadersBuilder( + Stores.persistentTimestampedWindowStoreWithHeaders(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + WindowHeadersStoreWriterProcessor::new); + } + + private void startStreamsWithWindowNonHeadersStore() throws Exception { + // A plain (non-WithHeaders) timestamped window store: the headers-aware query types are unsupported here. + startStreams( + Stores.timestampedWindowStoreBuilder( + Stores.persistentTimestampedWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()), + WindowPlainStoreWriterProcessor::new); + } + + private void startStreamsWithWindowPlainSupplierStore() throws Exception { + // A WithHeaders window builder over a plain (non-timestamped) window supplier: entries come back + // with timestamp = -1, which cannot be represented as a ReadOnlyRecord. + startStreams( + Stores.timestampedWindowStoreWithHeadersBuilder( + Stores.persistentWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + WindowHeadersStoreWriterProcessor::new); } private void assertUnknownQueryTypeAgainstNonHeadersStore(final Query query) throws Exception { - startStreamsWithNonHeadersStore(); + startStreamsWithKeyValueNonHeadersStore(); + assertUnknownQueryType(query); + } + + private void assertUnknownQueryType(final Query query) { produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); final StateQueryRequest request = @@ -379,7 +502,7 @@ private void assertUnknownQueryTypeAgainstNonHeadersStore(final Query que assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); } - private ReadOnlyRecord query(final int key) { + private ReadOnlyRecord keyQuery(final int key) { final StateQueryRequest> request = inStore(STORE_NAME) .withQuery(TimestampedKeyWithHeadersQuery.withKey(key)) @@ -394,7 +517,7 @@ private ReadOnlyRecord query(final int key) { return onlyResult == null ? null : onlyResult.getResult(); } - private ReadOnlyRecord querySkipCache(final int key) { + private ReadOnlyRecord keyQuerySkipCache(final int key) { // skipCache forwards the query past the record cache to the persistent store. Use the default // (unbounded) position bound on purpose: the persistent layer may legitimately be empty (nothing // flushed), so bounding on the input position would never be satisfied. @@ -420,6 +543,25 @@ private List> rangeQuery(final TimestampedRangeW return records; } + private List, String>> windowKeyQuery(final int key, + final Instant timeFrom, + final Instant timeTo) { + final StateQueryRequest, String>> request = + inStore(STORE_NAME) + .withQuery(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + key, timeFrom, timeTo)) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult, String>> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + final List, String>> records = new ArrayList<>(); + try (ReadOnlyRecordIterator, String> iterator = result.getOnlyPartitionResult().getResult()) { + while (iterator.hasNext()) { + records.add(iterator.next()); + } + } + return records; + } + private static List keys(final List> records) { return records.stream().map(ReadOnlyRecord::key).collect(Collectors.toList()); } @@ -467,7 +609,7 @@ private final void produceDataToTopicWithHeaders(final String topic, } } - private static class HeadersStoreWriterProcessor implements Processor { + private static class KeyValueHeadersStoreWriterProcessor implements Processor { private ProcessorContext context; private TimestampedKeyValueStoreWithHeaders store; @@ -490,7 +632,7 @@ public void process(final Record record) { } } - private static class PlainStoreWriterProcessor implements Processor { + private static class KeyValuePlainStoreWriterProcessor implements Processor { private ProcessorContext context; private TimestampedKeyValueStore store; @@ -508,4 +650,49 @@ public void process(final Record record) { context.forward(record); } } + + private static class WindowHeadersStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedWindowStoreWithHeaders store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + // The record timestamp doubles as the window start; a null value tombstones that window. + if (record.value() == null) { + store.put(record.key(), null, record.timestamp()); + } else { + store.put( + record.key(), + ValueTimestampHeaders.make(record.value(), record.timestamp(), record.headers()), + record.timestamp()); + } + context.forward(record); + } + } + + private static class WindowPlainStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedWindowStore store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + store.put( + record.key(), + ValueAndTimestamp.make(record.value(), record.timestamp()), + record.timestamp()); + context.forward(record); + } + } } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java index eb2c6a65f8ebb..0003b88dee0f7 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java @@ -38,9 +38,16 @@ import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; import org.apache.kafka.streams.processor.internals.ProcessorStateManager; import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; +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.TimestampedWindowKeyWithHeadersQuery; +import org.apache.kafka.streams.query.WindowKeyQuery; import org.apache.kafka.streams.state.KeyValueIterator; import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.streams.state.WindowStore; +import org.apache.kafka.streams.state.WindowStoreIterator; import org.apache.kafka.test.InternalMockProcessorContext; import org.apache.kafka.test.KeyValueIteratorStub; import org.apache.kafka.test.MockRecordCollector; @@ -49,12 +56,16 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import java.time.Instant; +import java.util.Iterator; import java.util.List; +import java.util.Optional; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -484,4 +495,54 @@ public void shouldUseHeadersFromValueToDeserializeKeyInBackwardFetchRange() { verify(keyDeserializer).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); } + + @Test + public void shouldForwardWindowKeyQueryBoundsForTimestampedWindowKeyWithHeadersQuery() { + setUp(); + store.init(context, store); + + final Instant timeFrom = Instant.ofEpochMilli(5); + final Instant timeTo = Instant.ofEpochMilli(100); + final WindowKeyQuery rawQuery = forwardedRawWindowKeyQuery( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange(KEY, timeFrom, timeTo)); + + // The typed query is translated into a raw byte-level WindowKeyQuery with the serialized key + // and the same window-start range, forwarded to the wrapped store. + assertEquals(KEY_BYTES, rawQuery.getKey()); + assertEquals(Optional.of(timeFrom), rawQuery.getTimeFrom()); + assertEquals(Optional.of(timeTo), rawQuery.getTimeTo()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private WindowKeyQuery forwardedRawWindowKeyQuery(final Query query) { + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult(windowStoreIterator(List.of()))); + store.query(query, PositionBound.unbounded(), new QueryConfig(false)); + final ArgumentCaptor captor = ArgumentCaptor.forClass(WindowKeyQuery.class); + verify(innerStoreMock).query(captor.capture(), any(PositionBound.class), any(QueryConfig.class)); + return captor.getValue(); + } + + private static WindowStoreIterator windowStoreIterator(final List> data) { + final Iterator> iterator = data.iterator(); + return new WindowStoreIterator<>() { + @Override + public void close() { } + + @Override + public Long peekNextKey() { + throw new UnsupportedOperationException(); + } + + @Override + public boolean hasNext() { + return iterator.hasNext(); + } + + @Override + public KeyValue next() { + return iterator.next(); + } + }; + } } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java index 13994dd36af45..a775839f4a49a 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java @@ -18,6 +18,7 @@ import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.query.FailureReason; import org.apache.kafka.streams.query.PositionBound; @@ -39,6 +40,7 @@ import java.time.Instant; import java.util.Properties; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -87,17 +89,31 @@ public void tearDown() { } @Test - public void shouldReturnUnknownQueryTypeForWindowKeyQuery() { + public void shouldHandleWindowKeyQuery() { + // KIP-1356 (window key query only): the native time-ordered window header store special-cases + // WindowKeyQuery to the inherited RocksDBTimeOrderedWindowStore handling (StoreQueryUtils), + // returning the raw stored header-format bytes (previously UNKNOWN_QUERY_TYPE); the metered store + // does the header-aware deserialization. This matches the adapter build path. + final Bytes key = new Bytes("test-key".getBytes()); + final byte[] storedBytes = "headers+timestamp+value".getBytes(); + final long windowStart = 1_000L; + windowStore.put(key, storedBytes, windowStart); + final WindowKeyQuery query = WindowKeyQuery.withKeyAndWindowStartRange( - new Bytes("test-key".getBytes()), + key, Instant.ofEpochMilli(0), - Instant.ofEpochMilli(Long.MAX_VALUE) + Instant.ofEpochMilli(RETENTION_PERIOD) ); final QueryResult> result = windowStore.query(query, PositionBound.unbounded(), new QueryConfig(false)); - assertFalse(result.isSuccess()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getFailureReason()); + assertTrue(result.isSuccess(), "Expected WindowKeyQuery to succeed"); + try (WindowStoreIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext(), "Expected the stored entry in the window key result"); + final KeyValue keyValue = iterator.next(); + assertEquals(windowStart, keyValue.key); + assertArrayEquals(storedBytes, keyValue.value, "Expected the raw stored bytes to be returned"); + } assertNotNull(result.getPosition()); } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java index 554f5559865a2..3a0b999390598 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java @@ -18,6 +18,7 @@ import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.query.FailureReason; import org.apache.kafka.streams.query.PositionBound; @@ -39,6 +40,7 @@ import java.time.Instant; import java.util.Properties; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -91,24 +93,31 @@ public void tearDown() { } @Test - public void shouldReturnUnknownQueryTypeForWindowKeyQuery() { + public void shouldHandleWindowKeyQuery() { + // KIP-1356 (window key query only): the native window header store special-cases WindowKeyQuery + // to the inherited RocksDBWindowStore handling (StoreQueryUtils), returning the raw stored + // header-format bytes (previously UNKNOWN_QUERY_TYPE); the metered store does the header-aware + // deserialization. This matches the adapter build path. + final Bytes key = new Bytes("test-key".getBytes()); + final byte[] storedBytes = "headers+timestamp+value".getBytes(); + final long windowStart = 1_000L; + windowStore.put(key, storedBytes, windowStart); + final WindowKeyQuery query = WindowKeyQuery.withKeyAndWindowStartRange( - new Bytes("test-key".getBytes()), + key, Instant.ofEpochMilli(0), - Instant.ofEpochMilli(Long.MAX_VALUE) - ); - final PositionBound positionBound = PositionBound.unbounded(); - final QueryConfig config = new QueryConfig(false); - - final QueryResult> result = windowStore.query(query, positionBound, config); - - // Verify: Window store with headers currently returns UNKNOWN_QUERY_TYPE - assertFalse(result.isSuccess(), "Expected query to fail with unknown query type"); - assertEquals( - FailureReason.UNKNOWN_QUERY_TYPE, - result.getFailureReason(), - "Expected UNKNOWN_QUERY_TYPE failure reason" + Instant.ofEpochMilli(RETENTION_PERIOD) ); + final QueryResult> result = + windowStore.query(query, PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess(), "Expected WindowKeyQuery to succeed"); + try (WindowStoreIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext(), "Expected the stored entry in the window key result"); + final KeyValue keyValue = iterator.next(); + assertEquals(windowStart, keyValue.key); + assertArrayEquals(storedBytes, keyValue.value, "Expected the raw stored bytes to be returned"); + } assertNotNull(result.getPosition(), "Expected position to be set"); } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java index 67216bca41157..991b513ef87e0 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java @@ -17,11 +17,34 @@ package org.apache.kafka.streams.state.internals; +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.common.utils.Bytes; import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.internals.LogContext; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.errors.StreamsException; +import org.apache.kafka.streams.kstream.Windowed; import org.apache.kafka.streams.processor.StateStore; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.processor.internals.MockStreamsMetrics; +import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; +import org.apache.kafka.streams.query.PositionBound; +import org.apache.kafka.streams.query.QueryConfig; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; +import org.apache.kafka.streams.query.WindowKeyQuery; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; +import org.apache.kafka.streams.state.ValueAndTimestamp; +import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.streams.state.WindowBytesStoreSupplier; +import org.apache.kafka.streams.state.WindowStore; +import org.apache.kafka.streams.state.WindowStoreIterator; +import org.apache.kafka.test.InternalMockProcessorContext; +import org.apache.kafka.test.TestUtils; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -31,17 +54,28 @@ import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.STRICT_STUBS) public class TimestampedWindowStoreWithHeadersBuilderTest { + private enum StoreType { NATIVE, ADAPTER, PLAIN_ADAPTER } + @Nested class BuilderTests { private static final String STORE_NAME = "name"; @@ -207,4 +241,256 @@ public void shouldThrowNullPointerIfInnerIsNull() { assertThrows(NullPointerException.class, () -> new TimestampedWindowStoreWithHeadersBuilder<>(null, Serdes.String(), Serdes.String(), new MockTime())); } } + + /** + * End-to-end query behavior of {@link TimestampedWindowKeyWithHeadersQuery} against a real store + * built by the builder over each supported build path. Caching is disabled: a window key query reads + * the underlying store directly (it never consults the cache), so writes must be store-served. + */ + @Nested + class QueryTests { + private static final String STORE_NAME = "test-window-store"; + private static final String METRICS_SCOPE = "metrics-scope"; + private static final long RETENTION = 60_000L; + private static final long SEGMENT_INTERVAL = 30_000L; + private static final long WINDOW_SIZE = 10_000L; + private static final long WINDOW_START = 1_000L; + + @Mock + private WindowBytesStoreSupplier supplier; + + private WindowStore innerStore(final StoreType storeType) { + switch (storeType) { + case NATIVE: + return new RocksDBTimestampedWindowStoreWithHeaders( + new RocksDBTimestampedSegmentedBytesStoreWithHeaders( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + case ADAPTER: + return new RocksDBTimestampedWindowStore( + new RocksDBTimestampedSegmentedBytesStore( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + case PLAIN_ADAPTER: + return new RocksDBWindowStore( + new RocksDBSegmentedBytesStore( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + default: + throw new IllegalArgumentException("unknown store type: " + storeType); + } + } + + private TimestampedWindowStoreWithHeaders buildAndInitStore(final StoreType storeType) { + lenient().when(supplier.name()).thenReturn(STORE_NAME); + lenient().when(supplier.metricsScope()).thenReturn(METRICS_SCOPE); + lenient().when(supplier.windowSize()).thenReturn(WINDOW_SIZE); + lenient().when(supplier.get()).thenReturn(innerStore(storeType)); + + final TimestampedWindowStoreWithHeaders store = + new TimestampedWindowStoreWithHeadersBuilder<>(supplier, Serdes.String(), Serdes.String(), new MockTime()) + .withLoggingDisabled() + .withCachingDisabled() + .build(); + + final ThreadCache cache = new ThreadCache(new LogContext("test "), 0, new MockStreamsMetrics(new Metrics())); + final InternalMockProcessorContext context = new InternalMockProcessorContext<>( + TestUtils.tempDirectory(), Serdes.String(), Serdes.String(), null, cache); + context.setRecordContext(new ProcessorRecordContext(0L, 0L, 0, "topic", new RecordHeaders())); + store.init(context, store); + return store; + } + + private Headers headersWith(final String key, final String value) { + return new RecordHeaders().add(key, value.getBytes(StandardCharsets.UTF_8)); + } + + private QueryResult, String>> windowKeyQuery( + final TimestampedWindowStoreWithHeaders store) { + return store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + "k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)), + PositionBound.unbounded(), + new QueryConfig(false)); + } + + @Test + public void shouldReturnHeadersForTimestampedWindowKeyWithHeadersQueryOnHeaderPersistingStore() { + // The native header window store persists headers, so the record carries them; the window in + // the Windowed key comes from the stored window start, the timestamp from the stored value. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + final Headers headers = headersWith("h", "x"); + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headers), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess(), "Expected TimestampedWindowKeyWithHeadersQuery to succeed"); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord, String> record = iterator.next(); + assertEquals("k", record.key().key()); + assertEquals(WINDOW_START, record.key().window().start()); + assertEquals(WINDOW_START + WINDOW_SIZE, record.key().window().end()); + assertEquals("v", record.value()); + assertEquals(1_005L, record.timestamp()); + assertEquals(headers, record.headers()); + // The IQ result is a read-only snapshot: its headers are immutable. + assertThrows(IllegalStateException.class, () -> record.headers().add("new", new byte[0])); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldReturnEmptyHeadersForTimestampedWindowKeyWithHeadersQueryOnAdapterStore() { + // The timestamped adapter keeps the timestamp but drops headers on write, so the record comes + // back with empty (never null) headers while value and timestamp still round-trip. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.ADAPTER); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord, String> record = iterator.next(); + assertEquals("k", record.key().key()); + assertEquals(WINDOW_START, record.key().window().start()); + assertEquals("v", record.value()); + assertEquals(1_005L, record.timestamp()); + assertEquals(new RecordHeaders(), record.headers()); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForTimestampedWindowKeyWithHeadersQueryOnPlainSupplier() { + // A plain (non-timestamped) window supplier surfaces every entry with timestamp = -1, which + // cannot be represented as a ReadOnlyRecord. The query succeeds, but iterating throws. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.PLAIN_ADAPTER); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess(), "The query itself succeeds; the failure surfaces while iterating"); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertThrows(StreamsException.class, iterator::next, + "An entry with ts=-1 cannot be represented as a ReadOnlyRecord"); + } + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForNegativeStoredTimestampForTimestampedWindowKeyWithHeadersQuery() { + // A caller can store a negative timestamp directly on a native (header-persisting) store. + // The query succeeds, but the lazily-evaluated iterator throws while advancing (a ReadOnlyRecord + // timestamp must be non-negative). + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } + } finally { + store.close(); + } + } + + @Test + public void shouldCollectExecutionInfoForTimestampedWindowKeyWithHeadersQueryWhenRequested() { + // With execution info enabled, the result must carry both the wrapped (native) store's entry + // and the metered handler's entry. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + "k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)), + PositionBound.unbounded(), + new QueryConfig(true)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + final String info = String.join("\n", result.getExecutionInfo()); + assertTrue( + info.contains(RocksDBTimestampedWindowStoreWithHeaders.class.getName()) + && info.contains(MeteredTimestampedWindowStoreWithHeaders.class.getName()), + "execution info missing an entry: " + info); + } + } finally { + store.close(); + } + } + + @Test + public void shouldNotCollectExecutionInfoForTimestampedWindowKeyWithHeadersQueryWhenNotRequested() { + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(result.getExecutionInfo().isEmpty(), "Expected no execution info: " + result.getExecutionInfo()); + } + } finally { + store.close(); + } + } + + @Test + public void shouldReturnIdenticalResultsForNativeAndAdapterBuiltStores() { + // Build-path parity for the existing (header-stripped) WindowKeyQuery: KIP-1356 makes the native + // store serve it exactly as the adapter build already did. + final ValueTimestampHeaders value = ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")); + final TimestampedWindowStoreWithHeaders nativeStore = buildAndInitStore(StoreType.NATIVE); + final TimestampedWindowStoreWithHeaders adapterStore = buildAndInitStore(StoreType.ADAPTER); + try { + nativeStore.put("k", value, WINDOW_START); + adapterStore.put("k", value, WINDOW_START); + + assertEquals( + plainWindowKeyResults(nativeStore), + plainWindowKeyResults(adapterStore), + "WindowKeyQuery results should be identical across native and adapter build paths"); + } finally { + nativeStore.close(); + adapterStore.close(); + } + } + + // Drains the (header-stripped) plain WindowKeyQuery, which yields the window-start timestamp keyed + // to a ValueAndTimestamp -- used to compare native and adapter build paths. + private List>> plainWindowKeyResults( + final TimestampedWindowStoreWithHeaders store) { + final WindowKeyQuery> query = + WindowKeyQuery.withKeyAndWindowStartRange("k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)); + final List>> out = new ArrayList<>(); + try (WindowStoreIterator> iterator = + store.query(query, PositionBound.unbounded(), new QueryConfig(false)).getResult()) { + while (iterator.hasNext()) { + out.add(iterator.next()); + } + } + return out; + } + } }