Skip to content

Commit bb1202d

Browse files
alex268alex268
committed
Fixed error of stat collection for prepared statements
1 parent 517be81 commit bb1202d

File tree

11 files changed

+90
-72
lines changed

11 files changed

+90
-72
lines changed

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import tech.ydb.jdbc.context.YdbContext;
1010
import tech.ydb.jdbc.context.YdbValidator;
1111
import tech.ydb.jdbc.query.ExplainedQuery;
12+
import tech.ydb.jdbc.query.YdbQuery;
1213
import tech.ydb.table.query.Params;
1314
import tech.ydb.table.result.ResultSetReader;
1415

@@ -36,27 +37,30 @@ public interface YdbConnection extends Connection {
3637
/**
3738
* Explicitly execute query as a data query
3839
*
39-
* @param yql query to execute
40+
* @param query query to execute
41+
* @param yql YQL text to execute
4042
* @param params parameters for query
4143
* @param timeout timeout of operation
4244
* @param keepInCache flag to store query in server-side cache
4345
* @param validator handler for logging and warnings
4446
* @return list of result set
4547
* @throws SQLException if query cannot be executed
4648
*/
47-
List<ResultSetReader> executeDataQuery(String yql, YdbValidator validator,
49+
List<ResultSetReader> executeDataQuery(YdbQuery query, String yql, YdbValidator validator,
4850
int timeout, boolean keepInCache, Params params) throws SQLException;
4951

5052
/**
5153
* Explicitly execute query as a scan query
5254
*
53-
* @param yql query to execute
55+
* @param query query to execute
56+
* @param yql YQL text to execute
5457
* @param params parameters for query
5558
* @param validator handler for logging and warnings
5659
* @return single result set with rows
5760
* @throws SQLException if query cannot be executed
5861
*/
59-
ResultSetReader executeScanQuery(String yql, YdbValidator validator, Params params) throws SQLException;
62+
ResultSetReader executeScanQuery(YdbQuery query, String yql, YdbValidator validator, Params params)
63+
throws SQLException;
6064

6165
/**
6266
* Explicitly explain this query

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import tech.ydb.core.UnexpectedResultException;
1010
import tech.ydb.jdbc.exception.ExceptionFactory;
1111
import tech.ydb.jdbc.query.QueryType;
12+
import tech.ydb.jdbc.query.YdbQuery;
1213
import tech.ydb.table.Session;
1314
import tech.ydb.table.TableClient;
1415
import tech.ydb.table.query.Params;
@@ -50,15 +51,18 @@ public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yq
5051
}
5152

5253
@Override
53-
public ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
54-
throws SQLException {
54+
public ResultSetReader executeScanQuery(
55+
YdbContext ctx, YdbValidator validator, YdbQuery query, String yql, Params params
56+
) throws SQLException {
5557
ensureOpened();
5658

5759
Collection<ResultSetReader> resultSets = new LinkedBlockingQueue<>();
5860
Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
5961
ExecuteScanQuerySettings settings = ExecuteScanQuerySettings.newBuilder()
6062
.withRequestTimeout(scanQueryTimeout)
6163
.build();
64+
65+
ctx.traceQuery(query, yql);
6266
try (Session session = createNewTableSession(validator)) {
6367
validator.execute(QueryType.SCAN_QUERY + " >>\n" + yql,
6468
() -> session.executeScanQuery(yql, params, settings).start(resultSets::add));

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import tech.ydb.jdbc.exception.ExceptionFactory;
1818
import tech.ydb.jdbc.query.ExplainedQuery;
1919
import tech.ydb.jdbc.query.QueryType;
20+
import tech.ydb.jdbc.query.YdbQuery;
2021
import tech.ydb.query.QueryClient;
2122
import tech.ydb.query.QuerySession;
2223
import tech.ydb.query.QueryStream;
@@ -196,7 +197,8 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException
196197

197198
@Override
198199
public List<ResultSetReader> executeDataQuery(
199-
YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
200+
YdbContext ctx, YdbValidator validator, YdbQuery query,
201+
String yql, long timeout, boolean keepInCache, Params params
200202
) throws SQLException {
201203
ensureOpened();
202204

@@ -206,6 +208,7 @@ public List<ResultSetReader> executeDataQuery(
206208
}
207209
final ExecuteQuerySettings settings = builder.build();
208210

211+
ctx.traceQuery(query, yql);
209212
if (tx == null) {
210213
tx = createNewQuerySession(validator).createNewTransaction(txMode);
211214
}

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

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
import tech.ydb.core.Status;
77
import tech.ydb.jdbc.common.FixedResultSetFactory;
8-
import tech.ydb.jdbc.query.YdbQuery;
98
import tech.ydb.table.result.ResultSetReader;
109

1110
/**
@@ -18,6 +17,7 @@ public class QueryStat {
1817
private static final FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
1918
.addTextColumn("sql")
2019
.addBooleanColumn("is_fullscan")
20+
.addBooleanColumn("is_error")
2121
.addLongColumn("executed")
2222
.addTextColumn("yql")
2323
.addTextColumn("ast")
@@ -31,23 +31,26 @@ public class QueryStat {
3131
private final String plan;
3232
private final LongAdder usage;
3333
private final boolean isFullScan;
34+
private final boolean isError;
3435

35-
public QueryStat(YdbQuery query, String ast, String plan) {
36-
this.originSQL = query.getOriginQuery();
37-
this.preparedYQL = query.getPreparedYql();
36+
public QueryStat(String sql, String yql, String ast, String plan) {
37+
this.originSQL = sql;
38+
this.preparedYQL = yql;
3839
this.ast = ast;
3940
this.plan = plan;
4041
this.usage = new LongAdder();
4142
this.isFullScan = plan.contains("\"Node Type\":\"TableFullScan\"");
43+
this.isError = false;
4244
}
4345

44-
public QueryStat(YdbQuery query, Status error) {
45-
this.originSQL = query.getOriginQuery();
46-
this.preparedYQL = query.getPreparedYql();
47-
this.ast = error.toString();
46+
public QueryStat(String sql, String yql, Status error) {
47+
this.originSQL = sql;
48+
this.preparedYQL = yql;
49+
this.ast = null;
4850
this.plan = error.toString();
4951
this.usage = new LongAdder();
5052
this.isFullScan = false;
53+
this.isError = true;
5154
}
5255

5356
public long getUsageCounter() {
@@ -74,6 +77,10 @@ public boolean isFullScan() {
7477
return isFullScan;
7578
}
7679

80+
public boolean isError() {
81+
return isError;
82+
}
83+
7784
public void incrementUsage() {
7885
this.usage.increment();
7986
}
@@ -84,6 +91,7 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
8491
builder.newRow()
8592
.withTextValue("sql", stat.originSQL)
8693
.withBoolValue("is_fullscan", stat.isFullScan)
94+
.withBoolValue("is_error", stat.isError)
8795
.withLongValue("executed", stat.usage.longValue())
8896
.withTextValue("yql", stat.preparedYQL)
8997
.withTextValue("ast", stat.ast)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import tech.ydb.jdbc.YdbConst;
1111
import tech.ydb.jdbc.query.ExplainedQuery;
1212
import tech.ydb.jdbc.query.QueryType;
13+
import tech.ydb.jdbc.query.YdbQuery;
1314
import tech.ydb.table.Session;
1415
import tech.ydb.table.query.DataQueryResult;
1516
import tech.ydb.table.query.ExplainDataQueryResult;
@@ -168,7 +169,8 @@ public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator
168169

169170
@Override
170171
public List<ResultSetReader> executeDataQuery(
171-
YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
172+
YdbContext ctx, YdbValidator validator, YdbQuery query,
173+
String yql, long timeout, boolean keepInCache, Params params
172174
) throws SQLException {
173175
ensureOpened();
174176

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

Lines changed: 32 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
import java.sql.SQLDataException;
44
import java.sql.SQLException;
55
import java.time.Duration;
6+
import java.util.ArrayList;
67
import java.util.Collection;
78
import java.util.Collections;
89
import java.util.Comparator;
10+
import java.util.List;
911
import java.util.Map;
10-
import java.util.Set;
11-
import java.util.TreeSet;
1212
import java.util.concurrent.Executors;
1313
import java.util.concurrent.atomic.AtomicInteger;
1414
import java.util.logging.Level;
@@ -191,18 +191,13 @@ public Collection<QueryStat> getQueryStats() {
191191
if (queryStatesCache == null) {
192192
return Collections.emptyList();
193193
}
194-
Set<QueryStat> sortedByUsage = new TreeSet<>(Comparator.comparingLong(QueryStat::getUsageCounter).reversed());
195-
sortedByUsage.addAll(queryStatesCache.asMap().values());
196-
return sortedByUsage;
197-
}
198-
199-
public void traceQueryExecution(YdbQuery query) {
200-
if (queryStatesCache != null) {
201-
QueryStat stat = queryStatesCache.getIfPresent(query.getOriginQuery());
202-
if (stat != null) {
203-
stat.incrementUsage();
204-
}
205-
}
194+
List<QueryStat> sorted = new ArrayList<>(queryStatesCache.asMap().values());
195+
Collections.sort(sorted,
196+
Comparator
197+
.comparingLong(QueryStat::getUsageCounter).reversed()
198+
.thenComparing(QueryStat::getPreparedYQL)
199+
);
200+
return sorted;
206201
}
207202

208203
public void register() {
@@ -308,26 +303,33 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
308303
queriesCache.put(sql, cached);
309304
}
310305

311-
if (queryStatesCache != null) {
312-
QueryStat stat = queryStatesCache.getIfPresent(sql);
313-
if (stat == null) {
314-
final String preparedYQL = cached.getPreparedYql();
315-
final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
316-
Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
317-
session -> session.explainDataQuery(preparedYQL, settings)
318-
).join();
319306

320-
if (res.isSuccess()) {
321-
ExplainDataQueryResult exp = res.getValue();
322-
stat = new QueryStat(cached, exp.getQueryAst(), exp.getQueryPlan());
323-
} else {
324-
stat = new QueryStat(cached, res.getStatus());
325-
}
326-
queryStatesCache.put(sql, stat);
307+
return cached;
308+
}
309+
310+
public void traceQuery(YdbQuery query, String yql) {
311+
if (queryStatesCache == null) {
312+
return;
313+
}
314+
315+
QueryStat stat = queryStatesCache.getIfPresent(yql);
316+
if (stat == null) {
317+
final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
318+
Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
319+
session -> session.explainDataQuery(yql, settings)
320+
).join();
321+
322+
if (res.isSuccess()) {
323+
ExplainDataQueryResult exp = res.getValue();
324+
stat = new QueryStat(query.getOriginQuery(), yql, exp.getQueryAst(), exp.getQueryPlan());
325+
} else {
326+
stat = new QueryStat(query.getOriginQuery(), yql, res.getStatus());
327327
}
328+
329+
queryStatesCache.put(yql, stat);
328330
}
329331

330-
return cached;
332+
stat.incrementUsage();
331333
}
332334

333335
public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) throws SQLException {

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import tech.ydb.jdbc.YdbConst;
77
import tech.ydb.jdbc.query.ExplainedQuery;
8+
import tech.ydb.jdbc.query.YdbQuery;
89
import tech.ydb.table.query.Params;
910
import tech.ydb.table.result.ResultSetReader;
1011

@@ -34,10 +35,10 @@ default void ensureOpened() throws SQLException {
3435

3536
void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException;
3637

37-
List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, String yql,
38+
List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, String yql,
3839
long timeout, boolean poolable, Params params) throws SQLException;
3940

40-
ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
41+
ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, String yql, Params params)
4142
throws SQLException;
4243

4344
ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, String yql)

‎jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,14 +189,16 @@ protected List<YdbResult> executeExplainQuery(YdbQuery query) throws SQLExceptio
189189
return Collections.singletonList(new YdbResult(new YdbResultSetImpl(this, result)));
190190
}
191191

192-
protected List<YdbResult> executeScanQuery(String yql, Params params) throws SQLException {
193-
ResultSetReader result = connection.executeScanQuery(yql, validator, params);
192+
protected List<YdbResult> executeScanQuery(YdbQuery query, String yql, Params params) throws SQLException {
193+
connection.getCtx().traceQuery(query, yql);
194+
ResultSetReader result = connection.executeScanQuery(query, yql, validator, params);
194195
return Collections.singletonList(new YdbResult(new YdbResultSetImpl(this, result)));
195196
}
196197

197198
protected List<YdbResult> executeDataQuery(YdbQuery query, String yql, Params params) throws SQLException {
199+
connection.getCtx().traceQuery(query, yql);
198200
List<ResultSetReader> resultSets = connection
199-
.executeDataQuery(yql, validator, getQueryTimeout(), isPoolable(), params);
201+
.executeDataQuery(query, yql, validator, getQueryTimeout(), isPoolable(), params);
200202

201203
List<YdbResult> results = new ArrayList<>();
202204
int idx = 0;

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -205,13 +205,14 @@ public void executeSchemeQuery(String yql, YdbValidator validator) throws SQLExc
205205
}
206206

207207
@Override
208-
public List<ResultSetReader> executeDataQuery(String yql, YdbValidator validator,
208+
public List<ResultSetReader> executeDataQuery(YdbQuery query, String yql, YdbValidator validator,
209209
int timeout, boolean poolable, Params params) throws SQLException {
210-
return executor.executeDataQuery(ctx, validator, yql, timeout, poolable, params);
210+
return executor.executeDataQuery(ctx, validator, query, yql, timeout, poolable, params);
211211
}
212212

213213
@Override
214-
public ResultSetReader executeScanQuery(String yql, YdbValidator validator, Params params) throws SQLException {
214+
public ResultSetReader executeScanQuery(YdbQuery query, String yql, YdbValidator validator, Params params)
215+
throws SQLException {
215216
executor.ensureOpened();
216217

217218
if (executor.isInsideTransaction()) {
@@ -227,7 +228,7 @@ public ResultSetReader executeScanQuery(String yql, YdbValidator validator, Para
227228
}
228229
}
229230

230-
return executor.executeScanQuery(ctx, validator, yql, params);
231+
return executor.executeScanQuery(ctx, validator, query, yql, params);
231232
}
232233

233234
@Override

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ public int[] executeBatch() throws SQLException {
9191

9292
try {
9393
for (Params prm: prepared.getBatchParams()) {
94-
getConnection().getCtx().traceQueryExecution(query);
9594
executeDataQuery(query, prepared.getQueryText(prm), prm);
9695
}
9796
} finally {
@@ -126,13 +125,12 @@ public boolean execute() throws SQLException {
126125
List<YdbResult> newState = null;
127126

128127
Params prms = prepared.getCurrentParams();
129-
getConnection().getCtx().traceQueryExecution(query);
130128
switch (query.getType()) {
131129
case DATA_QUERY:
132130
newState = executeDataQuery(query, prepared.getQueryText(prms), prms);
133131
break;
134132
case SCAN_QUERY:
135-
newState = executeScanQuery(prepared.getQueryText(prms), prms);
133+
newState = executeScanQuery(query, prepared.getQueryText(prms), prms);
136134
break;
137135
default:
138136
throw new IllegalStateException("Internal error. Unsupported query type " + query.getType());
@@ -146,7 +144,7 @@ public boolean execute() throws SQLException {
146144
public YdbResultSet executeScanQuery() throws SQLException {
147145
cleanState();
148146
Params prms = prepared.getCurrentParams();
149-
List<YdbResult> state = executeScanQuery(prepared.getQueryText(prms), prms);
147+
List<YdbResult> state = executeScanQuery(query, prepared.getQueryText(prms), prms);
150148
prepared.clearParameters();
151149
updateState(state);
152150
return getResultSet();

0 commit comments

Comments
 (0)