Skip to content

Commit

Permalink
Added option to FullScan detecting (#58)
Browse files Browse the repository at this point in the history
  • Loading branch information
@alex268
alex268 authored Aug 12, 2024
2 parents e46c943 + efbaf8a commit 16e24da
Show file tree
Hide file tree
Showing 9 changed files with 354 additions and 4 deletions.
95 changes: 95 additions & 0 deletions jdbc/src/main/java/tech/ydb/jdbc/context/QueryStat.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
package tech.ydb.jdbc.context;

import java.util.Collection;
import java.util.concurrent.atomic.LongAdder;

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 FixedResultSetFactory STATS_RS_FACTORY = FixedResultSetFactory.newBuilder()
.addTextColumn("sql")
.addBooleanColumn("is_fullscan")
.addLongColumn("executed")
.addTextColumn("yql")
.addTextColumn("ast")
.addTextColumn("plan")
.build();

private final String originSQL;
private final String preparedYQL;

private final String ast;
private final String plan;
private final LongAdder usage;
private final boolean isFullScan;

public QueryStat(YdbQuery query, String ast, String plan) {
this.originSQL = query.getOriginQuery();
this.preparedYQL = query.getPreparedYql();
this.ast = ast;
this.plan = plan;
this.usage = new LongAdder();
this.isFullScan = plan.contains("\"Node Type\":\"TableFullScan\"");
}

public QueryStat(YdbQuery query, Status error) {
this.originSQL = query.getOriginQuery();
this.preparedYQL = query.getPreparedYql();
this.ast = error.toString();
this.plan = error.toString();
this.usage = new LongAdder();
this.isFullScan = false;
}

public long getUsageCounter() {
return usage.longValue();
}

public String getOriginSQL() {
return originSQL;
}

public String getPreparedYQL() {
return preparedYQL;
}

public String getAat() {
return ast;
}

public String getPlan() {
return plan;
}

public boolean isFullScan() {
return isFullScan;
}

public void incrementUsage() {
this.usage.increment();
}

public static ResultSetReader toResultSetReader(Collection<QueryStat> stats) {
FixedResultSetFactory.ResultSetBuilder builder = STATS_RS_FACTORY.createResultSet();
for (QueryStat stat: stats) {
builder.newRow()
.withTextValue("sql", stat.originSQL)
.withBoolValue("is_fullscan", stat.isFullScan)
.withLongValue("executed", stat.usage.longValue())
.withTextValue("yql", stat.preparedYQL)
.withTextValue("ast", stat.ast)
.withTextValue("plan", stat.plan)
.build();
}
return builder.build();
}
}
59 changes: 57 additions & 2 deletions jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,12 @@
import java.sql.SQLDataException;
import java.sql.SQLException;
import java.time.Duration;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
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;
Expand Down Expand Up @@ -39,8 +44,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;
Expand Down Expand Up @@ -68,6 +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, Map<String, Type>> queryParamsCache;

private final boolean autoResizeSessionPool;
Expand Down Expand Up @@ -98,8 +106,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;
}
}
Expand Down Expand Up @@ -169,6 +183,28 @@ public boolean hasConnections() {
return connectionsCount.get() > 0;
}

public boolean queryStatsEnabled() {
return queryStatesCache != null;
}

public Collection<QueryStat> getQueryStats() {
if (queryStatesCache == null) {
return Collections.emptyList();
}
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 void register() {
int actual = connectionsCount.incrementAndGet();
int maxSize = tableClient.sessionPoolStats().getMaxSize();
Expand Down Expand Up @@ -272,6 +308,25 @@ 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;
}

Expand All @@ -288,8 +343,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);
}
Expand Down
3 changes: 3 additions & 0 deletions jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ public int[] executeBatch() throws SQLException {

try {
for (Params prm: prepared.getBatchParams()) {
getConnection().getCtx().traceQueryExecution(query);
executeDataQuery(query, prepared.getQueryText(prm), prm);
}
} finally {
Expand Down Expand Up @@ -123,7 +124,9 @@ public boolean execute() throws SQLException {
clearBatch();

List<YdbResult> newState = null;

Params prms = prepared.getCurrentParams();
getConnection().getCtx().traceQueryExecution(query);
switch (query.getType()) {
case DATA_QUERY:
newState = executeDataQuery(query, prepared.getQueryText(prms), prms);
Expand Down
18 changes: 17 additions & 1 deletion jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,16 @@
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

import tech.ydb.jdbc.YdbConnection;
import tech.ydb.jdbc.YdbConst;
import tech.ydb.jdbc.YdbResultSet;
import tech.ydb.jdbc.context.QueryStat;
import tech.ydb.jdbc.context.YdbContext;
import tech.ydb.jdbc.query.YdbQuery;
import tech.ydb.table.query.Params;

Expand Down Expand Up @@ -79,7 +82,20 @@ public int executeUpdate(String sql) throws SQLException {
public boolean execute(String sql) throws SQLException {
cleanState();

YdbQuery query = getConnection().getCtx().parseYdbQuery(sql);
YdbContext ctx = getConnection().getCtx();
YdbQuery query;

if (ctx.queryStatsEnabled()) {
if (sql != null && QueryStat.QUERY.equalsIgnoreCase(sql.trim())) {
YdbResultSet rs = new YdbResultSetImpl(this, QueryStat.toResultSetReader(ctx.getQueryStats()));
return updateState(Collections.singletonList(new YdbResult(rs)));
}
query = ctx.findOrParseYdbQuery(sql);
ctx.traceQueryExecution(query);
} else {
query = ctx.parseYdbQuery(sql);
}

List<YdbResult> newState = null;
switch (query.getType()) {
case SCHEME_QUERY:
Expand Down
1 change: 0 additions & 1 deletion jdbc/src/main/java/tech/ydb/jdbc/settings/YdbClientProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,5 +81,4 @@ public boolean applyToTableClient(TableClient.Builder table, QueryClient.Builder
query.sessionPoolMaxSize(maxSize).sessionPoolMinSize(minSize);
return false;
}

}
9 changes: 9 additions & 0 deletions jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ public class YdbConfig {
static final YdbProperty<Boolean> USE_QUERY_SERVICE = YdbProperty.bool("useQueryService",
"Use QueryService intead of TableService", false
);
static final YdbProperty<Boolean> FULLSCAN_DETECTOR_ENABLED = YdbProperty.bool(
"jdbcFullScanDetector", "Enable analizator for collecting query stats", false
);


private final String url;
Expand All @@ -53,6 +56,7 @@ public class YdbConfig {
private final boolean isCacheConnectionsInDriver;
private final int preparedStatementsCacheSize;
private final boolean useQueryService;
private final boolean fullScanDetectorEnabled;

private YdbConfig(
String url, String safeUrl, String connectionString, String username, String password, Properties props
Expand All @@ -66,6 +70,7 @@ private YdbConfig(
this.isCacheConnectionsInDriver = CACHE_CONNECTIONS_IN_DRIVER.readValue(props).getValue();
this.preparedStatementsCacheSize = Math.max(0, PREPARED_STATEMENT_CACHE_SIZE.readValue(props).getValue());
this.useQueryService = USE_QUERY_SERVICE.readValue(props).getValue();
this.fullScanDetectorEnabled = FULLSCAN_DETECTOR_ENABLED.readValue(props).getValue();
}

public Properties getSafeProps() {
Expand Down Expand Up @@ -96,6 +101,10 @@ public boolean isUseQueryService() {
return this.useQueryService;
}

public boolean isFullScanDetectorEnabled() {
return fullScanDetectorEnabled;
}

static boolean isSensetive(String key) {
return TOKEN_KEY.equalsIgnoreCase(key) || PASSWORD_KEY.equalsIgnoreCase(key);
}
Expand Down
83 changes: 83 additions & 0 deletions jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -925,4 +925,87 @@ public void testUnsupportedComplexTypes(String type) throws SQLException {
() -> statement.execute(sql));
}
}

@Test
public void testFullScanAnalyzer() throws SQLException {
try (Connection connection = jdbc.createCustomConnection("jdbcFullScanDetector", "true")) {
String selectAll = QUERIES.selectAllSQL();
String selectByKey = QUERIES.selectAllByKey("1");
String preparedSelect = QUERIES.selectAllByKey("?");

try (Statement st = connection.createStatement()) {
try (ResultSet rs = st.executeQuery(" print_JDBC_stats(); ")) {
Assertions.assertFalse(rs.next()); // not stats
}

try (ResultSet rs = st.executeQuery(selectAll)) {
Assertions.assertFalse(rs.next());
}

try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
Assertions.assertTrue(rs.next());
Assertions.assertEquals(selectAll, rs.getString("sql"));
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(1l, rs.getLong("executed"));

Assertions.assertFalse(rs.next());
}

try (ResultSet rs = st.executeQuery(selectAll)) {
Assertions.assertFalse(rs.next());
}
try (ResultSet rs = st.executeQuery(selectByKey)) {
Assertions.assertFalse(rs.next());
}

try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
Assertions.assertTrue(rs.next());
Assertions.assertEquals(selectAll, rs.getString("sql"));
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(2l, rs.getLong("executed"));

Assertions.assertTrue(rs.next());
Assertions.assertEquals(selectByKey, rs.getString("sql"));
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(1l, rs.getLong("executed"));

Assertions.assertFalse(rs.next());
}

try (PreparedStatement ps = connection.prepareStatement(preparedSelect)) {
ps.setLong(1, 1);
try (ResultSet rs = ps.executeQuery()) {
Assertions.assertFalse(rs.next());
}
ps.setLong(1, 2);
try (ResultSet rs = ps.executeQuery()) {
Assertions.assertFalse(rs.next());
}
ps.setLong(1, 3);
try (ResultSet rs = ps.executeQuery()) {
Assertions.assertFalse(rs.next());
}
}

try (ResultSet rs = st.executeQuery("Print_JDBC_stats();\n")) {
Assertions.assertTrue(rs.next());
Assertions.assertEquals(preparedSelect, rs.getString("sql"));
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(3l, rs.getLong("executed"));

Assertions.assertTrue(rs.next());
Assertions.assertEquals(selectAll, rs.getString("sql"));
Assertions.assertEquals(true, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(2l, rs.getLong("executed"));

Assertions.assertTrue(rs.next());
Assertions.assertEquals(selectByKey, rs.getString("sql"));
Assertions.assertEquals(false, rs.getBoolean("is_fullscan"));
Assertions.assertEquals(1l, rs.getLong("executed"));

Assertions.assertFalse(rs.next());
}
}
}
}
}
Loading

0 comments on commit 16e24da

Please sign in to comment.