1
+ // Copyright (c) Microsoft. All rights reserved.
1
2
package com .microsoft .semantickernel .connectors .data .jdbc ;
2
3
3
4
import com .microsoft .semantickernel .data .recorddefinition .VectorStoreRecordDefinition ;
4
- import com . microsoft . semantickernel . exceptions . SKException ;
5
+ import edu . umd . cs . findbugs . annotations . SuppressFBWarnings ;
5
6
import reactor .core .publisher .Mono ;
7
+ import reactor .core .scheduler .Schedulers ;
6
8
7
9
import javax .annotation .Nonnull ;
8
10
import javax .annotation .Nullable ;
9
11
import java .sql .Connection ;
10
- import java .sql .ResultSet ;
11
- import java .sql .SQLException ;
12
- import java .util .ArrayList ;
13
12
import java .util .List ;
14
13
15
14
/**
@@ -27,14 +26,18 @@ public class JDBCVectorStore implements SQLVectorStore<JDBCVectorStoreRecordColl
27
26
* @param connection the connection
28
27
* @param options the options
29
28
*/
30
- public JDBCVectorStore (@ Nonnull Connection connection , @ Nullable JDBCVectorStoreOptions options ) {
29
+ @ SuppressFBWarnings ("EI_EXPOSE_REP2" )
30
+ public JDBCVectorStore (@ Nonnull Connection connection ,
31
+ @ Nullable JDBCVectorStoreOptions options ) {
31
32
this .connection = connection ;
32
33
this .options = options ;
33
34
34
35
if (this .options != null && this .options .getQueryProvider () != null ) {
35
36
this .queryProvider = this .options .getQueryProvider ();
36
37
} else {
37
- this .queryProvider = new JDBCVectorStoreDefaultQueryProvider (connection );
38
+ this .queryProvider = JDBCVectorStoreDefaultQueryProvider .builder ()
39
+ .withConnection (connection )
40
+ .build ();
38
41
}
39
42
}
40
43
@@ -57,30 +60,30 @@ public static Builder builder() {
57
60
*/
58
61
@ Override
59
62
public <Key , Record > JDBCVectorStoreRecordCollection <?> getCollection (
60
- @ Nonnull String collectionName ,
61
- @ Nonnull Class <Record > recordClass ,
62
- @ Nullable VectorStoreRecordDefinition recordDefinition ) {
63
+ @ Nonnull String collectionName ,
64
+ @ Nonnull Class <Record > recordClass ,
65
+ @ Nullable VectorStoreRecordDefinition recordDefinition ) {
63
66
64
67
if (this .options != null && this .options .getVectorStoreRecordCollectionFactory () != null ) {
65
68
return this .options .getVectorStoreRecordCollectionFactory ()
66
69
.createVectorStoreRecordCollection (
67
70
connection ,
68
71
collectionName ,
69
72
JDBCVectorStoreRecordCollectionOptions .<Record >builder ()
70
- .withRecordClass (recordClass )
71
- .withRecordDefinition (recordDefinition )
72
- .withQueryProvider (this .queryProvider )
73
- .build ());
74
- }
75
-
76
- return new JDBCVectorStoreRecordCollection <>(
77
- connection ,
78
- collectionName ,
79
- JDBCVectorStoreRecordCollectionOptions .<Record >builder ()
80
73
.withRecordClass (recordClass )
81
74
.withRecordDefinition (recordDefinition )
82
75
.withQueryProvider (this .queryProvider )
83
76
.build ());
77
+ }
78
+
79
+ return new JDBCVectorStoreRecordCollection <>(
80
+ connection ,
81
+ collectionName ,
82
+ JDBCVectorStoreRecordCollectionOptions .<Record >builder ()
83
+ .withRecordClass (recordClass )
84
+ .withRecordDefinition (recordDefinition )
85
+ .withQueryProvider (this .queryProvider )
86
+ .build ());
84
87
}
85
88
86
89
/**
@@ -90,30 +93,17 @@ public <Key, Record> JDBCVectorStoreRecordCollection<?> getCollection(
90
93
*/
91
94
@ Override
92
95
public Mono <List <String >> getCollectionNamesAsync () {
93
- return Mono .fromCallable (() -> {
94
- List <String > collectionNames = new ArrayList <>();
95
- try {
96
- ResultSet resultSet = queryProvider .getCollectionNames ();
97
- while (resultSet .next ()) {
98
- collectionNames .add (resultSet .getString (1 ));
99
- }
100
-
101
- return collectionNames ;
102
- } catch (SQLException e ) {
103
- throw new SKException ("Failed to get collection names." , e );
104
- }
105
- });
96
+ return Mono .fromCallable (queryProvider ::getCollectionNames )
97
+ .subscribeOn (Schedulers .boundedElastic ());
106
98
}
107
99
100
+ /**
101
+ * Prepares the vector store.
102
+ */
108
103
@ Override
109
104
public Mono <Void > prepareAsync () {
110
- return Mono .fromRunnable (() -> {
111
- try {
112
- queryProvider .prepareVectorStore ();
113
- } catch (SQLException e ) {
114
- throw new SKException ("Failed to prepare vector store." , e );
115
- }
116
- });
105
+ return Mono .fromRunnable (queryProvider ::prepareVectorStore )
106
+ .subscribeOn (Schedulers .boundedElastic ()).then ();
117
107
}
118
108
119
109
/**
@@ -129,6 +119,7 @@ public static class Builder {
129
119
* @param connection the connection
130
120
* @return the builder
131
121
*/
122
+ @ SuppressFBWarnings ("EI_EXPOSE_REP2" )
132
123
public Builder withConnection (Connection connection ) {
133
124
this .connection = connection ;
134
125
return this ;
0 commit comments