aliehsaeedii commented on code in PR #22666: URL: https://github.com/apache/kafka/pull/22666#discussion_r3490581018
########## streams/src/main/java/org/apache/kafka/streams/query/TimestampedKeyWithHeadersQuery.java: ########## @@ -0,0 +1,87 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.streams.query; + +import org.apache.kafka.common.annotation.InterfaceStability.Evolving; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; + +import java.util.Objects; + +/** + * Interactive query for retrieving a single record, including its record headers, based on its key + * from a {@link TimestampedKeyValueStoreWithHeaders}. + * + * <p>This is the headers-aware parallel of {@link TimestampedKeyQuery}: it returns a + * {@link ReadOnlyRecord} carrying the key, value, timestamp, and headers, whereas + * {@link TimestampedKeyQuery} returns a {@link org.apache.kafka.streams.state.ValueAndTimestamp} + * (value and timestamp only, no key or headers). + * + * <p>Headers are only returned when the queried store was built with a KIP-1271 + * {@code WithHeaders} supplier. Against a plain (non-headers) store, this query type is + * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}. + * + * @param <K> Type of keys + * @param <V> Type of values + */ +@Evolving +public final class TimestampedKeyWithHeadersQuery<K, V> implements Query<ReadOnlyRecord<K, V>> { Review Comment: Have you thought of a user who built a `WithHeaders` store over a legacy supplier and queries for headers? Do you please test it and either update the java doc or inform us in this thread about it? (consider with-caching/without caching) ########## streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java: ########## @@ -228,356 +230,142 @@ public void shouldWrapPlainKeyValueStoreAsHeadersStore() { assertInstanceOf(PlainToHeadersStoreAdapter.class, ((WrappedStateStore) store).wrapped()); } - @Test - public void shouldHandleKeyQueryOnInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - - try { - // Put data into the store - final Headers headers = new RecordHeaders(); - headers.add("key1", "value1".getBytes()); - store.put("test-key", ValueTimestampHeaders.make("test-value", 12345L, headers)); - - // Verify wrapper type for InMemoryKeyValueStore - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped, - "Expected wrapper to implement HeadersBytesStore for InMemoryKeyValueStore"); - - // Query at typed level - KeyQuery should return just the value - final KeyQuery<String, String> query = KeyQuery.withKey("test-key"); - final QueryResult<String> result = store.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - // Verify IQv2 query result - assertTrue(result.isSuccess(), "Expected query to succeed on InMemoryKeyValueStore"); - assertNotNull(result.getPosition(), "Expected position to be set"); - assertInstanceOf(String.class, result.getResult()); - assertEquals("test-value", result.getResult(), "KeyQuery should return just the value"); - } finally { - store.close(); + // ---------------------------------------------------------------------------------------------- + // IQv2 query handling for the built header store. + // + // The header store can be built three ways; each is exercised with caching on and off through a + // single helper that builds the real store chain (Metered -> [Caching] -> inner) with a real + // record cache, then writes and reads real data at the typed (metered) level: + // NATIVE -> RocksDBTimestampedStoreWithHeaders (persists headers) + // ADAPTER -> RocksDBTimestampedStore via TimestampedToHeadersStoreAdapter (drops headers) + // IN_MEMORY -> InMemoryKeyValueStore via the in-memory headers marker + // ---------------------------------------------------------------------------------------------- + + private enum StoreType { NATIVE, ADAPTER, IN_MEMORY } + + private KeyValueStore<Bytes, byte[]> innerStore(final StoreType storeType) { + switch (storeType) { + case NATIVE: return new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope"); + case ADAPTER: return new RocksDBTimestampedStore("test-store", "metrics-scope"); + case IN_MEMORY: return new InMemoryKeyValueStore("test-store"); + default: throw new IllegalArgumentException("unknown store type: " + storeType); } } - @Test - public void shouldHandleTimestampedKeyQueryOnInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); + private TimestampedKeyValueStoreWithHeaders<String, String> buildAndInitStore( + final StoreType storeType, + final boolean cachingEnabled) { + lenient().when(supplier.name()).thenReturn("test-store"); + lenient().when(supplier.metricsScope()).thenReturn("metricScope"); + lenient().when(supplier.get()).thenReturn(innerStore(storeType)); builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); + supplier, Serdes.String(), Serdes.String(), new MockTime()); + final TimestampedKeyValueStoreWithHeaders<String, String> store = + (cachingEnabled + ? builder.withLoggingDisabled().withCachingEnabled() + : builder.withLoggingDisabled().withCachingDisabled()) + .build(); - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); + final ThreadCache cache = new ThreadCache( + new LogContext("test "), + cachingEnabled ? 10 * 1024 * 1024 : 0, + new MockStreamsMetrics(new Metrics())); final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); + TestUtils.tempDirectory(), Serdes.String(), Serdes.String(), null, cache); + context.setRecordContext(new ProcessorRecordContext(0L, 0L, 0, "topic", new RecordHeaders())); store.init(context, store); - - try { - // Put data into the store - final Headers headers = new RecordHeaders(); - headers.add("key1", "value1".getBytes()); - store.put("test-key", ValueTimestampHeaders.make("test-value", 12345L, headers)); - - // Verify wrapper type for InMemoryKeyValueStore - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped, - "Expected wrapper to implement HeadersBytesStore for InMemoryKeyValueStore"); - - // Query at typed level - TimestampedKeyQuery should return value + timestamp - final TimestampedKeyQuery<String, String> query = TimestampedKeyQuery.withKey("test-key"); - final QueryResult<ValueAndTimestamp<String>> result = store.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - // Verify IQv2 query result - assertTrue(result.isSuccess(), "Expected query to succeed on InMemoryKeyValueStore"); - assertNotNull(result.getPosition(), "Expected position to be set"); - assertNotNull(result.getResult(), "Expected non-null result"); - assertInstanceOf(ValueAndTimestamp.class, result.getResult()); - assertEquals("test-value", result.getResult().value(), "TimestampedKeyQuery should return the value"); - assertEquals(12345L, result.getResult().timestamp(), "TimestampedKeyQuery should return the timestamp"); - } finally { - store.close(); - } + return store; } - @Test - public void shouldReturnPositionFromHeadersStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - - try { - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - final Position position = wrapped.getPosition(); - - // Verify: Position is returned (should be non-null) - assertNotNull(position, "Expected non-null position"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); - } finally { - store.close(); - } + private static Headers headersWith(final String key, final String value) { + return new RecordHeaders().add(key, value.getBytes(StandardCharsets.UTF_8)); } - @Test - public void shouldReturnPositionFromAdaptedTimestampedStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStore("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @CsvSource({"NATIVE, true", "NATIVE, false", "ADAPTER, true", "ADAPTER, false", "IN_MEMORY, true", "IN_MEMORY, false"}) + public void shouldHandleKeyQuery(final StoreType storeType, final boolean cachingEnabled) { + final TimestampedKeyValueStoreWithHeaders<String, String> store = buildAndInitStore(storeType, cachingEnabled); try { - // Verify adapter is used - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(TimestampedToHeadersStoreAdapter.class, wrapped); + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); - // Get position from adapter (should delegate to underlying store) - final Position position = wrapped.getPosition(); + final QueryResult<String> result = + store.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)); - assertNotNull(position, "Expected non-null position from adapter"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); + assertTrue(result.isSuccess(), "Expected KeyQuery to succeed"); + assertEquals("v", result.getResult(), "KeyQuery returns the value only"); + assertNotNull(result.getPosition(), "Expected position to be set"); } finally { store.close(); } } - @Test - public void shouldReturnPositionFromInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @CsvSource({"NATIVE, true", "NATIVE, false", "ADAPTER, true", "ADAPTER, false", "IN_MEMORY, true", "IN_MEMORY, false"}) + public void shouldHandleTimestampedKeyQuery(final StoreType storeType, final boolean cachingEnabled) { + final TimestampedKeyValueStoreWithHeaders<String, String> store = buildAndInitStore(storeType, cachingEnabled); try { - // Verify marker wrapper is used - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped); + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); - // Get position from marker (should delegate to InMemoryKeyValueStore) - final Position position = wrapped.getPosition(); + final QueryResult<ValueAndTimestamp<String>> result = + store.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)); - assertNotNull(position, "Expected non-null position from in-memory store"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); + assertTrue(result.isSuccess(), "Expected TimestampedKeyQuery to succeed"); + assertEquals("v", result.getResult().value()); + assertEquals(123L, result.getResult().timestamp()); + assertNotNull(result.getPosition(), "Expected position to be set"); } finally { store.close(); } } - @Test - public void shouldMaintainPositionAcrossOperationsOnHeadersStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void shouldHandleTimestampedKeyWithHeadersQueryOnNativeStore(final boolean cachingEnabled) { Review Comment: (Re-anchored from a prior pending comment.) The non-native behavior of `TimestampedKeyWithHeadersQuery` is now pinned down for the ADAPTER build path (`shouldReturnEmptyHeadersForTimestampedKeyWithHeadersQueryOnAdapterStore` asserts empty headers), which addresses most of my earlier concern. The IN_MEMORY `WithHeaders` build still has no `TimestampedKeyWithHeadersQuery` test; it also drops headers on write, so it should behave like the adapter (empty headers). Either add an IN_MEMORY case (it's nearly free given `buildAndInitStore` is already parameterized over `StoreType`) or, if you consider the adapter test sufficient coverage for the header-stripping behavior, it's reasonable to leave this as-is. ########## streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedKeyValueStoreBuilderWithHeadersTest.java: ########## @@ -228,356 +230,142 @@ public void shouldWrapPlainKeyValueStoreAsHeadersStore() { assertInstanceOf(PlainToHeadersStoreAdapter.class, ((WrappedStateStore) store).wrapped()); } - @Test - public void shouldHandleKeyQueryOnInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - - try { - // Put data into the store - final Headers headers = new RecordHeaders(); - headers.add("key1", "value1".getBytes()); - store.put("test-key", ValueTimestampHeaders.make("test-value", 12345L, headers)); - - // Verify wrapper type for InMemoryKeyValueStore - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped, - "Expected wrapper to implement HeadersBytesStore for InMemoryKeyValueStore"); - - // Query at typed level - KeyQuery should return just the value - final KeyQuery<String, String> query = KeyQuery.withKey("test-key"); - final QueryResult<String> result = store.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - // Verify IQv2 query result - assertTrue(result.isSuccess(), "Expected query to succeed on InMemoryKeyValueStore"); - assertNotNull(result.getPosition(), "Expected position to be set"); - assertInstanceOf(String.class, result.getResult()); - assertEquals("test-value", result.getResult(), "KeyQuery should return just the value"); - } finally { - store.close(); + // ---------------------------------------------------------------------------------------------- + // IQv2 query handling for the built header store. + // + // The header store can be built three ways; each is exercised with caching on and off through a + // single helper that builds the real store chain (Metered -> [Caching] -> inner) with a real + // record cache, then writes and reads real data at the typed (metered) level: + // NATIVE -> RocksDBTimestampedStoreWithHeaders (persists headers) + // ADAPTER -> RocksDBTimestampedStore via TimestampedToHeadersStoreAdapter (drops headers) + // IN_MEMORY -> InMemoryKeyValueStore via the in-memory headers marker + // ---------------------------------------------------------------------------------------------- + + private enum StoreType { NATIVE, ADAPTER, IN_MEMORY } + + private KeyValueStore<Bytes, byte[]> innerStore(final StoreType storeType) { + switch (storeType) { + case NATIVE: return new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope"); + case ADAPTER: return new RocksDBTimestampedStore("test-store", "metrics-scope"); + case IN_MEMORY: return new InMemoryKeyValueStore("test-store"); + default: throw new IllegalArgumentException("unknown store type: " + storeType); } } - @Test - public void shouldHandleTimestampedKeyQueryOnInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); + private TimestampedKeyValueStoreWithHeaders<String, String> buildAndInitStore( + final StoreType storeType, + final boolean cachingEnabled) { + lenient().when(supplier.name()).thenReturn("test-store"); + lenient().when(supplier.metricsScope()).thenReturn("metricScope"); + lenient().when(supplier.get()).thenReturn(innerStore(storeType)); builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); + supplier, Serdes.String(), Serdes.String(), new MockTime()); + final TimestampedKeyValueStoreWithHeaders<String, String> store = + (cachingEnabled + ? builder.withLoggingDisabled().withCachingEnabled() + : builder.withLoggingDisabled().withCachingDisabled()) + .build(); - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); + final ThreadCache cache = new ThreadCache( + new LogContext("test "), + cachingEnabled ? 10 * 1024 * 1024 : 0, + new MockStreamsMetrics(new Metrics())); final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); + TestUtils.tempDirectory(), Serdes.String(), Serdes.String(), null, cache); + context.setRecordContext(new ProcessorRecordContext(0L, 0L, 0, "topic", new RecordHeaders())); store.init(context, store); - - try { - // Put data into the store - final Headers headers = new RecordHeaders(); - headers.add("key1", "value1".getBytes()); - store.put("test-key", ValueTimestampHeaders.make("test-value", 12345L, headers)); - - // Verify wrapper type for InMemoryKeyValueStore - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped, - "Expected wrapper to implement HeadersBytesStore for InMemoryKeyValueStore"); - - // Query at typed level - TimestampedKeyQuery should return value + timestamp - final TimestampedKeyQuery<String, String> query = TimestampedKeyQuery.withKey("test-key"); - final QueryResult<ValueAndTimestamp<String>> result = store.query(query, PositionBound.unbounded(), new QueryConfig(false)); - - // Verify IQv2 query result - assertTrue(result.isSuccess(), "Expected query to succeed on InMemoryKeyValueStore"); - assertNotNull(result.getPosition(), "Expected position to be set"); - assertNotNull(result.getResult(), "Expected non-null result"); - assertInstanceOf(ValueAndTimestamp.class, result.getResult()); - assertEquals("test-value", result.getResult().value(), "TimestampedKeyQuery should return the value"); - assertEquals(12345L, result.getResult().timestamp(), "TimestampedKeyQuery should return the timestamp"); - } finally { - store.close(); - } + return store; } - @Test - public void shouldReturnPositionFromHeadersStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - - try { - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - final Position position = wrapped.getPosition(); - - // Verify: Position is returned (should be non-null) - assertNotNull(position, "Expected non-null position"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); - } finally { - store.close(); - } + private static Headers headersWith(final String key, final String value) { + return new RecordHeaders().add(key, value.getBytes(StandardCharsets.UTF_8)); } - @Test - public void shouldReturnPositionFromAdaptedTimestampedStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStore("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @CsvSource({"NATIVE, true", "NATIVE, false", "ADAPTER, true", "ADAPTER, false", "IN_MEMORY, true", "IN_MEMORY, false"}) + public void shouldHandleKeyQuery(final StoreType storeType, final boolean cachingEnabled) { + final TimestampedKeyValueStoreWithHeaders<String, String> store = buildAndInitStore(storeType, cachingEnabled); try { - // Verify adapter is used - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(TimestampedToHeadersStoreAdapter.class, wrapped); + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); - // Get position from adapter (should delegate to underlying store) - final Position position = wrapped.getPosition(); + final QueryResult<String> result = + store.query(KeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)); - assertNotNull(position, "Expected non-null position from adapter"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); + assertTrue(result.isSuccess(), "Expected KeyQuery to succeed"); + assertEquals("v", result.getResult(), "KeyQuery returns the value only"); + assertNotNull(result.getPosition(), "Expected position to be set"); } finally { store.close(); } } - @Test - public void shouldReturnPositionFromInMemoryStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("in-memory"); - when(supplier.get()).thenReturn(new InMemoryKeyValueStore("test-store")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @CsvSource({"NATIVE, true", "NATIVE, false", "ADAPTER, true", "ADAPTER, false", "IN_MEMORY, true", "IN_MEMORY, false"}) + public void shouldHandleTimestampedKeyQuery(final StoreType storeType, final boolean cachingEnabled) { + final TimestampedKeyValueStoreWithHeaders<String, String> store = buildAndInitStore(storeType, cachingEnabled); try { - // Verify marker wrapper is used - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - assertInstanceOf(HeadersBytesStore.class, wrapped); + store.put("k", ValueTimestampHeaders.make("v", 123L, headersWith("h", "x"))); - // Get position from marker (should delegate to InMemoryKeyValueStore) - final Position position = wrapped.getPosition(); + final QueryResult<ValueAndTimestamp<String>> result = + store.query(TimestampedKeyQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)); - assertNotNull(position, "Expected non-null position from in-memory store"); - assertTrue(position.isEmpty(), "Expected position to be empty initially"); + assertTrue(result.isSuccess(), "Expected TimestampedKeyQuery to succeed"); + assertEquals("v", result.getResult().value()); + assertEquals(123L, result.getResult().timestamp()); + assertNotNull(result.getPosition(), "Expected position to be set"); } finally { store.close(); } } - @Test - public void shouldMaintainPositionAcrossOperationsOnHeadersStore() { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope")); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store = builder - .withLoggingDisabled() - .withCachingDisabled() - .build(); - - final File dir = TestUtils.tempDirectory(); - final Properties props = StreamsTestUtils.getStreamsConfig(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - new StreamsConfig(props) - ); - store.init(context, store); - + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void shouldHandleTimestampedKeyWithHeadersQueryOnNativeStore(final boolean cachingEnabled) { + // Only the native store persists headers; the adapter and in-memory builds drop them on write. + final TimestampedKeyValueStoreWithHeaders<String, String> store = buildAndInitStore(StoreType.NATIVE, cachingEnabled); try { - final StateStore wrapped = ((WrappedStateStore) store).wrapped(); - - // Get initial position - final Position initialPosition = wrapped.getPosition(); - assertNotNull(initialPosition, "Expected non-null initial position"); - - // Put some data - store.put("key1", ValueTimestampHeaders.make("value1", 100L, new RecordHeaders())); - store.put("key2", ValueTimestampHeaders.make("value2", 200L, new RecordHeaders())); - - // Get position after puts - final Position afterPutPosition = wrapped.getPosition(); - assertNotNull(afterPutPosition, "Expected non-null position after puts"); - - // Position object should be the same instance (stores maintain a single position) - // The position content might be updated internally by the context + final Headers headers = headersWith("h", "x"); + store.put("k", ValueTimestampHeaders.make("v", 123L, headers)); + + final QueryResult<ReadOnlyRecord<String, String>> result = + store.query(TimestampedKeyWithHeadersQuery.withKey("k"), PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess(), "Expected TimestampedKeyWithHeadersQuery to succeed"); + final ReadOnlyRecord<String, String> record = result.getResult(); + assertEquals("k", record.key()); + assertEquals("v", record.value()); + assertEquals(123L, record.timestamp()); + // Headers must round-trip on both the persistent path and the cache path. + assertEquals(headers, record.headers()); + assertNotNull(result.getPosition(), "Expected position to be set"); } finally { store.close(); } } - private static ThreadCache mockCacheHit() { - final ThreadCache cache = mock(ThreadCache.class); - final LRUCacheEntry entry = mock(LRUCacheEntry.class); - final byte[] entryValue = "mockEntryValue".getBytes(StandardCharsets.UTF_8); - lenient().when(entry.value()).thenReturn(entryValue); - lenient().when(cache.get(any(String.class), any(Bytes.class))).thenReturn(entry); - return cache; - } - - private TimestampedKeyValueStoreWithHeaders<String, String> headersStoreMaybeWithCache(final boolean cachingEnabled) { - when(supplier.name()).thenReturn("test-store"); - when(supplier.metricsScope()).thenReturn("metricScope"); - when(supplier.get()).thenReturn(new RocksDBTimestampedStoreWithHeaders("test-store", "metrics-scope")); - - final File dir = TestUtils.tempDirectory(); - final ThreadCache cache = mockCacheHit(); - final InternalMockProcessorContext<String, String> context = new InternalMockProcessorContext<>( - dir, - Serdes.String(), - Serdes.String(), - null, - cache - ); - - builder = new TimestampedKeyValueStoreBuilderWithHeaders<>( - supplier, - Serdes.String(), - Serdes.String(), - new MockTime() - ); - - final TimestampedKeyValueStoreWithHeaders<String, String> store; - if (cachingEnabled) { - store = builder.withLoggingDisabled() - .withCachingEnabled() - .build(); - } else { - store = builder.withLoggingDisabled() - .withCachingDisabled() - .build(); - } - - store.init(context, store); - return store; - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void shouldReturnUnknownQueryTypeForKeyQueryOnHeadersStore(final boolean cachingEnabled) { - final TimestampedKeyValueStoreWithHeaders<String, String> store = headersStoreMaybeWithCache(cachingEnabled); - + @Test + public void shouldReturnEmptyHeadersForTimestampedKeyWithHeadersQueryOnAdapterStore() { Review Comment: This test deliberately pins the adapter case to `caching=false`, and the comment sidesteps the caching-enabled case. But that case is the interesting one and it isn't covered anywhere. Trace the layering for an adapter-built store: `Metered -> Caching -> adapter -> plain store`. On write, the full serialized `ValueTimestampHeaders` (with headers) lands in the record cache, which sits *above* the adapter; the adapter only strips headers when it flushes down to the plain store. So with an adapter store + caching enabled, the *same* `TimestampedKeyWithHeadersQuery` returns the real headers before flush (cache hit) but empty headers after flush (read through the adapter) -- a flush-timing-dependent answer for one query. Worth either adding an `ADAPTER, true` case that documents/asserts this, or reconsidering whether returning headers from cache on a store that can't persist them is the behavior we want. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
