microsoft/semantic-kernel
1
0
Fork 0
mirror of https://github.com/microsoft/semantic-kernel.git synced 2026-10-12 02:19:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Java: Add batch operations to SQL memory stores (#2960)

Browse Source 
### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

Fix https://github.com/microsoft/semantic-kernel/issues/2619

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [ ] The code builds clean without any errors or warnings
- [ ] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [ ] All unit tests pass, and I have added new tests where possible
- [ ] I didn't break anyone 😄

---------

Co-authored-by: Luigi Montoya <[email protected]>
Co-authored-by: joe-braley <[email protected]>
Co-authored-by: Luigi96 <[email protected]>
Co-authored-by: Mark Wallace <[email protected]>
Co-authored-by: John Oliver <[email protected]>
Co-authored-by: David Grieve <[email protected]>
This commit is contained in:
Milder Hernandez
2023-09-28 18:49:14 +01:00
committed by GitHub
co-authored by Luigi Montoya joe-braley Luigi96 Mark Wallace John Oliver David Grieve
parent 428dadb200
commit e972ee8b4c
8 changed files with 346 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
java/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStoreTest.java
+8 -8
View File
@@ -314,16 +314,16 @@ public class JDBCMemoryStoreTest {
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
// Act
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
@@ -655,27 +655,27 @@ public class JDBCMemoryStoreTest {
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
collectionNum += numCollections;
// Act
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(initialSize + numCollections, collections.size());
// Act
for (String collection : collections) {
this.db.deleteCollectionAsync(collection).block();
db.deleteCollectionAsync(collection).block();
}
// Assert
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(0, collections.size());
}
java/connectors/semantickernel-connectors-memory-jdbc/pom.xml
+1
View File
@@ -9,6 +9,7 @@
</parent>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<name>Semantic Kernel JDBC Memory Connector</name>
<dependencies>
<dependency>
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnector.java
+142 -40
View File
@@ -12,6 +12,7 @@ import java.sql.Statement;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import javax.annotation.Nullable;
import reactor.core.publisher.Mono;
@@ -81,8 +82,8 @@ public class JDBCConnector implements SQLConnector, Closeable {
"CREATE TABLE IF NOT EXISTS "
+ TABLE_NAME
+ " ("
+ "collection TEXT, "
+ "key TEXT, "
+ "collection TEXT NOT NULL, "
+ "key TEXT NOT NULL, "
+ "metadata TEXT, "
+ "embedding TEXT, "
+ "timestamp TEXT, "
@@ -135,41 +136,7 @@ public class JDBCConnector implements SQLConnector, Closeable {
.then();
}
public Mono<Void> updateAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp) {
return Mono.fromRunnable(
() -> {
String query =
"UPDATE "
+ TABLE_NAME
+ " SET metadata = ?, embedding = ?, timestamp = ?"
+ " WHERE collection = ?"
+ " AND key = ?";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, metadata != null ? metadata : "");
statement.setString(2, embedding != null ? embedding : "");
statement.setString(3, formatDatetime(timestamp));
statement.setString(4, collection);
statement.setString(5, key != null && !key.isEmpty() ? key : null);
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"UPDATE\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
}
public Mono<Void> insertOrIgnoreAsync(
public Mono<String> upsertAsync(
String collection,
String key,
String metadata,
@@ -178,7 +145,7 @@ public class JDBCConnector implements SQLConnector, Closeable {
return Mono.fromRunnable(
() -> {
String query =
"INSERT OR IGNORE INTO "
"INSERT OR REPLACE INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?)";
@@ -193,12 +160,53 @@ public class JDBCConnector implements SQLConnector, Closeable {
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"INSERT OR IGNORE INTO\" failed",
"\"INSERT OR REPLACE INTO\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
.thenReturn(key);
}
@Override
public Mono<Collection<String>> upsertBatchAsync(
String collection, Collection<DatabaseEntry> records) {
Collection<String> keys = new ArrayList<>();
return Mono.fromRunnable(
() -> {
String query =
"INSERT OR REPLACE INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?)";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
for (DatabaseEntry entry : records) {
statement.setString(1, collection);
statement.setString(2, entry.getKey());
statement.setString(
3,
entry.getMetadata() != null ? entry.getMetadata() : "");
statement.setString(
4,
entry.getEmbedding() != null
? entry.getEmbedding()
: "");
statement.setString(5, formatDatetime(entry.getTimestamp()));
statement.addBatch();
keys.add(entry.getKey());
}
statement.executeBatch();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"INSERT OR REPLACE INTO\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.thenReturn(keys);
}
public Mono<Boolean> doesCollectionExistsAsync(String collectionName) {
@@ -304,6 +312,77 @@ public class JDBCConnector implements SQLConnector, Closeable {
.subscribeOn(Schedulers.boundedElastic());
}
enum BatchOperation {
SELECT,
DELETE
}
protected String batchQuery(BatchOperation operation, Collection<String> keys) {
String queryPrefix;
switch (operation) {
case SELECT:
queryPrefix = "SELECT * FROM " + TABLE_NAME + " WHERE collection = ?";
break;
case DELETE:
queryPrefix = "DELETE FROM " + TABLE_NAME + " WHERE collection = ?";
break;
default:
throw new IllegalArgumentException("Invalid batch operation");
}
StringBuilder queryBuilder = new StringBuilder(queryPrefix);
queryBuilder.append(" AND key IN (");
// Add placeholders for each key
for (int i = 0; i < keys.size(); i++) {
queryBuilder.append("?");
if (i < keys.size() - 1) {
queryBuilder.append(",");
}
}
queryBuilder.append(")");
return queryBuilder.toString();
}
@Override
public Mono<Collection<DatabaseEntry>> readBatchAsync(
String collectionName, Collection<String> keys) {
return Mono.defer(
() -> {
Collection<DatabaseEntry> entries = new ArrayList<>();
String query = batchQuery(BatchOperation.SELECT, keys);
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
int index = 2;
for (String key : keys) {
statement.setString(index++, key);
}
ResultSet resultSet = statement.executeQuery();
while (resultSet.next()) {
String key = resultSet.getString("key");
String metadata = resultSet.getString("metadata");
String embedding = resultSet.getString("embedding");
String timestamp = resultSet.getString("timestamp");
ZonedDateTime zonedDateTime = parseDatetime(timestamp);
entries.add(
new DatabaseEntry(
key, metadata, embedding, zonedDateTime));
}
} catch (SQLException e) {
return Mono.error(
new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"SELECT * FROM\" failed",
e));
}
return Mono.just(entries);
})
.subscribeOn(Schedulers.boundedElastic());
}
public Mono<Void> deleteCollectionAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
@@ -356,6 +435,29 @@ public class JDBCConnector implements SQLConnector, Closeable {
.then();
}
public Mono<Void> deleteBatchAsync(String collectionName, Collection<String> keys) {
return Mono.fromRunnable(
() -> {
String query = batchQuery(BatchOperation.DELETE, keys);
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
int index = 2;
for (String key : keys) {
statement.setString(index++, key);
}
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"DELETE FROM\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
}
public Mono<Void> deleteEmptyAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStore.java
+51 -44
View File
@@ -12,7 +12,6 @@ import java.util.*;
import java.util.stream.Collectors;
import javax.annotation.CheckReturnValue;
import javax.annotation.Nonnull;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
@@ -67,25 +66,42 @@ public class JDBCMemoryStore implements SQLMemoryStore {
.then(internalUpsertAsync(collectionName, record));
}
protected DatabaseEntry memoryRecordToDatabaseEntry(MemoryRecord record) {
try {
return new DatabaseEntry(
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
protected MemoryRecord databaseEntryToMemoryRecord(DatabaseEntry entry, boolean withEmbedding) {
try {
Embedding embedding =
withEmbedding
? new ObjectMapper().readValue(entry.getEmbedding(), Embedding.class)
: Embedding.empty();
return MemoryRecord.fromJsonMetadata(
entry.getMetadata(), embedding, entry.getKey(), entry.getTimestamp());
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"Error deserializing database entry",
e);
}
}
private Mono<String> internalUpsertAsync(String collectionName, MemoryRecord record) {
try {
Mono<Void> update =
this.dbConnector.updateAsync(
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
Mono<Void> insert =
this.dbConnector.insertOrIgnoreAsync(
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
return update.then(insert).then(Mono.just(record.getMetadata().getId()));
return this.dbConnector.upsertAsync(
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
@@ -113,9 +129,11 @@ public class JDBCMemoryStore implements SQLMemoryStore {
sink.next(exists);
})
.then(
Flux.fromIterable(records)
.concatMap(record -> internalUpsertAsync(collectionName, record))
.collect(Collectors.toCollection(ArrayList::new)));
dbConnector.upsertBatchAsync(
collectionName,
records.stream()
.map(this::memoryRecordToDatabaseEntry)
.collect(Collectors.toList())));
}
@Override
@@ -126,22 +144,6 @@ public class JDBCMemoryStore implements SQLMemoryStore {
return this.internalGetAsync(collectionName, key, withEmbedding);
}
private MemoryRecord memoryRecordFromDatabaseEntry(DatabaseEntry entry, boolean withEmbedding) {
try {
Embedding embedding =
withEmbedding
? new ObjectMapper().readValue(entry.getEmbedding(), Embedding.class)
: Embedding.empty();
return MemoryRecord.fromJsonMetadata(
entry.getMetadata(), embedding, entry.getKey(), entry.getTimestamp());
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"Error deserializing database entry",
e);
}
}
private Mono<MemoryRecord> internalGetAsync(
String collectionName, String key, boolean withEmbedding) {
Objects.requireNonNull(collectionName);
@@ -157,7 +159,7 @@ public class JDBCMemoryStore implements SQLMemoryStore {
return entry.map(
databaseEntry ->
memoryRecordFromDatabaseEntry(
databaseEntryToMemoryRecord(
databaseEntry, withEmbedding));
});
}
@@ -169,9 +171,16 @@ public class JDBCMemoryStore implements SQLMemoryStore {
boolean withEmbeddings) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.concatMap(key -> internalGetAsync(collectionName, key, withEmbeddings))
.collect(Collectors.toCollection(ArrayList::new));
return this.dbConnector
.readBatchAsync(collectionName, keys)
.flatMap(
databaseEntries -> {
List<MemoryRecord> records = new ArrayList<>();
for (DatabaseEntry entry : databaseEntries) {
records.add(databaseEntryToMemoryRecord(entry, withEmbeddings));
}
return Mono.just(records);
});
}
@Override
@@ -186,9 +195,7 @@ public class JDBCMemoryStore implements SQLMemoryStore {
@Nonnull String collectionName, @Nonnull Collection<String> keys) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.concatMap(key -> this.dbConnector.deleteAsync(collectionName, key))
.then();
return this.dbConnector.deleteBatchAsync(collectionName, keys);
}
@Override
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnector.java
+29 -19
View File
@@ -2,6 +2,7 @@
package com.microsoft.semantickernel.connectors.memory.jdbc;
import java.time.ZonedDateTime;
import java.util.Collection;
import java.util.List;
import reactor.core.publisher.Mono;
@@ -23,24 +24,7 @@ public interface SQLConnector {
Mono<Void> createCollectionAsync(String collectionName);
/**
* Asynchronously updates an entry.
*
* @param collection The name of the collection containing the entry.
* @param key The key identifying the entry to update.
* @param metadata The metadata associated with the entry.
* @param embedding The embedding data associated with the entry.
* @param timestamp The timestamp of the update.
* @return A Mono representing the completion of the update operation.
*/
Mono<Void> updateAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp);
/**
* Asynchronously inserts or ignores an entry.
* Asynchronously inserts or updates an entry.
*
* @param collection The name of the collection to insert the entry into.
* @param key The key identifying the entry.
@@ -49,13 +33,21 @@ public interface SQLConnector {
* @param timestamp The timestamp of the insertion.
* @return A Mono representing the completion of the insertion operation.
*/
Mono<Void> insertOrIgnoreAsync(
Mono<String> upsertAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp);
/**
* Asynchronously inserts or updates a batch of entries.
*
* @param records
* @return A Mono emitting a list of keys identifying the entries.
*/
Mono<Collection<String>> upsertBatchAsync(String collection, Collection<DatabaseEntry> records);
/**
* Asynchronously checks if a collection exists.
*
@@ -88,6 +80,15 @@ public interface SQLConnector {
*/
Mono<DatabaseEntry> readAsync(String collectionName, String key);
/**
* Asynchronously reads a batch of entries within a collection.
*
* @param collectionName The name of the collection to read from.
* @param keys The keys identifying the entries to read.
* @return A Mono emitting a list of DatabaseEntry objects representing the entries.
*/
Mono<Collection<DatabaseEntry>> readBatchAsync(String collectionName, Collection<String> keys);
/**
* Asynchronously deletes a collection.
*
@@ -105,6 +106,15 @@ public interface SQLConnector {
*/
Mono<Void> deleteAsync(String collectionName, String key);
/**
* Asynchronously deletes a batch of entries within a collection.
*
* @param collectionName The name of the collection to delete the entries from.
* @param keys The keys identifying the entries to delete.
* @return A Mono representing the completion of the entry deletion operation.
*/
Mono<Void> deleteBatchAsync(String collectionName, Collection<String> keys);
/**
* Asynchronously deletes all empty entries within a collection.
*
java/connectors/semantickernel-connectors-memory-jdbc/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnectorTest.java
+35
View File
@@ -0,0 +1,35 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
import static org.junit.jupiter.api.Assertions.*;
import java.util.ArrayList;
import java.util.Collection;
import org.junit.jupiter.api.Test;
public class JDBCConnectorTest {
@Test
void testBatchQuery() {
Collection<String> keys = new ArrayList<>();
keys.add("key1");
keys.add("key2");
keys.add("key3");
JDBCConnector connector = new JDBCConnector(null);
String read = connector.batchQuery(JDBCConnector.BatchOperation.SELECT, keys);
String delete = connector.batchQuery(JDBCConnector.BatchOperation.DELETE, keys);
assertNotNull(read);
assertNotNull(delete);
assertEquals(
"SELECT * FROM "
+ JDBCConnector.TABLE_NAME
+ " WHERE collection = ? AND key IN (?,?,?)",
read);
assertEquals(
"DELETE FROM "
+ JDBCConnector.TABLE_NAME
+ " WHERE collection = ? AND key IN (?,?,?)",
delete);
}
}
java/connectors/semantickernel-connectors-memory-postgresql/src/main/java/com/microsoft/semantickernel/connectors/memory/postgresql/PostgreSQLConnector.java
+66 -9
View File
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.postgresql;
import com.microsoft.semantickernel.connectors.memory.jdbc.DatabaseEntry;
import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCConnector;
import com.microsoft.semantickernel.connectors.memory.jdbc.SQLConnector;
import com.microsoft.semantickernel.connectors.memory.jdbc.SQLConnectorException;
@@ -9,6 +10,8 @@ import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.Collection;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
@@ -32,8 +35,8 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
"CREATE TABLE IF NOT EXISTS "
+ TABLE_NAME
+ " ("
+ "collection TEXT, "
+ "key TEXT, "
+ "collection TEXT NOT NULL, "
+ "key TEXT NOT NULL, "
+ "metadata TEXT, "
+ "embedding TEXT, "
+ "timestamp TEXT, "
@@ -67,7 +70,7 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
}
@Override
public Mono<Void> insertOrIgnoreAsync(
public Mono<String> upsertAsync(
String collection,
String key,
String metadata,
@@ -79,15 +82,22 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
"INSERT INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?)"
+ " ON CONFLICT (collection, key) DO NOTHING";
+ " VALUES (?, ?, ?, ?, ?) ON CONFLICT (collection,"
+ " key) DO UPDATE SET metadata = ?, embedding = ?,"
+ " timestamp = ?";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
String metadataString = metadata != null ? metadata : "";
String embeddingString = embedding != null ? embedding : "";
String timestampString = formatDatetime(timestamp);
statement.setString(1, collection);
statement.setString(2, key);
statement.setString(3, metadata != null ? metadata : "");
statement.setString(4, embedding != null ? embedding : "");
statement.setString(5, formatDatetime(timestamp));
statement.setString(3, metadataString);
statement.setString(4, embeddingString);
statement.setString(5, timestampString);
statement.setString(6, metadataString);
statement.setString(7, embeddingString);
statement.setString(8, timestampString);
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
@@ -97,6 +107,53 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
.thenReturn(key);
}
@Override
public Mono<Collection<String>> upsertBatchAsync(
String collection, Collection<DatabaseEntry> records) {
Collection<String> keys = new ArrayList<>();
return Mono.fromRunnable(
() -> {
String query =
"INSERT INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?) ON CONFLICT (collection,"
+ " key) DO UPDATE SET metadata = ?, embedding = ?,"
+ " timestamp = ?";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
for (DatabaseEntry entry : records) {
String metadataString =
entry.getMetadata() != null ? entry.getMetadata() : "";
String embeddingString =
entry.getEmbedding() != null
? entry.getEmbedding()
: "";
String timestampString = formatDatetime(entry.getTimestamp());
statement.setString(1, collection);
statement.setString(2, entry.getKey());
statement.setString(3, metadataString);
statement.setString(4, embeddingString);
statement.setString(5, timestampString);
statement.setString(6, metadataString);
statement.setString(7, embeddingString);
statement.setString(8, timestampString);
statement.addBatch();
keys.add(entry.getKey());
}
statement.executeBatch();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"INSERT INTO\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.thenReturn(keys);
}
}
java/connectors/semantickernel-connectors-memory-sqlite/src/test/java/com/microsoft/semantickernel/connectors/memory/sqlite/SQLiteMemoryStoreTest.java
+14 -14
View File
@@ -308,16 +308,16 @@ public class SQLiteMemoryStoreTest {
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
// Act
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
@@ -404,7 +404,7 @@ public class SQLiteMemoryStoreTest {
String collection = setUpNearestMatches();
// Act
float threshold = -1f;
float threshold = -1;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(collection, compareEmbedding, topN, threshold, false)
.block();
@@ -626,16 +626,16 @@ public class SQLiteMemoryStoreTest {
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
List<String> keys = new ArrayList<>();
for (String key : db.upsertBatchAsync(collection, records).block()) {
keys.add(key);
}
Collection<String> keys = db.upsertBatchAsync(collection, records).block();
assertNotNull(keys);
// Act
db.removeBatchAsync(collection, keys).block();
// Assert
for (MemoryRecord result : db.getBatchAsync(collection, keys, true).block()) {
Collection<MemoryRecord> results = db.getBatchAsync(collection, keys, false).block();
assertNotNull(results);
for (MemoryRecord result : results) {
assertNull(result);
}
}
@@ -649,27 +649,27 @@ public class SQLiteMemoryStoreTest {
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
collectionNum += numCollections;
// Act
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(initialSize + numCollections, collections.size());
// Act
for (String collection : collections) {
this.db.deleteCollectionAsync(collection).block();
db.deleteCollectionAsync(collection).block();
}
// Assert
collections = this.db.getCollectionsAsync().block();
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(0, collections.size());
}