Added collection of stats to YdbContext

ydb-platform · Aug 12, 2024 · 4e581be · 4e581be
1 parent 15d5c63
commit 4e581be
Showing 1 changed file with 32 additions and 2 deletions.
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
@@ -39,8 +39,10 @@
 import tech.ydb.table.description.TableColumn;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.impl.PooledTableClient;
+import tech.ydb.table.query.ExplainDataQueryResult;
 import tech.ydb.table.rpc.grpc.GrpcTableRpc;
 import tech.ydb.table.settings.DescribeTableSettings;
+import tech.ydb.table.settings.ExplainDataQuerySettings;
 import tech.ydb.table.settings.PrepareDataQuerySettings;
 import tech.ydb.table.settings.RequestSettings;
 import tech.ydb.table.values.Type;
@@ -68,6 +70,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, Map<String, Type>> queryParamsCache;
 
     private final boolean autoResizeSessionPool;
@@ -98,8 +101,14 @@ private YdbContext(
         if (cacheSize > 0) {
             queriesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
             queryParamsCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
+            if (config.isFullScanDetectorEnabled()) {
+                queryStatesCache = CacheBuilder.newBuilder().maximumSize(cacheSize).build();
+            } else {
+                queryStatesCache = null;
+            }
         } else {
             queriesCache = null;
+            queryStatesCache = null;
             queryParamsCache = null;
         }
     }
@@ -270,6 +279,27 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
         if (cached == null) {
             cached = parseYdbQuery(sql);
             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);
+                }
+
+                stat.incrementUsage();
+            }
         }
 
         return cached;
@@ -288,8 +318,8 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
                 ).join();
 
                 if (result.isSuccess()) {
-                    TableDescription d = result.getValue();
-                    types = result.getValue().getColumns().stream()
+                    TableDescription descrtiption = result.getValue();
+                    types = descrtiption.getColumns().stream()
                             .collect(Collectors.toMap(TableColumn::getName, TableColumn::getType));
                     queryParamsCache.put(query.getOriginQuery(), types);
                 }