Skip to content

Commit 8f2932b

Browse files
alex268alex268
committed
Added support of stats query via prepareStatement
1 parent c923a4d commit 8f2932b

File tree

6 files changed

+109
-103
lines changed

6 files changed

+109
-103
lines changed

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
* @author Aleksandr Gorshenin
1313
*/
1414
public class QueryStat {
15-
public static final String PRINT_QUERY = "print_jdbc_stats();";
16-
public static final String RESET_QUERY = "reset_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")
@@ -101,4 +101,12 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
101101
}
102102
return builder.build();
103103
}
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+
}
104112
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,12 @@ public void traceQuery(YdbQuery query, String yql) {
339339
}
340340

341341
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+
342348
if (query.getYqlBatcher() != null && mode == YdbPrepareMode.AUTO) {
343349
Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
344350
if (types == null) {

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

Lines changed: 16 additions & 1 deletion
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;
@@ -196,7 +198,20 @@ protected List<YdbResult> executeScanQuery(YdbQuery query, String yql, Params pa
196198
}
197199

198200
protected List<YdbResult> executeDataQuery(YdbQuery query, String yql, Params params) throws SQLException {
199-
connection.getCtx().traceQuery(query, yql);
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);
200215
List<ResultSetReader> resultSets = connection
201216
.executeDataQuery(query, yql, validator, getQueryTimeout(), isPoolable(), params);
202217

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

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,13 @@
55
import java.sql.Statement;
66
import java.util.ArrayList;
77
import java.util.Arrays;
8-
import java.util.Collections;
98
import java.util.List;
109
import java.util.logging.Level;
1110
import java.util.logging.Logger;
1211

1312
import tech.ydb.jdbc.YdbConnection;
1413
import tech.ydb.jdbc.YdbConst;
1514
import tech.ydb.jdbc.YdbResultSet;
16-
import tech.ydb.jdbc.context.QueryStat;
17-
import tech.ydb.jdbc.context.YdbContext;
1815
import tech.ydb.jdbc.query.YdbQuery;
1916
import tech.ydb.table.query.Params;
2017

@@ -82,19 +79,7 @@ public int executeUpdate(String sql) throws SQLException {
8279
public boolean execute(String sql) throws SQLException {
8380
cleanState();
8481

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

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

Lines changed: 39 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -969,7 +969,6 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
969969
.assertNoRows();
970970
}
971971

972-
// scheme queries don't collect stats
973972
ExceptionAssert.ydbException("Cannot find table", () -> st.execute(wrongQuery));
974973

975974
try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
@@ -1090,12 +1089,10 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
10901089
String preparedSelectByColumn = QUERIES.selectAllByColumnValue("c_Text", "?");
10911090

10921091
try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
1093-
try (Statement st = connection.createStatement()) {
1094-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1095-
sa.check(rs)
1096-
.assertMetaColumns()
1097-
.assertNoRows();
1098-
}
1092+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1093+
sa.check(ps.executeQuery())
1094+
.assertMetaColumns()
1095+
.assertNoRows();
10991096
}
11001097

11011098
try (PreparedStatement ps = connection.prepareStatement(preparedSelectByKey)) {
@@ -1106,18 +1103,16 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11061103
ps.execute();
11071104
}
11081105

1109-
try (Statement st = connection.createStatement()) {
1110-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1111-
TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
1106+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1107+
TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
11121108

1113-
check.nextRow(
1114-
sa.sql("select * from ydb_connection_test where key = ?"),
1115-
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1116-
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1117-
).assertAll();
1109+
check.nextRow(
1110+
sa.sql("select * from ydb_connection_test where key = ?"),
1111+
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1112+
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1113+
).assertAll();
11181114

1119-
check.assertNoRows();
1120-
}
1115+
check.assertNoRows();
11211116
}
11221117

11231118
try (PreparedStatement ps = connection.prepareStatement(preparedSelectByColumn)) {
@@ -1128,38 +1123,39 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11281123
ps.execute();
11291124
}
11301125

1131-
try (Statement st = connection.createStatement()) {
1132-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1133-
TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
1126+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1127+
TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
11341128

1135-
check.nextRow(
1136-
sa.sql("select * from ydb_connection_test where key = ?"),
1137-
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1138-
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1139-
).assertAll();
1129+
check.nextRow(
1130+
sa.sql("select * from ydb_connection_test where key = ?"),
1131+
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1132+
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1133+
).assertAll();
11401134

1141-
check.nextRow(
1142-
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1143-
sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
1144-
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1145-
).assertAll();
1135+
check.nextRow(
1136+
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1137+
sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
1138+
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1139+
).assertAll();
11461140

1147-
check.nextRow(
1148-
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1149-
sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
1150-
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1151-
).assertAll();
1141+
check.nextRow(
1142+
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1143+
sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
1144+
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1145+
).assertAll();
11521146

1153-
check.assertNoRows();
1154-
}
1147+
check.assertNoRows();
1148+
}
11551149

1156-
Assertions.assertFalse(st.execute("reset_JDBC_stats();"));
1157-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1158-
sa.check(rs)
1159-
.assertMetaColumns()
1160-
.assertNoRows();
1150+
try (PreparedStatement ps = connection.prepareStatement("reset_JDBC_stats();")) {
1151+
Assertions.assertFalse(ps.execute());
1152+
}
1153+
1154+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1155+
sa.check(ps.executeQuery())
1156+
.assertMetaColumns()
1157+
.assertNoRows();
11611158
}
11621159
}
11631160
}
11641161
}
1165-
}

‎jdbc/src/test/java/tech/ydb/jdbc/impl/YdbQueryConnectionImplTest.java

Lines changed: 37 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
import java.util.Map;
1212
import java.util.Properties;
1313

14-
import org.junit.Assert;
1514
import org.junit.jupiter.api.AfterAll;
1615
import org.junit.jupiter.api.AfterEach;
1716
import org.junit.jupiter.api.Assertions;
@@ -987,7 +986,7 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
987986
check.assertNoRows();
988987
}
989988

990-
Assert.assertFalse(st.execute("reset_jdbc_stats();\n"));
989+
Assertions.assertFalse(st.execute("reset_jdbc_stats();\n"));
991990
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
992991
sa.check(rs)
993992
.assertMetaColumns()
@@ -1075,7 +1074,7 @@ public void fullScanAnalyzerStatementTest() throws SQLException {
10751074
check.assertNoRows();
10761075
}
10771076

1078-
Assert.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
1077+
Assertions.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
10791078
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
10801079
sa.check(rs)
10811080
.assertMetaColumns()
@@ -1093,12 +1092,10 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
10931092
String preparedSelectByColumn = QUERIES.selectAllByColumnValue("c_Text", "?");
10941093

10951094
try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
1096-
try (Statement st = connection.createStatement()) {
1097-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1098-
sa.check(rs)
1095+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1096+
sa.check(ps.executeQuery())
10991097
.assertMetaColumns()
11001098
.assertNoRows();
1101-
}
11021099
}
11031100

11041101
try (PreparedStatement ps = connection.prepareStatement(preparedSelectByKey)) {
@@ -1109,18 +1106,16 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11091106
ps.execute();
11101107
}
11111108

1112-
try (Statement st = connection.createStatement()) {
1113-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1114-
TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
1109+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1110+
TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
11151111

1116-
check.nextRow(
1117-
sa.sql("select * from ydb_connection_test where key = ?"),
1118-
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1119-
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1120-
).assertAll();
1112+
check.nextRow(
1113+
sa.sql("select * from ydb_connection_test where key = ?"),
1114+
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1115+
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1116+
).assertAll();
11211117

1122-
check.assertNoRows();
1123-
}
1118+
check.assertNoRows();
11241119
}
11251120

11261121
try (PreparedStatement ps = connection.prepareStatement(preparedSelectByColumn)) {
@@ -1131,37 +1126,38 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11311126
ps.execute();
11321127
}
11331128

1134-
try (Statement st = connection.createStatement()) {
1135-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1136-
TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
1129+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1130+
TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
11371131

1138-
check.nextRow(
1139-
sa.sql("select * from ydb_connection_test where key = ?"),
1140-
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1141-
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1142-
).assertAll();
1132+
check.nextRow(
1133+
sa.sql("select * from ydb_connection_test where key = ?"),
1134+
sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
1135+
sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
1136+
).assertAll();
11431137

1144-
check.nextRow(
1145-
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1146-
sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
1147-
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1148-
).assertAll();
1138+
check.nextRow(
1139+
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1140+
sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
1141+
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1142+
).assertAll();
11491143

1150-
check.nextRow(
1151-
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1152-
sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
1153-
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1154-
).assertAll();
1144+
check.nextRow(
1145+
sa.sql("select * from ydb_connection_test where c_Text = ?"),
1146+
sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
1147+
sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
1148+
).assertAll();
11551149

1156-
check.assertNoRows();
1157-
}
1150+
check.assertNoRows();
1151+
}
11581152

1159-
Assert.assertFalse(st.execute("reset_JDBC_stats();"));
1160-
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1161-
sa.check(rs)
1153+
try (PreparedStatement ps = connection.prepareStatement("reset_JDBC_stats();")) {
1154+
Assertions.assertFalse(ps.execute());
1155+
}
1156+
1157+
try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
1158+
sa.check(ps.executeQuery())
11621159
.assertMetaColumns()
11631160
.assertNoRows();
1164-
}
11651161
}
11661162
}
11671163
}

0 commit comments

Comments
 (0)