Skip to content

Commit e40880f

Browse files
skyflow-bhartiskyflow-bharti
authored
Merge pull request #298 from skyflowapi/SK-2726-Fix-client-initialisation
SK-2726-fix-client-initialisation
2 parents d252eba + c22c341 commit e40880f

2 files changed

Lines changed: 149 additions & 18 deletions

File tree

‎v3/src/main/java/com/skyflow/VaultClient.java‎

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
package com.skyflow;
22

3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import java.util.concurrent.TimeUnit;
6+
37
import com.skyflow.config.Credentials;
48
import com.skyflow.config.VaultConfig;
59
import com.skyflow.enums.UpsertType;
@@ -19,19 +23,17 @@
1923
import com.skyflow.utils.Utils;
2024
import com.skyflow.utils.logger.LogUtil;
2125
import com.skyflow.utils.validations.Validations;
26+
import com.skyflow.vault.data.DeleteTokensRequest;
2227
import com.skyflow.vault.data.DetokenizeRequest;
2328
import com.skyflow.vault.data.InsertRecord;
24-
import com.skyflow.vault.data.DeleteTokensRequest;
25-
import com.skyflow.vault.data.TokenizeRequest;
2629
import com.skyflow.vault.data.TokenizeRecord;
30+
import com.skyflow.vault.data.TokenizeRequest;
2731
import io.github.cdimascio.dotenv.Dotenv;
2832
import io.github.cdimascio.dotenv.DotenvException;
33+
import okhttp3.ConnectionPool;
2934
import okhttp3.OkHttpClient;
3035
import okhttp3.Request;
3136

32-
import java.util.ArrayList;
33-
import java.util.List;
34-
3537

3638
public class VaultClient {
3739
private final VaultConfig vaultConfig;
@@ -41,6 +43,8 @@ public class VaultClient {
4143
private Credentials finalCredentials;
4244
private String token;
4345
private String apiKey;
46+
private OkHttpClient sharedHttpClient = null;
47+
private String currentVaultURL = null;
4448

4549
protected VaultClient(VaultConfig vaultConfig, Credentials credentials) throws SkyflowException {
4650
super();
@@ -78,10 +82,12 @@ protected void setBearerToken() throws SkyflowException {
7882
} else {
7983
LogUtil.printInfoLog(InfoLogs.REUSE_BEARER_TOKEN.getLog());
8084
}
81-
updateExecutorInHTTP(); // update executor
82-
this.apiClient = this.apiClientBuilder.build();
83-
}
8485

86+
if (apiClient == null) {
87+
updateExecutorInHTTP();
88+
this.apiClient = this.apiClientBuilder.build();
89+
}
90+
}
8591
private void updateVaultURL() throws SkyflowException {
8692
// Fetch vaultURL from ENV
8793
String vaultURL = Utils.getEnvVaultURL();
@@ -96,6 +102,10 @@ private void updateVaultURL() throws SkyflowException {
96102
vaultURL = Utils.getVaultURL(this.vaultConfig.getClusterId(), this.vaultConfig.getEnv());
97103
}
98104
this.apiClientBuilder.url(vaultURL);
105+
if (!vaultURL.equals(this.currentVaultURL)) {
106+
this.currentVaultURL = vaultURL;
107+
this.apiClient = null;
108+
}
99109
}
100110

101111
private void prioritiseCredentials() throws SkyflowException {
@@ -132,16 +142,18 @@ private void prioritiseCredentials() throws SkyflowException {
132142
}
133143

134144
protected void updateExecutorInHTTP() {
135-
OkHttpClient httpClient = new OkHttpClient.Builder()
136-
.addInterceptor(chain -> {
137-
Request original = chain.request();
138-
Request requestWithAuth = original.newBuilder()
139-
.header("Authorization", "Bearer " + this.token)
140-
.build();
141-
return chain.proceed(requestWithAuth);
142-
})
143-
.build();
144-
apiClientBuilder.httpClient(httpClient);
145+
if (sharedHttpClient == null) {
146+
sharedHttpClient = new OkHttpClient.Builder()
147+
.connectionPool(new ConnectionPool(10, 1, TimeUnit.MINUTES))
148+
.addInterceptor(chain -> {
149+
Request requestWithAuth = chain.request().newBuilder()
150+
.header("Authorization", "Bearer " + this.token)
151+
.build();
152+
return chain.proceed(requestWithAuth);
153+
})
154+
.build();
155+
apiClientBuilder.httpClient(sharedHttpClient);
156+
}
145157
}
146158

147159
protected V1InsertRequest getBulkInsertRequestBody(com.skyflow.vault.data.InsertRequest request, VaultConfig config) {

‎v3/src/test/java/com/skyflow/VaultClientTests.java‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
import static org.junit.Assert.*;
44
import com.skyflow.config.Credentials;
55
import com.skyflow.config.VaultConfig;
6+
import okhttp3.ConnectionPool;
7+
import okhttp3.OkHttpClient;
8+
import java.util.concurrent.TimeUnit;
69
import com.skyflow.enums.Env;
710
import com.skyflow.enums.UpsertType;
811
import com.skyflow.errors.ErrorCode;
@@ -477,6 +480,71 @@ public void testUpdateExecutorInHTTP() throws Exception {
477480
Assert.assertNotNull(recordsApi);
478481
}
479482

483+
@Test
484+
public void testConnectionPoolMaxIdleConnections() throws Exception {
485+
Credentials credentials = new Credentials();
486+
credentials.setApiKey("sky-ab123-abcd1234cdef1234abcd4321cdef4321");
487+
VaultConfig config = new VaultConfig();
488+
config.setVaultId(vaultID);
489+
config.setClusterId(clusterID);
490+
config.setEnv(Env.PROD);
491+
config.setCredentials(credentials);
492+
493+
VaultClient client = new VaultClient(config, credentials);
494+
client.updateExecutorInHTTP();
495+
496+
OkHttpClient httpClient = (OkHttpClient) getPrivateField(client, "sharedHttpClient");
497+
Assert.assertNotNull(httpClient);
498+
499+
ConnectionPool pool = httpClient.connectionPool();
500+
Assert.assertNotNull(pool);
501+
502+
Object delegate = getPrivateField(pool, "delegate");
503+
int maxIdleConnections = (int) getPrivateField(delegate, "maxIdleConnections");
504+
Assert.assertEquals(10, maxIdleConnections);
505+
}
506+
507+
@Test
508+
public void testConnectionPoolKeepAliveDuration() throws Exception {
509+
Credentials credentials = new Credentials();
510+
credentials.setApiKey("sky-ab123-abcd1234cdef1234abcd4321cdef4321");
511+
VaultConfig config = new VaultConfig();
512+
config.setVaultId(vaultID);
513+
config.setClusterId(clusterID);
514+
config.setEnv(Env.PROD);
515+
config.setCredentials(credentials);
516+
517+
VaultClient client = new VaultClient(config, credentials);
518+
client.updateExecutorInHTTP();
519+
520+
OkHttpClient httpClient = (OkHttpClient) getPrivateField(client, "sharedHttpClient");
521+
ConnectionPool pool = httpClient.connectionPool();
522+
523+
Object delegate = getPrivateField(pool, "delegate");
524+
long keepAliveDurationNs = (long) getPrivateField(delegate, "keepAliveDurationNs");
525+
Assert.assertEquals(TimeUnit.MINUTES.toNanos(1), keepAliveDurationNs);
526+
}
527+
528+
@Test
529+
public void testSharedHttpClientIsSingleton() throws Exception {
530+
Credentials credentials = new Credentials();
531+
credentials.setApiKey("sky-ab123-abcd1234cdef1234abcd4321cdef4321");
532+
VaultConfig config = new VaultConfig();
533+
config.setVaultId(vaultID);
534+
config.setClusterId(clusterID);
535+
config.setEnv(Env.PROD);
536+
config.setCredentials(credentials);
537+
538+
VaultClient client = new VaultClient(config, credentials);
539+
client.updateExecutorInHTTP();
540+
OkHttpClient first = (OkHttpClient) getPrivateField(client, "sharedHttpClient");
541+
542+
client.updateExecutorInHTTP();
543+
OkHttpClient second = (OkHttpClient) getPrivateField(client, "sharedHttpClient");
544+
545+
Assert.assertSame("sharedHttpClient must not be rebuilt on repeated calls", first, second);
546+
}
547+
480548
// Helper methods for reflection field access
481549
private Object getPrivateField(Object obj, String fieldName) throws Exception {
482550
java.lang.reflect.Field field = obj.getClass().getDeclaredField(fieldName);
@@ -533,6 +601,57 @@ public void testPrioritiseCredentialsThrowsWhenNoCredentials() throws Exception
533601
}
534602
}
535603

604+
@Test
605+
public void testApiClientNotReinitializedOnSubsequentCalls() throws Exception {
606+
Credentials credentials = new Credentials();
607+
credentials.setApiKey("sky-ab123-abcd1234cdef1234abcd4321cdef4321");
608+
VaultConfig config = new VaultConfig();
609+
config.setVaultId(vaultID);
610+
config.setClusterId(clusterID);
611+
config.setEnv(Env.PROD);
612+
config.setCredentials(credentials);
613+
614+
VaultClient client = new VaultClient(config, credentials);
615+
616+
client.setBearerToken();
617+
Object apiClientFirst = getPrivateField(client, "apiClient");
618+
Assert.assertNotNull(apiClientFirst);
619+
620+
client.setBearerToken();
621+
Object apiClientSecond = getPrivateField(client, "apiClient");
622+
623+
Assert.assertSame("apiClient must not be rebuilt on repeated setBearerToken calls", apiClientFirst, apiClientSecond);
624+
}
625+
626+
@Test
627+
public void testTokenUpdatedWithoutReinitializingApiClient() throws Exception {
628+
Credentials credentials = new Credentials();
629+
credentials.setApiKey("sky-ab123-abcd1234cdef1234abcd4321cdef4321");
630+
VaultConfig config = new VaultConfig();
631+
config.setVaultId(vaultID);
632+
config.setClusterId(clusterID);
633+
config.setEnv(Env.PROD);
634+
config.setCredentials(credentials);
635+
636+
VaultClient client = new VaultClient(config, credentials);
637+
client.setBearerToken();
638+
Object apiClientBefore = getPrivateField(client, "apiClient");
639+
String tokenBefore = (String) getPrivateField(client, "token");
640+
641+
// Rotate to a different API key
642+
Credentials newCredentials = new Credentials();
643+
newCredentials.setApiKey("sky-ab123-1234567890abcdef1234567890abcdef");
644+
config.setCredentials(newCredentials);
645+
646+
client.setBearerToken();
647+
Object apiClientAfter = getPrivateField(client, "apiClient");
648+
String tokenAfter = (String) getPrivateField(client, "token");
649+
650+
Assert.assertSame("apiClient must not be rebuilt after credential rotation", apiClientBefore, apiClientAfter);
651+
Assert.assertNotEquals("token must reflect the new credential", tokenBefore, tokenAfter);
652+
Assert.assertEquals("sky-ab123-1234567890abcdef1234567890abcdef", tokenAfter);
653+
}
654+
536655
@Test
537656
public void testPrioritiseCredentialsSetsFinalCredentialsFromSysCredentials() throws Exception {
538657
VaultConfig config = new VaultConfig();

0 commit comments

Comments
 (0)