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..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 @@ -30,25 +30,34 @@ 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.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; 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; 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; @@ -63,24 +72,29 @@ import java.io.IOException; import java.time.Duration; +import java.time.Instant; +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 { @@ -93,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; @@ -133,7 +152,7 @@ public void afterTest() { @Test public void shouldHandleTimestampedKeyWithHeadersQuery() throws Exception { - startStreams(); + startStreamsWithKeyValueHeadersStore(); // key 1 has headers, key 2 has empty headers, key 3 is tombstoned (null value) produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, @@ -144,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 @@ -171,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(); - startStreams(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()); @@ -188,60 +207,302 @@ 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())); + public void shouldFailWithUnknownQueryTypeForKeyQueryAgainstNonHeadersStore() throws Exception { + assertUnknownQueryTypeAgainstNonHeadersStore(TimestampedKeyWithHeadersQuery.withKey(1)); + } - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + @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. + startStreamsWithKeyValueHeadersStore(); - produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); + // 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(), + 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 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()); // 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: 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: 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")); + } - final StateQueryRequest> request = + @Test + public void shouldFailWithUnknownQueryTypeForRangeQueryAgainstNonHeadersStore() throws Exception { + assertUnknownQueryTypeAgainstNonHeadersStore(TimestampedRangeWithHeadersQuery.withNoBounds()); + } + + @Test + public void shouldThrowForTimestampedRangeWithHeadersQueryOnPlainSupplier() throws Exception { + // The query succeeds, but iterating throws because timestamp = -1 cannot be a ReadOnlyRecord. + startStreamsWithKeyValuePlainSupplierStore(); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, + KeyValue.pair(1, "one"), KeyValue.pair(2, "two")); + + final StateQueryRequest> request = inStore(STORE_NAME) - .withQuery(TimestampedKeyWithHeadersQuery.withKey(1)) + .withQuery(TimestampedRangeWithHeadersQuery.withNoBounds()) .withPositionBound(PositionBound.at(inputPosition)); - final StateQueryResult> result = + final StateQueryResult> result = IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); - assertTrue(result.getOnlyPartitionResult().isFailure()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); + final QueryResult> onlyResult = result.getOnlyPartitionResult(); + assertTrue(onlyResult.isSuccess()); + try (ReadOnlyRecordIterator iterator = onlyResult.getResult()) { + assertThrows(StreamsException.class, iterator::next); + } } - private void startStreams() throws Exception { + @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(); + 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 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). - startStreams(false); + startStreamsWithKeyValueHeadersStore(false); } - private void startStreams(final boolean cachingEnabled) throws Exception { - final StreamsBuilder builder = new StreamsBuilder(); + private void startStreamsWithKeyValueHeadersStore(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(), + KeyValueHeadersStoreWriterProcessor::new); + } - kafkaStreams = new KafkaStreams(builder.build(), props()); - IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + 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()), + KeyValuePlainStoreWriterProcessor::new); + } + + 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( + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + 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 { + startStreamsWithKeyValueNonHeadersStore(); + assertUnknownQueryType(query); + } + + private void assertUnknownQueryType(final Query query) { + 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) { + private ReadOnlyRecord keyQuery(final int key) { final StateQueryRequest> request = inStore(STORE_NAME) .withQuery(TimestampedKeyWithHeadersQuery.withKey(key)) @@ -256,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. @@ -268,6 +529,47 @@ 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 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()); + } + + 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(); @@ -307,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; @@ -319,14 +621,18 @@ 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); } } - private static class PlainStoreWriterProcessor implements Processor { + private static class KeyValuePlainStoreWriterProcessor implements Processor { private ProcessorContext context; private TimestampedKeyValueStore store; @@ -344,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/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..65762c72f57ac --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedRangeWithHeadersQuery.java @@ -0,0 +1,156 @@ +/* + * 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.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 +@InterfaceAudience.Public +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/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/ReadOnlyRecordIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java new file mode 100644 index 0000000000000..e186095b1db31 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyRecordIterator.java @@ -0,0 +1,46 @@ +/* + * 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.InterfaceAudience; +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}. + *

+ * 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 + */ +@InterfaceAudience.Public +@Evolving +public interface ReadOnlyRecordIterator extends Iterator>, Closeable { + + @Override + void close(); +} 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..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 @@ -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; @@ -54,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; @@ -106,6 +110,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) ) ); @@ -451,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, @@ -460,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()) { @@ -505,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()) { @@ -542,6 +546,40 @@ 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; + + final RangeQuery rawRangeQuery = + rawRangeQuery(typedQuery.lowerBound(), typedQuery.upperBound(), typedQuery.resultOrder()); + + 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 @@ -639,6 +677,7 @@ private MeteredTimestampedKeyValueStoreWithHeadersQueryIterator( this.startNs = time.nanoseconds(); this.startTimestampMs = time.milliseconds(); this.returnPlainValue = returnPlainValue; + numOpenIterators.increment(); openIterators.add(this); } @@ -683,6 +722,7 @@ public void close() { final long duration = time.nanoseconds() - startNs; sensor.record(duration); iteratorDurationSensor.record(duration); + numOpenIterators.decrement(); openIterators.remove(this); } } @@ -696,6 +736,91 @@ 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(); + numOpenIterators.increment(); + 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); + numOpenIterators.decrement(); + 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/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/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; - } - } 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(); 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..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 @@ -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,39 @@ public void shouldGetRangeFromInnerStoreAndRecordRangeMetric() { assertTrue((Double) metric.metricValue() > 0); } + @Test + public void shouldPropagateBoundsAndDescendingOrderForTimestampedRangeWithHeadersQuery() { + setUp(); + init(); + final RangeQuery rawQuery = forwardedRawRangeQuery( + TimestampedRangeWithHeadersQuery.withLowerBound("a").withDescendingKeys()); + assertEquals(ResultOrder.DESCENDING, rawQuery.resultOrder()); + assertEquals(Bytes.wrap("a".getBytes()), rawQuery.getLowerBound().get()); + assertFalse(rawQuery.getUpperBound().isPresent()); + } + + @Test + public void shouldPropagateBoundsAndAscendingOrderForTimestampedRangeWithHeadersQuery() { + setUp(); + init(); + final RangeQuery rawQuery = forwardedRawRangeQuery( + TimestampedRangeWithHeadersQuery.withUpperBound("z").withAscendingKeys()); + assertEquals(ResultOrder.ASCENDING, rawQuery.resultOrder()); + assertFalse(rawQuery.getLowerBound().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 public void shouldGetAllFromInnerStoreAndRecordAllMetric() { setUp(); 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/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/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/TimestampedKeyValueStoreBuilderWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java index cccc1174e53b9..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 @@ -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,24 +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) { @@ -615,4 +605,205 @@ public void shouldReturnIdenticalResultsForNativeAndAdapterBuiltStores(final boo 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; + } } 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; + } + } }