Skip to content

Commit c923a4d

Browse files
alex268alex268
committed
Added method to reset stats
1 parent 7736d9e commit c923a4d

File tree

5 files changed

+70
-15
lines changed

5 files changed

+70
-15
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@
1212
* @author Aleksandr Gorshenin
1313
*/
1414
public class QueryStat {
15-
public static final String QUERY = "print_jdbc_stats();";
15+
public static final String PRINT_QUERY = "print_jdbc_stats();";
16+
public static final String RESET_QUERY = "reset_jdbc_stats();";
1617

1718
private static final FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
1819
.addTextColumn("sql")

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

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -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,14 +184,20 @@ public boolean hasConnections() {
184184
}
185185

186186
public boolean queryStatsEnabled() {
187-
return queryStatesCache != null;
187+
return statsCache != null;
188+
}
189+
190+
public void resetQueryStats() {
191+
if (statsCache != null) {
192+
statsCache.invalidateAll();
193+
}
188194
}
189195

190196
public Collection<QueryStat> getQueryStats() {
191-
if (queryStatesCache == null) {
197+
if (statsCache == null) {
192198
return Collections.emptyList();
193199
}
194-
List<QueryStat> sorted = new ArrayList<>(queryStatesCache.asMap().values());
200+
List<QueryStat> sorted = new ArrayList<>(statsCache.asMap().values());
195201
Collections.sort(sorted,
196202
Comparator
197203
.comparingLong(QueryStat::getUsageCounter).reversed()
@@ -308,11 +314,11 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
308314
}
309315

310316
public void traceQuery(YdbQuery query, String yql) {
311-
if (queryStatesCache == null) {
317+
if (statsCache == null) {
312318
return;
313319
}
314320

315-
QueryStat stat = queryStatesCache.getIfPresent(yql);
321+
QueryStat stat = statsCache.getIfPresent(yql);
316322
if (stat == null) {
317323
final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
318324
Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
@@ -326,7 +332,7 @@ public void traceQuery(YdbQuery query, String yql) {
326332
stat = new QueryStat(query.getOriginQuery(), yql, res.getStatus());
327333
}
328334

329-
queryStatesCache.put(yql, stat);
335+
statsCache.put(yql, stat);
330336
}
331337

332338
stat.incrementUsage();

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,15 @@ public boolean execute(String sql) throws SQLException {
8383
cleanState();
8484

8585
YdbContext ctx = getConnection().getCtx();
86-
if (ctx.queryStatsEnabled() && sql != null && QueryStat.QUERY.equalsIgnoreCase(sql.trim())) {
87-
YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
88-
return updateState(Collections.singletonList(new YdbResult(rs)));
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+
}
8995
}
9096

9197
YdbQuery query = ctx.parseYdbQuery(sql);

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -983,6 +983,13 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
983983

984984
check.assertNoRows();
985985
}
986+
987+
Assertions.assertFalse(st.execute("reset_jdbc_stats();\n"));
988+
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
989+
sa.check(rs)
990+
.assertMetaColumns()
991+
.assertNoRows();
992+
}
986993
}
987994
}
988995
}
@@ -1064,6 +1071,13 @@ public void fullScanAnalyzerStatementTest() throws SQLException {
10641071

10651072
check.assertNoRows();
10661073
}
1074+
1075+
Assertions.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
1076+
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1077+
sa.check(rs)
1078+
.assertMetaColumns()
1079+
.assertNoRows();
1080+
}
10671081
}
10681082
}
10691083
}
@@ -1138,7 +1152,14 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11381152

11391153
check.assertNoRows();
11401154
}
1155+
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();
11411161
}
11421162
}
11431163
}
11441164
}
1165+
}

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

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

14+
import org.junit.Assert;
1415
import org.junit.jupiter.api.AfterAll;
1516
import org.junit.jupiter.api.AfterEach;
1617
import org.junit.jupiter.api.Assertions;
@@ -972,7 +973,6 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
972973
.assertNoRows();
973974
}
974975

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

978978
try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
@@ -986,6 +986,13 @@ public void fullScanAnalyzerSchemeWrongQueryTest() throws SQLException {
986986

987987
check.assertNoRows();
988988
}
989+
990+
Assert.assertFalse(st.execute("reset_jdbc_stats();\n"));
991+
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
992+
sa.check(rs)
993+
.assertMetaColumns()
994+
.assertNoRows();
995+
}
989996
}
990997
}
991998
}
@@ -1067,6 +1074,13 @@ public void fullScanAnalyzerStatementTest() throws SQLException {
10671074

10681075
check.assertNoRows();
10691076
}
1077+
1078+
Assert.assertFalse(st.execute("\t\treSet_jdbc_statS();"));
1079+
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1080+
sa.check(rs)
1081+
.assertMetaColumns()
1082+
.assertNoRows();
1083+
}
10701084
}
10711085
}
10721086
}
@@ -1141,6 +1155,13 @@ public void fullScanAnalyzerPreparedStatementTest() throws SQLException {
11411155

11421156
check.assertNoRows();
11431157
}
1158+
1159+
Assert.assertFalse(st.execute("reset_JDBC_stats();"));
1160+
try (ResultSet rs = st.executeQuery("print_JDBC_stats();")) {
1161+
sa.check(rs)
1162+
.assertMetaColumns()
1163+
.assertNoRows();
1164+
}
11441165
}
11451166
}
11461167
}

0 commit comments

Comments
 (0)