Add configuration of headers through CLI

4 files changed, 34 insertions, 1 deletions

diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliArguments.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliArguments.java
index 63b438134d6..06c994b12b6 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliArguments.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliArguments.java
@@ -16,6 +16,9 @@ import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.file.Path;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalInt;
 
@@ -38,6 +41,7 @@ class CliArguments {
     private static final String PRIVATE_KEY_OPTION = "private-key";
     private static final String CA_CERTIFICATES_OPTION = "ca-certificates";
     private static final String DISABLE_SSL_HOSTNAME_VERIFICATION_OPTION = "disable-ssl-hostname-verification";
+    private static final String HEADER_OPTION = "header";
 
     private final CommandLine arguments;
 
@@ -89,6 +93,24 @@ class CliArguments {
             .orElseThrow(() -> new CliArgumentsException("Feed file must be specified"));
     }
 
+    Map<String, String> headers() throws CliArgumentsException {
+        String[] rawArguments = arguments.getOptionValues(HEADER_OPTION);
+        if (rawArguments == null) return Collections.emptyMap();
+        Map<String, String> headers = new HashMap<>();
+        for (String rawArgument : rawArguments) {
+            if (rawArgument.startsWith("\"") || rawArgument.startsWith("'")) {
+                rawArgument = rawArgument.substring(1);
+            }
+            if (rawArgument.endsWith("\"") || rawArgument.endsWith("'")) {
+                rawArgument = rawArgument.substring(0, rawArgument.length() - 1);
+            }
+            int colonIndex = rawArgument.indexOf(':');
+            if (colonIndex == -1) throw new CliArgumentsException("Invalid header: '" + rawArgument + "'");
+            headers.put(rawArgument.substring(0, colonIndex), rawArgument.substring(colonIndex + 1).trim());
+        }
+        return Collections.unmodifiableMap(headers);
+    }
+
     boolean sslHostnameVerificationDisabled() { return has(DISABLE_SSL_HOSTNAME_VERIFICATION_OPTION); }
 
     private OptionalInt intValue(String option) throws CliArgumentsException {
@@ -126,6 +148,10 @@ class CliArguments {
                         .type(URL.class)
                         .build())
                 .addOption(Option.builder()
+                        .longOpt(HEADER_OPTION)
+                        .hasArgs()
+                        .build())
+                .addOption(Option.builder()
                         .longOpt(FILE_OPTION)
                         .type(File.class)
                         .hasArg()
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliClient.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliClient.java
index 83a0c650318..84a29f9c3e4 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliClient.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/CliClient.java
@@ -68,6 +68,7 @@ class CliClient {
             }
             builder.setSslContext(sslContextBuilder.build());
         }
+        cliArgs.headers().forEach(builder::addRequestHeader);
         return builder.build();
     }
 
diff --git a/vespa-feed-client/src/test/java/ai/vespa/feed/client/CliArgumentsTest.java b/vespa-feed-client/src/test/java/ai/vespa/feed/client/CliArgumentsTest.java
index b8dda66bf96..be479d294d5 100644
--- a/vespa-feed-client/src/test/java/ai/vespa/feed/client/CliArgumentsTest.java
+++ b/vespa-feed-client/src/test/java/ai/vespa/feed/client/CliArgumentsTest.java
@@ -20,7 +20,8 @@ class CliArgumentsTest {
         CliArguments args = CliArguments.fromRawArgs(new String[]{
                 "--endpoint=https://vespa.ai:4443/", "--file=feed.json", "--connections=10",
                 "--max-streams-per-connection=128", "--certificate=cert.pem", "--private-key=key.pem",
-                "--ca-certificates=ca-certs.pem", "--disable-ssl-hostname-verification"});
+                "--ca-certificates=ca-certs.pem", "--disable-ssl-hostname-verification",
+                "--header=\"My-Header: my-value\"", "--header", "Another-Header: another-value"});
         assertEquals(URI.create("https://vespa.ai:4443/"), args.endpoint());
         assertEquals(Paths.get("feed.json"), args.inputFile());
         assertEquals(10, args.connections().getAsInt());
@@ -30,6 +31,10 @@ class CliArgumentsTest {
         assertEquals(Paths.get("ca-certs.pem"), args.caCertificates().get());
         assertTrue(args.sslHostnameVerificationDisabled());
         assertFalse(args.helpSpecified());
+        assertFalse(args.versionSpecified());
+        assertEquals(2, args.headers().size());
+        assertEquals("my-value", args.headers().get("My-Header"));
+        assertEquals("another-value", args.headers().get("Another-Header"));
     }
 
     @Test
diff --git a/vespa-feed-client/src/test/resources/help.txt b/vespa-feed-client/src/test/resources/help.txt
index d0017003a11..8ad153bc0e0 100644
--- a/vespa-feed-client/src/test/resources/help.txt
+++ b/vespa-feed-client/src/test/resources/help.txt
@@ -6,6 +6,7 @@ Vespa feed client
     --disable-ssl-hostname-verification
     --endpoint <arg>
     --file <arg>
+    --header <arg>
     --help
     --max-streams-per-connection <arg>
     --private-key <arg>