Skip to content

Commit 8538b4a

Browse files
clunclun
committed
Adding vector
1 parent e5d2280 commit 8538b4a

File tree

17 files changed

+319
-1089
lines changed

17 files changed

+319
-1089
lines changed

‎astra-sdk-devops/src/main/java/com/dtsx/astra/sdk/db/DatabaseClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ private void downloadSecureConnectBundle(Datacenter dc, String destination) {
208208
Utils.downloadFile(dc.getSecureBundleUrl(), destination);
209209
LOGGER.info("+ Downloading SCB to : {}", destination);
210210
} else {
211-
LOGGER.info("+ SCB {} already available.", destination);
211+
LOGGER.debug("+ SCB already available ({}) ", destination);
212212
}
213213
}
214214

‎astra-sdk-devops/src/main/java/com/dtsx/astra/sdk/streaming/domain/CdcDefinition.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public class CdcDefinition {
4848
/** Cdc. */
4949
private String codStatus;
5050

51-
@JsonFormat(shape = JsonFormat.Shape.STRING, pattern = "yyyy-MM-dd'T'HH:mm:ss'Z'")
51+
@JsonFormat(shape = JsonFormat.Shape.STRING, pattern = "yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'")
5252
private Date createdAt;
5353

5454
@JsonFormat(shape = JsonFormat.Shape.STRING, pattern = "yyyy-MM-dd'T'HH:mm:ss'Z'")

‎astra-sdk-devops/src/main/java/com/dtsx/astra/sdk/utils/TestUtils.java‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,17 @@
77
import com.dtsx.astra.sdk.db.domain.Database;
88
import com.dtsx.astra.sdk.db.domain.DatabaseCreationRequest;
99
import com.dtsx.astra.sdk.db.domain.DatabaseStatusType;
10+
import org.apache.hc.client5.http.classic.methods.HttpGet;
11+
import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
12+
import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse;
13+
import org.apache.hc.client5.http.impl.classic.HttpClients;
14+
import org.apache.hc.core5.http.HttpHeaders;
15+
import org.slf4j.Logger;
16+
import org.slf4j.LoggerFactory;
1017

18+
import java.io.IOException;
1119
import java.util.List;
20+
import java.util.Optional;
1221
import java.util.stream.Collectors;
1322

1423
/**
@@ -122,4 +131,85 @@ public static void terminateDatabaseByName(AstraDbClient devopsDbCli, String dbN
122131
}
123132
}
124133

134+
/**
135+
* Read Token for tests.
136+
*
137+
* @return
138+
* token for test or error
139+
*/
140+
public static String readToken() {
141+
String token = null;
142+
if (AstraRc.isDefaultConfigFileExists()) {
143+
token = new AstraRc()
144+
.getSectionKey(AstraRc.ASTRARC_DEFAULT, AstraRc.ASTRA_DB_APPLICATION_TOKEN)
145+
.orElse(null);
146+
}
147+
return Utils.readEnvVariable(AstraRc.ASTRA_DB_APPLICATION_TOKEN).orElse(token);
148+
}
149+
150+
/**
151+
* Database name.
152+
*
153+
* @param db
154+
* database name
155+
*/
156+
private static void resumeDb(Database db) {
157+
try(CloseableHttpClient httpClient = HttpClients.createDefault()) {
158+
HttpGet request = new HttpGet(ApiLocator
159+
.getApiRestEndpoint(db.getId(), db.getInfo().getRegion()) +
160+
"/v2/schemas/keyspace");
161+
request.setHeader(HttpHeaders.CONTENT_TYPE, "application/json");
162+
request.setHeader("X-Cassandra-Token", readToken());
163+
request.setHeader("Content-Type", "application/json");
164+
httpClient.execute(request).close();
165+
} catch (IOException e) {
166+
throw new IllegalStateException("Cannot resume DB", e);
167+
}
168+
}
169+
170+
/**
171+
* Logger for the class.
172+
*/
173+
static Logger logger = LoggerFactory.getLogger(TestUtils.class);
174+
175+
/**
176+
* Create DB if not exist
177+
*
178+
* @param dbName
179+
* database name
180+
* @param keyspace
181+
* keyspace name
182+
* @return
183+
* database client
184+
*/
185+
public static String setupDatabase(String dbName, String keyspace)
186+
throws InterruptedException {
187+
AstraDbClient astraDb = new AstraDbClient(readToken());
188+
Optional<Database> optDb = astraDb.findByName(dbName).findAny();
189+
String dbId = null;
190+
if (!optDb.isPresent()) {
191+
logger.info("Creating database '{}' as it does not exist, this operation takes about 90s, please wait... ", dbName);
192+
dbId = astraDb.create(DatabaseCreationRequest
193+
.builder().name(dbName)
194+
.keyspace(keyspace)
195+
.cloudRegion(TEST_REGION)
196+
.withVector()
197+
.build());
198+
} else {
199+
dbId = optDb.get().getId();
200+
DatabaseClient dbClient = astraDb.database(dbId);
201+
if (optDb.get().getStatus().equals(DatabaseStatusType.HIBERNATED)) {
202+
logger.info("Resume DB {} as HIBERNATED ", dbName);
203+
resumeDb(optDb.get());
204+
waitForDbStatus(dbClient, DatabaseStatusType.ACTIVE, 500);
205+
}
206+
if (!dbClient.keyspaces().findAll().contains(keyspace)) {
207+
dbClient.keyspaces().create(keyspace);
208+
}
209+
}
210+
TestUtils.waitForDbStatus(astraDb.database(dbId), DatabaseStatusType.ACTIVE, 500);
211+
return dbId;
212+
}
213+
214+
125215
}

‎astra-sdk-devops/src/test/java/com/dtsx/astra/sdk/db/DatabasesClientTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public void failInitializationsWithInvalidParamsTest() {
3434
@DisplayName("Create DB and test existence")
3535
public void shouldCreateServerlessDb() {
3636
String dbId;
37-
Optional<Database> optDb = getDatabasesClient().databaseByName(SDK_TEST_DB_NAME).find();
37+
Optional<Database> optDb = getDatabasesClient().findByName(SDK_TEST_DB_NAME).findFirst();
3838
if (!optDb.isPresent()) {
3939
dbId = getDatabasesClient().create(DatabaseCreationRequest
4040
.builder()
@@ -139,7 +139,7 @@ public void shouldFindDatabaseByNameTest() {
139139
@DisplayName("Find database by id")
140140
public void shouldFindDatabaseByIdTest() {
141141
// --> Getting a valid id
142-
Assertions.assertTrue(getDatabasesClient().databaseByName(SDK_TEST_DB_NAME).exist());
142+
Assertions.assertTrue(getDatabasesClient().findByName(SDK_TEST_DB_NAME).findFirst().isPresent());
143143
String dbId = getDatabasesClient().databaseByName(SDK_TEST_DB_NAME).get().getId();
144144
// <---
145145
Assertions.assertThrows(IllegalArgumentException.class, () -> getDatabasesClient().database(""));

‎astra-sdk-vector/pom.xml‎

Lines changed: 17 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -13,29 +13,18 @@
1313
</parent>
1414

1515
<properties>
16-
<openai-java.version>0.15.0</openai-java.version>
17-
<langchain4j.version>0.22.0</langchain4j.version>
18-
<maven.compiler.target>17</maven.compiler.target>
19-
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
16+
<openai-java.version>0.16.0</openai-java.version>
2017
</properties>
2118

2219
<dependencies>
23-
<dependency>
24-
<groupId>org.projectlombok</groupId>
25-
<artifactId>lombok</artifactId>
26-
<scope>provided</scope>
27-
</dependency>
2820

2921
<dependency>
30-
<groupId>dev.langchain4j</groupId>
31-
<artifactId>langchain4j</artifactId>
32-
<version>${langchain4j.version}</version>
22+
<groupId>org.slf4j</groupId>
23+
<artifactId>slf4j-api</artifactId>
3324
</dependency>
34-
3525
<dependency>
36-
<groupId>com.theokanning.openai-gpt3-java</groupId>
37-
<artifactId>service</artifactId>
38-
<version>${openai-java.version}</version>
26+
<groupId>ch.qos.logback</groupId>
27+
<artifactId>logback-classic</artifactId>
3928
</dependency>
4029

4130
<dependency>
@@ -50,28 +39,27 @@
5039
</exclusions>
5140
</dependency>
5241

42+
<dependency>
43+
<groupId>org.projectlombok</groupId>
44+
<artifactId>lombok</artifactId>
45+
<scope>provided</scope>
46+
</dependency>
47+
5348
<dependency>
5449
<groupId>com.datastax.oss</groupId>
5550
<artifactId>java-driver-query-builder</artifactId>
5651
</dependency>
5752

53+
<!-- Testing with OpenAI -->
5854
<dependency>
59-
<groupId>org.junit.jupiter</groupId>
60-
<artifactId>junit-jupiter-engine</artifactId>
55+
<groupId>com.theokanning.openai-gpt3-java</groupId>
56+
<artifactId>service</artifactId>
57+
<version>${openai-java.version}</version>
6158
<scope>test</scope>
6259
</dependency>
6360
<dependency>
64-
<groupId>org.slf4j</groupId>
65-
<artifactId>slf4j-api</artifactId>
66-
</dependency>
67-
<dependency>
68-
<groupId>ch.qos.logback</groupId>
69-
<artifactId>logback-classic</artifactId>
70-
</dependency>
71-
<dependency>
72-
<groupId>org.checkerframework</groupId>
73-
<artifactId>checker-qual</artifactId>
74-
<version>3.33.0</version>
61+
<groupId>org.junit.jupiter</groupId>
62+
<artifactId>junit-jupiter-engine</artifactId>
7563
<scope>test</scope>
7664
</dependency>
7765
</dependencies>

‎astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/cassio/MetadataVectorCassandraTable.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
99
import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
1010
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
11+
import lombok.Builder;
1112
import lombok.Data;
1213
import lombok.Getter;
1314
import lombok.extern.slf4j.Slf4j;
@@ -58,6 +59,7 @@ public MetadataVectorCassandraTable(CqlSession session, String keyspaceName, Str
5859
super(session, keyspaceName, tableName);
5960
this.vectorDimension = vectorDimension;
6061
this.similarityMetric = metric;
62+
createSchema();
6163
}
6264

6365
/**
@@ -208,7 +210,7 @@ public static class Record implements Serializable {
208210
* Default Constructor
209211
*/
210212
public Record() {
211-
this.rowId = Uuids.timeBased().toString();
213+
this(Uuids.timeBased().toString(), null);
212214
}
213215

214216
/**
@@ -217,12 +219,22 @@ public Record() {
217219
* @param vector current vector.
218220
*/
219221
public Record(List<Float> vector) {
220-
this();
222+
this(Uuids.timeBased().toString(), vector);
223+
}
224+
225+
/**
226+
* Create a record with a vector.
227+
* @param rowId identifier for the row
228+
* @param vector current vector.
229+
*/
230+
public Record(String rowId, List<Float> vector) {
231+
this.rowId = rowId;
221232
this.vector = vector;
222233
}
223234

224235
/**
225-
* Dynamic insert
236+
* Build insert statement dynamically.
237+
*
226238
* @param keyspaceName
227239
* keyspace name
228240
* @param tableName

0 commit comments

Comments
 (0)