Support dynamic queries instead

2 files changed, 18 insertions, 1 deletions

diff --git a/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java b/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
index 2055cd6d74a..a9f36da9f97 100644
--- a/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
+++ b/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
@@ -26,11 +26,13 @@ import java.io.UncheckedIOException;
 import java.net.URI;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.logging.Logger;
 
 import static java.util.Objects.requireNonNull;
@@ -67,10 +69,14 @@ public abstract class AbstractHttpClient implements HttpClient {
 
         Throwable thrown = null;
         for (URI host : builder.hosts) {
+            Query query = builder.query;
+            for (Supplier<Query> dynamic : builder.dynamicQuery)
+                query = query.set(dynamic.get().lastEntries());
+
             ClassicHttpRequest request = ClassicRequestBuilder.create(builder.method.name())
                                                               .setUri(HttpURL.from(host)
                                                                              .appendPath(builder.path)
-                                                                             .appendQuery(builder.query)
+                                                                             .appendQuery(query)
                                                                              .asURI())
                                                               .build();
             request.setEntity(builder.entity);
@@ -143,6 +149,7 @@ public abstract class AbstractHttpClient implements HttpClient {
         private final HostStrategy hosts;
         private HttpURL.Path path = Path.empty();
         private HttpURL.Query query = Query.empty();
+        private List<Supplier<Query>> dynamicQuery = new ArrayList<>();
         private HttpEntity entity;
         private RequestConfig config = HttpClient.defaultRequestConfig;
         private ResponseVerifier verifier = HttpClient.throwOnError;
@@ -203,6 +210,12 @@ public abstract class AbstractHttpClient implements HttpClient {
         }
 
         @Override
+        public HttpClient.RequestBuilder parameters(Supplier<Query> query) {
+            dynamicQuery.add(query);
+            return this;
+        }
+
+        @Override
         public RequestBuilder timeout(Duration timeout) {
             return config(RequestConfig.copy(config)
                                        .setResponseTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS)
diff --git a/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java b/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
index f5805ce5b94..e5a8ebcc8b3 100644
--- a/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
+++ b/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
@@ -26,6 +26,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.IntStream;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -99,6 +100,9 @@ public interface HttpClient extends Closeable {
         /** Appends all parameters from the given query. */
         RequestBuilder parameters(Query query);
 
+        /** Sets all parameters from the given query dynamically, when creating retried requests. */
+        RequestBuilder parameters(Supplier<Query> query);
+
         /** Overrides the default socket read timeout of the request. {@code Duration.ZERO} gives infinite timeout. */
         RequestBuilder timeout(Duration timeout);