1 files changed, 22 insertions, 13 deletions
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);