Skip to content

Commit 132fe80

Browse files
clunclun
committed
refectoring and alignment
1 parent dc37d47 commit 132fe80

File tree

14 files changed

+1486
-126
lines changed

14 files changed

+1486
-126
lines changed

‎astra-db-client/pom.xml‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
<groupId>com.datastax.astra</groupId>
3737
<artifactId>astra-sdk-pulsar</artifactId>
3838
</exclusion>
39-
<!--
4039
<exclusion>
4140
<groupId>com.datastax.stargate</groupId>
4241
<artifactId>stargate-sdk-document</artifactId>
@@ -49,7 +48,14 @@
4948
<groupId>com.datastax.stargate</groupId>
5049
<artifactId>stargate-sdk-grpc</artifactId>
5150
</exclusion>
52-
-->
51+
<exclusion>
52+
<groupId>com.datastax.stargate</groupId>
53+
<artifactId>stargate-sdk-grpc</artifactId>
54+
</exclusion>
55+
<exclusion>
56+
<groupId>com.datastax.stargate</groupId>
57+
<artifactId>stargate-sdk-cql</artifactId>
58+
</exclusion>
5359
</exclusions>
5460
</dependency>
5561

‎astra-db-client/src/main/java/com/dtsx/astra/sdk/AstraDB.java‎

Lines changed: 89 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package com.dtsx.astra.sdk;
22

3-
import com.datastax.astra.sdk.AstraClient;
4-
import com.datastax.oss.driver.api.core.CqlSession;
53
import com.dtsx.astra.sdk.db.AstraDBOpsClient;
64
import com.dtsx.astra.sdk.db.domain.Database;
75
import com.dtsx.astra.sdk.db.exception.DatabaseNotFoundException;
@@ -11,15 +9,13 @@
119
import io.stargate.sdk.api.SimpleTokenProvider;
1210
import io.stargate.sdk.http.ServiceHttp;
1311
import io.stargate.sdk.json.ApiClient;
14-
import io.stargate.sdk.json.CollectionClient;
15-
import io.stargate.sdk.json.CollectionRepository;
1612
import io.stargate.sdk.json.NamespaceClient;
1713
import io.stargate.sdk.json.domain.CollectionDefinition;
14+
import io.stargate.sdk.json.domain.SimilarityMetric;
1815
import lombok.Getter;
1916
import lombok.NonNull;
2017
import lombok.extern.slf4j.Slf4j;
2118

22-
import java.util.Objects;
2319
import java.util.Optional;
2420
import java.util.UUID;
2521
import java.util.stream.Stream;
@@ -59,13 +55,26 @@ public class AstraDB {
5955
* api endpoint
6056
*/
6157
public AstraDB(String token, String apiEndpoint) {
62-
Objects.requireNonNull(token, "token");
63-
Objects.requireNonNull(apiEndpoint, "apiEndpoint");
64-
// Fixing api
58+
this(token, apiEndpoint, AstraDBClient.DEFAULT_KEYSPACE);
59+
}
60+
61+
/**
62+
* Initialization with endpoint and apikey.
63+
*
64+
* @param token
65+
* api token
66+
* @param apiEndpoint
67+
* api endpoint
68+
* @param keyspace
69+
* keyspace
70+
*/
71+
public AstraDB(@NonNull String token, @NonNull String apiEndpoint, @NonNull String keyspace) {
72+
// Support for apiEndpoint with or without /api/json
6573
if (apiEndpoint.endsWith("com")) {
6674
apiEndpoint = apiEndpoint + "/api/json";
6775
}
6876
this.apiEndpoint = apiEndpoint;
77+
6978
// Finding Environment based on apiEndpoint (looping to devops)
7079
if (apiEndpoint.contains(AstraEnvironment.PROD.getAppsSuffix())) {
7180
this.env = AstraEnvironment.PROD;
@@ -82,7 +91,7 @@ public AstraDB(String token, String apiEndpoint) {
8291
jsonDeploy.addDatacenterTokenProvider("default", new SimpleTokenProvider(token));
8392
jsonDeploy.addDatacenterServices("default", new ServiceHttp("json", apiEndpoint, apiEndpoint));
8493
this.apiClient = new ApiClient(jsonDeploy);
85-
this.nsClient = apiClient.namespace(AstraDBClient.DEFAULT_KEYSPACE);
94+
this.nsClient = apiClient.namespace(keyspace);
8695
}
8796

8897
/**
@@ -94,7 +103,7 @@ public AstraDB(String token, String apiEndpoint) {
94103
* database identifier
95104
*/
96105
public AstraDB(@NonNull String token, @NonNull UUID databaseId) {
97-
this(token, databaseId, null, AstraEnvironment.PROD);
106+
this(token, databaseId, null, AstraEnvironment.PROD, AstraDBClient.DEFAULT_KEYSPACE);
98107
}
99108

100109
/**
@@ -104,15 +113,32 @@ public AstraDB(@NonNull String token, @NonNull UUID databaseId) {
104113
* token
105114
* @param databaseId
106115
* database identifier
107-
* @param env
108-
* environment
116+
* @param keyspace
117+
* database keyspace
109118
*/
110-
public AstraDB(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEnvironment env) {
111-
this(token, databaseId, null, env);
119+
public AstraDB(@NonNull String token, @NonNull UUID databaseId, @NonNull String keyspace) {
120+
this(token, databaseId, null, AstraEnvironment.PROD, keyspace);
121+
}
122+
123+
/**
124+
* Full constructor.
125+
*
126+
* @param token
127+
* token
128+
* @param databaseId
129+
* database identifier
130+
* @param region
131+
* database region
132+
* @param keyspace
133+
* keyspace
134+
*/
135+
public AstraDB(@NonNull String token, @NonNull UUID databaseId, @NonNull String region, @NonNull String keyspace) {
136+
this(token, databaseId, region, AstraEnvironment.PROD, keyspace);
112137
}
113138

114139
/**
115140
* Accessing the database with id an region.
141+
*
116142
* @param token
117143
* astra token
118144
* @param databaseId
@@ -122,7 +148,7 @@ public AstraDB(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEn
122148
* @param env
123149
* environment
124150
*/
125-
public AstraDB(@NonNull String token, @NonNull UUID databaseId, String region, @NonNull AstraEnvironment env) {
151+
public AstraDB(@NonNull String token, @NonNull UUID databaseId, String region, @NonNull AstraEnvironment env, String keyspace) {
126152
this.env = env;
127153
Database db = new AstraDBOpsClient(token, env)
128154
.findById(databaseId.toString())
@@ -135,18 +161,10 @@ public AstraDB(@NonNull String token, @NonNull UUID databaseId, String region, @
135161
jsonDeploy.addDatacenterTokenProvider("default", new SimpleTokenProvider(token));
136162
jsonDeploy.addDatacenterServices("default", new ServiceHttp("json", apiEndpoint, apiEndpoint));
137163
this.apiClient = new ApiClient(jsonDeploy);
138-
139-
// will inherit 'default_keyspace' from the database
140-
this.nsClient = apiClient.namespace(db.getInfo().getKeyspace());
141-
}
142-
143-
/**
144-
* Gets apiEndpoint
145-
*
146-
* @return value of apiEndpoint
147-
*/
148-
public String getApiEndpoint() {
149-
return apiEndpoint;
164+
if (keyspace == null) {
165+
keyspace = db.getInfo().getKeyspace();
166+
}
167+
this.nsClient = apiClient.namespace(keyspace);
150168
}
151169

152170
// --------------------------
@@ -211,8 +229,8 @@ public void deleteCollection(String name) {
211229
* @return
212230
* json vector store
213231
*/
214-
public CollectionClient createCollection(String name) {
215-
return nsClient.createCollection(name);
232+
public AstraDBCollection createCollection(String name) {
233+
return new AstraDBCollection(nsClient.createCollection(name));
216234
}
217235

218236
/**
@@ -227,8 +245,8 @@ public CollectionClient createCollection(String name) {
227245
* @return
228246
* json vector store
229247
*/
230-
public <DOC> CollectionRepository<DOC> createCollection(String name, Class<DOC> clazz) {
231-
return nsClient.createCollection(name, clazz);
248+
public <DOC> AstraDBRepository<DOC> createCollection(String name, Class<DOC> clazz) {
249+
return new AstraDBRepository<DOC>(nsClient.createCollection(name, clazz));
232250
}
233251

234252
/**
@@ -241,8 +259,26 @@ public <DOC> CollectionRepository<DOC> createCollection(String name, Class<DOC>
241259
* @return
242260
* json vector store
243261
*/
244-
public CollectionClient createCollection(String name, int vectorDimension) {
245-
return nsClient.createCollection(name, vectorDimension);
262+
public AstraDBCollection createCollection(String name, int vectorDimension) {
263+
return new AstraDBCollection(nsClient.createCollection(name, vectorDimension));
264+
}
265+
266+
/**
267+
* Create the minimal store.
268+
*
269+
* @param name
270+
* store name
271+
* @param vectorDimension
272+
* dimension
273+
* @return
274+
* json vector store
275+
*/
276+
public AstraDBCollection createCollection(String name, int vectorDimension, SimilarityMetric metric) {
277+
return new AstraDBCollection(nsClient.createCollection(CollectionDefinition
278+
.builder()
279+
.name(name)
280+
.vector(vectorDimension, metric)
281+
.build()));
246282
}
247283

248284
/**
@@ -259,8 +295,8 @@ public CollectionClient createCollection(String name, int vectorDimension) {
259295
* @param <T>
260296
* object type
261297
*/
262-
public <T> CollectionRepository<T> createCollection(String name, int vectorDimension, Class<T> bean) {
263-
return nsClient.createCollection(name, vectorDimension, bean);
298+
public <T> AstraDBRepository<T> createCollection(String name, int vectorDimension, Class<T> bean) {
299+
return new AstraDBRepository(nsClient.createCollection(name, vectorDimension, bean));
264300
}
265301

266302

@@ -272,8 +308,8 @@ public <T> CollectionRepository<T> createCollection(String name, int vectorDimen
272308
* @return
273309
* json vector store
274310
*/
275-
public CollectionClient createCollection(CollectionDefinition def) {
276-
return nsClient.createCollection(def);
311+
public AstraDBCollection createCollection(CollectionDefinition def) {
312+
return new AstraDBCollection(nsClient.createCollection(def));
277313
}
278314

279315
/**
@@ -288,8 +324,8 @@ public CollectionClient createCollection(CollectionDefinition def) {
288324
* @return
289325
* json vector store
290326
*/
291-
public <DOC> CollectionRepository<DOC> createCollection(CollectionDefinition def, Class<DOC> clazz) {
292-
return nsClient.createCollection(def, clazz);
327+
public <DOC> AstraDBRepository<DOC> createCollection(CollectionDefinition def, Class<DOC> clazz) {
328+
return new AstraDBRepository<DOC>(nsClient.createCollection(def, clazz));
293329
}
294330

295331
// --------------------
@@ -304,8 +340,8 @@ public <DOC> CollectionRepository<DOC> createCollection(CollectionDefinition def
304340
* @return
305341
* storeName client
306342
*/
307-
public CollectionClient collection(@NonNull String storeName) {
308-
return nsClient.collection(storeName);
343+
public AstraDBCollection collection(@NonNull String storeName) {
344+
return new AstraDBCollection(nsClient.collection(storeName));
309345
}
310346

311347
/**
@@ -320,8 +356,18 @@ public CollectionClient collection(@NonNull String storeName) {
320356
* @param <T>
321357
* type of the bean in use
322358
*/
323-
public <T> CollectionRepository<T> collectionRepository(@NonNull String storeName, Class<T> clazz) {
324-
return nsClient.collectionRepository(storeName, clazz);
359+
public <T> AstraDBRepository<T> collectionRepository(@NonNull String storeName, Class<T> clazz) {
360+
return new AstraDBRepository<T>(nsClient.collectionRepository(storeName, clazz));
361+
}
362+
363+
/**
364+
* Access the low level Stargate Namespace resource operation.
365+
*
366+
* @return
367+
* raw namespace client
368+
*/
369+
public NamespaceClient getRawNamespaceClient() {
370+
return nsClient;
325371
}
326372

327373
}

‎astra-db-client/src/main/java/com/dtsx/astra/sdk/AstraDBClient.java‎

Lines changed: 2 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package com.dtsx.astra.sdk;
22

3-
import com.datastax.astra.sdk.AstraClient;
4-
import com.datastax.oss.driver.api.core.CqlSession;
53
import com.dtsx.astra.sdk.db.AstraDBOpsClient;
64
import com.dtsx.astra.sdk.db.DbOpsClient;
75
import com.dtsx.astra.sdk.db.domain.CloudProviderType;
@@ -295,7 +293,7 @@ public AstraDB database(@NonNull String databaseName) {
295293
* database client
296294
*/
297295
public AstraDB database(UUID databaseId) {
298-
return new AstraDB(token, databaseId, env);
296+
return new AstraDB(token, databaseId, null, env, AstraDBClient.DEFAULT_KEYSPACE);
299297
}
300298

301299
/**
@@ -386,55 +384,13 @@ public ApiClient getApiClient(@NonNull UUID databaseId) {
386384
return database(databaseId).getApiClient();
387385
}
388386

389-
Map<String, CqlSession> cqlSessions = new HashMap<>();
390-
391-
/**
392-
* Access the cqlSession for the database.
393-
*
394-
* @param databaseName
395-
* current database
396-
* @return
397-
* current cqlSession.
398-
*/
399-
public CqlSession getCqlSession(String databaseName) {
400-
// initialized if not present
401-
cqlSessions.computeIfAbsent(databaseName, k -> {
402-
List<Database> dbs = findDatabaseByName(databaseName).collect(Collectors.toList());
403-
if (dbs.isEmpty()) {
404-
throw new DatabaseNotFoundException(databaseName);
405-
}
406-
if (dbs.size()>1) {
407-
throw new IllegalStateException("More than one database exists with the same name, use its uuid.");
408-
}
409-
Database db = dbs.get(0);
410-
return AstraClient.builder()
411-
.env(env)
412-
.withDatabaseRegion(db.getInfo().getRegion())
413-
.withDatabaseId(db.getId())
414-
.disableCrossRegionFailOver()
415-
.enableCql()
416-
.enableDownloadSecureConnectBundle()
417-
.build().cqlSession();
418-
});
419-
return cqlSessions.get(databaseName);
420-
}
421-
422-
// create a runtime shutdown hook to close each session
423-
{
424-
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
425-
cqlSessions.values().stream()
426-
.filter(session -> !session.isClosed())
427-
.forEach(CqlSession::close);
428-
}));
429-
}
430-
431387
/**
432388
* Access the devops client.
433389
*
434390
* @return
435391
* devops client.
436392
*/
437-
public AstraDBOpsClient getAstraDBOps() {
393+
public AstraDBOpsClient getRawDevopsApiClient() {
438394
return this.devopsDbClient;
439395
}
440396

0 commit comments

Comments
 (0)