Add support to custom headers

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## main
+* Add custom headers support
+* Expose http client
+```
+var client = new Client(cfg);
+var headers = client.getDefaultHttpHeaders();
+headers.put("h1", "v1");
+
+var httpClient = client.getHttpClient();
+```
+
+
 ## v0.4.1-alpha
 * Adding support to new java versions.
 * Upgrading `apache arrow` version to `9.0.0`.

rai-sdk/src/main/java/com/relationalai/Client.java

@@ -77,15 +77,13 @@ public Client(Config cfg) {
             this.host = cfg.host;
         if (cfg.port != null)
             this.port = Integer.parseInt(cfg.port);
+        this.httpClient = HttpClient.newBuilder().build();
         this.credentials = cfg.credentials;
         this.setAccessTokenHandler(new DefaultAccessTokenHandler());
     }
 
     // Returns the current `HttpClient` instance, creating one if necessarry.
-    HttpClient getHttpClient() {
-        if (this.httpClient == null) {
-            this.httpClient = HttpClient.newBuilder().build();
-        }
+    public HttpClient getHttpClient() {
         return this.httpClient;
     }
 
@@ -95,6 +93,11 @@ public Client setHttpClient(HttpClient httpClient) {
         return this;
     }
 
+    // Returns the default http headers
+    public Map<String, String> getDefaultHttpHeaders() {
+        return defaultHeaders;
+    }
+
     public void setAccessTokenHandler(AccessTokenHandler handler) {
         this.accessTokenHandler = handler;
     }

rai-sdk/src/test/java/com/relationalai/UnitTest.java

@@ -19,6 +19,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
+import java.util.Map;
 import java.util.UUID;
 import java.util.function.Predicate;
 
@@ -35,17 +36,24 @@ Config getConfig() throws IOException {
 
         var cfg = String.format(
                 "[default]\nregion=us-east\nport=443\nscheme=https\nclient_id=%s\nclient_secret=%s\nclient_credentials_url=%s",
-                System.getenv("CLIENT_ID"),
-                System.getenv("CLIENT_SECRET"),
-                System.getenv("CLIENT_CREDENTIALS_URL")
+                getenv("CLIENT_ID"),
+                getenv("CLIENT_SECRET"),
+                getenv("CLIENT_CREDENTIALS_URL")
         );
         var stream = new ByteArrayInputStream(cfg.getBytes());
         return Config.loadConfig(stream);
     }
     // Returns a new client object constructed from default config settings.
     Client createClient() throws IOException {
         var cfg = getConfig();
-        return new Client(cfg);
+        var customHeaders = (Map<String, String>) Json.deserialize(getenv("CUSTOM_HEADERS", "{}"), Map.class);
+
+        var testClient = new Client(cfg);
+        var httpHeaders = testClient.getDefaultHttpHeaders();
+        for (var header : customHeaders.entrySet()) {
+            httpHeaders.put(header.getKey(), header.getValue());
+        }
+        return testClient;
     }
 
     // Ensure that the test database exists.
@@ -92,4 +100,12 @@ static Relation findRelation(Relation[] relations, String colName) {
         }
         return null;
     }
+
+    String getenv(String name, String defaultValue) {
+        return System.getenv(name) == null ? defaultValue : System.getenv(name);
+    }
+
+    String getenv(String name) {
+        return getenv(name, null);
+    }
 }