Skip to content

Commit efbaf8a

Browse files
alex268alex268
committed
Added impelemntation and tests of query analyzer
1 parent 4e581be commit efbaf8a

File tree

8 files changed

+332
-20
lines changed

8 files changed

+332
-20
lines changed

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

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: 43 additions & 18 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;
@@ -178,6 +183,28 @@ public boolean hasConnections() {
178183
return connectionsCount.get() > 0;
179184
}
180185

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+
181208
public void register() {
182209
int actual = connectionsCount.incrementAndGet();
183210
int maxSize = tableClient.sessionPoolStats().getMaxSize();
@@ -279,26 +306,24 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
279306
if (cached == null) {
280307
cached = parseYdbQuery(sql);
281308
queriesCache.put(sql, cached);
309+
}
282310

283-
if (queryStatesCache != null) {
284-
QueryStat stat = queryStatesCache.getIfPresent(sql);
285-
if (stat == null) {
286-
final String preparedYQL = cached.getPreparedYql();
287-
final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
288-
Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
289-
session -> session.explainDataQuery(preparedYQL, settings)
290-
).join();
291-
292-
if (res.isSuccess()) {
293-
ExplainDataQueryResult exp = res.getValue();
294-
stat = new QueryStat(cached, exp.getQueryAst(), exp.getQueryPlan());
295-
} else {
296-
stat = new QueryStat(cached, res.getStatus());
297-
}
298-
queryStatesCache.put(sql, stat);
299-
}
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();
300319

301-
stat.incrementUsage();
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);
302327
}
303328
}
304329

‎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/YdbConfig.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public class YdbConfig {
4141
"Use QueryService intead of TableService", false
4242
);
4343
static final YdbProperty<Boolean> FULLSCAN_DETECTOR_ENABLED = YdbProperty.bool(
44-
"jdbc.ydb.fullscan_analyze", "Enable analizator for collecting query stats", false
44+
"jdbcFullScanDetector", "Enable analizator for collecting query stats", false
4545
);
4646

4747

‎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)