Added auto batch mode

alex268 · alex268 · commit c9c179833be7 · 2024-07-17T10:46:02.000+01:00
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
@@ -8,10 +8,12 @@
 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;
@@ -34,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;
@@ -271,17 +276,44 @@ public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
     }
 
     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.isPrepareDataQueries()) {
             return new InMemoryQuery(query, queryOptions.isDeclareJdbcParameters());
         }
 
-        String yql = query.getPreparedYql();
-        PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings());
+        // 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()
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/YdbQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/YdbQuery.java
@@ -15,15 +15,21 @@ public class YdbQuery {
     private final String originQuery;
     private final String preparedYQL;
     private final List<QueryStatement> statements;
+    private final YqlBatcher batch;
 
     private final QueryType type;
     private final boolean isPlainYQL;
 
     YdbQuery(String originQuery, String preparedYQL, List<QueryStatement> stats, QueryType type) {
+        this(originQuery, preparedYQL, stats, null, type);
+    }
+
+    YdbQuery(String originQuery, String preparedYQL, List<QueryStatement> stats, YqlBatcher batch, QueryType type) {
         this.originQuery = originQuery;
         this.preparedYQL = preparedYQL;
         this.statements = stats;
         this.type = type;
+        this.batch = batch;
 
         boolean hasJdbcParamters = false;
         for (QueryStatement st: statements) {
@@ -36,6 +42,10 @@ public QueryType getType() {
         return type;
     }
 
+    public YqlBatcher getYqlBatcher() {
+        return batch;
+    }
+
     public boolean isPlainYQL() {
         return isPlainYQL;
     }
@@ -52,43 +62,6 @@ public List<QueryStatement> getStatements() {
         return statements;
     }
 
-//    public String withParams(Params params) throws SQLException {
-//        if (isPlainYQL) {
-//            return preparedYQL;
-//        }
-//
-//        if (params == null) {
-//            if (isAutoDeclare) {
-//                int paramCount = statements.stream().mapToInt(st -> st.getParams().size()).sum();
-//                // Comment in place where must be declare section
-//                return "-- DECLARE " + paramCount + " PARAMETERS\n" + preparedYQL;
-//            }
-//            return preparedYQL;
-//        }
-//
-//        StringBuilder yql = new StringBuilder();
-//        Map<String, Value<?>> values = params.values();
-//        for (QueryStatement st: statements) {
-//            for (ParamDescription prm: st.getParams()) {
-//                if (!values.containsKey(prm.name())) {
-//                    throw new SQLDataException(YdbConst.MISSING_VALUE_FOR_PARAMETER + prm);
-//                }
-//
-//                if (isAutoDeclare) {
-//                    String prmType = values.get(prm.name()).getType().toString();
-//                    yql.append("DECLARE ")
-//                            .append(prm.name())
-//                            .append(" AS ")
-//                            .append(prmType)
-//                            .append(";\n");
-//                }
-//            }
-//        }
-//
-//        yql.append(preparedYQL);
-//        return yql.toString();
-//    }
-
     public static YdbQuery parseQuery(String query, YdbQueryProperties opts) throws SQLException {
         YdbQueryParser parser = new YdbQueryParser(opts.isDetectQueryType(), opts.isDetectJdbcParameters());
         String preparedYQL = parser.parseSQL(query);
@@ -97,7 +70,9 @@ public static YdbQuery parseQuery(String query, YdbQueryProperties opts) throws
         if (type == null) {
             type = parser.detectQueryType();
         }
-
+        if (parser.getYqlBatcher().isValidBatch()) {
+            return new YdbQuery(query, preparedYQL, parser.getStatements(), parser.getYqlBatcher(), type);
+        }
         return new YdbQuery(query, preparedYQL, parser.getStatements(), type);
     }
 }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java
@@ -18,6 +18,7 @@
 import tech.ydb.jdbc.query.ParamDescription;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
 import tech.ydb.jdbc.query.YdbQuery;
+import tech.ydb.jdbc.query.YqlBatcher;
 import tech.ydb.table.query.Params;
 import tech.ydb.table.values.ListType;
 import tech.ydb.table.values.ListValue;
@@ -211,4 +212,47 @@ public static BatchedQuery tryCreateBatched(YdbQuery query, Map<String, Type> ty
         StructType itemType = (StructType) innerType;
         return new BatchedQuery(query.getPreparedYql(), listName, itemType);
     }
+
+    public static BatchedQuery createAutoBatched(YqlBatcher batcher, Map<String, Type> tableColumns) {
+        StringBuilder sb = new StringBuilder();
+        Map<String, Type> structTypes = new HashMap<>();
+
+        sb.append("DECLARE $batch AS List<Struct<");
+        int idx = 1;
+        for (String column: batcher.getColumns()) {
+            Type type = tableColumns.get(column);
+            if (type == null) {
+                return null;
+            }
+            if (idx > 1) {
+                sb.append(", ");
+            }
+            sb.append("p").append(idx).append(":").append(type.toString());
+            structTypes.put("p" + idx, type);
+            idx++;
+        }
+        sb.append(">>;\n");
+
+        if (batcher.isInsert()) {
+            sb.append("INSERT ");
+        }
+        if (batcher.isUpsert()) {
+            sb.append("UPSERT ");
+        }
+
+        sb.append("INTO `").append(batcher.getTableName()).append("` SELECT ");
+
+        idx = 1;
+        for (String column: batcher.getColumns()) {
+            if (idx > 1) {
+                sb.append(", ");
+            }
+            sb.append("p").append(idx).append(" AS `").append(column).append("`");
+            idx++;
+        }
+
+        sb.append(" FROM AS_TABLE($batch);");
+
+        return new BatchedQuery(sb.toString(), "$batch", StructType.of(structTypes));
+    }
 }
diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java
@@ -112,6 +112,10 @@ public void executeUpdateTest(SqlQueries.JdbcQuery query) throws SQLException {
     @ParameterizedTest(name = "with {0}")
     @EnumSource(SqlQueries.JdbcQuery.class)
     public void executeWithMissingParameter(SqlQueries.JdbcQuery query) throws SQLException {
+        if (query == SqlQueries.JdbcQuery.STANDART) {
+            // TODO: auto batching don't check missing parameter
+            return;
+        }
         String sql = TEST_TABLE.upsertOne(query, "c_Text", "Text");
 
         try (PreparedStatement statement = jdbc.connection().prepareStatement(sql)) {
