Fixed usePrefixPath for preparing of YQL statements

Author: alex268

jdbc/src/main/java/tech/ydb/jdbc/YdbConst.java

@@ -91,6 +91,7 @@ public final class YdbConst {
     public static final String CANNOT_LOAD_DATA_FROM_IS = "Unable to load data from input stream: ";
     public static final String CANNOT_LOAD_DATA_FROM_READER = "Unable to load data from reader: ";
     public static final String STATEMENT_IS_NOT_A_BATCH = "Statement cannot be executed as batch statement: ";
+    public static final String UNABLE_PREPARE_STATEMENT = "Cannot prepare statement: ";
     public static final String MULTI_TYPES_IN_ONE_QUERY = "Query cannot contain expressions with different types: ";
     public static final String SCAN_QUERY_INSIDE_TRANSACTION = "Scan query cannot be executed inside active "
             + "transaction. This behavior may be changed by property scanQueryTxMode";

jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java

@@ -24,6 +24,7 @@
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.jdbc.YdbConst;
 import tech.ydb.jdbc.YdbPrepareMode;
+import tech.ydb.jdbc.YdbTracer;
 import tech.ydb.jdbc.exception.ExceptionFactory;
 import tech.ydb.jdbc.query.QueryType;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
@@ -45,6 +46,7 @@
 import tech.ydb.table.TableClient;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.impl.PooledTableClient;
+import tech.ydb.table.query.DataQuery;
 import tech.ydb.table.query.ExplainDataQueryResult;
 import tech.ydb.table.rpc.grpc.GrpcTableRpc;
 import tech.ydb.table.settings.DescribeTableSettings;
@@ -393,11 +395,20 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
             String tablePath = joined(getPrefixPath(), query.getYqlBatcher().getTableName());
             TableDescription description = tableDescribeCache.getIfPresent(tablePath);
             if (description == null) {
+                YdbTracer tracer = config.isTxTracedEnabled() ? YdbTracer.current() : null;
+                if (tracer != null) {
+                    tracer.trace("--> describe table");
+                    tracer.traceRequest(tablePath);
+                }
                 DescribeTableSettings settings = withDefaultTimeout(new DescribeTableSettings());
                 Result<TableDescription> result = retryCtx.supplyResult(
                         session -> session.describeTable(tablePath, settings)
                 ).join();
 
+                if (tracer != null) {
+                    tracer.trace("<-- " + result.getStatus());
+                }
+
                 if (result.isSuccess()) {
                     description = result.getValue();
                     tableDescribeCache.put(query.getOriginQuery(), description);
@@ -426,32 +437,46 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
         }
 
         // try to prepare data query
-        try {
-            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.getOriginQuery(), types);
+        Map<String, Type> types = queryParamsCache.getIfPresent(query.getOriginQuery());
+        if (types == null) {
+            String yql = prefixPragma + query.getPreparedYql();
+            YdbTracer tracer = config.isTxTracedEnabled() ? YdbTracer.current() : null;
+            if (tracer != null) {
+                tracer.trace("--> prepare data query");
+                tracer.traceRequest(yql);
             }
 
-            boolean requireBatch = mode == YdbPrepareMode.DATA_QUERY_BATCH;
-            if (requireBatch || (mode == YdbPrepareMode.AUTO && queryOptions.isDetectBatchQueries())) {
-                BatchedQuery params = BatchedQuery.tryCreateBatched(query, types);
-                if (params != null) {
-                    return params;
-                }
+            PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings());
+            Result<DataQuery> result = retryCtx.supplyResult(
+                    session -> session.prepareDataQuery(yql, settings)
+            ).join();
 
-                if (requireBatch) {
-                    throw new SQLDataException(YdbConst.STATEMENT_IS_NOT_A_BATCH + query.getOriginQuery());
+            if (tracer != null) {
+                tracer.trace("<-- " + result.getStatus());
+            }
+            if (!result.isSuccess()) {
+                if (tracer != null) {
+                    tracer.close();
                 }
+                throw ExceptionFactory.createException("Cannot prepare data query: " + result.getStatus(),
+                        new UnexpectedResultException("Unexpected status", result.getStatus()));
+            }
+
+            types = result.getValue().types();
+            queryParamsCache.put(query.getOriginQuery(), types);
+        }
+
+        boolean requireBatch = mode == YdbPrepareMode.DATA_QUERY_BATCH;
+        if (requireBatch || (mode == YdbPrepareMode.AUTO && queryOptions.isDetectBatchQueries())) {
+            BatchedQuery params = BatchedQuery.tryCreateBatched(query, types);
+            if (params != null) {
+                return params;
+            }
+
+            if (requireBatch) {
+                throw new SQLDataException(YdbConst.STATEMENT_IS_NOT_A_BATCH + query.getOriginQuery());
             }
-            return new PreparedQuery(query, types);
-        } catch (UnexpectedResultException ex) {
-            throw ExceptionFactory.createException("Cannot prepare data query: " + ex.getMessage(), ex);
         }
+        return new PreparedQuery(query, types);
     }
 }

jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTablesTest.java

@@ -102,9 +102,9 @@ public void defaultModeTest() throws SQLException {
             }
 
             // read all
-            try (Statement st = connection.createStatement()) {
+            try (PreparedStatement ps = connection.prepareStatement(SELECT_ALL)) {
                 int readed = 0;
-                try (ResultSet rs = st.executeQuery(SELECT_ALL)) {
+                try (ResultSet rs = ps.executeQuery()) {
                     while (rs.next()) {
                         readed++;
                         Assertions.assertEquals(readed, rs.getInt("id"));