Skip to content

Commit 16e24da

Browse files
authored
Added option to FullScan detecting (#58)
2 parents e46c943 + efbaf8a commit 16e24da

File tree

9 files changed

+354
-4
lines changed

9 files changed

+354
-4
lines changed
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package tech.ydb.jdbc.context;
2+
3+
import java.util.Collection;
4+
import java.util.concurrent.atomic.LongAdder;
5+
6+
import tech.ydb.core.Status;
7+
import tech.ydb.jdbc.common.FixedResultSetFactory;
8+
import tech.ydb.jdbc.query.YdbQuery;
9+
import tech.ydb.table.result.ResultSetReader;
10+
11+
/**
12+
*
13+
* @author Aleksandr Gorshenin
14+
*/
15+
public class QueryStat {
16+
public static final String QUERY = "print_jdbc_stats();";
17+
18+
private static final FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
19+
.addTextColumn("sql")
20+
.addBooleanColumn("is_fullscan")
21+
.addLongColumn("executed")
22+
.addTextColumn("yql")
23+
.addTextColumn("ast")
24+
.addTextColumn("plan")
25+
.build();
26+
27+
private final String originSQL;
28+
private final String preparedYQL;
29+
30+
private final String ast;
31+
private final String plan;
32+
private final LongAdder usage;
33+
private final boolean isFullScan;
34+
35+
public QueryStat(YdbQuery query, String ast, String plan) {
36+
this.originSQL = query.getOriginQuery();
37+
this.preparedYQL = query.getPreparedYql();
38+
this.ast = ast;
39+
this.plan = plan;
40+
this.usage = new LongAdder();
41+
this.isFullScan = plan.contains("\"Node Type\":\"TableFullScan\"");
42+
}
43+
44+
public QueryStat(YdbQuery query, Status error) {
45+
this.originSQL = query.getOriginQuery();
46+
this.preparedYQL = query.getPreparedYql();
47+
this.ast = error.toString();
48+
this.plan = error.toString();
49+
this.usage = new LongAdder();
50+
this.isFullScan = false;
51+
}
52+
53+
public long getUsageCounter() {
54+
return usage.longValue();
55+
}
56+
57+
public String getOriginSQL() {
58+
return originSQL;
59+
}
60+
61+
public String getPreparedYQL() {
62+
return preparedYQL;
63+
}
64+
65+
public String getAat() {
66+
return ast;
67+
}
68+
69+
public String getPlan() {
70+
return plan;
71+
}
72+
73+
public boolean isFullScan() {
74+
return isFullScan;
75+
}
76+
77+
public void incrementUsage() {
78+
this.usage.increment();
79+
}
80+
81+
public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
82+
FixedResultSetFactory.ResultSetBuilder builder = STATS_RS_FACTORY.createResultSet();
83+
for (QueryStat stat: stats) {
84+
builder.newRow()
85+
.withTextValue("sql", stat.originSQL)
86+
.withBoolValue("is_fullscan", stat.isFullScan)
87+
.withLongValue("executed", stat.usage.longValue())
88+
.withTextValue("yql", stat.preparedYQL)
89+
.withTextValue("ast", stat.ast)
90+
.withTextValue("plan", stat.plan)
91+
.build();
92+
}
93+
return builder.build();
94+
}
95+
}

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

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@
33
import java.sql.SQLDataException;
44
import java.sql.SQLException;
55
import java.time.Duration;
6+
import java.util.Collection;
7+
import java.util.Collections;
8+
import java.util.Comparator;
69
import java.util.Map;
10+
import java.util.Set;
11+
import java.util.TreeSet;
712
import java.util.concurrent.Executors;
813
import java.util.concurrent.atomic.AtomicInteger;
914
import java.util.logging.Level;
@@ -39,8 +44,10 @@
3944
import tech.ydb.table.description.TableColumn;
4045
import tech.ydb.table.description.TableDescription;
4146
import tech.ydb.table.impl.PooledTableClient;
47+
import tech.ydb.table.query.ExplainDataQueryResult;
4248
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
4349
import tech.ydb.table.settings.DescribeTableSettings;
50+
import tech.ydb.table.settings.ExplainDataQuerySettings;
4451
import tech.ydb.table.settings.PrepareDataQuerySettings;
4552
import tech.ydb.table.settings.RequestSettings;
4653
import tech.ydb.table.values.Type;
@@ -68,6 +75,7 @@ public class YdbContext implements AutoCloseable {
6875
private final SessionRetryContext retryCtx;
6976

7077
private final Cache<String, YdbQuery> queriesCache;
78+
private final Cache<String, QueryStat> queryStatesCache;
7179
private final Cache<String, Map<String, Type>> queryParamsCache;
7280

7381
private final boolean autoResizeSessionPool;
@@ -98,8 +106,14 @@ private YdbContext(
98106
if (cacheSize > 0) {
99107
queriesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
100108
queryParamsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
109+
if (config.isFullScanDetectorEnabled()) {
110+
queryStatesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
111+
} else {
112+
queryStatesCache = null;
113+
}
101114
} else {
102115
queriesCache = null;
116+
queryStatesCache = null;
103117
queryParamsCache = null;
104118
}
105119
}
@@ -169,6 +183,28 @@ public boolean hasConnections() {
169183
return connectionsCount.get() > 0;
170184
}
171185

186+
public boolean queryStatsEnabled() {
187+
return queryStatesCache != null;
188+
}
189+
190+
public Collection<QueryStat> getQueryStats() {
191+
if (queryStatesCache == null) {
192+
return Collections.emptyList();
193+
}
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+
}
206+
}
207+
172208
public void register() {
173209
int actual = connectionsCount.incrementAndGet();
174210
int maxSize = tableClient.sessionPoolStats().getMaxSize();
@@ -272,6 +308,25 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
272308
queriesCache.put(sql, cached);
273309
}
274310

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();
319+
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);
327+
}
328+
}
329+
275330
return cached;
276331
}
277332

@@ -288,8 +343,8 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
288343
).join();
289344

290345
if (result.isSuccess()) {
291-
TableDescription d = result.getValue();
292-
types = result.getValue().getColumns().stream()
346+
TableDescription descrtiption = result.getValue();
347+
types = descrtiption.getColumns().stream()
293348
.collect(Collectors.toMap(TableColumn::getName, TableColumn::getType));
294349
queryParamsCache.put(query.getOriginQuery(), types);
295350
}

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

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

9292
try {
9393
for (Params prm: prepared.getBatchParams()) {
94+
getConnection().getCtx().traceQueryExecution(query);
9495
executeDataQuery(query, prepared.getQueryText(prm), prm);
9596
}
9697
} finally {
@@ -123,7 +124,9 @@ public boolean execute() throws SQLException {
123124
clearBatch();
124125

125126
List<YdbResult> newState = null;
127+
126128
Params prms = prepared.getCurrentParams();
129+
getConnection().getCtx().traceQueryExecution(query);
127130
switch (query.getType()) {
128131
case DATA_QUERY:
129132
newState = executeDataQuery(query, prepared.getQueryText(prms), prms);

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,16 @@
55
import java.sql.Statement;
66
import java.util.ArrayList;
77
import java.util.Arrays;
8+
import java.util.Collections;
89
import java.util.List;
910
import java.util.logging.Level;
1011
import java.util.logging.Logger;
1112

1213
import tech.ydb.jdbc.YdbConnection;
1314
import tech.ydb.jdbc.YdbConst;
1415
import tech.ydb.jdbc.YdbResultSet;
16+
import tech.ydb.jdbc.context.QueryStat;
17+
import tech.ydb.jdbc.context.YdbContext;
1518
import tech.ydb.jdbc.query.YdbQuery;
1619
import tech.ydb.table.query.Params;
1720

@@ -79,7 +82,20 @@ public int executeUpdate(String sql) throws SQLException {
7982
public boolean execute(String sql) throws SQLException {
8083
cleanState();
8184

82-
YdbQuery query = getConnection().getCtx().parseYdbQuery(sql);
85+
YdbContext ctx = getConnection().getCtx();
86+
YdbQuery query;
87+
88+
if (ctx.queryStatsEnabled()) {
89+
if (sql != null && QueryStat.QUERY.equalsIgnoreCase(sql.trim())) {
90+
YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
91+
return updateState(Collections.singletonList(new YdbResult(rs)));
92+
}
93+
query = ctx.findOrParseYdbQuery(sql);
94+
ctx.traceQueryExecution(query);
95+
} else {
96+
query = ctx.parseYdbQuery(sql);
97+
}
98+
8399
List<YdbResult> newState = null;
84100
switch (query.getType()) {
85101
case SCHEME_QUERY:

jdbc/src/main/java/tech/ydb/jdbc/settings/YdbClientProperties.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,5 +81,4 @@ public boolean applyToTableClient(TableClient.Builder table, QueryClient.Builder
8181
query.sessionPoolMaxSize(maxSize).sessionPoolMinSize(minSize);
8282
return false;
8383
}
84-
8584
}

jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ public class YdbConfig {
4040
static final YdbProperty<Boolean> USE_QUERY_SERVICE = YdbProperty.bool("useQueryService",
4141
"Use QueryService intead of TableService", false
4242
);
43+
static final YdbProperty<Boolean> FULLSCAN_DETECTOR_ENABLED = YdbProperty.bool(
44+
"jdbcFullScanDetector", "Enable analizator for collecting query stats", false
45+
);
4346

4447

4548
private final String url;
@@ -53,6 +56,7 @@ public class YdbConfig {
5356
private final boolean isCacheConnectionsInDriver;
5457
private final int preparedStatementsCacheSize;
5558
private final boolean useQueryService;
59+
private final boolean fullScanDetectorEnabled;
5660

5761
private YdbConfig(
5862
String url, String safeUrl, String connectionString, String username, String password, Properties props
@@ -66,6 +70,7 @@ private YdbConfig(
6670
this.isCacheConnectionsInDriver = CACHE_CONNECTIONS_IN_DRIVER.readValue(props).getValue();
6771
this.preparedStatementsCacheSize = Math.max(0, PREPARED_STATEMENT_CACHE_SIZE.readValue(props).getValue());
6872
this.useQueryService = USE_QUERY_SERVICE.readValue(props).getValue();
73+
this.fullScanDetectorEnabled = FULLSCAN_DETECTOR_ENABLED.readValue(props).getValue();
6974
}
7075

7176
public Properties getSafeProps() {
@@ -96,6 +101,10 @@ public boolean isUseQueryService() {
96101
return this.useQueryService;
97102
}
98103

104+
public boolean isFullScanDetectorEnabled() {
105+
return fullScanDetectorEnabled;
106+
}
107+
99108
static boolean isSensetive(String key) {
100109
return TOKEN_KEY.equalsIgnoreCase(key) || PASSWORD_KEY.equalsIgnoreCase(key);
101110
}

jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -925,4 +925,87 @@ public void testUnsupportedComplexTypes(String type) throws SQLException {
925925
() -> statement.execute(sql));
926926
}
927927
}
928+
929+
@Test
930+
public void testFullScanAnalyzer() throws SQLException {
931+
try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
932+
String selectAll = QUERIES.selectAllSQL();
933+
String selectByKey = QUERIES.selectAllByKey("1");
934+
String preparedSelect = QUERIES.selectAllByKey("?");
935+
936+
try (Statement st = connection.createStatement()) {
937+
try (ResultSet rs = st.executeQuery(" print_JDBC_stats(); ")) {
938+
Assertions.assertFalse(rs.next()); // not stats
939+
}
940+
941+
try (ResultSet rs = st.executeQuery(selectAll)) {
942+
Assertions.assertFalse(rs.next());
943+
}
944+
945+
try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
946+
Assertions.assertTrue(rs.next());
947+
Assertions.assertEquals(selectAll, rs.getString("sql"));
948+
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
949+
Assertions.assertEquals(1l, rs.getLong("executed"));
950+
951+
Assertions.assertFalse(rs.next());
952+
}
953+
954+
try (ResultSet rs = st.executeQuery(selectAll)) {
955+
Assertions.assertFalse(rs.next());
956+
}
957+
try (ResultSet rs = st.executeQuery(selectByKey)) {
958+
Assertions.assertFalse(rs.next());
959+
}
960+
961+
try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
962+
Assertions.assertTrue(rs.next());
963+
Assertions.assertEquals(selectAll, rs.getString("sql"));
964+
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
965+
Assertions.assertEquals(2l, rs.getLong("executed"));
966+
967+
Assertions.assertTrue(rs.next());
968+
Assertions.assertEquals(selectByKey, rs.getString("sql"));
969+
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
970+
Assertions.assertEquals(1l, rs.getLong("executed"));
971+
972+
Assertions.assertFalse(rs.next());
973+
}
974+
975+
try (PreparedStatement ps = connection.prepareStatement(preparedSelect)) {
976+
ps.setLong(1, 1);
977+
try (ResultSet rs = ps.executeQuery()) {
978+
Assertions.assertFalse(rs.next());
979+
}
980+
ps.setLong(1, 2);
981+
try (ResultSet rs = ps.executeQuery()) {
982+
Assertions.assertFalse(rs.next());
983+
}
984+
ps.setLong(1, 3);
985+
try (ResultSet rs = ps.executeQuery()) {
986+
Assertions.assertFalse(rs.next());
987+
}
988+
}
989+
990+
try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
991+
Assertions.assertTrue(rs.next());
992+
Assertions.assertEquals(preparedSelect, rs.getString("sql"));
993+
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
994+
Assertions.assertEquals(3l, rs.getLong("executed"));
995+
996+
Assertions.assertTrue(rs.next());
997+
Assertions.assertEquals(selectAll, rs.getString("sql"));
998+
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
999+
Assertions.assertEquals(2l, rs.getLong("executed"));
1000+
1001+
Assertions.assertTrue(rs.next());
1002+
Assertions.assertEquals(selectByKey, rs.getString("sql"));
1003+
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
1004+
Assertions.assertEquals(1l, rs.getLong("executed"));
1005+
1006+
Assertions.assertFalse(rs.next());
1007+
}
1008+
}
1009+
}
1010+
}
9281011
}

0 commit comments

Comments
 (0)