Skip to content

Commit 1a1d31a

Browse files
clunclun
committed
Update compose
1 parent 31901d3 commit 1a1d31a

File tree

21 files changed

+685
-158
lines changed

21 files changed

+685
-158
lines changed

‎astra-db-java/src/main/java/com/datastax/astra/client/Collection.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@
8484

8585
import static com.datastax.astra.client.exception.DataApiException.ERROR_CODE_INTERRUPTED;
8686
import static com.datastax.astra.client.exception.DataApiException.ERROR_CODE_TIMEOUT;
87+
import static com.datastax.astra.internal.http.RetryHttpClient.HEADER_EMBEDDING_SERVICE_API_KEY;
8788
import static com.datastax.astra.internal.utils.AnsiUtils.cyan;
8889
import static com.datastax.astra.internal.utils.AnsiUtils.green;
8990
import static com.datastax.astra.internal.utils.AnsiUtils.magenta;
@@ -949,6 +950,12 @@ private Callable<InsertManyResult> getInsertManyResultCallable(List<? extends T>
949950
Command insertMany = new Command("insertMany")
950951
.withDocuments(documents.subList(start, end))
951952
.withOptions(new Document().append(INPUT_ORDERED, options.isOrdered()));
953+
if (options.getMaxTimeMS() != null) {
954+
insertMany.setMaxTimeMS(options.getMaxTimeMS());
955+
}
956+
if (options.getEmbeddingServiceApiKey() != null) {
957+
insertMany.withHeader(HEADER_EMBEDDING_SERVICE_API_KEY, options.getEmbeddingServiceApiKey());
958+
}
952959
return new InsertManyResult(runCommand(insertMany)
953960
.getStatusKeyAsList(RESULT_INSERTED_IDS, Object.class)
954961
.stream()
@@ -1044,6 +1051,12 @@ public Optional<T> findOne(Filter filter, FindOneOptions options) {
10441051
.withProjection(options.getProjection())
10451052
.withOptions(new Document()
10461053
.appendIfNotNull(INPUT_INCLUDE_SIMILARITY, options.getIncludeSimilarity()));
1054+
if (options.getMaxTimeMS() != null) {
1055+
findOne.setMaxTimeMS(options.getMaxTimeMS());
1056+
}
1057+
if (options.getEmbeddingServiceApiKey() != null) {
1058+
findOne.withHeader(HEADER_EMBEDDING_SERVICE_API_KEY, options.getEmbeddingServiceApiKey());
1059+
}
10471060
return Optional.ofNullable(
10481061
runCommand(findOne)
10491062
.getData().getDocument()

‎astra-db-java/src/main/java/com/datastax/astra/client/DataAPIClients.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@
2020
* #L%
2121
*/
2222

23-
import com.datastax.astra.internal.command.LoggingCommandObserver;
2423
import com.datastax.astra.client.admin.DataAPIDatabaseAdmin;
25-
import com.datastax.astra.internal.auth.TokenProviderStargate;
24+
import com.datastax.astra.internal.auth.TokenProviderStargateV2;
25+
import com.datastax.astra.internal.command.LoggingCommandObserver;
2626

2727
import static com.datastax.astra.client.admin.AstraDBAdmin.DEFAULT_NAMESPACE;
2828

@@ -75,7 +75,7 @@ private DataAPIClients() {}
7575
*/
7676
public static DataAPIClient createForLocal() {
7777
return new DataAPIClient(
78-
new TokenProviderStargate().getToken(),
78+
new TokenProviderStargateV2().getToken(),
7979
DataAPIOptions.builder().withDestination(DataAPIOptions.DataAPIDestination.CASSANDRA).build());
8080
}
8181

‎astra-db-java/src/main/java/com/datastax/astra/client/model/CollectionOptions.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ public static class Service {
164164
private Authentication authentication;
165165

166166
/** Free form parameters. */
167-
private Map<String, Parameters> parameters;
167+
private Map<String, Object> parameters;
168168

169169
/** Default constructor. */
170170
public Service() {
@@ -372,6 +372,24 @@ public CollectionOptionsBuilder vectorize(String provider, String modeName) {
372372
return this;
373373
}
374374

375+
/**
376+
* Enable Vectorization within the collection.
377+
*
378+
* @param provider
379+
* provider Name (LLM)
380+
* @param modeName
381+
* mode name
382+
* @param parameters
383+
* expected parameters for vectorize
384+
* @return
385+
* self reference
386+
*/
387+
public CollectionOptionsBuilder vectorize(String provider, String modeName, Map<String, Object> parameters) {
388+
vectorize(provider, modeName);
389+
getVector().getService().setParameters(parameters);
390+
return this;
391+
}
392+
375393
/**
376394
* Build the output.
377395
*

‎astra-db-java/src/main/java/com/datastax/astra/client/model/Command.java

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@ public class Command implements Serializable {
4747
/** Command payload.*/
4848
protected Document payload = new Document();
4949

50+
/** Extra headers. */
51+
protected Map<String, String> headers;
52+
53+
/** Maximum response time in milliseconds. */
54+
protected Long maxTimeMS;
55+
5056
/**
5157
* Default constructor.
5258
*/
@@ -90,6 +96,37 @@ public Command withFilter(Filter filter) {
9096
return this;
9197
}
9298

99+
/**
100+
* Builder pattern, update filter.
101+
*
102+
* @param maxTimeMS
103+
* maximum response time in milliseconds
104+
* @return
105+
* self-reference
106+
*/
107+
public Command withMaxTimeMS(long maxTimeMS) {
108+
this.maxTimeMS = maxTimeMS;
109+
return this;
110+
}
111+
112+
/**
113+
* Builder pattern, update filter.
114+
*
115+
* @param key
116+
* name of the extra header
117+
* @param value
118+
* value of the extra header
119+
* @return
120+
* self-reference
121+
*/
122+
public Command withHeader(String key, String value) {
123+
if (headers == null) {
124+
headers = new LinkedHashMap<>();
125+
}
126+
headers.put(key, value);
127+
return this;
128+
}
129+
93130
/**
94131
* Builder pattern, update replacement document.
95132
*

‎astra-db-java/src/main/java/com/datastax/astra/client/model/FindOneOptions.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,16 @@ public class FindOneOptions {
4848
*/
4949
private Boolean includeSimilarity;
5050

51+
/**
52+
* When use some form of vectorization, the key to use for embedding
53+
*/
54+
private String embeddingServiceApiKey;
55+
56+
/**
57+
* Maximum amount of time to execute a query
58+
*/
59+
private Long maxTimeMS;
60+
5161
/**
5262
* Default constructor.
5363
*/
@@ -123,6 +133,32 @@ public FindOneOptions includeSimilarity() {
123133
return this;
124134
}
125135

136+
/**
137+
* Setter for timeout.
138+
*
139+
* @param maxTimeMS
140+
* timeout value
141+
* @return
142+
* insert many options
143+
*/
144+
public FindOneOptions maxTimeMS(long maxTimeMS) {
145+
this.maxTimeMS = maxTimeMS;
146+
return this;
147+
}
148+
149+
/**
150+
* Setter for timeout.
151+
*
152+
* @param apiKey
153+
* embedding service Api keu
154+
* @return
155+
* insert many options
156+
*/
157+
public FindOneOptions embeddingServiceApiKey(String apiKey) {
158+
this.embeddingServiceApiKey = apiKey;
159+
return this;
160+
}
161+
126162
/**
127163
* Builder for creating {@link FindOneAndUpdateOptions} instances with a fluent API.
128164
*/
@@ -186,5 +222,29 @@ public static FindOneOptions includeSimilarity() {
186222
return new FindOneOptions().includeSimilarity();
187223
}
188224

225+
/**
226+
* Setter maxTimeMS
227+
*
228+
* @param maxTimeMS
229+
* timeout value
230+
* @return
231+
* insert many options
232+
*/
233+
public static InsertManyOptions maxTimeMS(long maxTimeMS) {
234+
return new InsertManyOptions().maxTimeMS(maxTimeMS);
235+
}
236+
237+
/**
238+
* Setter for Api Key
239+
*
240+
* @param apiKey
241+
* embedding service Api keu
242+
* @return
243+
* insert many options
244+
*/
245+
public static InsertManyOptions embeddingServiceApiKey(String apiKey) {
246+
return new InsertManyOptions().embeddingServiceApiKey(apiKey);
247+
}
248+
189249
}
190250
}

‎astra-db-java/src/main/java/com/datastax/astra/client/model/InsertManyOptions.java

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,17 @@ public class InsertManyOptions {
5050
private int timeout = DataAPIOptions.DEFAULT_REQUEST_TIMEOUT_SECONDS * 1000;
5151

5252
/**
53-
* Populate inserMany options
53+
* When use some form of vectorization, the key to use for embedding
54+
*/
55+
private String embeddingServiceApiKey;
56+
57+
/**
58+
* Maximum amount of time to execute a query
59+
*/
60+
private Long maxTimeMS;
61+
62+
/**
63+
* Populate insertMany options
5464
*/
5565
public InsertManyOptions() {
5666
// left blank, jackson serialization
@@ -109,6 +119,32 @@ public InsertManyOptions timeout(int timeout) {
109119
return this;
110120
}
111121

122+
/**
123+
* Setter for timeout.
124+
*
125+
* @param maxTimeMS
126+
* timeout value
127+
* @return
128+
* insert many options
129+
*/
130+
public InsertManyOptions maxTimeMS(long maxTimeMS) {
131+
this.maxTimeMS = maxTimeMS;
132+
return this;
133+
}
134+
135+
/**
136+
* Setter for timeout.
137+
*
138+
* @param apiKey
139+
* embedding service Api keu
140+
* @return
141+
* insert many options
142+
*/
143+
public InsertManyOptions embeddingServiceApiKey(String apiKey) {
144+
this.embeddingServiceApiKey = apiKey;
145+
return this;
146+
}
147+
112148
/**
113149
* Builder for creating {@link FindOneAndUpdateOptions} instances with a fluent API.
114150
*/
@@ -161,6 +197,30 @@ public static InsertManyOptions timeout(int timeout) {
161197
return new InsertManyOptions().timeout(timeout);
162198
}
163199

200+
/**
201+
* Setter maxTimeMS
202+
*
203+
* @param maxTimeMS
204+
* timeout value
205+
* @return
206+
* insert many options
207+
*/
208+
public static InsertManyOptions maxTimeMS(long maxTimeMS) {
209+
return new InsertManyOptions().maxTimeMS(maxTimeMS);
210+
}
211+
212+
/**
213+
* Setter for Api Key
214+
*
215+
* @param apiKey
216+
* embedding service Api keu
217+
* @return
218+
* insert many options
219+
*/
220+
public static InsertManyOptions embeddingServiceApiKey(String apiKey) {
221+
return new InsertManyOptions().embeddingServiceApiKey(apiKey);
222+
}
223+
164224
}
165225

166226
}

‎astra-db-java/src/main/java/com/datastax/astra/internal/auth/TokenProviderStargateV2.java

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package com.datastax.astra.internal.auth;
2+
3+
/*-
4+
* #%L
5+
* Data API Java Client
6+
* --
7+
* Copyright (C) 2024 DataStax
8+
* --
9+
* Licensed under the Apache License, Version 2.0
10+
* You may not use this file except in compliance with the License.
11+
* You may obtain a copy of the License at
12+
*
13+
* http://www.apache.org/licenses/LICENSE-2.0
14+
*
15+
* Unless required by applicable law or agreed to in writing, software
16+
* distributed under the License is distributed on an "AS IS" BASIS,
17+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18+
* See the License for the specific language governing permissions and
19+
* limitations under the License.
20+
* #L%
21+
*/
22+
23+
import com.datastax.astra.internal.http.RetryHttpClient;
24+
import com.datastax.astra.internal.utils.Assert;
25+
26+
import java.util.Arrays;
27+
import java.util.Base64;
28+
29+
/**
30+
* Creating a token with base64 encoded credentials.
31+
*
32+
* @author Cedrick LUNVEN (@clunven)
33+
*/
34+
public class TokenProviderStargateV2 implements TokenProvider {
35+
36+
/** Simple Client. */
37+
public static final RetryHttpClient HTTP_CLIENT = new RetryHttpClient();
38+
39+
/** Default username for Cassandra. */
40+
public static final String DEFAULT_USERNAME = "cassandra";
41+
42+
/** Default password for Cassandra. */
43+
public static final String DEFAULT_CREDENTIALS = "cassandra";
44+
45+
/** Storing an authentication token to speed up queries. */
46+
private final String token;
47+
48+
/** Using defaults settings. */
49+
public TokenProviderStargateV2() {
50+
this(DEFAULT_USERNAME, DEFAULT_CREDENTIALS);
51+
}
52+
53+
/**
54+
* Full-fledged constructor.
55+
*
56+
* @param username
57+
* username
58+
* @param password
59+
* password
60+
*/
61+
public TokenProviderStargateV2(String username, String password) {
62+
Assert.hasLength(username, "username");
63+
Assert.hasLength(password, "password");
64+
Base64.Encoder encoder = Base64.getEncoder();
65+
this.token = "Cassandra:" +
66+
Base64.getEncoder().encodeToString(username.getBytes()) + ":" +
67+
Base64.getEncoder().encodeToString(username.getBytes());
68+
}
69+
70+
/**
71+
* Generate or renew authentication token.
72+
*
73+
* @return String
74+
*/
75+
@Override
76+
public String getToken() {
77+
return token;
78+
}
79+
80+
}

0 commit comments

Comments
 (0)