Query execution interceptor

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

@@ -4,6 +4,7 @@
 import java.sql.Statement;
 
 import tech.ydb.jdbc.context.YdbValidator;
+import tech.ydb.table.query.stats.QueryStatsCollectionMode;
 
 public interface YdbStatement extends Statement {
     /**
@@ -58,4 +59,8 @@ public interface YdbStatement extends Statement {
 
     @Override
     int getMaxRows();
+
+    void setStatsCollectionMode(QueryStatsCollectionMode mode);
+
+    QueryStatsCollectionMode getStatsCollectionMode();
 }

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

@@ -32,6 +32,7 @@
 import tech.ydb.query.settings.CommitTransactionSettings;
 import tech.ydb.query.settings.ExecuteQuerySettings;
 import tech.ydb.query.settings.QueryExecMode;
+import tech.ydb.query.settings.QueryStatsMode;
 import tech.ydb.query.settings.RollbackTransactionSettings;
 import tech.ydb.query.tools.QueryReader;
 import tech.ydb.table.query.Params;
@@ -233,6 +234,7 @@ protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query
 
         int timeout = statement.getQueryTimeout();
         ExecuteQuerySettings.Builder settings = ExecuteQuerySettings.newBuilder();
+        settings = settings.withStatsMode(QueryStatsMode.valueOf(statement.getStatsCollectionMode().name()));
         if (timeout > 0) {
             settings = settings.withRequestTimeout(timeout, TimeUnit.SECONDS);
         }
@@ -300,7 +302,10 @@ public void onClose(Status status, Throwable th) {
             for (int idx = 0; idx < readers.length; idx++) {
                 readers[idx] = new YdbResultSetMemory(types, statement, result.getResultSet(idx));
             }
-            return updateCurrentResult(new YdbQueryResultStatic(query, readers));
+
+            YdbQueryResultStatic queryResult = new YdbQueryResultStatic(query, readers);
+            queryResult.setQueryStats(result.getQueryInfo().getStats());
+            return updateCurrentResult(queryResult);
         } finally {
             if (!localTx.isActive()) {
                 if (tx.compareAndSet(localTx, null)) {

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

@@ -66,7 +66,7 @@ public String getPreparedYQL() {
         return preparedYQL;
     }
 
-    public String getAat() {
+    public String getAst() {
         return ast;
     }
 

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

@@ -173,6 +173,8 @@ private ExecuteDataQuerySettings dataQuerySettings(YdbStatement statement) {
             settings = settings.disableQueryCache();
         }
 
+        settings = settings.setCollectStats(statement.getStatsCollectionMode());
+
         return settings;
     }
 
@@ -231,7 +233,9 @@ protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query
                 readers[idx] = new YdbResultSetMemory(types, statement, rs);
             }
 
-            return updateCurrentResult(new YdbQueryResultStatic(query, readers));
+            YdbQueryResultStatic queryResult = new YdbQueryResultStatic(query, readers);
+            queryResult.setQueryStats(result.getQueryStats());
+            return updateCurrentResult(queryResult);
         } catch (SQLException | RuntimeException ex) {
             updateState(tx.withRollback(session));
             throw ex;

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

@@ -3,6 +3,8 @@
 import java.sql.SQLException;
 import java.time.Duration;
 import java.util.Collection;
+import java.util.Iterator;
+import java.util.ServiceLoader;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
@@ -23,6 +25,7 @@
 import tech.ydb.jdbc.settings.YdbConnectionProperties;
 import tech.ydb.jdbc.settings.YdbOperationProperties;
 import tech.ydb.jdbc.settings.YdbQueryProperties;
+import tech.ydb.jdbc.spi.YDBQueryExtensionService;
 import tech.ydb.query.QueryClient;
 import tech.ydb.query.impl.QueryClientImpl;
 import tech.ydb.scheme.SchemeClient;
@@ -58,6 +61,8 @@ public class YdbContext implements AutoCloseable {
     private final boolean autoResizeSessionPool;
     private final AtomicInteger connectionsCount = new AtomicInteger();
 
+    private YDBQueryExtensionService queryExtensionService = null;
+
     private YdbContext(
             YdbConfig config,
             YdbOperationProperties operationProperties,
@@ -96,6 +101,13 @@ private YdbContext(
             this.cache = new YdbCache(this,
                     queryProperties, config.getPreparedStatementsCachecSize(), config.isFullScanDetectorEnabled());
         }
+
+        Iterator<YDBQueryExtensionService> extLoaderIterator = ServiceLoader
+                .load(YDBQueryExtensionService.class)
+                .iterator();
+        if (extLoaderIterator.hasNext()) {
+            queryExtensionService = extLoaderIterator.next();
+        }
     }
 
     public YdbTypes getTypes() {
@@ -312,4 +324,8 @@ public YdbQuery parseYdbQuery(QueryKey key) throws SQLException {
     public YdbPreparedQuery prepareYdbQuery(YdbQuery query, YdbPrepareMode mode) throws SQLException {
         return cache.prepareYdbQuery(query, mode);
     }
+
+    public YDBQueryExtensionService getQueryExtensionService() {
+        return queryExtensionService;
+    }
 }

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbQueryResultStatic.java

@@ -4,6 +4,7 @@
 import java.sql.SQLException;
 
 import tech.ydb.jdbc.YdbResultSet;
+import tech.ydb.jdbc.query.UnifiedQueryStats;
 import tech.ydb.jdbc.query.YdbQuery;
 
 /**
@@ -12,6 +13,7 @@
  */
 public class YdbQueryResultStatic extends YdbQueryResultBase {
     private final YdbResultSet[] rs;
+    private UnifiedQueryStats queryStats;
 
     public YdbQueryResultStatic(YdbQuery query, YdbResultSet... rs) {
         super(query, rs != null ? rs.length : 0);
@@ -33,4 +35,20 @@ protected void closeResultSet(int index) throws SQLException {
         }
         rs[index].close();
     }
+
+    public UnifiedQueryStats getQueryStats() {
+        return queryStats;
+    }
+
+    public void setQueryStats(tech.ydb.table.query.stats.QueryStats src) {
+        if (src != null) {
+            queryStats = new UnifiedQueryStats(src);
+        }
+    }
+
+    public void setQueryStats(tech.ydb.query.result.QueryStats src) {
+        if (src != null) {
+            queryStats = new UnifiedQueryStats(src);
+        }
+    }
 }

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementBase.java

@@ -19,7 +19,9 @@
 import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.jdbc.settings.FakeTxMode;
 import tech.ydb.jdbc.settings.YdbOperationProperties;
+import tech.ydb.jdbc.spi.YDBQueryExtensionService;
 import tech.ydb.table.query.Params;
+import tech.ydb.table.query.stats.QueryStatsCollectionMode;
 import tech.ydb.table.result.ResultSetReader;
 import tech.ydb.table.values.ListValue;
 
@@ -41,6 +43,7 @@ public abstract class YdbStatementBase implements YdbStatement {
     private int queryTimeout;
     private boolean isPoolable;
     private boolean isClosed = false;
+    private QueryStatsCollectionMode statsMode = QueryStatsCollectionMode.NONE;
 
     /** @see Statement#getMaxRows() */
     private int maxRows = 0; // no limit
@@ -203,7 +206,20 @@ protected YdbQueryResult executeDataQuery(YdbQuery query, String yql, Params par
         }
 
         ctx.traceQueryByFullScanDetector(query, yql);
-        return connection.getExecutor().executeDataQuery(this, query, yql, params);
+
+        YDBQueryExtensionService ext = ctx.getQueryExtensionService();
+        if (ext != null) {
+            ext.dataQueryPreExecute(ctx, this, query, yql, params);
+        }
+
+        YdbQueryResult result;
+        result = connection.getExecutor().executeDataQuery(this, query, yql, params);
+
+        if (ext != null) {
+            ext.dataQueryPostExecute(ctx, this, query, yql, params, result);
+        }
+
+        return result;
     }
 
     protected YdbQueryResult executeSchemeQuery(YdbQuery query) throws SQLException {
@@ -319,4 +335,14 @@ public int getFetchSize() {
     public int getResultSetConcurrency() {
         return ResultSet.CONCUR_READ_ONLY;
     }
+
+    @Override
+    public QueryStatsCollectionMode getStatsCollectionMode() {
+        return statsMode;
+    }
+
+    @Override
+    public void setStatsCollectionMode(QueryStatsCollectionMode mode) {
+        this.statsMode = mode == null ? QueryStatsCollectionMode.NONE : mode;
+    }
 }