Skip to content

Commit cf2be50

Browse files
alex268alex268
committed
Update implementation of SPI extention
1 parent 08ff4ab commit cf2be50

12 files changed

+246
-392
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/YdbStatement.java‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import java.sql.Statement;
55

66
import tech.ydb.jdbc.context.YdbValidator;
7-
import tech.ydb.table.query.stats.QueryStatsCollectionMode;
87

98
public interface YdbStatement extends Statement {
109
/**
@@ -59,8 +58,4 @@ public interface YdbStatement extends Statement {
5958

6059
@Override
6160
int getMaxRows();
62-
63-
void setStatsCollectionMode(QueryStatsCollectionMode mode);
64-
65-
QueryStatsCollectionMode getStatsCollectionMode();
6661
}

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

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,15 @@
1515
import tech.ydb.jdbc.YdbQueryResult;
1616
import tech.ydb.jdbc.YdbStatement;
1717
import tech.ydb.jdbc.YdbTracer;
18+
import tech.ydb.jdbc.exception.YdbStatusable;
1819
import tech.ydb.jdbc.impl.YdbQueryResultExplain;
1920
import tech.ydb.jdbc.impl.YdbQueryResultReader;
2021
import tech.ydb.jdbc.impl.YdbQueryResultStatic;
2122
import tech.ydb.jdbc.impl.YdbResultSetMemory;
2223
import tech.ydb.jdbc.query.QueryType;
2324
import tech.ydb.jdbc.query.YdbQuery;
2425
import tech.ydb.jdbc.settings.YdbOperationProperties;
26+
import tech.ydb.jdbc.spi.YdbQueryExtentionService;
2527
import tech.ydb.query.QueryClient;
2628
import tech.ydb.query.QuerySession;
2729
import tech.ydb.query.QueryStream;
@@ -31,7 +33,6 @@
3133
import tech.ydb.query.settings.CommitTransactionSettings;
3234
import tech.ydb.query.settings.ExecuteQuerySettings;
3335
import tech.ydb.query.settings.QueryExecMode;
34-
import tech.ydb.query.settings.QueryStatsMode;
3536
import tech.ydb.query.settings.RollbackTransactionSettings;
3637
import tech.ydb.query.tools.QueryReader;
3738
import tech.ydb.table.query.Params;
@@ -44,6 +45,7 @@ public class QueryServiceExecutor extends BaseYdbExecutor {
4445
private final Duration sessionTimeout;
4546
private final QueryClient queryClient;
4647
private final boolean useStreamResultSet;
48+
private final YdbQueryExtentionService querySpi;
4749

4850
private int transactionLevel;
4951
private boolean isReadOnly;
@@ -59,6 +61,7 @@ public QueryServiceExecutor(YdbContext ctx) throws SQLException {
5961
this.sessionTimeout = options.getSessionTimeout();
6062
this.queryClient = ctx.getQueryClient();
6163
this.useStreamResultSet = options.getUseStreamResultSets();
64+
this.querySpi = ctx.getQuerySpi();
6265

6366
this.transactionLevel = options.getTransactionLevel();
6467
this.isAutoCommit = options.isAutoCommit();
@@ -225,17 +228,21 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException
225228
}
226229

227230
@Override
228-
public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, String preparedYql, Params params)
229-
throws SQLException {
231+
public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, YdbQuery query, String preparedYql,
232+
Params params) throws SQLException {
230233
ensureOpened();
231234

232235
YdbValidator validator = statement.getValidator();
236+
237+
String yql = prefixPragma + preparedYql;
238+
YdbQueryExtentionService.QueryCall spi = querySpi.newDataQuery(statement, query, yql);
239+
233240
int timeout = statement.getQueryTimeout();
234241
ExecuteQuerySettings.Builder settings = ExecuteQuerySettings.newBuilder();
235-
settings = settings.withStatsMode(QueryStatsMode.valueOf(statement.getStatsCollectionMode().name()));
236242
if (timeout > 0) {
237243
settings = settings.withRequestTimeout(timeout, TimeUnit.SECONDS);
238244
}
245+
settings = spi.prepareQuerySettings(settings);
239246

240247
QueryTransaction nextTx = tx.get();
241248
while (nextTx == null) {
@@ -246,11 +253,9 @@ public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, String
246253
}
247254
}
248255

249-
final QueryTransaction localTx = nextTx;
256+
QueryTransaction localTx = nextTx;
250257
YdbTracer tracer = statement.getConnection().getCtx().getTracer();
251258

252-
String yql = prefixPragma + preparedYql;
253-
254259
try {
255260
tracer.trace("--> data query");
256261
tracer.query(yql);
@@ -266,8 +271,19 @@ public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, String
266271
readers[idx] = new YdbResultSetMemory(types, statement, result.getResultSet(idx));
267272
}
268273

269-
// queryResult.setQueryStats(result.getQueryInfo().getStats());
274+
if (result.getQueryInfo().hasStats()) {
275+
spi.onQueryStats(result.getQueryInfo().getStats());
276+
}
277+
278+
spi.onQueryResult(Status.SUCCESS, null);
270279
return readers;
280+
} catch (SQLException | RuntimeException ex) {
281+
if (ex instanceof YdbStatusable) {
282+
spi.onQueryResult(((YdbStatusable) ex).getStatus(), null);
283+
} else {
284+
spi.onQueryResult(null, ex);
285+
}
286+
throw ex;
271287
} finally {
272288
if (!localTx.isActive()) {
273289
if (tx.compareAndSet(localTx, null)) {
@@ -289,16 +305,20 @@ public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, S
289305
ensureOpened();
290306

291307
if (!useStreamResultSet) {
292-
YdbResultSetMemory[] readers = executeInMemoryQuery(statement, preparedYql, params);
308+
YdbResultSetMemory[] readers = executeInMemoryQuery(statement, query, preparedYql, params);
293309
return updateCurrentResult(new YdbQueryResultStatic(query, readers));
294310
}
295311

296312
YdbValidator validator = statement.getValidator();
313+
String yql = prefixPragma + preparedYql;
314+
YdbQueryExtentionService.QueryCall spi = querySpi.newDataQuery(statement, query, yql);
315+
297316
int timeout = statement.getQueryTimeout();
298317
ExecuteQuerySettings.Builder settings = ExecuteQuerySettings.newBuilder();
299318
if (timeout > 0) {
300319
settings = settings.withRequestTimeout(timeout, TimeUnit.SECONDS);
301320
}
321+
settings = spi.prepareQuerySettings(settings);
302322

303323
QueryTransaction nextTx = tx.get();
304324
while (nextTx == null) {
@@ -309,18 +329,16 @@ public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, S
309329
}
310330
}
311331

312-
final QueryTransaction localTx = nextTx;
332+
QueryTransaction localTx = nextTx;
313333
YdbTracer tracer = statement.getConnection().getCtx().getTracer();
314-
315-
String yql = prefixPragma + preparedYql;
316-
317334
tracer.trace("--> stream query");
318335
tracer.query(yql);
319336
String msg = "STREAM_QUERY >>\n" + yql;
320337

321338
YdbQueryResultReader reader = new YdbQueryResultReader(types, statement, query) {
322339
@Override
323340
public void onClose(Status status, Throwable th) {
341+
spi.onQueryResult(status, th);
324342
if (th != null) {
325343
tracer.trace("<-- " + th.getMessage());
326344
}
@@ -344,7 +362,7 @@ public void onClose(Status status, Throwable th) {
344362

345363
settings = settings.withGrpcFlowControl(reader);
346364
QueryStream stream = localTx.createQuery(yql, isAutoCommit, params, settings.build());
347-
validator.execute(msg, tracer, () -> reader.load(validator, stream));
365+
validator.execute(msg, tracer, () -> reader.load(validator, stream, spi::onQueryStats));
348366
return updateCurrentResult(reader);
349367
}
350368

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

Lines changed: 35 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,20 @@
55
import java.sql.SQLFeatureNotSupportedException;
66
import java.time.Duration;
77

8+
import tech.ydb.core.Status;
89
import tech.ydb.jdbc.YdbConst;
910
import tech.ydb.jdbc.YdbQueryResult;
1011
import tech.ydb.jdbc.YdbStatement;
1112
import tech.ydb.jdbc.YdbTracer;
13+
import tech.ydb.jdbc.exception.YdbStatusable;
1214
import tech.ydb.jdbc.impl.YdbQueryResultExplain;
1315
import tech.ydb.jdbc.impl.YdbQueryResultStatic;
1416
import tech.ydb.jdbc.impl.YdbResultSetMemory;
1517
import tech.ydb.jdbc.query.QueryType;
1618
import tech.ydb.jdbc.query.YdbQuery;
1719
import tech.ydb.jdbc.settings.YdbOperationProperties;
20+
import tech.ydb.jdbc.spi.YdbQueryExtentionService;
21+
import tech.ydb.query.result.QueryStats;
1822
import tech.ydb.table.Session;
1923
import tech.ydb.table.query.DataQueryResult;
2024
import tech.ydb.table.query.ExplainDataQueryResult;
@@ -33,13 +37,15 @@
3337
*/
3438
public class TableServiceExecutor extends BaseYdbExecutor {
3539
private final boolean failOnTruncatedResult;
40+
private final YdbQueryExtentionService querySpi;
3641
private volatile TxState tx;
3742

3843
public TableServiceExecutor(YdbContext ctx) throws SQLException {
3944
super(ctx);
4045
YdbOperationProperties options = ctx.getOperationProperties();
4146
this.tx = createTx(options.getTransactionLevel(), options.isAutoCommit());
4247
this.failOnTruncatedResult = options.isFailOnTruncatedResult();
48+
this.querySpi = ctx.getQuerySpi();
4349
}
4450

4551
@Override
@@ -172,8 +178,6 @@ private ExecuteDataQuerySettings dataQuerySettings(YdbStatement statement) {
172178
settings = settings.disableQueryCache();
173179
}
174180

175-
settings = settings.setCollectStats(statement.getStatsCollectionMode());
176-
177181
return settings;
178182
}
179183

@@ -202,8 +206,18 @@ public YdbQueryResult executeExplainQuery(YdbStatement statement, YdbQuery query
202206
}
203207
}
204208

205-
private DataQueryResult executeTableQuery(YdbValidator validator, YdbTracer tracer, String yql,
206-
ExecuteDataQuerySettings settings, Params prms) throws SQLException {
209+
@Override
210+
public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, YdbQuery query, String preparedYql,
211+
Params params) throws SQLException {
212+
ensureOpened();
213+
214+
YdbValidator validator = statement.getValidator();
215+
String yql = prefixPragma + preparedYql;
216+
YdbQueryExtentionService.QueryCall spi = querySpi.newDataQuery(statement, query, yql);
217+
218+
YdbTracer tracer = statement.getConnection().getCtx().getTracer();
219+
ExecuteDataQuerySettings settings = spi.prepareDataQuerySettings(dataQuerySettings(statement));
220+
207221
Session session = tx.getSession(validator);
208222
try {
209223
tracer.trace("--> data query");
@@ -212,10 +226,14 @@ private DataQueryResult executeTableQuery(YdbValidator validator, YdbTracer trac
212226
DataQueryResult result = validator.call(
213227
QueryType.DATA_QUERY + " >>\n" + yql,
214228
tracer,
215-
() -> session.executeDataQuery(yql, tx.txControl(), prms, settings)
229+
() -> session.executeDataQuery(yql, tx.txControl(), params, settings)
216230
);
217231
updateState(tx.withDataQuery(session, result.getTxId()));
218232

233+
if (result.hasQueryStats()) {
234+
spi.onQueryStats(new QueryStats(result.getRawQueryStats()));
235+
}
236+
219237
if (failOnTruncatedResult) {
220238
for (int idx = 0; idx < result.getResultSetCount(); idx += 1) {
221239
ResultSetReader rs = result.getResultSet(idx);
@@ -226,30 +244,22 @@ private DataQueryResult executeTableQuery(YdbValidator validator, YdbTracer trac
226244
}
227245
}
228246

229-
return result;
230-
} catch (SQLException | RuntimeException ex) {
231-
updateState(tx.withRollback(session));
232-
throw ex;
233-
}
234-
}
235-
236-
@Override
237-
public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, String preparedYql, Params params)
238-
throws SQLException {
239-
ensureOpened();
240-
241-
YdbValidator validator = statement.getValidator();
242-
String yql = prefixPragma + preparedYql;
243-
YdbTracer tracer = statement.getConnection().getCtx().getTracer();
244-
245-
try {
246-
DataQueryResult result = executeTableQuery(validator, tracer, yql, dataQuerySettings(statement), params);
247247
YdbResultSetMemory[] readers = new YdbResultSetMemory[result.getResultSetCount()];
248248
for (int idx = 0; idx < result.getResultSetCount(); idx += 1) {
249249
readers[idx] = new YdbResultSetMemory(types, statement, result.getResultSet(idx));
250250
}
251-
// queryResult.setQueryStats(result.getQueryStats());
251+
252+
spi.onQueryResult(Status.SUCCESS, null);
252253
return readers;
254+
} catch (SQLException | RuntimeException ex) {
255+
if (ex instanceof YdbStatusable) {
256+
spi.onQueryResult(((YdbStatusable) ex).getStatus(), null);
257+
} else {
258+
spi.onQueryResult(null, ex);
259+
}
260+
261+
updateState(tx.withRollback(session));
262+
throw ex;
253263
} finally {
254264
if (tx.isInsideTransaction()) {
255265
tracer.setId(tx.txID());
@@ -262,7 +272,7 @@ public YdbResultSetMemory[] executeInMemoryQuery(YdbStatement statement, String
262272
@Override
263273
public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params)
264274
throws SQLException {
265-
YdbResultSetMemory[] readers = executeInMemoryQuery(statement, preparedYql, params);
275+
YdbResultSetMemory[] readers = executeInMemoryQuery(statement, query, preparedYql, params);
266276
return updateCurrentResult(new YdbQueryResultStatic(query, readers));
267277
}
268278

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

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@
33
import java.sql.SQLException;
44
import java.time.Duration;
55
import java.util.Collection;
6-
import java.util.Iterator;
7-
import java.util.ServiceLoader;
86
import java.util.concurrent.Executors;
97
import java.util.concurrent.atomic.AtomicInteger;
108
import java.util.logging.Level;
@@ -26,7 +24,7 @@
2624
import tech.ydb.jdbc.settings.YdbConnectionProperties;
2725
import tech.ydb.jdbc.settings.YdbOperationProperties;
2826
import tech.ydb.jdbc.settings.YdbQueryProperties;
29-
import tech.ydb.jdbc.spi.YDBQueryExtensionService;
27+
import tech.ydb.jdbc.spi.YdbQueryExtentionService;
3028
import tech.ydb.query.QueryClient;
3129
import tech.ydb.query.impl.QueryClientImpl;
3230
import tech.ydb.scheme.SchemeClient;
@@ -62,7 +60,7 @@ public class YdbContext implements AutoCloseable {
6260
private final boolean autoResizeSessionPool;
6361
private final AtomicInteger connectionsCount = new AtomicInteger();
6462

65-
private YDBQueryExtensionService queryExtensionService = null;
63+
private final YdbQueryExtentionService querySpi;
6664

6765
private YdbContext(
6866
YdbConfig config,
@@ -103,12 +101,7 @@ private YdbContext(
103101
queryProperties, config.getPreparedStatementsCachecSize(), config.isFullScanDetectorEnabled());
104102
}
105103

106-
Iterator<YDBQueryExtensionService> extLoaderIterator = ServiceLoader
107-
.load(YDBQueryExtensionService.class)
108-
.iterator();
109-
if (extLoaderIterator.hasNext()) {
110-
queryExtensionService = extLoaderIterator.next();
111-
}
104+
this.querySpi = YdbServiceLoader.loadQuerySpi();
112105
}
113106

114107
public YdbTypes getTypes() {
@@ -127,6 +120,10 @@ public YdbTracer getTracer() {
127120
return config.isTxTracedEnabled() ? YdbTracer.current() : YdbTracerNone.DISABLED;
128121
}
129122

123+
public YdbQueryExtentionService getQuerySpi() {
124+
return querySpi;
125+
}
126+
130127
static String joined(String path1, String path2) {
131128
return path1.endsWith("/") || path2.startsWith("/") ? path1 + path2 : path1 + "/" + path2;
132129
}
@@ -325,8 +322,4 @@ public YdbQuery parseYdbQuery(QueryKey key) throws SQLException {
325322
public YdbPreparedQuery prepareYdbQuery(YdbQuery query, YdbPrepareMode mode) throws SQLException {
326323
return cache.prepareYdbQuery(query, mode);
327324
}
328-
329-
public YDBQueryExtensionService getQueryExtensionService() {
330-
return queryExtensionService;
331-
}
332325
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ public interface YdbExecutor {
3535
YdbQueryResult executeExplainQuery(YdbStatement st, YdbQuery query) throws SQLException;
3636
YdbQueryResult executeScanQuery(YdbStatement st, YdbQuery query, String yql, Params prms) throws SQLException;
3737
YdbQueryResult executeDataQuery(YdbStatement st, YdbQuery query, String yql, Params prms) throws SQLException;
38-
YdbResultSetMemory[] executeInMemoryQuery(YdbStatement st, String yql, Params prms) throws SQLException;
38+
YdbResultSetMemory[] executeInMemoryQuery(YdbStatement st, YdbQuery query, String yql, Params prms)
39+
throws SQLException;
3940

4041
void commit(YdbContext ctx, YdbValidator validator) throws SQLException;
4142
void rollback(YdbContext ctx, YdbValidator validator) throws SQLException;

0 commit comments

Comments
 (0)