88import java .util .concurrent .atomic .AtomicInteger ;
99import java .util .logging .Level ;
1010import java .util .logging .Logger ;
11+ import java .util .stream .Collectors ;
1112
1213import com .google .common .cache .Cache ;
1314import com .google .common .cache .CacheBuilder ;
1415
16+ import tech .ydb .core .Result ;
1517import tech .ydb .core .UnexpectedResultException ;
1618import tech .ydb .core .grpc .GrpcTransport ;
1719import tech .ydb .core .grpc .GrpcTransportBuilder ;
1820import tech .ydb .core .settings .BaseRequestSettings ;
1921import tech .ydb .jdbc .YdbConst ;
2022import tech .ydb .jdbc .YdbPrepareMode ;
2123import tech .ydb .jdbc .exception .ExceptionFactory ;
22- import tech .ydb .jdbc .query .JdbcParams ;
23- import tech .ydb .jdbc .query .JdbcQueryLexer ;
24+ import tech .ydb .jdbc .query .YdbPreparedQuery ;
2425import tech .ydb .jdbc .query .YdbQuery ;
25- import tech .ydb .jdbc .query .YdbQueryBuilder ;
26- import tech .ydb .jdbc .query .params .BatchedParams ;
27- import tech .ydb .jdbc .query .params .InMemoryParams ;
28- import tech .ydb .jdbc .query .params .PreparedParams ;
26+ import tech .ydb .jdbc .query .params .BatchedQuery ;
27+ import tech .ydb .jdbc .query .params .InMemoryQuery ;
28+ import tech .ydb .jdbc .query .params .PreparedQuery ;
2929import tech .ydb .jdbc .settings .YdbClientProperties ;
3030import tech .ydb .jdbc .settings .YdbConfig ;
3131import tech .ydb .jdbc .settings .YdbConnectionProperties ;
3636import tech .ydb .scheme .SchemeClient ;
3737import tech .ydb .table .SessionRetryContext ;
3838import tech .ydb .table .TableClient ;
39+ import tech .ydb .table .description .TableColumn ;
40+ import tech .ydb .table .description .TableDescription ;
3941import tech .ydb .table .impl .PooledTableClient ;
4042import tech .ydb .table .rpc .grpc .GrpcTableRpc ;
43+ import tech .ydb .table .settings .DescribeTableSettings ;
4144import tech .ydb .table .settings .PrepareDataQuerySettings ;
4245import tech .ydb .table .settings .RequestSettings ;
4346import tech .ydb .table .values .Type ;
@@ -255,9 +258,7 @@ public <T extends BaseRequestSettings.BaseBuilder<T>> T withRequestTimeout(T bui
255258 }
256259
257260 public YdbQuery parseYdbQuery (String sql ) throws SQLException {
258- YdbQueryBuilder builder = new YdbQueryBuilder (sql , queryOptions .getForcedQueryType ());
259- JdbcQueryLexer .buildQuery (builder , queryOptions );
260- return builder .build (queryOptions );
261+ return YdbQuery .parseQuery (sql , queryOptions );
261262 }
262263
263264 public YdbQuery findOrParseYdbQuery (String sql ) throws SQLException {
@@ -274,37 +275,64 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
274275 return cached ;
275276 }
276277
277- public JdbcParams findOrCreateJdbcParams (YdbQuery query , YdbPrepareMode mode ) throws SQLException {
278- if (query .hasIndexesParameters ()
278+ public YdbPreparedQuery findOrPrepareParams (YdbQuery query , YdbPrepareMode mode ) throws SQLException {
279+ if (query .getYqlBatcher () != null && mode == YdbPrepareMode .AUTO ) {
280+ Map <String , Type > types = queryParamsCache .getIfPresent (query .getOriginQuery ());
281+ if (types == null ) {
282+ String tableName = query .getYqlBatcher ().getTableName ();
283+ String tablePath = tableName .startsWith ("/" ) ? tableName : getDatabase () + "/" + tableName ;
284+
285+ DescribeTableSettings settings = withDefaultTimeout (new DescribeTableSettings ());
286+ Result <TableDescription > result = retryCtx .supplyResult (
287+ session -> session .describeTable (tablePath , settings )
288+ ).join ();
289+
290+ if (result .isSuccess ()) {
291+ TableDescription d = result .getValue ();
292+ types = result .getValue ().getColumns ().stream ()
293+ .collect (Collectors .toMap (TableColumn ::getName , TableColumn ::getType ));
294+ queryParamsCache .put (query .getOriginQuery (), types );
295+ }
296+ }
297+ if (types != null ) {
298+ BatchedQuery params = BatchedQuery .createAutoBatched (query .getYqlBatcher (), types );
299+ if (params != null ) {
300+ return params ;
301+ }
302+ }
303+ }
304+
305+ if (!query .isPlainYQL ()
279306 || mode == YdbPrepareMode .IN_MEMORY
280- || !queryOptions .iPrepareDataQueries ()) {
281- return new InMemoryParams (query . getIndexesParameters ());
307+ || !queryOptions .isPrepareDataQueries ()) {
308+ return new InMemoryQuery (query , queryOptions . isDeclareJdbcParameters ());
282309 }
283310
284- String yql = query .getYqlQuery (null );
285- PrepareDataQuerySettings settings = withDefaultTimeout (new PrepareDataQuerySettings ());
311+ // try to prepare data query
286312 try {
287- Map <String , Type > types = queryParamsCache .getIfPresent (query .originSQL ());
313+ Map <String , Type > types = queryParamsCache .getIfPresent (query .getOriginQuery ());
288314 if (types == null ) {
315+ String yql = query .getPreparedYql ();
316+ PrepareDataQuerySettings settings = withDefaultTimeout (new PrepareDataQuerySettings ());
289317 types = retryCtx .supplyResult (session -> session .prepareDataQuery (yql , settings ))
290318 .join ()
291319 .getValue ()
292320 .types ();
293- queryParamsCache .put (query .originSQL (), types );
321+ queryParamsCache .put (query .getOriginQuery (), types );
294322 }
295323
296324 boolean requireBatch = mode == YdbPrepareMode .DATA_QUERY_BATCH ;
297325 if (requireBatch || (mode == YdbPrepareMode .AUTO && queryOptions .isDetectBatchQueries ())) {
298- BatchedParams params = BatchedParams .tryCreateBatched (types );
326+ BatchedQuery params = BatchedQuery .tryCreateBatched (query , types );
299327 if (params != null ) {
300328 return params ;
301329 }
302330
303331 if (requireBatch ) {
304- throw new SQLDataException (YdbConst .STATEMENT_IS_NOT_A_BATCH + query .originSQL ());
332+ throw new SQLDataException (YdbConst .STATEMENT_IS_NOT_A_BATCH + query .getOriginQuery ());
305333 }
306334 }
307- return new PreparedParams ( types );
335+ return new PreparedQuery ( query , types );
308336 } catch (UnexpectedResultException ex ) {
309337 throw ExceptionFactory .createException ("Cannot prepare data query: " + ex .getMessage (), ex );
310338 }
0 commit comments