Skip to content

Commit 21dbf86

Browse files
alex268alex268
authored
Fixed full scan detector (#61)
2 parents 517be81 + 8f2932b commit 21dbf86

17 files changed

+627
-272
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: 3 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

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

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,20 @@
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
/**
1211
*
1312
* @author Aleksandr Gorshenin
1413
*/
1514
public class QueryStat {
16-
public static final String QUERY = "print_jdbc_stats();";
15+
private static final String PRINT_QUERY = "print_jdbc_stats();";
16+
private static final String RESET_QUERY = "reset_jdbc_stats();";
1717

1818
private static final FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
1919
.addTextColumn("sql")
2020
.addBooleanColumn("is_fullscan")
21+
.addBooleanColumn("is_error")
2122
.addLongColumn("executed")
2223
.addTextColumn("yql")
2324
.addTextColumn("ast")
@@ -31,23 +32,26 @@ public class QueryStat {
3132
private final String plan;
3233
private final LongAdder usage;
3334
private final boolean isFullScan;
35+
private final boolean isError;
3436

35-
public QueryStat(YdbQuery query, String ast, String plan) {
36-
this.originSQL = query.getOriginQuery();
37-
this.preparedYQL = query.getPreparedYql();
37+
public QueryStat(String sql, String yql, String ast, String plan) {
38+
this.originSQL = sql;
39+
this.preparedYQL = yql;
3840
this.ast = ast;
3941
this.plan = plan;
4042
this.usage = new LongAdder();
4143
this.isFullScan = plan.contains("\"Node Type\":\"TableFullScan\"");
44+
this.isError = false;
4245
}
4346

44-
public QueryStat(YdbQuery query, Status error) {
45-
this.originSQL = query.getOriginQuery();
46-
this.preparedYQL = query.getPreparedYql();
47-
this.ast = error.toString();
47+
public QueryStat(String sql, String yql, Status error) {
48+
this.originSQL = sql;
49+
this.preparedYQL = yql;
50+
this.ast = null;
4851
this.plan = error.toString();
4952
this.usage = new LongAdder();
5053
this.isFullScan = false;
54+
this.isError = true;
5155
}
5256

5357
public long getUsageCounter() {
@@ -74,6 +78,10 @@ public boolean isFullScan() {
7478
return isFullScan;
7579
}
7680

81+
public boolean isError() {
82+
return isError;
83+
}
84+
7785
public void incrementUsage() {
7886
this.usage.increment();
7987
}
@@ -84,6 +92,7 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
8492
builder.newRow()
8593
.withTextValue("sql", stat.originSQL)
8694
.withBoolValue("is_fullscan", stat.isFullScan)
95+
.withBoolValue("is_error", stat.isError)
8796
.withLongValue("executed", stat.usage.longValue())
8897
.withTextValue("yql", stat.preparedYQL)
8998
.withTextValue("ast", stat.ast)
@@ -92,4 +101,12 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
92101
}
93102
return builder.build();
94103
}
104+
105+
public static boolean isPrint(String sql) {
106+
return sql != null && PRINT_QUERY.equalsIgnoreCase(sql.trim());
107+
}
108+
109+
public static boolean isReset(String sql) {
110+
return sql != null && RESET_QUERY.equalsIgnoreCase(sql.trim());
111+
}
95112
}

‎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: 49 additions & 35 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;
@@ -75,7 +75,7 @@ public class YdbContext implements AutoCloseable {
7575
private final SessionRetryContext retryCtx;
7676

7777
private final Cache<String, YdbQuery> queriesCache;
78-
private final Cache<String, QueryStat> queryStatesCache;
78+
private final Cache<String, QueryStat> statsCache;
7979
private final Cache<String, Map<String, Type>> queryParamsCache;
8080

8181
private final boolean autoResizeSessionPool;
@@ -107,13 +107,13 @@ private YdbContext(
107107
queriesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
108108
queryParamsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
109109
if (config.isFullScanDetectorEnabled()) {
110-
queryStatesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
110+
statsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
111111
} else {
112-
queryStatesCache = null;
112+
statsCache = null;
113113
}
114114
} else {
115115
queriesCache = null;
116-
queryStatesCache = null;
116+
statsCache = null;
117117
queryParamsCache = null;
118118
}
119119
}
@@ -184,25 +184,26 @@ public boolean hasConnections() {
184184
}
185185

186186
public boolean queryStatsEnabled() {
187-
return queryStatesCache != null;
187+
return statsCache != null;
188188
}
189189

190-
public Collection<QueryStat> getQueryStats() {
191-
if (queryStatesCache == null) {
192-
return Collections.emptyList();
190+
public void resetQueryStats() {
191+
if (statsCache != null) {
192+
statsCache.invalidateAll();
193193
}
194-
Set<QueryStat> sortedByUsage = new TreeSet<>(Comparator.comparingLong(QueryStat::getUsageCounter).reversed());
195-
sortedByUsage.addAll(queryStatesCache.asMap().values());
196-
return sortedByUsage;
197194
}
198195

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-
}
196+
public Collection<QueryStat> getQueryStats() {
197+
if (statsCache == null) {
198+
return Collections.emptyList();
205199
}
200+
List<QueryStat> sorted = new ArrayList<>(statsCache.asMap().values());
201+
Collections.sort(sorted,
202+
Comparator
203+
.comparingLong(QueryStat::getUsageCounter).reversed()
204+
.thenComparing(QueryStat::getPreparedYQL)
205+
);
206+
return sorted;
206207
}
207208

208209
public void register() {
@@ -308,29 +309,42 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
308309
queriesCache.put(sql, cached);
309310
}
310311

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();
319312

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);
313+
return cached;
314+
}
315+
316+
public void traceQuery(YdbQuery query, String yql) {
317+
if (statsCache == null) {
318+
return;
319+
}
320+
321+
QueryStat stat = statsCache.getIfPresent(yql);
322+
if (stat == null) {
323+
final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
324+
Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
325+
session -> session.explainDataQuery(yql, settings)
326+
).join();
327+
328+
if (res.isSuccess()) {
329+
ExplainDataQueryResult exp = res.getValue();
330+
stat = new QueryStat(query.getOriginQuery(), yql, exp.getQueryAst(), exp.getQueryPlan());
331+
} else {
332+
stat = new QueryStat(query.getOriginQuery(), yql, res.getStatus());
327333
}
334+
335+
statsCache.put(yql, stat);
328336
}
329337

330-
return cached;
338+
stat.incrementUsage();
331339
}
332340

333341
public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) throws SQLException {
342+
if (statsCache != null) {
343+
if (QueryStat.isPrint(query.getOriginQuery()) || QueryStat.isReset(query.getOriginQuery())) {
344+
return new InMemoryQuery(query, queryOptions.isDeclareJdbcParameters());
345+
}
346+
}
347+
334348
if (query.getYqlBatcher() != null && mode == YdbPrepareMode.AUTO) {
335349
Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
336350
if (types == null) {

‎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: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
import tech.ydb.jdbc.YdbResultSet;
1717
import tech.ydb.jdbc.YdbStatement;
1818
import tech.ydb.jdbc.common.FixedResultSetFactory;
19+
import tech.ydb.jdbc.context.QueryStat;
20+
import tech.ydb.jdbc.context.YdbContext;
1921
import tech.ydb.jdbc.context.YdbValidator;
2022
import tech.ydb.jdbc.query.ExplainedQuery;
2123
import tech.ydb.jdbc.query.QueryStatement;
@@ -189,14 +191,29 @@ protected List<YdbResult> executeExplainQuery(YdbQuery query) throws SQLExceptio
189191
return Collections.singletonList(new YdbResult(new YdbResultSetImpl(this, result)));
190192
}
191193

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

197200
protected List<YdbResult> executeDataQuery(YdbQuery query, String yql, Params params) throws SQLException {
201+
YdbContext ctx = connection.getCtx();
202+
203+
if (ctx.queryStatsEnabled()) {
204+
if (QueryStat.isPrint(yql)) {
205+
YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
206+
return Collections.singletonList(new YdbResult(rs));
207+
}
208+
if (QueryStat.isReset(yql)) {
209+
getConnection().getCtx().resetQueryStats();
210+
return null;
211+
}
212+
}
213+
214+
ctx.traceQuery(query, yql);
198215
List<ResultSetReader> resultSets = connection
199-
.executeDataQuery(yql, validator, getQueryTimeout(), isPoolable(), params);
216+
.executeDataQuery(query, yql, validator, getQueryTimeout(), isPoolable(), params);
200217

201218
List<YdbResult> results = new ArrayList<>();
202219
int idx = 0;

0 commit comments

Comments
 (0)