KAFKA-9487: Follow-up PR of Kafka-9445

streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java

@@ -1160,27 +1160,25 @@ public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
 
 
     /**
-     * @deprecated since 2.5 release; use {@link #store(StoreQueryParams)}  instead
+     * @deprecated since 2.5 release; use {@link #store(StoreQueryParameters)}  instead
      */
     @Deprecated
     public <T> T store(final String storeName, final QueryableStoreType<T> queryableStoreType) {
-        return store(StoreQueryParams.fromNameAndType(storeName, queryableStoreType));
+        return store(StoreQueryParameters.fromNameAndType(storeName, queryableStoreType));
     }
 
     /**
-     * Get a facade wrapping the local {@link StateStore} instances with the provided {@link StoreQueryParams}.
-     * StoreQueryParams need required parameters to be set, which are {@code storeName} and if
-     * type is accepted by the provided {@link QueryableStoreType#accepts(StateStore) queryableStoreType}.
+     * Get a facade wrapping the local {@link StateStore} instances with the provided {@link StoreQueryParameters}.
      * The returned object can be used to query the {@link StateStore} instances.
      *
-     * @param storeQueryParams   to set the optional parameters to fetch type of stores user wants to fetch when a key is queried
+     * @param storeQueryParameters   the parameters used to fetch a queryable store
      * @return A facade wrapping the local {@link StateStore} instances
      * @throws InvalidStateStoreException if Kafka Streams is (re-)initializing or a store with {@code storeName} and
      * {@code queryableStoreType} doesn't exist
      */
-    public <T> T store(final StoreQueryParams<T> storeQueryParams) {
+    public <T> T store(final StoreQueryParameters<T> storeQueryParameters) {
         validateIsRunningOrRebalancing();
-        return queryableStoreProvider.getStore(storeQueryParams);
+        return queryableStoreProvider.getStore(storeQueryParameters);
     }
 
     /**

streams/src/main/java/org/apache/kafka/streams/StoreQueryParams.java → streams/src/main/java/org/apache/kafka/streams/StoreQueryParameters.java

@@ -21,60 +21,51 @@
 import java.util.Objects;
 
 /**
- * Represents all the query options that a user can provide to state what kind of stores it is expecting.
- * The options would be whether a user would want to enable/disable stale stores
- * or whether it knows the list of partitions that it specifically wants to fetch.
- * If this information is not provided the default behavior is to fetch the stores for all the partitions
- * available on that instance for that particular store name.
- * It contains a partition, which for a point queries can be populated from the {@link KeyQueryMetadata}.
+ * {@code StoreQueryParameters} allows you to pass a variety of parameters when fetching a store for interactive query.
  */
-public class StoreQueryParams<T> {
+public class StoreQueryParameters<T> {
 
     private Integer partition;
     private boolean staleStores;
     private final String storeName;
     private final QueryableStoreType<T> queryableStoreType;
 
-    private StoreQueryParams(final String storeName, final QueryableStoreType<T>  queryableStoreType) {
+    private StoreQueryParameters(final String storeName, final QueryableStoreType<T>  queryableStoreType, final Integer partition, final boolean staleStores) {
         this.storeName = storeName;
         this.queryableStoreType = queryableStoreType;
+        this.partition = partition;
+        this.staleStores = staleStores;
     }
 
-    public static <T> StoreQueryParams<T> fromNameAndType(final String storeName,
+    public static <T> StoreQueryParameters<T> fromNameAndType(final String storeName,
                                                           final QueryableStoreType<T>  queryableStoreType) {
-        return new StoreQueryParams<T>(storeName, queryableStoreType);
+        return new StoreQueryParameters<T>(storeName, queryableStoreType, null, false);
     }
 
     /**
      * Set a specific partition that should be queried exclusively.
      *
-     * @param partition   The specific integer partition to be fetched from the stores list by using {@link StoreQueryParams}.
+     * @param partition   The specific integer partition to be fetched from the stores list by using {@link StoreQueryParameters}.
      *
-     * @return String storeName
+     * @return StoreQueryParameters a new {@code StoreQueryParameters} instance configured with the specified partition
      */
-    public StoreQueryParams<T> withPartition(final Integer partition) {
-        final StoreQueryParams<T> storeQueryParams = StoreQueryParams.fromNameAndType(this.storeName(), this.queryableStoreType());
-        storeQueryParams.partition = partition;
-        storeQueryParams.staleStores = this.staleStores;
-        return storeQueryParams;
+    public StoreQueryParameters<T> withPartition(final Integer partition) {
+        return new StoreQueryParameters<T>(storeName, queryableStoreType, partition, staleStores);
     }
 
     /**
      * Enable querying of stale state stores, i.e., allow to query active tasks during restore as well as standby tasks.
      *
-     * @return String storeName
+     * @return StoreQueryParameters a new {@code StoreQueryParameters} instance configured with serving from stale stores enabled
      */
-    public StoreQueryParams<T> enableStaleStores() {
-        final StoreQueryParams<T> storeQueryParams = StoreQueryParams.fromNameAndType(this.storeName(), this.queryableStoreType());
-        storeQueryParams.partition = this.partition;
-        storeQueryParams.staleStores = true;
-        return storeQueryParams;
+    public StoreQueryParameters<T> enableStaleStores() {
+        return new StoreQueryParameters<T>(storeName, queryableStoreType, partition, true);
     }
 
     /**
-     * Get the store name for which key is queried by the user.
+     * Get the name of the state store that should be queried.
      *
-     * @return String storeName
+     * @return String state store name
      */
     public String storeName() {
         return storeName;
@@ -83,16 +74,16 @@ public String storeName() {
     /**
      * Get the queryable store type for which key is queried by the user.
      *
-     * @return QueryableStoreType queryableStoreType
+     * @return QueryableStoreType type of queryable store
      */
     public QueryableStoreType<T> queryableStoreType() {
         return queryableStoreType;
     }
 
     /**
-     * Get the partition to be used to fetch list of stores.
+     * Get the store partition that will be queried.
      * If the method returns {@code null}, it would mean that no specific partition has been requested,
-     * so all the local partitions for the store will be returned.
+     * so all the local partitions for the store will be queried.
      *
      * @return Integer partition
      */
@@ -111,19 +102,19 @@ public boolean staleStoresEnabled() {
 
     @Override
     public boolean equals(final Object obj) {
-        if (!(obj instanceof StoreQueryParams)) {
+        if (!(obj instanceof StoreQueryParameters)) {
             return false;
         }
-        final StoreQueryParams storeQueryParams = (StoreQueryParams) obj;
-        return Objects.equals(storeQueryParams.partition, partition)
-                && Objects.equals(storeQueryParams.staleStores, staleStores)
-                && Objects.equals(storeQueryParams.storeName, storeName)
-                && Objects.equals(storeQueryParams.queryableStoreType, queryableStoreType);
+        final StoreQueryParameters storeQueryParameters = (StoreQueryParameters) obj;
+        return Objects.equals(storeQueryParameters.partition, partition)
+                && Objects.equals(storeQueryParameters.staleStores, staleStores)
+                && Objects.equals(storeQueryParameters.storeName, storeName)
+                && Objects.equals(storeQueryParameters.queryableStoreType, queryableStoreType);
     }
 
     @Override
     public String toString() {
-        return "StoreQueryParams {" +
+        return "StoreQueryParameters {" +
                 "partition=" + partition +
                 ", staleStores=" + staleStores +
                 ", storeName=" + storeName +

streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java

@@ -1149,6 +1149,10 @@ public Map<String, List<String>> stateStoreNameToSourceTopics() {
         return results;
     }
 
+    public Collection<String> sourceTopicsForStore(final String storeName) {
+        return maybeDecorateInternalSourceTopics(stateStoreNameToSourceTopics.get(storeName));
+    }
+
     public synchronized Collection<Set<String>> copartitionGroups() {
         // compute transitive closures of copartitionGroups to relieve registering code to know all members
         // of a copartitionGroup at the same time

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -1153,6 +1153,10 @@ private void updateThreadMetadata(final Map<TaskId, Task> activeTasks,
             standbyTasksMetadata);
     }
 
+    public Map<TaskId, Task> activeTaskMap() {
+        return taskManager.activeTaskMap();
+    }
+
     public List<Task> activeTasks() {
         return taskManager.activeTaskIterable();
     }

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.PartitionInfo;
@@ -109,7 +110,7 @@ public synchronized Collection<StreamsMetadata> getAllMetadataForStore(final Str
             return allMetadata;
         }
 
-        final List<String> sourceTopics = builder.stateStoreNameToSourceTopics().get(storeName);
+        final Collection<String> sourceTopics = builder.sourceTopicsForStore(storeName);
         if (sourceTopics == null) {
             return Collections.emptyList();
         }
@@ -408,7 +409,7 @@ private <K> StreamsMetadata getStreamsMetadataForKey(final String storeName,
     }
 
     private SourceTopicsInfo getSourceTopicsInfo(final String storeName) {
-        final List<String> sourceTopics = builder.stateStoreNameToSourceTopics().get(storeName);
+        final List<String> sourceTopics = builder.sourceTopicsForStore(storeName).stream().collect(Collectors.toList());
         if (sourceTopics == null || sourceTopics.isEmpty()) {
             return null;
         }

streams/src/main/java/org/apache/kafka/streams/state/internals/QueryableStoreProvider.java

@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals;
 
-import org.apache.kafka.streams.StoreQueryParams;
+import org.apache.kafka.streams.StoreQueryParameters;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.state.QueryableStoreType;
@@ -42,29 +42,29 @@ public QueryableStoreProvider(final List<StreamThreadStateStoreProvider> storePr
      * Get a composite object wrapping the instances of the {@link StateStore} with the provided
      * storeName and {@link QueryableStoreType}
      *
-     * @param storeQueryParams       if stateStoresEnabled is used i.e. staleStoresEnabled is true, include standbys and recovering stores;
+     * @param storeQueryParameters       if stateStoresEnabled is used i.e. staleStoresEnabled is true, include standbys and recovering stores;
      *                                        if stateStoresDisabled i.e. staleStoresEnabled is false, only include running actives;
      *                                        if partition is null then it fetches all local partitions on the instance;
      *                                        if partition is set then it fetches a specific partition.
      * @param <T>                The expected type of the returned store
      * @return A composite object that wraps the store instances.
      */
-    public <T> T getStore(final StoreQueryParams<T> storeQueryParams) {
-        final String storeName = storeQueryParams.storeName();
-        final QueryableStoreType<T> queryableStoreType = storeQueryParams.queryableStoreType();
+    public <T> T getStore(final StoreQueryParameters<T> storeQueryParameters) {
+        final String storeName = storeQueryParameters.storeName();
+        final QueryableStoreType<T> queryableStoreType = storeQueryParameters.queryableStoreType();
         final List<T> globalStore = globalStoreProvider.stores(storeName, queryableStoreType);
         if (!globalStore.isEmpty()) {
             return queryableStoreType.create(globalStoreProvider, storeName);
         }
         final List<T> allStores = new ArrayList<>();
         for (final StreamThreadStateStoreProvider storeProvider : storeProviders) {
-            allStores.addAll(storeProvider.stores(storeQueryParams));
+            allStores.addAll(storeProvider.stores(storeQueryParameters));
         }
         if (allStores.isEmpty()) {
             throw new InvalidStateStoreException("The state store, " + storeName + ", may have migrated to another instance.");
         }
         return queryableStoreType.create(
-            new WrappingStoreProvider(storeProviders, storeQueryParams),
+            new WrappingStoreProvider(storeProviders, storeQueryParameters),
             storeName
         );
     }

streams/src/main/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProvider.java

@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals;
 
-import org.apache.kafka.streams.StoreQueryParams;
+import org.apache.kafka.streams.StoreQueryParameters;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.TaskId;
@@ -47,7 +47,7 @@ public StreamThreadStateStoreProvider(final StreamThread streamThread,
     }
 
     @SuppressWarnings("unchecked")
-    public <T> List<T> stores(final StoreQueryParams storeQueryParams) {
+    public <T> List<T> stores(final StoreQueryParameters storeQueryParams) {
         final String storeName = storeQueryParams.storeName();
         final QueryableStoreType<T> queryableStoreType = storeQueryParams.queryableStoreType();
         final TaskId keyTaskId = createKeyTaskId(storeName, storeQueryParams.partition());
@@ -56,34 +56,47 @@ public <T> List<T> stores(final StoreQueryParams storeQueryParams) {
         }
         final StreamThread.State state = streamThread.state();
         if (storeQueryParams.staleStoresEnabled() ? state.isAlive() : state == StreamThread.State.RUNNING) {
-            final Iterable<? extends Task> tasks = storeQueryParams.staleStoresEnabled() ? streamThread.allTasks().values() : streamThread.activeTasks();
+            final Map<TaskId, ? extends Task> tasks = storeQueryParams.staleStoresEnabled() ? streamThread.allTasks() : streamThread.activeTaskMap();
             final List<T> stores = new ArrayList<>();
-            for (final Task task : tasks) {
-                if (keyTaskId != null && !keyTaskId.equals(task.id())) {
-                    continue;
+            if (keyTaskId != null) {
+                final T store = validateAndListStores(tasks.get(keyTaskId).getStore(storeName), queryableStoreType, storeName, keyTaskId);
+                if (store != null) {
+                    return Collections.singletonList(store);
                 }
-                final StateStore store = task.getStore(storeName);
-                if (store != null && queryableStoreType.accepts(store)) {
-                    if (!store.isOpen()) {
-                        throw new InvalidStateStoreException(
-                            "Cannot get state store " + storeName + " for task " + task +
-                                " because the store is not open. " +
-                                "The state store may have migrated to another instances.");
-                    }
-                    if (store instanceof TimestampedKeyValueStore && queryableStoreType instanceof QueryableStoreTypes.KeyValueStoreType) {
-                        stores.add((T) new ReadOnlyKeyValueStoreFacade<>((TimestampedKeyValueStore<Object, Object>) store));
-                    } else if (store instanceof TimestampedWindowStore && queryableStoreType instanceof QueryableStoreTypes.WindowStoreType) {
-                        stores.add((T) new ReadOnlyWindowStoreFacade<>((TimestampedWindowStore<Object, Object>) store));
-                    } else {
-                        stores.add((T) store);
+            } else {
+                for (final Task streamTask : tasks.values()) {
+                    final T store = validateAndListStores(streamTask.getStore(storeName), queryableStoreType, storeName, streamTask.id());
+                    if (store != null) {
+                        stores.add(store);
                     }
                 }
             }
             return stores;
         } else {
             throw new InvalidStateStoreException("Cannot get state store " + storeName + " because the stream thread is " +
-                                                     state + ", not RUNNING" +
-                                                     (storeQueryParams.staleStoresEnabled() ? " or REBALANCING" : ""));
+                                                    state + ", not RUNNING" +
+                                                    (storeQueryParams.staleStoresEnabled() ? " or REBALANCING" : ""));
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private <T> T validateAndListStores(final StateStore store, final QueryableStoreType<T> queryableStoreType, final String storeName, final TaskId taskId) {
+        if (store != null && queryableStoreType.accepts(store)) {
+            if (!store.isOpen()) {
+                throw new InvalidStateStoreException(
+                        "Cannot get state store " + storeName + " for task " + taskId +
+                            " because the store is not open. " +
+                            "The state store may have migrated to another instances.");
+            }
+            if (store instanceof TimestampedKeyValueStore && queryableStoreType instanceof QueryableStoreTypes.KeyValueStoreType) {
+                return (T) new ReadOnlyKeyValueStoreFacade<>((TimestampedKeyValueStore<Object, Object>) store);
+            } else if (store instanceof TimestampedWindowStore && queryableStoreType instanceof QueryableStoreTypes.WindowStoreType) {
+                return (T) new ReadOnlyWindowStoreFacade<>((TimestampedWindowStore<Object, Object>) store);
+            } else {
+                return (T) store;
+            }
+        } else {
+            return null;
         }
     }
 

streams/src/main/java/org/apache/kafka/streams/state/internals/WrappingStoreProvider.java

@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals;
 
-import org.apache.kafka.streams.StoreQueryParams;
+import org.apache.kafka.streams.StoreQueryParameters;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.state.QueryableStoreType;
 
@@ -29,25 +29,25 @@
 public class WrappingStoreProvider implements StateStoreProvider {
 
     private final List<StreamThreadStateStoreProvider> storeProviders;
-    private StoreQueryParams storeQueryParams;
+    private StoreQueryParameters storeQueryParameters;
 
     WrappingStoreProvider(final List<StreamThreadStateStoreProvider> storeProviders,
-                          final StoreQueryParams storeQueryParams) {
+                          final StoreQueryParameters storeQueryParameters) {
         this.storeProviders = storeProviders;
-        this.storeQueryParams = storeQueryParams;
+        this.storeQueryParameters = storeQueryParameters;
     }
 
     //visible for testing
-    public void setStoreQueryParams(final StoreQueryParams storeQueryParams) {
-        this.storeQueryParams = storeQueryParams;
+    public void setStoreQueryParameters(final StoreQueryParameters storeQueryParameters) {
+        this.storeQueryParameters = storeQueryParameters;
     }
 
     @Override
     public <T> List<T> stores(final String storeName,
                               final QueryableStoreType<T> queryableStoreType) {
         final List<T> allStores = new ArrayList<>();
         for (final StreamThreadStateStoreProvider provider : storeProviders) {
-            final List<T> stores = provider.stores(storeQueryParams);
+            final List<T> stores = provider.stores(storeQueryParameters);
             allStores.addAll(stores);
         }
         if (allStores.isEmpty()) {