Skip to content

Commit

Permalink
Updated yql query parser with support of RETURNING, OFFSET/LIMIT and …
Browse files Browse the repository at this point in the history 
…auto batching (#56)
  • Loading branch information
@alex268
alex268 authored Jul 17, 2024
2 parents fc82d87 + c9c1798 commit f340f22
Show file tree
Hide file tree
Showing 34 changed files with 1,787 additions and 875 deletions.
17 changes: 8 additions & 9 deletions jdbc/src/main/java/tech/ydb/jdbc/YdbConnection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
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;

Expand All @@ -36,47 +35,47 @@ public interface YdbConnection extends Connection {
/**
* Explicitly execute query as a schema query
*
* @param query query (DDL) to execute
* @param yql query (DDL) to execute
* @param validator handler for logging and warnings
* @throws SQLException if query cannot be executed
*/
void executeSchemeQuery(YdbQuery query, YdbValidator validator) throws SQLException;
void executeSchemeQuery(String yql, YdbValidator validator) throws SQLException;

/**
* Explicitly execute query as a data query
*
* @param query query to execute
* @param yql query 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(YdbQuery query, YdbValidator validator,
List<ResultSetReader> executeDataQuery(String yql, YdbValidator validator,
int timeout, boolean keepInCache, Params params) throws SQLException;

/**
* Explicitly execute query as a scan query
*
* @param query query to execute
* @param yql query 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(YdbQuery query, YdbValidator validator, Params params) throws SQLException;
ResultSetReader executeScanQuery(String yql, YdbValidator validator, Params params) throws SQLException;

/**
* Explicitly explain this query
*
* @param query query to explain
* @param yql query to explain
* @param validator handler for logging and warnings
* @return list of result set of two string columns: {@link YdbConst#EXPLAIN_COLUMN_AST}
* and {@link YdbConst#EXPLAIN_COLUMN_PLAN}
* @throws SQLException if query cannot be explained
*/
ExplainedQuery executeExplainQuery(YdbQuery query, YdbValidator validator) throws SQLException;
ExplainedQuery executeExplainQuery(String yql, YdbValidator validator) throws SQLException;

@Override
YdbDatabaseMetaData getMetaData() throws SQLException;
Expand Down
8 changes: 2 additions & 6 deletions jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
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;
Expand Down Expand Up @@ -42,22 +41,19 @@ protected Session createNewTableSession(YdbValidator validator) throws SQLExcept
}

@Override
public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, YdbQuery query) throws SQLException {
public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException {
// Scheme query does not affect transactions or result sets
ExecuteSchemeQuerySettings settings = ctx.withDefaultTimeout(new ExecuteSchemeQuerySettings());
final String yql = query.getYqlQuery(null);

try (Session session = createNewTableSession(validator)) {
validator.execute(QueryType.SCHEME_QUERY + " >>\n" + yql, () -> session.executeSchemeQuery(yql, settings));
}
}

@Override
public ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, Params params)
public ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
throws SQLException {
ensureOpened();

String yql = query.getYqlQuery(params);
Collection<ResultSetReader> resultSets = new LinkedBlockingQueue<>();
Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
ExecuteScanQuerySettings settings = ExecuteScanQuerySettings.newBuilder()
Expand Down
11 changes: 3 additions & 8 deletions jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
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;
Expand Down Expand Up @@ -197,11 +196,10 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException

@Override
public List<ResultSetReader> executeDataQuery(
YdbContext ctx, YdbValidator validator, YdbQuery query, long timeout, boolean keepInCache, Params params
YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
) throws SQLException {
ensureOpened();

final String yql = query.getYqlQuery(params);
ExecuteQuerySettings.Builder builder = ExecuteQuerySettings.newBuilder();
if (timeout > 0) {
builder = builder.withRequestTimeout(timeout, TimeUnit.SECONDS);
Expand Down Expand Up @@ -229,13 +227,11 @@ public List<ResultSetReader> executeDataQuery(
}

@Override
public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, YdbQuery query) throws SQLException {
public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException {
ensureOpened();

// Scheme query does not affect transactions or result sets
ExecuteQuerySettings settings = ctx.withRequestTimeout(ExecuteQuerySettings.newBuilder()).build();
final String yql = query.getYqlQuery(null);

try (QuerySession session = createNewQuerySession(validator)) {
validator.call(QueryType.SCHEME_QUERY + " >>\n" + yql, () -> session
.createQuery(yql, TxMode.NONE, Params.empty(), settings)
Expand All @@ -245,15 +241,14 @@ public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, YdbQuery
}

@Override
public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, YdbQuery query)
public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, String yql)
throws SQLException {
ensureOpened();

// Scheme query does not affect transactions or result sets
ExecuteQuerySettings settings = ctx.withRequestTimeout(ExecuteQuerySettings.newBuilder())
.withExecMode(QueryExecMode.EXPLAIN)
.build();
final String yql = query.getYqlQuery(null);

try (QuerySession session = createNewQuerySession(validator)) {
QueryInfo res = validator.call(QueryType.EXPLAIN_QUERY + " >>\n" + yql, () -> session
Expand Down
7 changes: 2 additions & 5 deletions jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
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;
Expand Down Expand Up @@ -155,11 +154,10 @@ private ExecuteDataQuerySettings dataQuerySettings(long timeout, boolean keepInC
}

@Override
public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, YdbQuery query)
public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, String yql)
throws SQLException {
ensureOpened();

String yql = query.getYqlQuery(null);
ExplainDataQuerySettings settings = ctx.withDefaultTimeout(new ExplainDataQuerySettings());
try (Session session = createNewTableSession(validator)) {
String msg = QueryType.EXPLAIN_QUERY + " >>\n" + yql;
Expand All @@ -170,11 +168,10 @@ public ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator

@Override
public List<ResultSetReader> executeDataQuery(
YdbContext ctx, YdbValidator validator, YdbQuery query, long timeout, boolean keepInCache, Params params
YdbContext ctx, YdbValidator validator, String yql, long timeout, boolean keepInCache, Params params
) throws SQLException {
ensureOpened();

final String yql = query.getYqlQuery(params);
final Session session = tx.getSession(validator);
try {
DataQueryResult result = validator.call(
Expand Down
68 changes: 48 additions & 20 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 @@ -8,24 +8,24 @@
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;

import tech.ydb.core.Result;
import tech.ydb.core.UnexpectedResultException;
import tech.ydb.core.grpc.GrpcTransport;
import tech.ydb.core.grpc.GrpcTransportBuilder;
import tech.ydb.core.settings.BaseRequestSettings;
import tech.ydb.jdbc.YdbConst;
import tech.ydb.jdbc.YdbPrepareMode;
import tech.ydb.jdbc.exception.ExceptionFactory;
import tech.ydb.jdbc.query.JdbcParams;
import tech.ydb.jdbc.query.JdbcQueryLexer;
import tech.ydb.jdbc.query.YdbPreparedQuery;
import tech.ydb.jdbc.query.YdbQuery;
import tech.ydb.jdbc.query.YdbQueryBuilder;
import tech.ydb.jdbc.query.params.BatchedParams;
import tech.ydb.jdbc.query.params.InMemoryParams;
import tech.ydb.jdbc.query.params.PreparedParams;
import tech.ydb.jdbc.query.params.BatchedQuery;
import tech.ydb.jdbc.query.params.InMemoryQuery;
import tech.ydb.jdbc.query.params.PreparedQuery;
import tech.ydb.jdbc.settings.YdbClientProperties;
import tech.ydb.jdbc.settings.YdbConfig;
import tech.ydb.jdbc.settings.YdbConnectionProperties;
Expand All @@ -36,8 +36,11 @@
import tech.ydb.scheme.SchemeClient;
import tech.ydb.table.SessionRetryContext;
import tech.ydb.table.TableClient;
import tech.ydb.table.description.TableColumn;
import tech.ydb.table.description.TableDescription;
import tech.ydb.table.impl.PooledTableClient;
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
import tech.ydb.table.settings.DescribeTableSettings;
import tech.ydb.table.settings.PrepareDataQuerySettings;
import tech.ydb.table.settings.RequestSettings;
import tech.ydb.table.values.Type;
Expand Down Expand Up @@ -255,9 +258,7 @@ public <T extends BaseRequestSettings.BaseBuilder<T>> T withRequestTimeout(T bui
}

public YdbQuery parseYdbQuery(String sql) throws SQLException {
YdbQueryBuilder builder = new YdbQueryBuilder(sql, queryOptions.getForcedQueryType());
JdbcQueryLexer.buildQuery(builder, queryOptions);
return builder.build(queryOptions);
return YdbQuery.parseQuery(sql, queryOptions);
}

public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
Expand All @@ -274,37 +275,64 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
return cached;
}

public JdbcParams findOrCreateJdbcParams(YdbQuery query, YdbPrepareMode mode) throws SQLException {
if (query.hasIndexesParameters()
public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) throws SQLException {
if (query.getYqlBatcher() != null && mode == YdbPrepareMode.AUTO) {
Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
if (types == null) {
String tableName = query.getYqlBatcher().getTableName();
String tablePath = tableName.startsWith("/") ? tableName : getDatabase() + "/" + tableName;

DescribeTableSettings settings = withDefaultTimeout(new DescribeTableSettings());
Result<TableDescription> result = retryCtx.supplyResult(
session -> session.describeTable(tablePath, settings)
).join();

if (result.isSuccess()) {
TableDescription d = result.getValue();
types = result.getValue().getColumns().stream()
.collect(Collectors.toMap(TableColumn::getName, TableColumn::getType));
queryParamsCache.put(query.getOriginQuery(), types);
}
}
if (types != null) {
BatchedQuery params = BatchedQuery.createAutoBatched(query.getYqlBatcher(), types);
if (params != null) {
return params;
}
}
}

if (!query.isPlainYQL()
|| mode == YdbPrepareMode.IN_MEMORY
|| !queryOptions.iPrepareDataQueries()) {
return new InMemoryParams(query.getIndexesParameters());
|| !queryOptions.isPrepareDataQueries()) {
return new InMemoryQuery(query, queryOptions.isDeclareJdbcParameters());
}

String yql = query.getYqlQuery(null);
PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings());
// try to prepare data query
try {
Map<String, Type> types = queryParamsCache.getIfPresent(query.originSQL());
Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
if (types == null) {
String yql = query.getPreparedYql();
PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings());
types = retryCtx.supplyResult(session -> session.prepareDataQuery(yql, settings))
.join()
.getValue()
.types();
queryParamsCache.put(query.originSQL(), types);
queryParamsCache.put(query.getOriginQuery(), types);
}

boolean requireBatch = mode == YdbPrepareMode.DATA_QUERY_BATCH;
if (requireBatch || (mode == YdbPrepareMode.AUTO && queryOptions.isDetectBatchQueries())) {
BatchedParams params = BatchedParams.tryCreateBatched(types);
BatchedQuery params = BatchedQuery.tryCreateBatched(query, types);
if (params != null) {
return params;
}

if (requireBatch) {
throw new SQLDataException(YdbConst.STATEMENT_IS_NOT_A_BATCH + query.originSQL());
throw new SQLDataException(YdbConst.STATEMENT_IS_NOT_A_BATCH + query.getOriginQuery());
}
}
return new PreparedParams(types);
return new PreparedQuery(query, types);
} catch (UnexpectedResultException ex) {
throw ExceptionFactory.createException("Cannot prepare data query: " + ex.getMessage(), ex);
}
Expand Down
9 changes: 4 additions & 5 deletions jdbc/src/main/java/tech/ydb/jdbc/context/YdbExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

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;

Expand Down Expand Up @@ -33,15 +32,15 @@ default void ensureOpened() throws SQLException {
void setReadOnly(boolean readOnly) throws SQLException;
void setAutoCommit(boolean autoCommit) throws SQLException;

void executeSchemeQuery(YdbContext ctx, YdbValidator validator, YdbQuery query) throws SQLException;
void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException;

List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, YdbQuery query,
List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, String yql,
long timeout, boolean poolable, Params params) throws SQLException;

ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, Params params)
ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, String yql, Params params)
throws SQLException;

ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, YdbQuery query)
ExplainedQuery executeExplainQuery(YdbContext ctx, YdbValidator validator, String yql)
throws SQLException;

void commit(YdbContext ctx, YdbValidator validator) throws SQLException;
Expand Down
Loading

0 comments on commit f340f22

Please sign in to comment.