Skip to content

Commit 1ea06a4

Browse files
alex268alex268
committed
Added tests for Query spi
1 parent cf2be50 commit 1ea06a4

File tree

5 files changed

+349
-6
lines changed

5 files changed

+349
-6
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/context/YdbServiceLoader.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
package tech.ydb.jdbc.context;
22

3+
import java.sql.SQLException;
34
import java.util.ArrayList;
45
import java.util.List;
56
import java.util.ServiceLoader;
6-
import java.util.stream.Collectors;
77

88
import tech.ydb.core.Status;
99
import tech.ydb.jdbc.YdbStatement;
@@ -46,9 +46,18 @@ private static class ProxySpi implements YdbQueryExtentionService {
4646
}
4747

4848
@Override
49-
public QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) {
50-
List<QueryCall> proxed = spis.stream().map(spi -> newDataQuery(statement, query, yql))
51-
.collect(Collectors.toList());
49+
public QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) throws SQLException {
50+
List<QueryCall> proxed = new ArrayList<>();
51+
try {
52+
for (YdbQueryExtentionService spi: spis) {
53+
proxed.add(spi.newDataQuery(statement, query, yql));
54+
}
55+
} catch (SQLException | RuntimeException th) {
56+
for (QueryCall call: proxed) {
57+
call.onQueryResult(null, th);
58+
}
59+
throw th;
60+
}
5261

5362
return new QueryCall() {
5463
@Override

‎jdbc/src/main/java/tech/ydb/jdbc/spi/YdbQueryExtentionService.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package tech.ydb.jdbc.spi;
22

3+
import java.sql.SQLException;
4+
35
import tech.ydb.core.Status;
46
import tech.ydb.jdbc.YdbStatement;
57
import tech.ydb.jdbc.query.YdbQuery;
@@ -66,6 +68,7 @@ default void onQueryResult(Status status, Throwable th) {
6668
* @param query Internal query information
6769
* @param yql Prepared YQL query, might be different for different parameters
6870
* @return current query handler
71+
* @throws java.sql.SQLException if SPI rejected query execution
6972
*/
70-
QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql);
73+
QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) throws SQLException;
7174
}

‎jdbc/src/test/java/tech/ydb/jdbc/context/QuerySpiTestLoader.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package tech.ydb.jdbc.context;
2+
3+
import java.io.ByteArrayInputStream;
4+
import java.io.IOException;
5+
import java.io.InputStream;
6+
import java.net.URL;
7+
import java.net.URLClassLoader;
8+
import java.net.URLConnection;
9+
import java.net.URLStreamHandler;
10+
import java.util.Collections;
11+
import java.util.Enumeration;
12+
import java.util.stream.Collectors;
13+
import java.util.stream.Stream;
14+
15+
import tech.ydb.jdbc.spi.YdbQueryExtentionService;
16+
17+
/**
18+
*
19+
* @author Aleksandr Gorshenin
20+
*/
21+
public class QuerySpiTestLoader extends URLClassLoader {
22+
private final static Class<?> SPI = YdbQueryExtentionService.class;
23+
private Class<?>[] classes;
24+
25+
public QuerySpiTestLoader(ClassLoader prev, Class<?>... implementingClasses) {
26+
super(new URL[0], prev);
27+
this.classes = implementingClasses;
28+
}
29+
30+
@Override
31+
public Enumeration<URL> getResources(String name) throws IOException {
32+
if (name.equals("META-INF/services/" + SPI.getName())) {
33+
if (classes == null) {
34+
return Collections.emptyEnumeration();
35+
}
36+
URL url = new URL("mock", "junit", 1234, "/service", new URLStreamHandler() {
37+
@Override
38+
protected URLConnection openConnection(URL u) {
39+
return new URLConnection(u) {
40+
@Override
41+
public void connect() { }
42+
43+
@Override
44+
public InputStream getInputStream() throws IOException {
45+
return new ByteArrayInputStream(Stream.of(classes)
46+
.map(Class::getName)
47+
.collect(Collectors.joining("\n"))
48+
.getBytes());
49+
}
50+
};
51+
}
52+
});
53+
54+
return new Enumeration<URL>() {
55+
boolean hasNext = true;
56+
57+
@Override
58+
public boolean hasMoreElements() {
59+
return hasNext;
60+
}
61+
62+
@Override
63+
public URL nextElement() {
64+
hasNext = false;
65+
return url;
66+
}
67+
};
68+
}
69+
return super.getResources(name);
70+
}
71+
}

‎jdbc/src/test/java/tech/ydb/jdbc/context/YdbDriverQuerySpiTest.java‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
package tech.ydb.jdbc.context;
2+
3+
import java.sql.Connection;
4+
import java.sql.DriverManager;
5+
import java.sql.PreparedStatement;
6+
import java.sql.ResultSet;
7+
import java.sql.SQLException;
8+
import java.sql.Statement;
9+
import java.util.concurrent.ConcurrentLinkedQueue;
10+
import java.util.concurrent.atomic.AtomicLong;
11+
12+
import org.junit.jupiter.api.Assertions;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.extension.RegisterExtension;
15+
import org.junit.jupiter.params.ParameterizedTest;
16+
import org.junit.jupiter.params.provider.ValueSource;
17+
18+
import tech.ydb.core.Status;
19+
import tech.ydb.core.StatusCode;
20+
import tech.ydb.jdbc.YdbStatement;
21+
import tech.ydb.jdbc.impl.helper.ExceptionAssert;
22+
import tech.ydb.jdbc.impl.helper.JdbcUrlHelper;
23+
import tech.ydb.jdbc.query.YdbQuery;
24+
import tech.ydb.jdbc.spi.YdbQueryExtentionService;
25+
import tech.ydb.query.result.QueryStats;
26+
import tech.ydb.query.settings.ExecuteQuerySettings;
27+
import tech.ydb.query.settings.QueryStatsMode;
28+
import tech.ydb.table.query.stats.QueryStatsCollectionMode;
29+
import tech.ydb.table.settings.ExecuteDataQuerySettings;
30+
import tech.ydb.test.junit5.YdbHelperExtension;
31+
32+
/**
33+
*
34+
* @author Aleksandr Gorshenin
35+
*/
36+
public class YdbDriverQuerySpiTest {
37+
38+
@RegisterExtension
39+
private static final YdbHelperExtension ydb = new YdbHelperExtension();
40+
41+
private static final JdbcUrlHelper jdbcURL = new JdbcUrlHelper(ydb)
42+
.withArg("cacheConnectionsInDriver", "false");
43+
44+
@BeforeEach
45+
public void clean() {
46+
EmptiSpi.COUNT.set(0);
47+
FullStatsSpi.QUEUE.clear();
48+
}
49+
50+
@ParameterizedTest
51+
@ValueSource(strings = {"true", "false"})
52+
public void defaultUsageTest(String useQS) throws SQLException {
53+
ClassLoader prev = Thread.currentThread().getContextClassLoader();
54+
Thread.currentThread().setContextClassLoader(new QuerySpiTestLoader(prev, EmptiSpi.class));
55+
56+
try (Connection conn = DriverManager.getConnection(jdbcURL.withArg("useQueryService", useQS).build())) {
57+
Assertions.assertEquals(0, EmptiSpi.COUNT.get());
58+
59+
try (ResultSet rs = conn.createStatement().executeQuery("SELECT 1 + 2")) {
60+
Assertions.assertTrue(rs.next());
61+
Assertions.assertFalse(rs.next());
62+
}
63+
64+
Assertions.assertEquals(1, EmptiSpi.COUNT.get());
65+
66+
try (PreparedStatement ps = conn.prepareStatement("SELECT ? + ?")) {
67+
ps.setInt(1, 1);
68+
ps.setInt(2, 2);
69+
Assertions.assertTrue(ps.execute());
70+
71+
Assertions.assertEquals(2, EmptiSpi.COUNT.get());
72+
73+
ps.setInt(1, 2);
74+
ps.setInt(2, 3);
75+
ps.addBatch();
76+
ps.setLong(1, 2);
77+
ps.setLong(2, 3);
78+
ps.addBatch();
79+
80+
Assertions.assertEquals(2, ps.executeBatch().length);
81+
Assertions.assertEquals(4, EmptiSpi.COUNT.get());
82+
}
83+
84+
try (Statement st = conn.createStatement()) {
85+
ExceptionAssert.ydbException("code = GENERIC_ERROR", () -> st.executeQuery("SELECT 1 + 'test'u"));
86+
Assertions.assertEquals(5, EmptiSpi.COUNT.get());
87+
}
88+
} finally {
89+
Thread.currentThread().setContextClassLoader(prev);
90+
}
91+
}
92+
93+
@ParameterizedTest
94+
@ValueSource(strings = {"true", "false"})
95+
public void enableStatsModeTest(String useQS) throws SQLException {
96+
ClassLoader prev = Thread.currentThread().getContextClassLoader();
97+
Thread.currentThread().setContextClassLoader(new QuerySpiTestLoader(prev, FullStatsSpi.class));
98+
99+
try (Connection conn = DriverManager.getConnection(jdbcURL.withArg("useQueryService", useQS).build())) {
100+
Assertions.assertTrue(FullStatsSpi.QUEUE.isEmpty());
101+
102+
try (ResultSet rs = conn.createStatement().executeQuery("SELECT 1 + 2")) {
103+
Assertions.assertTrue(rs.next());
104+
Assertions.assertFalse(rs.next());
105+
}
106+
107+
Assertions.assertFalse(FullStatsSpi.QUEUE.isEmpty());
108+
Assertions.assertNotNull(FullStatsSpi.QUEUE.poll().stats);
109+
Assertions.assertTrue(FullStatsSpi.QUEUE.isEmpty());
110+
111+
try (PreparedStatement ps = conn.prepareStatement("SELECT ? + ?")) {
112+
ps.setInt(1, 1);
113+
ps.setInt(2, 2);
114+
Assertions.assertTrue(ps.execute());
115+
116+
Assertions.assertEquals(1, FullStatsSpi.QUEUE.size());
117+
FullStatsSpi.Record record = FullStatsSpi.QUEUE.poll();
118+
Assertions.assertNotNull(record.stats);
119+
Assertions.assertEquals(Status.SUCCESS, record.status);
120+
Assertions.assertNull(record.th);
121+
122+
ps.setInt(1, 2);
123+
ps.setInt(2, 3);
124+
ps.addBatch();
125+
ps.setLong(1, 2);
126+
ps.setLong(2, 3);
127+
ps.addBatch();
128+
129+
Assertions.assertEquals(2, ps.executeBatch().length);
130+
131+
Assertions.assertEquals(2, FullStatsSpi.QUEUE.size());
132+
133+
record = FullStatsSpi.QUEUE.poll();
134+
Assertions.assertNotNull(record.stats);
135+
Assertions.assertEquals(Status.SUCCESS, record.status);
136+
Assertions.assertNull(record.th);
137+
138+
record = FullStatsSpi.QUEUE.poll();
139+
Assertions.assertNotNull(record.stats);
140+
Assertions.assertEquals(Status.SUCCESS, record.status);
141+
Assertions.assertNull(record.th);
142+
}
143+
144+
try (Statement st = conn.createStatement()) {
145+
ExceptionAssert.ydbException("code = GENERIC_ERROR", () -> st.executeQuery("SELECT 1 + 'test'u"));
146+
147+
Assertions.assertEquals(1, FullStatsSpi.QUEUE.size());
148+
FullStatsSpi.Record record = FullStatsSpi.QUEUE.poll();
149+
Assertions.assertNull(record.stats);
150+
Assertions.assertEquals(StatusCode.GENERIC_ERROR, record.status.getCode());
151+
Assertions.assertNull(record.th);
152+
}
153+
} finally {
154+
Thread.currentThread().setContextClassLoader(prev);
155+
}
156+
}
157+
158+
@ParameterizedTest
159+
@ValueSource(strings = {"true", "false"})
160+
public void validateQueryTest(String useQS) throws SQLException {
161+
ClassLoader prev = Thread.currentThread().getContextClassLoader();
162+
Thread.currentThread().setContextClassLoader(new QuerySpiTestLoader(prev, FullStatsSpi.class, ValidateSpi.class,
163+
EmptiSpi.class));
164+
165+
try (Connection conn = DriverManager.getConnection(jdbcURL.withArg("useQueryService", useQS).build())) {
166+
Assertions.assertTrue(FullStatsSpi.QUEUE.isEmpty());
167+
Assertions.assertEquals(0, EmptiSpi.COUNT.get());
168+
169+
try (ResultSet rs = conn.createStatement().executeQuery("SELECT 1 + 2")) {
170+
Assertions.assertTrue(rs.next());
171+
Assertions.assertFalse(rs.next());
172+
}
173+
174+
Assertions.assertEquals(1, FullStatsSpi.QUEUE.size());
175+
FullStatsSpi.Record record = FullStatsSpi.QUEUE.poll();
176+
Assertions.assertNotNull(record.stats);
177+
Assertions.assertEquals(Status.SUCCESS, record.status);
178+
Assertions.assertNull(record.th);
179+
Assertions.assertEquals(1, EmptiSpi.COUNT.get());
180+
181+
try (Statement st = conn.createStatement()) {
182+
ExceptionAssert.sqlException("INVALID QUERY", () -> st.executeQuery("SELECT 2 + 3"));
183+
}
184+
185+
Assertions.assertEquals(1, FullStatsSpi.QUEUE.size());
186+
record = FullStatsSpi.QUEUE.poll();
187+
Assertions.assertNull(record.stats);
188+
Assertions.assertNull(record.status);
189+
Assertions.assertNotNull(record.th);
190+
Assertions.assertEquals("INVALID QUERY", record.th.getMessage());
191+
192+
Assertions.assertEquals(1, EmptiSpi.COUNT.get());
193+
} finally {
194+
Thread.currentThread().setContextClassLoader(prev);
195+
}
196+
}
197+
198+
public static class EmptiSpi implements YdbQueryExtentionService {
199+
private static final AtomicLong COUNT = new AtomicLong(0);
200+
201+
@Override
202+
public QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) {
203+
COUNT.incrementAndGet();
204+
return new YdbQueryExtentionService.QueryCall() {
205+
};
206+
}
207+
}
208+
209+
public static class ValidateSpi implements YdbQueryExtentionService {
210+
211+
@Override
212+
public QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) throws SQLException {
213+
if ("SELECT 2 + 3".equals(query.getOriginQuery())) {
214+
throw new SQLException("INVALID QUERY");
215+
}
216+
217+
return new YdbQueryExtentionService.QueryCall() {
218+
};
219+
}
220+
}
221+
222+
public static class FullStatsSpi implements YdbQueryExtentionService {
223+
private static final ConcurrentLinkedQueue<Record> QUEUE = new ConcurrentLinkedQueue<>();
224+
225+
@Override
226+
public QueryCall newDataQuery(YdbStatement statement, YdbQuery query, String yql) {
227+
Record r = new Record();
228+
QUEUE.add(r);
229+
return r;
230+
}
231+
232+
static class Record implements YdbQueryExtentionService.QueryCall {
233+
234+
private QueryStats stats = null;
235+
private Status status = null;
236+
private Throwable th = null;
237+
238+
@Override
239+
public ExecuteQuerySettings.Builder prepareQuerySettings(ExecuteQuerySettings.Builder builder) {
240+
return builder.withStatsMode(QueryStatsMode.FULL);
241+
}
242+
243+
@Override
244+
public ExecuteDataQuerySettings prepareDataQuerySettings(ExecuteDataQuerySettings settings) {
245+
return settings.setCollectStats(QueryStatsCollectionMode.FULL);
246+
}
247+
248+
@Override
249+
public void onQueryStats(QueryStats stats) {
250+
this.stats = stats;
251+
}
252+
253+
@Override
254+
public void onQueryResult(Status status, Throwable th) {
255+
this.status = status;
256+
this.th = th;
257+
}
258+
}
259+
}
260+
}

‎jdbc/src/test/java/tech/ydb/jdbc/impl/helper/ExceptionAssert.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ public static void ydbException(String message, Executable exec) {
2222
"Invalid statement must throw YdbSQLException"
2323
);
2424
Assertions.assertTrue(ex.getMessage().contains(message),
25-
"YdbNonRetryableException '" + ex.getMessage() + "' doesn't contain message '" + message + "'");
25+
"YdbSQLException '" + ex.getMessage() + "' doesn't contain message '" + message + "'");
2626
}
2727

2828
public static void sqlDataException(String message, Executable exec) {

0 commit comments

Comments
 (0)