Skip to content

Commit 4d17c1f

Browse files
author
Milder Hernandez Cagua
committed
Add JDBC Vector Store
1 parent 21082b8 commit 4d17c1f

File tree

13 files changed

+1769
-0
lines changed

13 files changed

+1769
-0
lines changed

‎api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java‎

Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
package com.microsoft.semantickernel.tests.connectors.memory.jdbc;
2+
3+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollection;
4+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions;
5+
import com.microsoft.semantickernel.connectors.data.jdbc.MySQLVectorStoreQueryProvider;
6+
import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions;
7+
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
8+
import org.junit.jupiter.api.BeforeAll;
9+
import org.junit.jupiter.api.Test;
10+
import org.testcontainers.containers.MySQLContainer;
11+
import org.testcontainers.junit.jupiter.Container;
12+
import org.testcontainers.junit.jupiter.Testcontainers;
13+
14+
import javax.annotation.Nonnull;
15+
import java.sql.Connection;
16+
import java.sql.DriverManager;
17+
import java.sql.SQLException;
18+
import java.util.ArrayList;
19+
import java.util.Arrays;
20+
import java.util.List;
21+
22+
import static org.junit.jupiter.api.Assertions.assertEquals;
23+
import static org.junit.jupiter.api.Assertions.assertNotNull;
24+
import static org.junit.jupiter.api.Assertions.assertNull;
25+
26+
@Testcontainers
27+
public class JDBCVectorStoreRecordCollectionTest {
28+
@Container
29+
private static final MySQLContainer<?> CONTAINER = new MySQLContainer<>("mysql:5.7.34");
30+
private static final String MYSQL_USER = "test";
31+
private static final String MYSQL_PASSWORD = "test";
32+
private static Connection connection;
33+
@BeforeAll
34+
static void setup() throws SQLException {
35+
connection = DriverManager.getConnection(CONTAINER.getJdbcUrl(), MYSQL_USER, MYSQL_PASSWORD);
36+
}
37+
38+
private JDBCVectorStoreRecordCollection<Hotel> buildRecordCollection(@Nonnull String collectionName) {
39+
JDBCVectorStoreRecordCollection<Hotel> recordCollection = new JDBCVectorStoreRecordCollection<>(
40+
connection,
41+
collectionName,
42+
JDBCVectorStoreRecordCollectionOptions.<Hotel>builder()
43+
.withRecordClass(Hotel.class)
44+
.withQueryProvider(MySQLVectorStoreQueryProvider.builder()
45+
.withConnection(connection)
46+
.build())
47+
.build());
48+
49+
recordCollection.prepareAsync().block();
50+
recordCollection.createCollectionIfNotExistsAsync().block();
51+
return recordCollection;
52+
}
53+
54+
@Test
55+
public void buildRecordCollection() {
56+
assertNotNull(buildRecordCollection("buildTest"));
57+
}
58+
59+
private List<Hotel> getHotels() {
60+
return List.of(
61+
new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(1.0f, 2.0f, 3.0f), 4.0),
62+
new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(1.0f, 2.0f, 3.0f), 3.0),
63+
new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(1.0f, 2.0f, 3.0f), 5.0),
64+
new Hotel("id_4", "Hotel 4", 4, "Hotel 4 description", Arrays.asList(1.0f, 2.0f, 3.0f), 4.0),
65+
new Hotel("id_5", "Hotel 5", 5, "Hotel 5 description", Arrays.asList(1.0f, 2.0f, 3.0f), 5.0)
66+
);
67+
}
68+
69+
@Test
70+
public void upsertAndGetRecordAsync() {
71+
String collectionName = "upsertAndGetRecordAsync";
72+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
73+
74+
List<Hotel> hotels = getHotels();
75+
for (Hotel hotel : hotels) {
76+
recordStore.upsertAsync(hotel, null).block();
77+
}
78+
79+
for (Hotel hotel : hotels) {
80+
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), null).block();
81+
assertNotNull(retrievedHotel);
82+
assertEquals(hotel.getId(), retrievedHotel.getId());
83+
}
84+
}
85+
86+
@Test
87+
public void getBatchAsync() {
88+
String collectionName = "getBatchAsync";
89+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
90+
91+
List<Hotel> hotels = getHotels();
92+
for (Hotel hotel : hotels) {
93+
recordStore.upsertAsync(hotel, null).block();
94+
}
95+
96+
List<String> keys = new ArrayList<>();
97+
for (Hotel hotel : hotels) {
98+
keys.add(hotel.getId());
99+
}
100+
101+
List<Hotel> retrievedHotels = recordStore.getBatchAsync(keys, null).block();
102+
assertNotNull(retrievedHotels);
103+
assertEquals(hotels.size(), retrievedHotels.size());
104+
}
105+
106+
@Test
107+
public void upsertBatchAndGetBatchAsync() {
108+
String collectionName = "upsertBatchAndGetBatchAsync";
109+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
110+
111+
List<Hotel> hotels = getHotels();
112+
recordStore.upsertBatchAsync(hotels, null).block();
113+
114+
List<String> keys = new ArrayList<>();
115+
for (Hotel hotel : hotels) {
116+
keys.add(hotel.getId());
117+
}
118+
119+
List<Hotel> retrievedHotels = recordStore.getBatchAsync(keys, null).block();
120+
assertNotNull(retrievedHotels);
121+
assertEquals(hotels.size(), retrievedHotels.size());
122+
}
123+
124+
@Test
125+
public void insertAndReplaceAsync() {
126+
String collectionName = "insertAndReplaceAsync";
127+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
128+
129+
List<Hotel> hotels = getHotels();
130+
recordStore.upsertBatchAsync(hotels, null).block();
131+
recordStore.upsertBatchAsync(hotels, null).block();
132+
recordStore.upsertBatchAsync(hotels, null).block();
133+
134+
List<String> keys = new ArrayList<>();
135+
for (Hotel hotel : hotels) {
136+
keys.add(hotel.getId());
137+
}
138+
139+
List<Hotel> retrievedHotels = recordStore.getBatchAsync(keys, null).block();
140+
assertNotNull(retrievedHotels);
141+
assertEquals(hotels.size(), retrievedHotels.size());
142+
}
143+
144+
@Test
145+
public void deleteRecordAsync() {
146+
String collectionName = "deleteRecordAsync";
147+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
148+
149+
List<Hotel> hotels = getHotels();
150+
recordStore.upsertBatchAsync(hotels, null).block();
151+
152+
for (Hotel hotel : hotels) {
153+
recordStore.deleteAsync(hotel.getId(), null).block();
154+
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), null).block();
155+
assertNull(retrievedHotel);
156+
}
157+
}
158+
159+
@Test
160+
public void deleteBatchAsync() {
161+
String collectionName = "deleteBatchAsync";
162+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
163+
164+
List<Hotel> hotels = getHotels();
165+
recordStore.upsertBatchAsync(hotels, null).block();
166+
167+
List<String> keys = new ArrayList<>();
168+
for (Hotel hotel : hotels) {
169+
keys.add(hotel.getId());
170+
}
171+
172+
recordStore.deleteBatchAsync(keys, null).block();
173+
174+
for (String key : keys) {
175+
Hotel retrievedHotel = recordStore.getAsync(key, null).block();
176+
assertNull(retrievedHotel);
177+
}
178+
}
179+
180+
@Test
181+
public void getWithNoVectors() {
182+
String collectionName = "getWithNoVectors";
183+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
184+
185+
List<Hotel> hotels = getHotels();
186+
recordStore.upsertBatchAsync(hotels, null).block();
187+
188+
GetRecordOptions options = GetRecordOptions.builder()
189+
.includeVectors(false)
190+
.build();
191+
192+
for (Hotel hotel : hotels) {
193+
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), options).block();
194+
assertNotNull(retrievedHotel);
195+
assertEquals(hotel.getId(), retrievedHotel.getId());
196+
assertNull(retrievedHotel.getDescriptionEmbedding());
197+
}
198+
199+
options = GetRecordOptions.builder()
200+
.includeVectors(true)
201+
.build();
202+
203+
for (Hotel hotel : hotels) {
204+
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), options).block();
205+
assertNotNull(retrievedHotel);
206+
assertEquals(hotel.getId(), retrievedHotel.getId());
207+
assertNotNull(retrievedHotel.getDescriptionEmbedding());
208+
}
209+
}
210+
211+
@Test
212+
public void getBatchWithNoVectors() {
213+
String collectionName = "getBatchWithNoVectors";
214+
JDBCVectorStoreRecordCollection<Hotel> recordStore = buildRecordCollection(collectionName);
215+
216+
List<Hotel> hotels = getHotels();
217+
recordStore.upsertBatchAsync(hotels, null).block();
218+
219+
GetRecordOptions options = GetRecordOptions.builder()
220+
.includeVectors(false)
221+
.build();
222+
223+
List<String> keys = new ArrayList<>();
224+
for (Hotel hotel : hotels) {
225+
keys.add(hotel.getId());
226+
}
227+
228+
List<Hotel> retrievedHotels = recordStore.getBatchAsync(keys, options).block();
229+
assertNotNull(retrievedHotels);
230+
assertEquals(hotels.size(), retrievedHotels.size());
231+
232+
for (Hotel hotel : retrievedHotels) {
233+
assertNull(hotel.getDescriptionEmbedding());
234+
}
235+
236+
options = GetRecordOptions.builder()
237+
.includeVectors(true)
238+
.build();
239+
240+
retrievedHotels = recordStore.getBatchAsync(keys, options).block();
241+
assertNotNull(retrievedHotels);
242+
assertEquals(hotels.size(), retrievedHotels.size());
243+
244+
for (Hotel hotel : retrievedHotels) {
245+
assertNotNull(hotel.getDescriptionEmbedding());
246+
}
247+
}
248+
}

‎api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package com.microsoft.semantickernel.tests.connectors.memory.jdbc;
2+
3+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStore;
4+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreOptions;
5+
import com.microsoft.semantickernel.connectors.data.jdbc.MySQLVectorStoreQueryProvider;
6+
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
7+
import org.junit.jupiter.api.BeforeAll;
8+
import org.junit.jupiter.api.Test;
9+
import org.testcontainers.containers.MySQLContainer;
10+
import org.testcontainers.junit.jupiter.Container;
11+
import org.testcontainers.junit.jupiter.Testcontainers;
12+
13+
import java.sql.Connection;
14+
import java.sql.DriverManager;
15+
import java.sql.SQLException;
16+
import java.util.Arrays;
17+
import java.util.List;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertNotNull;
21+
import static org.junit.jupiter.api.Assertions.assertTrue;
22+
23+
@Testcontainers
24+
public class JDBCVectorStoreTest {
25+
@Container
26+
private static final MySQLContainer<?> CONTAINER = new MySQLContainer<>("mysql:5.7.34");
27+
private static final String MYSQL_USER = "test";
28+
private static final String MYSQL_PASSWORD = "test";
29+
private static Connection connection;
30+
31+
@BeforeAll
32+
static void setup() throws SQLException {
33+
connection = DriverManager.getConnection(CONTAINER.getJdbcUrl(), MYSQL_USER, MYSQL_PASSWORD);
34+
}
35+
36+
@Test
37+
public void getCollectionNamesAsync() {
38+
JDBCVectorStoreOptions options = JDBCVectorStoreOptions.builder()
39+
.withQueryProvider(MySQLVectorStoreQueryProvider.builder()
40+
.withConnection(connection)
41+
.build())
42+
.build();
43+
44+
JDBCVectorStore vectorStore = JDBCVectorStore.builder()
45+
.withConnection(connection)
46+
.withOptions(options)
47+
.build();
48+
49+
vectorStore.getCollectionNamesAsync().block();
50+
51+
List<String> collectionNames = Arrays.asList("collection1", "collection2", "collection3");
52+
53+
for (String collectionName : collectionNames) {
54+
vectorStore.getCollection(collectionName, Hotel.class, null).createCollectionAsync().block();
55+
}
56+
57+
List<String> retrievedCollectionNames = vectorStore.getCollectionNamesAsync().block();
58+
assertNotNull(retrievedCollectionNames);
59+
assertEquals(collectionNames.size(), retrievedCollectionNames.size());
60+
for (String collectionName : collectionNames) {
61+
assertTrue(retrievedCollectionNames.contains(collectionName));
62+
}
63+
}
64+
}

0 commit comments

Comments
 (0)