Add synthetic targets so that you can always use cluster.schema as source for both streaming and indexed.

- Make a SearchChainInvocationSpec proxy for all possible searchcluster.schema combinations.
- It will modify the query with the actual source to use, and restrict to the given schema.

5 files changed, 84 insertions, 46 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
index 570e76d0127..72184c5ea32 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
@@ -18,6 +18,7 @@ import com.yahoo.search.Result;
 import com.yahoo.search.Searcher;
 import com.yahoo.search.federation.selection.FederationTarget;
 import com.yahoo.search.federation.selection.TargetSelector;
+import com.yahoo.search.federation.sourceref.ModifyQueryAndResult;
 import com.yahoo.search.federation.sourceref.SearchChainInvocationSpec;
 import com.yahoo.search.federation.sourceref.SearchChainResolver;
 import com.yahoo.search.federation.sourceref.SingleTarget;
@@ -43,7 +44,6 @@ import java.time.Clock;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
@@ -155,15 +155,36 @@ public class FederationSearcher extends ForkingSearcher {
         return builder.build();
     }
 
+    private static class SearchChaininvocationProxy extends SearchChainInvocationSpec {
+        SearchChaininvocationProxy(ComponentId searchChainId, FederationOptions federationOptions, String schema) {
+            super(searchChainId, federationOptions, List.of(schema));
+        }
+
+        @Override
+        public void modifyTargetQuery(Query query) {
+            query.getModel().setSources(searchChainId.getName());
+            query.getModel().setRestrict(schemas.get(0));
+        }
+    }
+
     private static void addSearchChain(SearchChainResolver.Builder builder,
                                        FederationConfig.Target target,
                                        FederationConfig.Target.SearchChain searchChain)
     {
-        if (!target.id().equals(searchChain.searchChainId()))
-            throw new RuntimeException("Invalid federation config, " + target.id() + " != " + searchChain.searchChainId());
-
-        builder.addSearchChain(ComponentId.fromString(searchChain.searchChainId()),
-                               federationOptions(searchChain), searchChain.documentTypes());
+        String id = target.id();
+        if (!id.equals(searchChain.searchChainId()))
+            throw new RuntimeException("Invalid federation config, " + id + " != " + searchChain.searchChainId());
+
+        ComponentId searchChainId = ComponentId.fromString(id);
+        builder.addSearchChain(searchChainId, federationOptions(searchChain), searchChain.documentTypes());
+        // Here we make synthetic SearchChain proxies for all cluster.schema combinations possible
+        // Given a source on the form saerchcluster.schema will rewrite it to source=searchcluster and restrict to schema.
+        // TODO Consider solving this in the config model by making many synthetic search clusters
+        for (String schema : searchChain.documentTypes()) {
+            String virtualChainId = id + "." + schema;
+            builder.addSearchChain(ComponentId.fromString(virtualChainId),
+                    new SearchChaininvocationProxy(searchChainId, federationOptions(searchChain).setUseByDefault(false), schema));
+        }
     }
 
     private static void addSourceForProvider(SearchChainResolver.Builder builder, FederationConfig.Target target,
@@ -215,7 +236,6 @@ public class FederationSearcher extends ForkingSearcher {
 
     private void search(Query query, Execution execution, Target target, Result mergedResults) {
         mergeResult(query, target, mergedResults, search(query, execution, target).orElse(createSearchChainTimedOutResult(query, target)));
-
     }
 
     private void search(Query query, Execution execution, Collection<Target> targets, Result mergedResults) {
@@ -583,7 +603,7 @@ public class FederationSearcher extends ForkingSearcher {
         Collection<SearchChainInvocationSpec> prunedTargets = new ArrayList<>();
 
         for (SearchChainInvocationSpec target : targets) {
-            if (target.documentTypes.isEmpty() || documentTypeIntersectionIsNonEmpty(restrict, target))
+            if (target.schemas.isEmpty() || documentTypeIntersectionIsNonEmpty(restrict, target))
                 prunedTargets.add(target);
         }
 
@@ -591,7 +611,7 @@ public class FederationSearcher extends ForkingSearcher {
     }
 
     private boolean documentTypeIntersectionIsNonEmpty(Set<String> restrict, SearchChainInvocationSpec target) {
-        for (String documentType : target.documentTypes) {
+        for (String documentType : target.schemas) {
             if (restrict.contains(documentType))
                 return true;
         }
@@ -621,11 +641,9 @@ public class FederationSearcher extends ForkingSearcher {
     }
 
     /** A target for federation, containing a chain to which a federation query can be forwarded. */
-    static abstract class Target {
+    static abstract class Target implements ModifyQueryAndResult {
 
         abstract Chain<Searcher> getChain();
-        abstract void modifyTargetQuery(Query query);
-        abstract void modifyTargetResult(Result result);
 
         ComponentId getId() {
             return getChain().getId();
@@ -656,9 +674,9 @@ public class FederationSearcher extends ForkingSearcher {
         Chain<Searcher> getChain() { return chain; }
 
         @Override
-        void modifyTargetQuery(Query query) {}
+        public void modifyTargetQuery(Query query) { target.modifyTargetQuery(query); }
         @Override
-        void modifyTargetResult(Result result) {}
+        public void modifyTargetResult(Result result) { target.modifyTargetResult(result); }
 
         @Override
         public FederationOptions federationOptions() { return target.federationOptions; }
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java
new file mode 100644
index 00000000000..7a9cf1b561f
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java
@@ -0,0 +1,9 @@
+package com.yahoo.search.federation.sourceref;
+
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+
+public interface ModifyQueryAndResult {
+    void modifyTargetQuery(Query query);
+    void modifyTargetResult(Result result);
+}
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
index 32af08f95fc..9822c4ea09a 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
@@ -2,6 +2,8 @@
 package com.yahoo.search.federation.sourceref;
 
 import com.yahoo.component.ComponentId;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
 import com.yahoo.search.searchchain.model.federation.FederationOptions;
 
 import java.util.List;
@@ -13,7 +15,7 @@ import java.util.Objects;
  *
  * @author Tony Vaagenes
  */
-public class SearchChainInvocationSpec implements Cloneable {
+public class SearchChainInvocationSpec implements ModifyQueryAndResult, Cloneable {
 
     public final ComponentId searchChainId;
 
@@ -23,16 +25,20 @@ public class SearchChainInvocationSpec implements Cloneable {
     public final ComponentId provider;
 
     public final FederationOptions federationOptions;
-    public final List<String> documentTypes;
+    public final List<String> schemas;
 
-    SearchChainInvocationSpec(ComponentId searchChainId,
-                              ComponentId source, ComponentId provider, FederationOptions federationOptions,
-                              List<String> documentTypes) {
+    public SearchChainInvocationSpec(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas) {
+        this(searchChainId, null, null, federationOptions, schemas);
+    }
+
+
+    SearchChainInvocationSpec(ComponentId searchChainId, ComponentId source, ComponentId provider,
+                              FederationOptions federationOptions, List<String> schemas) {
         this.searchChainId = searchChainId;
         this.source = source;
         this.provider = provider;
         this.federationOptions = federationOptions;
-        this.documentTypes = List.copyOf(documentTypes);
+        this.schemas = List.copyOf(schemas);
     }
 
     @Override
@@ -49,13 +55,16 @@ public class SearchChainInvocationSpec implements Cloneable {
         if ( ! Objects.equals(this.source, other.source)) return false;
         if ( ! Objects.equals(this.provider, other.provider)) return false;
         if ( ! Objects.equals(this.federationOptions, other.federationOptions)) return false;
-        if ( ! Objects.equals(this.documentTypes, other.documentTypes)) return false;
+        if ( ! Objects.equals(this.schemas, other.schemas)) return false;
         return true;
     }
 
     @Override
     public int hashCode() { 
-        return Objects.hash(searchChainId, source, provider, federationOptions, documentTypes);
+        return Objects.hash(searchChainId, source, provider, federationOptions, schemas);
     }
 
+    public void modifyTargetQuery(Query query) { }
+    public void modifyTargetResult(Result result) {}
+
 }
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
index df91b968750..7dc65c819e4 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
@@ -46,6 +46,16 @@ public class SearchChainResolver {
 
     public static class Builder {
 
+        public interface InvocationSpecFactory {
+            SearchChainInvocationSpec create(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas);
+        }
+
+        private class DefaultInvocationSpecFactory implements InvocationSpecFactory {
+            public SearchChainInvocationSpec create(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas) {
+                return new SearchChainInvocationSpec(searchChainId, federationOptions, schemas);
+            }
+        }
+
         private final SortedSet<Target> defaultTargets = new TreeSet<>();
 
         private final ComponentRegistry<Target> targets = new ComponentRegistry<>() {
@@ -60,22 +70,22 @@ public class SearchChainResolver {
             return addSearchChain(searchChainId, federationOptions, List.of());
         }
 
-        public Builder addSearchChain(ComponentId searchChainId, List<String> documentTypes) {
-            return addSearchChain(searchChainId, new FederationOptions(), documentTypes);
+        public Builder addSearchChain(ComponentId searchChainId, List<String> schemas) {
+            return addSearchChain(searchChainId, new FederationOptions(), schemas);
         }
 
         public Builder addSearchChain(ComponentId searchChainId,
                                       FederationOptions federationOptions,
-                                      List<String> documentTypes) {
-            registerTarget(new SingleTarget(searchChainId,
-                                            new SearchChainInvocationSpec(searchChainId,
-                                                                          null,
-                                                                          null,
-                                                                          federationOptions,
-                                                                          documentTypes),
-                                            false));
+                                      List<String> schemas) {
+            addSearchChain(new SearchChainInvocationSpec(searchChainId, federationOptions, schemas));
             return this;
         }
+        private Builder addSearchChain(SearchChainInvocationSpec invocationSpec) {
+            return registerTarget(new SingleTarget(invocationSpec.searchChainId, invocationSpec, false));
+        }
+        public Builder addSearchChain(ComponentId id, SearchChainInvocationSpec invocationSpec) {
+            return registerTarget(new SingleTarget(id, invocationSpec, false));
+        }
 
         private Builder registerTarget(SingleTarget singleTarget) {
             targets.register(singleTarget.getId(), singleTarget);
@@ -87,9 +97,8 @@ public class SearchChainResolver {
 
         public Builder addSourceForProvider(ComponentId sourceId, ComponentId providerId, ComponentId searchChainId,
                                             boolean isDefaultProviderForSource, FederationOptions federationOptions,
-                                            List<String> documentTypes) {
-            SearchChainInvocationSpec searchChainInvocationSpec =
-                    new SearchChainInvocationSpec(searchChainId, sourceId, providerId, federationOptions, documentTypes);
+                                            List<String> schemas) {
+            var searchChainInvocationSpec = new SearchChainInvocationSpec(searchChainId, sourceId, providerId, federationOptions, schemas);
 
             SourcesTarget sourcesTarget = getOrRegisterSourceTarget(sourceId);
             sourcesTarget.addSource(providerId, searchChainInvocationSpec, isDefaultProviderForSource);
diff --git a/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java b/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
index bd39aed38fa..b1ce1fbd3ee 100644
--- a/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
+++ b/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
@@ -120,17 +120,10 @@ public class FederationResultTest {
             this.timeout = timeout;
         }
 
-        @Override
-        Chain<Searcher> getChain() { return chain; }
-
-        @Override
-        void modifyTargetQuery(Query query) { }
-
-        @Override
-        void modifyTargetResult(Result result) { }
-
-        @Override
-        public FederationOptions federationOptions() {
+        @Override Chain<Searcher> getChain() { return chain; }
+        @Override public void modifyTargetQuery(Query query) { }
+        @Override public void modifyTargetResult(Result result) { }
+        @Override public FederationOptions federationOptions() {
             return new FederationOptions(false, timeout, true);
         }