additionHeaders

Author: clun

astra-db-java/src/main/java/com/datastax/astra/client/DataAPIOptions.java

@@ -32,6 +32,7 @@
 
 import java.net.http.HttpClient;
 import java.time.Duration;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.TreeMap;
 
@@ -42,6 +43,9 @@
 @Getter
 public class DataAPIOptions {
 
+    /** Feature Flag Tables. */
+    public static final String HEADER_FEATURE_FLAG_TABLES = "Feature-Flag-tables";
+
     /** path for json api. */
     public static final String DEFAULT_VERSION = "v1";
 
@@ -75,6 +79,12 @@ public class DataAPIOptions {
     /** Embedding auth provider. */
     final EmbeddingHeadersProvider embeddingAuthProvider;
 
+    /** Add headers to db calls. */
+    final Map<String, Object> databaseAdditionalHeaders;
+
+    /** Add headers to admin calls. */
+    final Map<String, Object> adminAdditionalHeaders;
+
     /** Observers for the commands. */
     final Map<String, CommandObserver> observers;
 
@@ -103,6 +113,8 @@ private DataAPIOptions(DataAPIClientOptionsBuilder builder) {
         this.embeddingAuthProvider = builder.embeddingAuthProvider;
         this.httpClientOptions     = builder.httpClientOptions;
         this.observers             = builder.observers;
+        this.databaseAdditionalHeaders = builder.databaseAdditionalHeaders;
+        this.adminAdditionalHeaders    = builder.adminAdditionalHeaders;
     }
 
     /**
@@ -128,6 +140,12 @@ public static class DataAPIClientOptionsBuilder {
         /** The embedding service API key can be provided at top level. */
         private EmbeddingHeadersProvider embeddingAuthProvider;
 
+        /** Add headers to admin calls. */
+        final Map<String, Object> databaseAdditionalHeaders = new HashMap<>();
+
+        /** Add headers to admin calls. */
+        final Map<String, Object> adminAdditionalHeaders = new HashMap<>();
+
         /** Observers for the commands. */
         private final Map<String, CommandObserver> observers = new TreeMap<>();
 
@@ -295,6 +313,10 @@ public DataAPIClientOptionsBuilder withDestination(DataAPIDestination destinatio
             return this;
         }
 
+        // --------------------------------------------
+        // ----------------- HEADERS  -----------------
+        // --------------------------------------------
+
         /**
          * Builder pattern, update http connection Timeout
          *
@@ -346,10 +368,18 @@ public DataAPIClientOptionsBuilder addCaller(String name, String version) {
          *      self reference
          */
         public DataAPIClientOptionsBuilder addDatabaseAdditionalHeader(String key, Object value) {
-            httpClientOptions.addDbHeader(key, value);
+            databaseAdditionalHeaders.put(key, value);
             return this;
         }
 
+        public DataAPIClientOptionsBuilder enableFeatureFlagTables() {
+            return addDatabaseAdditionalHeader(HEADER_FEATURE_FLAG_TABLES, "true");
+        }
+
+        public DataAPIClientOptionsBuilder disableFeatureFlagTables() {
+            return addDatabaseAdditionalHeader(HEADER_FEATURE_FLAG_TABLES, null);
+        }
+
         /**
          * Add a header to the admin calls.
          *
@@ -360,13 +390,8 @@ public DataAPIClientOptionsBuilder addDatabaseAdditionalHeader(String key, Objec
          * @return
          *      self reference
          */
-        public DataAPIClientOptionsBuilder addAdminAdditionHeader(String key, Object value) {
-            httpClientOptions.addAdminHeader(key, value);
-            return this;
-        }
-
-        public DataAPIClientOptionsBuilder enableFeatureFlagTables() {
-            httpClientOptions.enableFeatureFlagTables();
+        public DataAPIClientOptionsBuilder addAdminAdditionalHeader(String key, Object value) {
+            adminAdditionalHeaders.put(key, value);
             return this;
         }
 

astra-db-java/src/main/java/com/datastax/astra/client/model/command/CommandOptions.java

@@ -26,10 +26,13 @@
 import com.datastax.astra.internal.command.CommandObserver;
 import lombok.Getter;
 
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Optional;
 
+import static com.datastax.astra.client.DataAPIOptions.HEADER_FEATURE_FLAG_TABLES;
+
 /**
  * Options that will be provided to all commands for this collection.
  *
@@ -59,6 +62,12 @@ public class CommandOptions<T extends CommandOptions<T>>{
      */
     protected EmbeddingHeadersProvider embeddingAuthProvider;
 
+    /** Add headers to db calls. */
+    protected Map<String, Object> databaseAdditionalHeaders = new HashMap<>();
+
+    /** Add headers to admin calls. */
+    protected Map<String, Object> adminAdditionalHeaders  = new HashMap<>();
+
     /**
      * Provide the token.
      *
@@ -200,6 +209,46 @@ public Optional<HttpClientOptions> getHttpClientOptions() {
         return Optional.ofNullable(httpClientOptions);
     }
 
+    /**
+     * Add a header to the db calls.
+     *
+     * @param key
+     *      key
+     * @param value
+     *      value
+     * @return
+     *      self reference
+     */
+    @SuppressWarnings("unchecked")
+    public T addDatabaseAdditionalHeader(String key, Object value) {
+        databaseAdditionalHeaders.put(key, value);
+        return (T) this;
+    }
+
+    public T enableFeatureFlagTables() {
+        return addDatabaseAdditionalHeader(HEADER_FEATURE_FLAG_TABLES, "true");
+    }
+
+    public T disableFeatureFlagTables() {
+        return addDatabaseAdditionalHeader(HEADER_FEATURE_FLAG_TABLES, null);
+    }
+
+    /**
+     * Add a header to the admin calls.
+     *
+     * @param key
+     *      key
+     * @param value
+     *      value
+     * @return
+     *      self reference
+     */
+    @SuppressWarnings("unchecked")
+    public T addAdminAdditionalHeader(String key, Object value) {
+        adminAdditionalHeaders.put(key, value);
+        return (T) this;
+    }
+
     /**
      * Default Constructor.
      */

astra-db-java/src/main/java/com/datastax/astra/client/model/http/HttpClientOptions.java

@@ -41,20 +41,11 @@ public class HttpClientOptions {
     // ----------------- HEADERS  -----------------
     // --------------------------------------------
 
-    /** Feature Flag Tables. */
-    public static final String HEADER_FEATURE_FLAG_TABLES = "Feature-Flag-tables";
-
     /** Default caller name. */
     public static final Caller DEFAULT_CALLER = new Caller("astra-db-java",
             DataAPIOptions.class.getPackage().getImplementationVersion() != null ?
             DataAPIOptions.class.getPackage().getImplementationVersion() : "dev");
 
-    /** Add headers to db calls. */
-    Map<String, Object> dbAdditionalHeaders = new HashMap<>();
-
-    /** Add headers to admin calls. */
-    Map<String, Object> adminAdditionalHeaders  = new HashMap<>();
-
     /** Add headers to db calls. */
     List<Caller> callers = new ArrayList<>();
 
@@ -70,44 +61,6 @@ public void addCaller(Caller caller) {
         callers.add(caller);
     }
 
-    /**
-     * Add a header to the db calls.
-     *
-     * @param key
-     *      key
-     * @param value
-     *      value
-     * @return
-     *      self reference
-     */
-    public HttpClientOptions addDbHeader(String key, Object value) {
-        dbAdditionalHeaders.put(key, value);
-        return this;
-    }
-
-    /**
-     * Add a header to the admin calls.
-     *
-     * @param key
-     *      key
-     * @param value
-     *      value
-     * @return
-     *      self reference
-     */
-    public HttpClientOptions addAdminHeader(String key, Object value) {
-        adminAdditionalHeaders.put(key, value);
-        return this;
-    }
-
-    /**
-     * Enable Feature Flag Tables.
-     */
-    public HttpClientOptions enableFeatureFlagTables() {
-        dbAdditionalHeaders.put(HEADER_FEATURE_FLAG_TABLES, true);
-        return this;
-    }
-
     // --------------------------------------------
     // ----------------- TIMEOUTS -----------------
     // --------------------------------------------

astra-db-java/src/test/java/com/datastax/astra/test/unit/DataApiOptionsTest.java

@@ -1,10 +1,12 @@
 package com.datastax.astra.test.unit;
 
+import com.datastax.astra.client.DataAPIDestination;
 import com.datastax.astra.client.DataAPIOptions;
 import com.datastax.astra.client.model.CollectionIdTypes;
 import com.datastax.astra.client.model.CollectionOptions;
 import com.datastax.astra.client.model.DeleteOneOptions;
 import com.datastax.astra.client.model.DeleteResult;
+import com.datastax.astra.client.model.http.HttpProxy;
 import com.datastax.astra.client.model.query.Filter;
 import com.datastax.astra.client.model.query.FilterOperator;
 import com.datastax.astra.client.model.FindOneAndDeleteOptions;
@@ -28,9 +30,11 @@
 import org.junit.jupiter.api.Test;
 
 import java.net.http.HttpClient;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 
+import static com.datastax.astra.client.DataAPIOptions.HEADER_FEATURE_FLAG_TABLES;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
@@ -39,12 +43,14 @@ class DataApiOptionsTest {
     @Test
     void shouldPopulateOptions() {
         DataAPIOptions options = DataAPIOptions.builder()
-                .withHttpProxy(new DataAPIOptions.HttpProxy("localhost", 8080))
+                .withHttpProxy(new HttpProxy("localhost", 8080))
                 .withApiVersion("v1")
                 .withHttpRedirect(HttpClient.Redirect.NORMAL)
-                .withHttpRetryCount(5)
-                .withHttpRetryDelayMillis(1000)
-                .withDestination(DataAPIOptions.DataAPIDestination.DSE)
+                .withHttpRetries(5, Duration.ofMillis(1000))
+                .withDestination(DataAPIDestination.DSE)
+                .enableFeatureFlagTables()
+                // equivalent to:
+                .addDatabaseAdditionalHeader(HEADER_FEATURE_FLAG_TABLES, "true")
                 .build();
         assertThat(options.getHttpClientOptions().getProxy().getHostname()).isEqualTo("localhost");
     }