Fixed full scan detector

jdbc/src/main/java/tech/ydb/jdbc/YdbConnection.java

@@ -9,6 +9,7 @@
 import tech.ydb.jdbc.context.YdbContext;
 import tech.ydb.jdbc.context.YdbValidator;
 import tech.ydb.jdbc.query.ExplainedQuery;
+import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.query.Params;
 import tech.ydb.table.result.ResultSetReader;
 
@@ -36,27 +37,30 @@ public interface YdbConnection extends Connection {
     /**
      * Explicitly execute query as a data query
      *
-     * @param yql query to execute
+     * @param query query to execute
+     * @param yql YQL text to execute
      * @param params parameters for query
      * @param timeout timeout of operation
      * @param keepInCache flag to store query in server-side cache
      * @param validator handler for logging and warnings
      * @return list of result set
      * @throws SQLException if query cannot be executed
      */
-    List<ResultSetReader> executeDataQuery(String yql, YdbValidator validator,
+    List<ResultSetReader> executeDataQuery(YdbQuery query, String yql, YdbValidator validator,
             int timeout, boolean keepInCache, Params params) throws SQLException;
 
     /**
      * Explicitly execute query as a scan query
      *
-     * @param yql query to execute
+     * @param query query to execute
+     * @param yql YQL text to execute
      * @param params parameters for query
      * @param validator handler for logging and warnings
      * @return single result set with rows
      * @throws SQLException if query cannot be executed
      */
-    ResultSetReader executeScanQuery(String yql, YdbValidator validator, Params params) throws SQLException;
+    ResultSetReader executeScanQuery(YdbQuery query, String yql, YdbValidator validator, Params params)
+            throws SQLException;
 
     /**
      * Explicitly explain this query

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

@@ -9,6 +9,7 @@
 import tech.ydb.core.UnexpectedResultException;
 import tech.ydb.jdbc.exception.ExceptionFactory;
 import tech.ydb.jdbc.query.QueryType;
+import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.Session;
 import tech.ydb.table.TableClient;
 import tech.ydb.table.query.Params;
@@ -50,15 +51,18 @@ public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yq
     }
 
     @Override
-    public ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
-            throws SQLException {
+    public ResultSetReader executeScanQuery(
+            YdbContext ctx, YdbValidator validator, YdbQuery query, String yql, Params params
+    ) throws SQLException {
         ensureOpened();
 
         Collection<ResultSetReader> resultSets = new LinkedBlockingQueue<>();
         Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
         ExecuteScanQuerySettings settings = ExecuteScanQuerySettings.newBuilder()
                 .withRequestTimeout(scanQueryTimeout)
                 .build();
+
+        ctx.traceQuery(query, yql);
         try (Session session = createNewTableSession(validator)) {
             validator.execute(QueryType.SCAN_QUERY + " >>\n" + yql,
                     () -> session.executeScanQuery(yql, params, settings).start(resultSets::add));

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

@@ -17,6 +17,7 @@
 import tech.ydb.jdbc.exception.ExceptionFactory;
 import tech.ydb.jdbc.query.ExplainedQuery;
 import tech.ydb.jdbc.query.QueryType;
+import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.query.QueryClient;
 import tech.ydb.query.QuerySession;
 import tech.ydb.query.QueryStream;
@@ -196,7 +197,8 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException
 
     @Override
     public List<ResultSetReader> executeDataQuery(
-            YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
+            YdbContext ctx, YdbValidator validator, YdbQuery query,
+            String yql, long timeout, boolean keepInCache, Params params
     ) throws SQLException {
         ensureOpened();
 

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

@@ -5,19 +5,20 @@
 
 import tech.ydb.core.Status;
 import tech.ydb.jdbc.common.FixedResultSetFactory;
-import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.result.ResultSetReader;
 
 /**
  *
  * @author Aleksandr Gorshenin
  */
 public class QueryStat {
-    public static final String QUERY = "print_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")
             .addBooleanColumn("is_fullscan")
+            .addBooleanColumn("is_error")
             .addLongColumn("executed")
             .addTextColumn("yql")
             .addTextColumn("ast")
@@ -31,23 +32,26 @@ public class QueryStat {
     private final String plan;
     private final LongAdder usage;
     private final boolean isFullScan;
+    private final boolean isError;
 
-    public QueryStat(YdbQuery query, String ast, String plan) {
-        this.originSQL = query.getOriginQuery();
-        this.preparedYQL = query.getPreparedYql();
+    public QueryStat(String sql, String yql, String ast, String plan) {
+        this.originSQL = sql;
+        this.preparedYQL = yql;
         this.ast = ast;
         this.plan = plan;
         this.usage = new LongAdder();
         this.isFullScan = plan.contains("\"Node Type\":\"TableFullScan\"");
+        this.isError = false;
     }
 
-    public QueryStat(YdbQuery query, Status error) {
-        this.originSQL = query.getOriginQuery();
-        this.preparedYQL = query.getPreparedYql();
-        this.ast = error.toString();
+    public QueryStat(String sql, String yql, Status error) {
+        this.originSQL = sql;
+        this.preparedYQL = yql;
+        this.ast = null;
         this.plan = error.toString();
         this.usage = new LongAdder();
         this.isFullScan = false;
+        this.isError = true;
     }
 
     public long getUsageCounter() {
@@ -74,6 +78,10 @@ public boolean isFullScan() {
         return isFullScan;
     }
 
+    public boolean isError() {
+        return isError;
+    }
+
     public void incrementUsage() {
         this.usage.increment();
     }
@@ -84,6 +92,7 @@ public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
             builder.newRow()
                     .withTextValue("sql", stat.originSQL)
                     .withBoolValue("is_fullscan", stat.isFullScan)
+                    .withBoolValue("is_error", stat.isError)
                     .withLongValue("executed", stat.usage.longValue())
                     .withTextValue("yql", stat.preparedYQL)
                     .withTextValue("ast", stat.ast)
@@ -92,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());
+    }
 }

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

@@ -10,6 +10,7 @@
 import tech.ydb.jdbc.YdbConst;
 import tech.ydb.jdbc.query.ExplainedQuery;
 import tech.ydb.jdbc.query.QueryType;
+import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.Session;
 import tech.ydb.table.query.DataQueryResult;
 import tech.ydb.table.query.ExplainDataQueryResult;
@@ -168,7 +169,8 @@ public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator
 
     @Override
     public List<ResultSetReader> executeDataQuery(
-            YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
+            YdbContext ctx, YdbValidator validator, YdbQuery query,
+            String yql, long timeout, boolean keepInCache, Params params
     ) throws SQLException {
         ensureOpened();
 

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

@@ -3,12 +3,12 @@
 import java.sql.SQLDataException;
 import java.sql.SQLException;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
-import java.util.Set;
-import java.util.TreeSet;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
@@ -75,7 +75,7 @@ public class YdbContext implements AutoCloseable {
     private final SessionRetryContext retryCtx;
 
     private final Cache<String, YdbQuery> queriesCache;
-    private final Cache<String, QueryStat> queryStatesCache;
+    private final Cache<String, QueryStat> statsCache;
     private final Cache<String, Map<String, Type>> queryParamsCache;
 
     private final boolean autoResizeSessionPool;
@@ -107,13 +107,13 @@ private YdbContext(
             queriesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
             queryParamsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
             if (config.isFullScanDetectorEnabled()) {
-                queryStatesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
+                statsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
             } else {
-                queryStatesCache = null;
+                statsCache = null;
             }
         } else {
             queriesCache = null;
-            queryStatesCache = null;
+            statsCache = null;
             queryParamsCache = null;
         }
     }
@@ -184,25 +184,26 @@ public boolean hasConnections() {
     }
 
     public boolean queryStatsEnabled() {
-        return queryStatesCache != null;
+        return statsCache != null;
     }
 
-    public Collection<QueryStat> getQueryStats() {
-        if (queryStatesCache == null) {
-            return Collections.emptyList();
+    public void resetQueryStats() {
+        if (statsCache != null) {
+            statsCache.invalidateAll();
         }
-        Set<QueryStat> sortedByUsage = new TreeSet<>(Comparator.comparingLong(QueryStat::getUsageCounter).reversed());
-        sortedByUsage.addAll(queryStatesCache.asMap().values());
-        return sortedByUsage;
     }
 
-    public void traceQueryExecution(YdbQuery query) {
-        if (queryStatesCache != null) {
-            QueryStat stat = queryStatesCache.getIfPresent(query.getOriginQuery());
-            if (stat != null) {
-                stat.incrementUsage();
-            }
+    public Collection<QueryStat> getQueryStats() {
+        if (statsCache == null) {
+            return Collections.emptyList();
         }
+        List<QueryStat> sorted = new ArrayList<>(statsCache.asMap().values());
+        Collections.sort(sorted,
+                Comparator
+                        .comparingLong(QueryStat::getUsageCounter).reversed()
+                        .thenComparing(QueryStat::getPreparedYQL)
+        );
+        return sorted;
     }
 
     public void register() {
@@ -308,29 +309,42 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
             queriesCache.put(sql, cached);
         }
 
-        if (queryStatesCache != null) {
-            QueryStat stat = queryStatesCache.getIfPresent(sql);
-            if (stat == null) {
-                final String preparedYQL = cached.getPreparedYql();
-                final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
-                Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
-                        session -> session.explainDataQuery(preparedYQL, settings)
-                ).join();
 
-                if (res.isSuccess()) {
-                    ExplainDataQueryResult exp = res.getValue();
-                    stat = new QueryStat(cached, exp.getQueryAst(), exp.getQueryPlan());
-                } else {
-                    stat = new QueryStat(cached, res.getStatus());
-                }
-                queryStatesCache.put(sql, stat);
+        return cached;
+    }
+
+    public void traceQuery(YdbQuery query, String yql) {
+        if (statsCache == null) {
+            return;
+        }
+
+        QueryStat stat = statsCache.getIfPresent(yql);
+        if (stat == null) {
+            final ExplainDataQuerySettings settings = withDefaultTimeout(new ExplainDataQuerySettings());
+            Result<ExplainDataQueryResult> res = retryCtx.supplyResult(
+                    session -> session.explainDataQuery(yql, settings)
+            ).join();
+
+            if (res.isSuccess()) {
+                ExplainDataQueryResult exp = res.getValue();
+                stat = new QueryStat(query.getOriginQuery(), yql, exp.getQueryAst(), exp.getQueryPlan());
+            } else {
+                stat = new QueryStat(query.getOriginQuery(), yql, res.getStatus());
             }
+
+            statsCache.put(yql, stat);
         }
 
-        return cached;
+        stat.incrementUsage();
     }
 
     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) {

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

@@ -5,6 +5,7 @@
 
 import tech.ydb.jdbc.YdbConst;
 import tech.ydb.jdbc.query.ExplainedQuery;
+import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.table.query.Params;
 import tech.ydb.table.result.ResultSetReader;
 
@@ -34,10 +35,10 @@ default void ensureOpened() throws SQLException {
 
     void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException;
 
-    List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, String yql,
+    List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, String yql,
             long timeout, boolean poolable, Params params) throws SQLException;
 
-    ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
+    ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, String yql, Params params)
             throws SQLException;
 
     ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, String yql)

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;
@@ -189,14 +191,29 @@ protected List<YdbResult> executeExplainQuery(YdbQuery query) throws SQLExceptio
         return Collections.singletonList(new YdbResult(new YdbResultSetImpl(this, result)));
     }
 
-    protected List<YdbResult> executeScanQuery(String yql, Params params) throws SQLException {
-        ResultSetReader result = connection.executeScanQuery(yql, validator, params);
+    protected List<YdbResult> executeScanQuery(YdbQuery query, String yql, Params params) throws SQLException {
+        connection.getCtx().traceQuery(query, yql);
+        ResultSetReader result = connection.executeScanQuery(query, yql, validator, params);
         return Collections.singletonList(new YdbResult(new YdbResultSetImpl(this, result)));
     }
 
     protected List<YdbResult> executeDataQuery(YdbQuery query, String yql, Params params) throws SQLException {
+        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(yql, validator, getQueryTimeout(), isPoolable(), params);
+                .executeDataQuery(query, yql, validator, getQueryTimeout(), isPoolable(), params);
 
         List<YdbResult> results = new ArrayList<>();
         int idx = 0;