diff --git a/jdbc/src/main/java/tech/ydb/jdbc/YdbStatement.java b/jdbc/src/main/java/tech/ydb/jdbc/YdbStatement.java index ed13055..4d6e73a 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/YdbStatement.java +++ b/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(); } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java index b113f4c..31d089b 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java +++ b/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)) { diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java index ccdc15e..8736a53 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java +++ b/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; } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java index 28fe73a..e90ee9b 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java +++ b/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; diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java index 7fe4885..764830b 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java +++ b/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 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; + } } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbQueryResultStatic.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbQueryResultStatic.java index f95a7f8..20cadf0 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbQueryResultStatic.java +++ b/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); + } + } } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementBase.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementBase.java index af961a0..7965048 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementBase.java +++ b/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; + } } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/UnifiedQueryStats.java b/jdbc/src/main/java/tech/ydb/jdbc/query/UnifiedQueryStats.java new file mode 100644 index 0000000..780e4ac --- /dev/null +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/UnifiedQueryStats.java @@ -0,0 +1,230 @@ +package tech.ydb.jdbc.query; + +import java.util.List; + +import static java.util.stream.Collectors.toList; + +/** + * Data class to wrap QueryStats of Table and Query services + */ +public class UnifiedQueryStats { + private final String queryPlan; + private final String queryAst; + private final long totalDurationUs; + private final long totalCpuTimeUs; + private final long processCpuTimeUs; + private final long compilationTimeUs; + private final long compilationCpuTimeUs; + private final boolean compilationIsCached; + private final List queryPhases; + + public UnifiedQueryStats(tech.ydb.query.result.QueryStats src) { + this.queryPlan = src.getQueryPlan(); + this.queryAst = src.getQueryAst(); + this.totalDurationUs = src.getTotalDurationUs(); + this.totalCpuTimeUs = src.getTotalCpuTimeUs(); + this.processCpuTimeUs = src.getProcessCpuTimeUs(); + this.compilationTimeUs = src.getCompilationStats().getDurationUs(); + this.compilationCpuTimeUs = src.getCompilationStats().getCpuTimeUs(); + this.compilationIsCached = src.getCompilationStats().isFromCache(); + this.queryPhases = src.getPhases().stream().map(PhaseStats::new).collect(toList()); + } + + public UnifiedQueryStats(tech.ydb.table.query.stats.QueryStats src) { + this.queryPlan = src.getQueryPlan(); + this.queryAst = src.getQueryAst(); + this.totalDurationUs = src.getTotalDurationUs(); + this.totalCpuTimeUs = src.getTotalCpuTimeUs(); + this.processCpuTimeUs = src.getProcessCpuTimeUs(); + this.compilationTimeUs = src.getCompilation().getDurationUs(); + this.compilationCpuTimeUs = src.getCompilation().getCpuTimeUs(); + this.compilationIsCached = src.getCompilation().getFromCache(); + this.queryPhases = src.getQueryPhasesList().stream().map(PhaseStats::new).collect(toList()); + } + + public String getQueryPlan() { + return queryPlan; + } + + public String getQueryAst() { + return queryAst; + } + + public long getTotalDurationUs() { + return totalDurationUs; + } + + public long getTotalCpuTimeUs() { + return totalCpuTimeUs; + } + + public long getProcessCpuTimeUs() { + return processCpuTimeUs; + } + + public long getCompilationTimeUs() { + return compilationTimeUs; + } + + public long getCompilationCpuTimeUs() { + return compilationCpuTimeUs; + } + + public boolean isCompilationIsCached() { + return compilationIsCached; + } + + public List getQueryPhases() { + return queryPhases; + } + + @Override + public String toString() { + return "UnifiedQueryStats{" + + "queryPlan='" + queryPlan + '\'' + + ", queryAst='" + queryAst + '\'' + + ", totalDurationUs=" + totalDurationUs + + ", totalCpuTimeUs=" + totalCpuTimeUs + + ", processCpuTimeUs=" + processCpuTimeUs + + ", compilationTimeUs=" + compilationTimeUs + + ", compilationCpuTimeUs=" + compilationCpuTimeUs + + ", compilationIsCached=" + compilationIsCached + + ", queryPhases=" + queryPhases + + '}'; + } + + public static class PhaseStats { + private final long durationUs; + private final long cpuTimeUs; + private final long affectedShards; + private final boolean isLiteralPhase; + private final List tableAccesses; + + public PhaseStats(tech.ydb.query.result.QueryStats.QueryPhase src) { + this.durationUs = src.getDurationUs(); + this.cpuTimeUs = src.getCpuTimeUs(); + this.affectedShards = src.getAffectedShards(); + this.isLiteralPhase = src.isLiteralPhase(); + this.tableAccesses = src.getTableAccesses().stream().map(TableAccess::new).collect(toList()); + } + + public PhaseStats(tech.ydb.table.query.stats.QueryPhaseStats src) { + this.durationUs = src.getDurationUs(); + this.cpuTimeUs = src.getCpuTimeUs(); + this.affectedShards = src.getAffectedShards(); + this.isLiteralPhase = src.getLiteralPhase(); + this.tableAccesses = src.getTableAccessList().stream().map(TableAccess::new).collect(toList()); + } + + public long getDurationUs() { + return durationUs; + } + + public long getCpuTimeUs() { + return cpuTimeUs; + } + + public long getAffectedShards() { + return affectedShards; + } + + public boolean isLiteralPhase() { + return isLiteralPhase; + } + + public List getTableAccesses() { + return tableAccesses; + } + + @Override + public String toString() { + return "PhaseStats{" + + "durationUs=" + durationUs + + ", cpuTimeUs=" + cpuTimeUs + + ", affectedShards=" + affectedShards + + ", isLiteralPhase=" + isLiteralPhase + + ", tableAccesses=" + tableAccesses + + '}'; + } + + public static class TableAccess { + private final String name; + private final long partitionsCount; + private final TableOperation reads; + private final TableOperation updates; + private final TableOperation deletes; + + public TableAccess(tech.ydb.query.result.QueryStats.TableAccess src) { + this.name = src.getTableName(); + this.partitionsCount = src.getPartitionsCount(); + this.reads = new TableOperation(src.getReads().getRows(), src.getReads().getBytes()); + this.updates = new TableOperation(src.getUpdates().getRows(), src.getUpdates().getBytes()); + this.deletes = new TableOperation(src.getDeletes().getRows(), src.getDeletes().getBytes()); + } + + public TableAccess(tech.ydb.table.query.stats.TableAccessStats src) { + this.name = src.getName(); + this.partitionsCount = src.getPartitionsCount(); + this.reads = new TableOperation(src.getReads().getRows(), src.getReads().getBytes()); + this.updates = new TableOperation(src.getUpdates().getRows(), src.getUpdates().getBytes()); + this.deletes = new TableOperation(src.getDeletes().getRows(), src.getDeletes().getBytes()); + } + + public String getName() { + return name; + } + + public long getPartitionsCount() { + return partitionsCount; + } + + public TableOperation getReads() { + return reads; + } + + public TableOperation getUpdates() { + return updates; + } + + public TableOperation getDeletes() { + return deletes; + } + + @Override + public String toString() { + return "TableAccess{" + + "name='" + name + '\'' + + ", partitionsCount=" + partitionsCount + + ", reads=" + reads + + ", updates=" + updates + + ", deletes=" + deletes + + '}'; + } + + public static class TableOperation { + private final long rows; + private final long bytes; + + public TableOperation(long rows, long bytes) { + this.rows = rows; + this.bytes = bytes; + } + + public long getRows() { + return rows; + } + + public long getBytes() { + return bytes; + } + + @Override + public String toString() { + return "