Added customization of BulkUpsert data

Author: alex268

table/src/main/java/tech/ydb/table/Session.java

@@ -11,6 +11,7 @@
 import tech.ydb.core.grpc.GrpcReadStream;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.description.TableOptionDescription;
+import tech.ydb.table.query.BulkUpsertData;
 import tech.ydb.table.query.DataQuery;
 import tech.ydb.table.query.DataQueryResult;
 import tech.ydb.table.query.ExplainDataQueryResult;
@@ -178,7 +179,15 @@ default CompletableFuture<Status> executeScanQuery(String query, Params params,
 
     CompletableFuture<Result<State>> keepAlive(KeepAliveSessionSettings settings);
 
-    CompletableFuture<Status> executeBulkUpsert(String tablePath, ListValue rows, BulkUpsertSettings settings);
+    default CompletableFuture<Status> executeBulkUpsert(String tablePath, ListValue rows, BulkUpsertSettings settings) {
+        return executeBulkUpsert(tablePath, new BulkUpsertData(rows), settings);
+    }
+
+    CompletableFuture<Status> executeBulkUpsert(String tablePath, BulkUpsertData data, BulkUpsertSettings settings);
+
+    default CompletableFuture<Status> executeBulkUpsert(String tablePath, BulkUpsertData data) {
+        return executeBulkUpsert(tablePath, data, new BulkUpsertSettings());
+    }
 
     default CompletableFuture<Status> createTable(String path, TableDescription tableDescriptions) {
         return createTable(path, tableDescriptions, new CreateTableSettings());

table/src/main/java/tech/ydb/table/impl/BaseSession.java

@@ -57,6 +57,7 @@
 import tech.ydb.table.description.TableIndex;
 import tech.ydb.table.description.TableOptionDescription;
 import tech.ydb.table.description.TableTtl;
+import tech.ydb.table.query.BulkUpsertData;
 import tech.ydb.table.query.DataQuery;
 import tech.ydb.table.query.DataQueryResult;
 import tech.ydb.table.query.ExplainDataQueryResult;
@@ -99,7 +100,6 @@
 import tech.ydb.table.transaction.Transaction;
 import tech.ydb.table.transaction.TxControl;
 import tech.ydb.table.values.ListType;
-import tech.ydb.table.values.ListValue;
 import tech.ydb.table.values.PrimitiveValue;
 import tech.ydb.table.values.StructValue;
 import tech.ydb.table.values.TupleValue;
@@ -1394,19 +1394,13 @@ public CompletableFuture<Result<State>> keepAlive(KeepAliveSessionSettings setti
     }
 
     @Override
-    public CompletableFuture<Status> executeBulkUpsert(String tablePath, ListValue rows, BulkUpsertSettings settings) {
-        ValueProtos.TypedValue typedRows = ValueProtos.TypedValue.newBuilder()
-                .setType(rows.getType().toPb())
-                .setValue(rows.toPb())
-                .build();
-
-        YdbTable.BulkUpsertRequest request = YdbTable.BulkUpsertRequest.newBuilder()
+    public CompletableFuture<Status> executeBulkUpsert(String tablePath, BulkUpsertData data, BulkUpsertSettings st) {
+        YdbTable.BulkUpsertRequest.Builder request = YdbTable.BulkUpsertRequest.newBuilder()
                 .setTable(tablePath)
-                .setRows(typedRows)
-                .setOperationParams(Operation.buildParams(settings.toOperationSettings()))
-                .build();
+                .setOperationParams(Operation.buildParams(st.toOperationSettings()));
 
-        return interceptStatus(rpc.bulkUpsert(request, makeOptions(settings).build()));
+        data.applyToRequest(request);
+        return interceptStatus(rpc.bulkUpsert(request.build(), makeOptions(st).build()));
     }
 
     private static State mapSessionStatus(YdbTable.KeepAliveResult result) {

table/src/main/java/tech/ydb/table/query/BulkUpsertData.java

@@ -0,0 +1,29 @@
+package tech.ydb.table.query;
+
+
+import tech.ydb.proto.ValueProtos;
+import tech.ydb.proto.table.YdbTable;
+import tech.ydb.table.values.ListValue;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class BulkUpsertData {
+    private final ValueProtos.TypedValue rows;
+
+    public BulkUpsertData(ListValue rows) {
+        this.rows = ValueProtos.TypedValue.newBuilder()
+                .setType(rows.getType().toPb())
+                .setValue(rows.toPb())
+                .build();
+    }
+
+    public BulkUpsertData(ValueProtos.TypedValue rows) {
+        this.rows = rows;
+    }
+
+    public void applyToRequest(YdbTable.BulkUpsertRequest.Builder builder) {
+        builder.setRows(rows);
+    }
+}

table/src/test/java/tech/ydb/table/SessionStub.java

@@ -10,6 +10,7 @@
 import tech.ydb.core.utils.FutureTools;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.description.TableOptionDescription;
+import tech.ydb.table.query.BulkUpsertData;
 import tech.ydb.table.query.DataQuery;
 import tech.ydb.table.query.DataQueryResult;
 import tech.ydb.table.query.ExplainDataQueryResult;
@@ -40,7 +41,6 @@
 import tech.ydb.table.transaction.TableTransaction;
 import tech.ydb.table.transaction.Transaction;
 import tech.ydb.table.transaction.TxControl;
-import tech.ydb.table.values.ListValue;
 
 
 /**
@@ -178,7 +178,7 @@ public CompletableFuture<Result<Session.State>> keepAlive(KeepAliveSessionSettin
     }
 
     @Override
-    public CompletableFuture<Status> executeBulkUpsert(String tablePath, ListValue rows, BulkUpsertSettings settings) {
+    public CompletableFuture<Status> executeBulkUpsert(String tablePath, BulkUpsertData rows, BulkUpsertSettings settings) {
         return notImplemented("bulkUpsert()");
     }
 

table/src/test/java/tech/ydb/table/integration/ReadTableTest.java

@@ -28,10 +28,10 @@
 import tech.ydb.table.SessionRetryContext;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.impl.SimpleTableClient;
+import tech.ydb.table.query.BulkUpsertData;
 import tech.ydb.table.query.ReadTablePart;
 import tech.ydb.table.result.ResultSetReader;
 import tech.ydb.table.rpc.grpc.GrpcTableRpc;
-import tech.ydb.table.settings.BulkUpsertSettings;
 import tech.ydb.table.settings.ReadTableSettings;
 import tech.ydb.table.values.ListType;
 import tech.ydb.table.values.PrimitiveType;
@@ -98,9 +98,9 @@ public static void prepareTable() {
             );
         }).collect(Collectors.toList());
 
-        retryCtx.supplyStatus(session -> session.executeBulkUpsert(
-                tablePath, ListType.of(batchType).newValue(batchData), new BulkUpsertSettings())
-        ).join().expectSuccess("bulk upsert problem in table " + tablePath);
+        retryCtx.supplyStatus(session -> session.executeBulkUpsert(tablePath,
+                new BulkUpsertData(ProtoValue.toTypedValue(ListType.of(batchType).newValue(batchData)))
+        )).join().expectSuccess("bulk upsert problem in table " + tablePath);
     }
 
     @AfterClass