Added support of stats query via prepareStatement

alex268 · Sep 2, 2024 · 8f2932b · 8f2932b
1 parent c923a4d
commit 8f2932b
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 103 deletions.
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java
@@ -12,8 +12,8 @@
  * @author Aleksandr Gorshenin
  */
 public class QueryStat {
- public static final String PRINT_QUERY = "print_jdbc_stats();";
- public static final String RESET_QUERY = "reset_jdbc_stats();";
+ private static final String PRINT_QUERY = "print_jdbc_stats();";
+ private static final String RESET_QUERY = "reset_jdbc_stats();";
 
  private static final FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
  .addTextColumn("sql")
@@ -101,4 +101,12 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
  }
  return builder.build();
  }
+
+ public static boolean isPrint(String sql) {
+ return sql != null && PRINT_QUERY.equalsIgnoreCase(sql.trim());
+ }
+
+ public static boolean isReset(String sql) {
+ return sql != null && RESET_QUERY.equalsIgnoreCase(sql.trim());
+ }
 }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
@@ -339,6 +339,12 @@ public void traceQuery(YdbQuery query, String yql) {
  }
 
  public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) throws SQLException {
+ if (statsCache != null) {
+ if (QueryStat.isPrint(query.getOriginQuery()) || QueryStat.isReset(query.getOriginQuery())) {
+ return new InMemoryQuery(query, queryOptions.isDeclareJdbcParameters());
+ }
+ }
+
  if (query.getYqlBatcher() != null && mode == YdbPrepareMode.AUTO) {
  Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
  if (types == null) {

diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java
@@ -16,6 +16,8 @@
 import tech.ydb.jdbc.YdbResultSet;
 import tech.ydb.jdbc.YdbStatement;
 import tech.ydb.jdbc.common.FixedResultSetFactory;
+import tech.ydb.jdbc.context.QueryStat;
+import tech.ydb.jdbc.context.YdbContext;
 import tech.ydb.jdbc.context.YdbValidator;
 import tech.ydb.jdbc.query.ExplainedQuery;
 import tech.ydb.jdbc.query.QueryStatement;
@@ -196,7 +198,20 @@ protected List<YdbResult> executeScanQuery(YdbQuery query, String yql, Params pa
  }
 
  protected List<YdbResult> executeDataQuery(YdbQuery query, String yql, Params params) throws SQLException {
- connection.getCtx().traceQuery(query, yql);
+ YdbContext ctx = connection.getCtx();
+
+ if (ctx.queryStatsEnabled()) {
+ if (QueryStat.isPrint(yql)) {
+ YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
+ return Collections.singletonList(new YdbResult(rs));
+ }
+ if (QueryStat.isReset(yql)) {
+ getConnection().getCtx().resetQueryStats();
+ return null;
+ }
+ }
+
+ ctx.traceQuery(query, yql);
  List<ResultSetReader> resultSets = connection
  .executeDataQuery(query, yql, validator, getQueryTimeout(), isPoolable(), params);
 

diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java
@@ -5,16 +5,13 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import tech.ydb.jdbc.YdbConnection;
 import tech.ydb.jdbc.YdbConst;
 import tech.ydb.jdbc.YdbResultSet;
-import tech.ydb.jdbc.context.QueryStat;
-import tech.ydb.jdbc.context.YdbContext;
 import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.query.Params;
 
@@ -82,19 +79,7 @@ public int executeUpdate(String sql) throws SQLException {
  public boolean execute(String sql) throws SQLException {
  cleanState();
 
- YdbContext ctx = getConnection().getCtx();
- if (ctx.queryStatsEnabled() && sql != null) {
- if (QueryStat.PRINT_QUERY.equalsIgnoreCase(sql.trim())) {
- YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
- return updateState(Collections.singletonList(new YdbResult(rs)));
- }
- if (QueryStat.RESET_QUERY.equalsIgnoreCase(sql.trim())) {
- getConnection().getCtx().resetQueryStats();
- return updateState(null);
- }
- }
-
- YdbQuery query = ctx.parseYdbQuery(sql);
+ YdbQuery query = getConnection().getCtx().parseYdbQuery(sql);
  List<YdbResult> newState = null;
  switch (query.getType()) {
  case SCHEME_QUERY:

diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java
@@ -969,7 +969,6 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
  .assertNoRows();
  }
 
- // scheme queries don't collect stats
  ExceptionAssert.ydbException("Cannot find table", () -> st.execute(wrongQuery));
 
  try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
@@ -1090,12 +1089,10 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  String preparedSelectByColumn = QUERIES.selectAllByColumnValue("c_Text", "?");
 
  try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- sa.check(rs)
- .assertMetaColumns()
- .assertNoRows();
- }
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ sa.check(ps.executeQuery())
+ .assertMetaColumns()
+ .assertNoRows();
  }
 
  try (PreparedStatement ps = connection.prepareStatement(preparedSelectByKey)) {
@@ -1106,18 +1103,16 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  ps.execute();
  }
 
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where key = ?"),
-  sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
-  sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where key = ?"),
+ sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
+ sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
- check.assertNoRows();
- }
+ check.assertNoRows();
  }
 
  try (PreparedStatement ps = connection.prepareStatement(preparedSelectByColumn)) {
@@ -1128,38 +1123,39 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  ps.execute();
  }
 
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where key = ?"),
-  sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
-  sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where key = ?"),
+ sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
+ sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where c_Text = ?"),
-  sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
-  sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where c_Text = ?"),
+ sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
+ sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where c_Text = ?"),
-  sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
-  sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where c_Text = ?"),
+ sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
+ sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.assertNoRows();
-  }
+ check.assertNoRows();
+ }
 
- Assertions.assertFalse(st.execute("reset_JDBC_stats();"));
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- sa.check(rs)
- .assertMetaColumns()
- .assertNoRows();
+ try (PreparedStatement ps = connection.prepareStatement("reset_JDBC_stats();")) {
+ Assertions.assertFalse(ps.execute());
+ }
+
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ sa.check(ps.executeQuery())
+ .assertMetaColumns()
+ .assertNoRows();
  }
  }
  }
 }
-}
diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbQueryConnectionImplTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbQueryConnectionImplTest.java
@@ -11,7 +11,6 @@
 import java.util.Map;
 import java.util.Properties;
 
-import org.junit.Assert;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
@@ -987,7 +986,7 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
  check.assertNoRows();
  }
 
- Assert.assertFalse(st.execute("reset_jdbc_stats();\n"));
+ Assertions.assertFalse(st.execute("reset_jdbc_stats();\n"));
  try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
  sa.check(rs)
  .assertMetaColumns()
@@ -1075,7 +1074,7 @@ public void fullScanAnalyzerStatementTest() throws SQLException {
  check.assertNoRows();
  }
 
- Assert.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
+ Assertions.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
  try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
  sa.check(rs)
  .assertMetaColumns()
@@ -1093,12 +1092,10 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  String preparedSelectByColumn = QUERIES.selectAllByColumnValue("c_Text", "?");
 
  try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- sa.check(rs)
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ sa.check(ps.executeQuery())
  .assertMetaColumns()
  .assertNoRows();
- }
  }
 
  try (PreparedStatement ps = connection.prepareStatement(preparedSelectByKey)) {
@@ -1109,18 +1106,16 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  ps.execute();
  }
 
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where key = ?"),
-  sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
-  sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where key = ?"),
+ sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
+ sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
- check.assertNoRows();
- }
+ check.assertNoRows();
  }
 
  try (PreparedStatement ps = connection.prepareStatement(preparedSelectByColumn)) {
@@ -1131,37 +1126,38 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
  ps.execute();
  }
 
- try (Statement st = connection.createStatement()) {
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- TableAssert.ResultSetAssert check = sa.check(rs).assertMetaColumns();
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ TableAssert.ResultSetAssert check = sa.check(ps.executeQuery()).assertMetaColumns();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where key = ?"),
-  sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
-  sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where key = ?"),
+ sa.yql("DECLARE $jp1 AS Int32;\nselect * from ydb_connection_test where key = $jp1"),
+ sa.isNotFullScan(), sa.isNotError(), sa.executed(2), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where c_Text = ?"),
-  sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
-  sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where c_Text = ?"),
+ sa.yql("DECLARE $jp1 AS Text;\nselect * from ydb_connection_test where c_Text = $jp1"),
+ sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.nextRow(
-  sa.sql("select * from ydb_connection_test where c_Text = ?"),
-  sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
-  sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
-  ).assertAll();
+ check.nextRow(
+ sa.sql("select * from ydb_connection_test where c_Text = ?"),
+ sa.yql("DECLARE $jp1 AS Text?;\nselect * from ydb_connection_test where c_Text = $jp1"),
+ sa.isFullScan(), sa.isNotError(), sa.executed(1), sa.hasAst(), sa.hasPlan()
+ ).assertAll();
 
-  check.assertNoRows();
-  }
+ check.assertNoRows();
+ }
 
- Assert.assertFalse(st.execute("reset_JDBC_stats();"));
- try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
- sa.check(rs)
+ try (PreparedStatement ps = connection.prepareStatement("reset_JDBC_stats();")) {
+ Assertions.assertFalse(ps.execute());
+ }
+
+ try (PreparedStatement ps = connection.prepareStatement("print_JDBC_stats();")) {
+ sa.check(ps.executeQuery())
  .assertMetaColumns()
  .assertNoRows();
- }
  }
  }
  }