Added implementation for BULK upserts

alex268 · alex268 · commit 94880ccd9f4f · 2024-09-12T18:16:58.000+01:00
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/YdbConnection.java b/jdbc/src/main/java/tech/ydb/jdbc/YdbConnection.java
@@ -38,12 +38,14 @@ public interface YdbConnection extends Connection {
     /**
      * Explicitly execute bulk upsert to the table
      *
+     * @param yql description of request
      * @param tablePath path to table
      * @param validator handler for logging and warnings
      * @param rows bulk rows
      * @throws SQLException if query cannot be executed
      */
-    void executeBulkUpsertQuery(String tablePath, YdbValidator validator, ListValue rows) throws SQLException;
+    void executeBulkUpsertQuery(String yql, String tablePath, YdbValidator validator, ListValue rows)
+            throws SQLException;
 
     /**
      * Explicitly execute query as a data query
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/YdbConst.java b/jdbc/src/main/java/tech/ydb/jdbc/YdbConst.java
@@ -65,7 +65,9 @@ public final class YdbConst {
     public static final String PARAMETER_NOT_FOUND = "Parameter not found: ";
     public static final String PARAMETER_TYPE_UNKNOWN = "Unable to convert sqlType %s to YDB type for parameter: %s";
     public static final String INVALID_ROW = "Current row index is out of bounds: ";
-    public static final String BULKS_UNSUPPORTED = "Bulk upserts are supported only in prepared statements";
+    public static final String BULKS_UNSUPPORTED = "BULK mode is available only for prepared statement with one UPSERT";
+    public static final String INVALID_BATCH_COLUMN = "Cannot prepared batch request: cannot find a column";
+    public static final String BULKS_DESCRIBE_ERROR = "Cannot parse BULK upsert: ";
     public static final String METADATA_RS_UNSUPPORTED_IN_PS = "ResultSet metadata is not supported " +
             "in prepared statements";
     public static final String CANNOT_UNWRAP_TO = "Cannot unwrap to ";
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java
@@ -38,10 +38,10 @@ public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yq
     }
 
     @Override
-    public void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String tablePath, ListValue rows)
+    public void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String yql, String tablePath, ListValue rows)
             throws SQLException {
         ensureOpened();
-        validator.execute(QueryType.BULK_QUERY + " >>\n" + tablePath,
+        validator.execute(QueryType.BULK_QUERY + " >>\n" + yql,
                 () -> retryCtx.supplyStatus(session -> session.executeBulkUpsert(tablePath, rows))
         );
     }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
@@ -30,6 +30,7 @@
 import tech.ydb.jdbc.query.YdbPreparedQuery;
 import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.jdbc.query.params.BatchedQuery;
+import tech.ydb.jdbc.query.params.BulkUpsertQuery;
 import tech.ydb.jdbc.query.params.InMemoryQuery;
 import tech.ydb.jdbc.query.params.PreparedQuery;
 import tech.ydb.jdbc.settings.YdbClientProperties;
@@ -350,16 +351,23 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
             }
         }
 
-        if (query.getType() == QueryType.EXPLAIN_QUERY || query.getType() == QueryType.SCHEME_QUERY) {
+        QueryType type = query.getType();
+
+        if (type == QueryType.BULK_QUERY) {
+            if (query.getYqlBatcher() == null || query.getYqlBatcher().isInsert()) {
+                throw new SQLException(YdbConst.BULKS_UNSUPPORTED);
+            }
+        }
+
+        if (type == QueryType.EXPLAIN_QUERY || type == QueryType.SCHEME_QUERY) {
             return new InMemoryQuery(query, queryOptions.isDeclareJdbcParameters());
         }
 
-        if (query.getYqlBatcher() != null && mode == YdbPrepareMode.AUTO) {
+        if (query.getYqlBatcher() != null && (mode == YdbPrepareMode.AUTO || type == QueryType.BULK_QUERY)) {
+            String tableName = query.getYqlBatcher().getTableName();
+            String tablePath = tableName.startsWith("/") ? tableName : getDatabase() + "/" + tableName;
             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)
@@ -370,8 +378,16 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
                     types = descrtiption.getColumns().stream()
                             .collect(Collectors.toMap(TableColumn::getName, TableColumn::getType));
                     queryParamsCache.put(query.getOriginQuery(), types);
+                } else {
+                    if (type == QueryType.BULK_QUERY) {
+                        throw new SQLException(YdbConst.BULKS_DESCRIBE_ERROR + result.getStatus());
+                    }
                 }
             }
+            if (type == QueryType.BULK_QUERY) {
+                return BulkUpsertQuery.build(tablePath, query.getYqlBatcher().getColumns(), types);
+            }
+
             if (types != null) {
                 BatchedQuery params = BatchedQuery.createAutoBatched(query.getYqlBatcher(), types);
                 if (params != null) {
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbExecutor.java
@@ -36,7 +36,7 @@ default void ensureOpened() throws SQLException {
 
     void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yql) throws SQLException;
 
-    void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String tablePath, ListValue rows)
+    void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String yql, String tablePath, ListValue rows)
             throws SQLException;
 
     List<ResultSetReader> executeDataQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, String yql,
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java
@@ -25,6 +25,7 @@
 import tech.ydb.jdbc.settings.YdbOperationProperties;
 import tech.ydb.table.query.Params;
 import tech.ydb.table.result.ResultSetReader;
+import tech.ydb.table.values.ListValue;
 
 /**
  *
@@ -167,6 +168,18 @@ protected boolean updateState(List<YdbResult> results) {
         return state.hasResultSets();
     }
 
+    protected List<YdbResult> executeBulkUpsert(YdbQuery query, String yql, String tablePath, ListValue rows)
+            throws SQLException {
+        connection.executeBulkUpsertQuery(yql, tablePath, validator, rows);
+
+        int expressionsCount = query.getStatements().isEmpty() ? 1 : query.getStatements().size();
+        List<YdbResult> results = new ArrayList<>();
+        for (int i = 0; i < expressionsCount; i++) {
+            results.add(HAS_UPDATED);
+        }
+        return results;
+    }
+
     protected List<YdbResult> executeSchemeQuery(YdbQuery query) throws SQLException {
         connection.executeSchemeQuery(query.getPreparedYql(), validator);
 
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java
@@ -235,7 +235,8 @@ public ResultSetReader executeScanQuery(YdbQuery query, String yql, YdbValidator
     }
 
     @Override
-    public void executeBulkUpsertQuery(String tablePath, YdbValidator validator, ListValue rows) throws SQLException {
+    public void executeBulkUpsertQuery(String yql, String tablePath, YdbValidator validator, ListValue rows)
+            throws SQLException {
         executor.ensureOpened();
 
         if (executor.isInsideTransaction()) {
@@ -251,7 +252,7 @@ public void executeBulkUpsertQuery(String tablePath, YdbValidator validator, Lis
             }
         }
 
-        executor.executeBulkUpsert(ctx, validator, tablePath, rows);
+        executor.executeBulkUpsert(ctx, validator, yql, tablePath, rows);
     }
 
     @Override
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java
@@ -32,8 +32,10 @@
 import tech.ydb.jdbc.YdbPreparedStatement;
 import tech.ydb.jdbc.YdbResultSet;
 import tech.ydb.jdbc.common.MappingSetters;
+import tech.ydb.jdbc.query.QueryType;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
 import tech.ydb.jdbc.query.YdbQuery;
+import tech.ydb.jdbc.query.params.BulkUpsertQuery;
 import tech.ydb.table.query.Params;
 import tech.ydb.table.values.Type;
 
@@ -90,8 +92,14 @@ public int[] executeBatch() throws SQLException {
         }
 
         try {
-            for (Params prm: prepared.getBatchParams()) {
-                executeDataQuery(query, prepared.getQueryText(prm), prm);
+            if (query.getType() == QueryType.BULK_QUERY && (prepared instanceof BulkUpsertQuery)) {
+                BulkUpsertQuery bulk = (BulkUpsertQuery) prepared;
+                String yql = bulk.getQueryText(null);
+                executeBulkUpsert(query, yql, bulk.getTablePath(), bulk.getBatchedBulk());
+            } else {
+                for (Params prm: prepared.getBatchParams()) {
+                    executeDataQuery(query, prepared.getQueryText(prm), prm);
+                }
             }
         } finally {
             clearBatch();
@@ -138,6 +146,17 @@ public boolean execute() throws SQLException {
             case EXPLAIN_QUERY:
                 newState = executeExplainQuery(query);
                 break;
+            case BULK_QUERY:
+                if (prepared instanceof BulkUpsertQuery) {
+                    BulkUpsertQuery bulk = (BulkUpsertQuery) prepared;
+                    String yql = bulk.getQueryText(null);
+                    newState = executeBulkUpsert(query, yql, bulk.getTablePath(), bulk.getCurrentBulk());
+                } else {
+                    throw new IllegalStateException(
+                            "Internal error. Incorrect class of bulk prepared query " + prepared.getClass()
+                    );
+                }
+                break;
             default:
                 throw new IllegalStateException("Internal error. Unsupported query type " + query.getType());
         }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbStatementImpl.java
@@ -94,6 +94,8 @@ public boolean execute(String sql) throws SQLException {
             case EXPLAIN_QUERY:
                 newState = executeExplainQuery(query);
                 break;
+            case BULK_QUERY:
+                throw new SQLException(YdbConst.BULKS_UNSUPPORTED);
             default:
                 throw new IllegalStateException("Internal error. Unsupported query type " + query.getType());
         }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/QueryType.java b/jdbc/src/main/java/tech/ydb/jdbc/query/QueryType.java
@@ -13,6 +13,6 @@ public enum QueryType {
     // EXPLAIN
     EXPLAIN_QUERY,
 
-    // BULK UPSERT
-    BULK_QUERY,
+    // BULK
+    BULK_QUERY;
 }
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
@@ -4,6 +4,7 @@
 import java.sql.SQLDataException;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -40,50 +41,22 @@ public class BatchedQuery implements YdbPreparedQuery {
     private final List<StructValue> batchList = new ArrayList<>();
     private final Map<String, Value<?>> currentValues = new HashMap<>();
 
-    private BatchedQuery(String yql, String listName, StructType structType) {
+    protected BatchedQuery(String yql, String listName, List<String> paramNames, Map<String, Type> types)
+            throws SQLException {
         this.yql = yql;
         this.batchParamName = listName;
         this.paramsByName = new HashMap<>();
-        this.params = new ParamDescription[structType.getMembersCount()];
+        this.params = new ParamDescription[paramNames.size()];
 
-        Map<String, Type> types = new HashMap<>();
-        for (int idx = 0; idx < structType.getMembersCount(); idx += 1) {
-            types.put(structType.getMemberName(idx), structType.getMemberType(idx));
-        }
-
-        // Firstly put all indexed params (p1, p2, ...,  pN) in correct places of paramNames
-        Set<String> indexedNames = new HashSet<>();
-        for (int idx = 0; idx < structType.getMembersCount(); idx += 1) {
-            String indexedName = YdbConst.INDEXED_PARAMETER_PREFIX + (1 + idx);
-            if (types.containsKey(indexedName)) {
-                String displayName = YdbConst.VARIABLE_PARAMETER_PREFIX + indexedName;
-                TypeDescription typeDesc = TypeDescription.of(types.get(indexedName));
-                ParamDescription paramDesc = new ParamDescription(indexedName, displayName, typeDesc);
-
-                params[idx] = paramDesc;
-                paramsByName.put(indexedName, paramDesc);
-                indexedNames.add(indexedName);
-            }
-        }
-
-        // Then put all others params in free places of paramNames in alphabetic order
-        Iterator<String> sortedIter = new TreeSet<>(types.keySet()).iterator();
-        for (int idx = 0; idx < params.length; idx += 1) {
-            if (params[idx] != null) {
-                continue;
-            }
-
-            String param = sortedIter.next();
-            while (indexedNames.contains(param)) {
-                param = sortedIter.next();
+        for (int idx = 0; idx < paramNames.size(); idx += 1) {
+            String name = paramNames.get(idx);
+            if (!types.containsKey(name)) {
+                throw new SQLException(YdbConst.INVALID_BATCH_COLUMN + name);
             }
-
-            String displayName = YdbConst.VARIABLE_PARAMETER_PREFIX + param;
-            TypeDescription typeDesc = TypeDescription.of(types.get(param));
-            ParamDescription paramDesc = new ParamDescription(param, displayName, typeDesc);
-
-            params[idx] = paramDesc;
-            paramsByName.put(param, paramDesc);
+            TypeDescription type = TypeDescription.of(types.get(name));
+            ParamDescription desc = new ParamDescription(name, YdbConst.VARIABLE_PARAMETER_PREFIX + name, type);
+            params[idx] = desc;
+            paramsByName.put(name, desc);
         }
     }
 
@@ -109,7 +82,7 @@ public void clearParameters() {
 
     @Override
     public void addBatch() throws SQLException {
-        batchList.add(validatedCurrentStruct());
+        batchList.add(getCurrentValues());
         currentValues.clear();
     }
 
@@ -118,7 +91,7 @@ public void clearBatch() {
         batchList.clear();
     }
 
-    private StructValue validatedCurrentStruct() throws SQLException {
+    protected StructValue getCurrentValues() throws SQLException {
         for (ParamDescription prm: params) {
             if (currentValues.containsKey(prm.name())) {
                 continue;
@@ -134,9 +107,13 @@ private StructValue validatedCurrentStruct() throws SQLException {
         return StructValue.of(currentValues);
     }
 
+    protected List<StructValue> getBatchedValues() {
+        return batchList;
+    }
+
     @Override
     public Params getCurrentParams() throws SQLException {
-        ListValue list = ListValue.of(validatedCurrentStruct());
+        ListValue list = ListValue.of(getCurrentValues());
         return Params.of(batchParamName, list);
     }
 
@@ -187,14 +164,14 @@ public TypeDescription getDescription(int index) throws SQLException {
         return params[index - 1].type();
     }
 
-    public static BatchedQuery tryCreateBatched(YdbQuery query, Map<String, Type> types) {
+    public static BatchedQuery tryCreateBatched(YdbQuery query, Map<String, Type> preparedTypes) throws SQLException {
         // Only single parameter
-        if (types.size() != 1) {
+        if (preparedTypes.size() != 1) {
             return null;
         }
 
-        String listName = types.keySet().iterator().next();
-        Type type = types.get(listName);
+        String listName = preparedTypes.keySet().iterator().next();
+        Type type = preparedTypes.get(listName);
 
         // Only list of values
         if (type.getKind() != Type.Kind.LIST) {
@@ -210,12 +187,46 @@ public static BatchedQuery tryCreateBatched(YdbQuery query, Map<String, Type> ty
         }
 
         StructType itemType = (StructType) innerType;
-        return new BatchedQuery(query.getPreparedYql(), listName, itemType);
+
+        String[] columns = new String[itemType.getMembersCount()];
+        Map<String, Type> types = new HashMap<>();
+        for (int idx = 0; idx < itemType.getMembersCount(); idx += 1) {
+            types.put(itemType.getMemberName(idx), itemType.getMemberType(idx));
+        }
+
+        // Firstly put all indexed params (p1, p2, ...,  pN) in correct places of paramNames
+        Set<String> indexedNames = new HashSet<>();
+        for (int idx = 0; idx < itemType.getMembersCount(); idx += 1) {
+            String indexedName = YdbConst.INDEXED_PARAMETER_PREFIX + (1 + idx);
+            if (types.containsKey(indexedName)) {
+                columns[idx] = indexedName;
+                indexedNames.add(indexedName);
+            }
+        }
+
+        // Then put all others params in free places of paramNames in alphabetic order
+        Iterator<String> sortedIter = new TreeSet<>(types.keySet()).iterator();
+        for (int idx = 0; idx < columns.length; idx += 1) {
+            if (columns[idx] != null) {
+                continue;
+            }
+
+            String param = sortedIter.next();
+            while (indexedNames.contains(param)) {
+                param = sortedIter.next();
+            }
+
+            columns[idx] = param;
+        }
+
+        return new BatchedQuery(query.getPreparedYql(), listName, Arrays.asList(columns), types);
     }
 
-    public static BatchedQuery createAutoBatched(YqlBatcher batcher, Map<String, Type> tableColumns) {
+    public static BatchedQuery createAutoBatched(YqlBatcher batcher, Map<String, Type> tableColumns)
+            throws SQLException {
         StringBuilder sb = new StringBuilder();
         Map<String, Type> structTypes = new HashMap<>();
+        List<String> columns = new ArrayList<>();
 
         sb.append("DECLARE $batch AS List<Struct<");
         int idx = 1;
@@ -229,6 +240,7 @@ public static BatchedQuery createAutoBatched(YqlBatcher batcher, Map<String, Typ
             }
             sb.append("p").append(idx).append(":").append(type.toString());
             structTypes.put("p" + idx, type);
+            columns.add("p" + idx);
             idx++;
         }
         sb.append(">>;\n");
@@ -253,6 +265,6 @@ public static BatchedQuery createAutoBatched(YqlBatcher batcher, Map<String, Typ
 
         sb.append(" FROM AS_TABLE($batch);");
 
-        return new BatchedQuery(sb.toString(), "$batch", StructType.of(structTypes));
+        return new BatchedQuery(sb.toString(), "$batch", columns, structTypes);
     }
 }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java
diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/helper/SqlQueries.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/helper/SqlQueries.java
diff --git a/jdbc/src/test/resources/sql/upsert/bulk.sql b/jdbc/src/test/resources/sql/upsert/bulk.sql

Original file line number	Diff line number	Diff line change
`@@ -38,10 +38,10 @@ public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, String yq`
`38`	`38`	`}`
`39`	`39`
`40`	`40`	`@Override`
`41`		`- public void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String tablePath, ListValue rows)`
	`41`	`+ public void executeBulkUpsert(YdbContext ctx, YdbValidator validator, String yql, String tablePath, ListValue rows)`
`42`	`42`	`throws SQLException {`
`43`	`43`	`ensureOpened();`
`44`		`- validator.execute(QueryType.BULK_QUERY + " >>\n" + tablePath,`
	`44`	`+ validator.execute(QueryType.BULK_QUERY + " >>\n" + yql,`
`45`	`45`	`() -> retryCtx.supplyStatus(session -> session.executeBulkUpsert(tablePath, rows))`
`46`	`46`	`);`
`47`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -235,7 +235,8 @@ public ResultSetReader executeScanQuery(YdbQuery query, String yql, YdbValidator`
`235`	`235`	`}`
`236`	`236`
`237`	`237`	`@Override`
`238`		`- public void executeBulkUpsertQuery(String tablePath, YdbValidator validator, ListValue rows) throws SQLException {`
	`238`	`+ public void executeBulkUpsertQuery(String yql, String tablePath, YdbValidator validator, ListValue rows)`
	`239`	`+ throws SQLException {`
`239`	`240`	`executor.ensureOpened();`
`240`	`241`
`241`	`242`	`if (executor.isInsideTransaction()) {`
`@@ -251,7 +252,7 @@ public void executeBulkUpsertQuery(String tablePath, YdbValidator validator, Lis`
`251`	`252`	`}`
`252`	`253`	`}`
`253`	`254`
`254`		`- executor.executeBulkUpsert(ctx, validator, tablePath, rows);`
	`255`	`+ executor.executeBulkUpsert(ctx, validator, yql, tablePath, rows);`
`255`	`256`	`}`
`256`	`257`
`257`	`258`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,8 @@ public boolean execute(String sql) throws SQLException {`
`94`	`94`	`case EXPLAIN_QUERY:`
`95`	`95`	`newState = executeExplainQuery(query);`
`96`	`96`	`break;`
	`97`	`+ case BULK_QUERY:`
	`98`	`+ throw new SQLException(YdbConst.BULKS_UNSUPPORTED);`
`97`	`99`	`default:`
`98`	`100`	`throw new IllegalStateException("Internal error. Unsupported query type " + query.getType());`
`99`	`101`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,6 @@ public enum QueryType {`
`13`	`13`	`// EXPLAIN`
`14`	`14`	`EXPLAIN_QUERY,`
`15`	`15`
`16`		`- // BULK UPSERT`
`17`		`- BULK_QUERY,`
	`16`	`+ // BULK`
	`17`	`+ BULK_QUERY;`
`18`	`18`	`}`