Added grpc flow control option to readTable,scanQuery and queryService

Author: alex268

query/src/main/java/tech/ydb/query/impl/SessionImpl.java

@@ -183,7 +183,7 @@ private static YdbQuery.StatsMode mapStatsMode(QueryStatsMode mode) {
     GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
             String query, YdbQuery.TransactionControl tx, Params prms, ExecuteQuerySettings settings
     ) {
-        YdbQuery.ExecuteQueryRequest.Builder requestBuilder = YdbQuery.ExecuteQueryRequest.newBuilder()
+        YdbQuery.ExecuteQueryRequest.Builder request = YdbQuery.ExecuteQueryRequest.newBuilder()
                 .setSessionId(sessionId)
                 .setExecMode(mapExecMode(settings.getExecMode()))
                 .setStatsMode(mapStatsMode(settings.getStatsMode()))
@@ -197,15 +197,19 @@ GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
 
         String resourcePool = settings.getResourcePool();
         if (resourcePool != null && !resourcePool.isEmpty()) {
-            requestBuilder.setPoolId(resourcePool);
+            request.setPoolId(resourcePool);
         }
 
         if (tx != null) {
-            requestBuilder.setTxControl(tx);
+            request.setTxControl(tx);
         }
 
-        YdbQuery.ExecuteQueryRequest request = requestBuilder.build();
-        return rpc.executeQuery(request, makeOptions(settings).build());
+        GrpcRequestSettings.Builder options = makeOptions(settings);
+        if (settings.getGrpcFlowControl() != null) {
+            options = options.withFlowControl(settings.getGrpcFlowControl());
+        }
+
+        return rpc.executeQuery(request.build(), options.build());
     }
 
     @Override

query/src/main/java/tech/ydb/query/settings/ExecuteQuerySettings.java

@@ -1,5 +1,6 @@
 package tech.ydb.query.settings;
 
+import tech.ydb.core.grpc.GrpcFlowControl;
 import tech.ydb.core.settings.BaseRequestSettings;
 
 /**
@@ -11,13 +12,15 @@ public class ExecuteQuerySettings extends BaseRequestSettings {
     private final QueryStatsMode statsMode;
     private final boolean concurrentResultSets;
     private final String resourcePool;
+    private final GrpcFlowControl flowControl;
 
     private ExecuteQuerySettings(Builder builder) {
         super(builder);
         this.execMode = builder.execMode;
         this.statsMode = builder.statsMode;
         this.concurrentResultSets = builder.concurrentResultSets;
         this.resourcePool = builder.resourcePool;
+        this.flowControl = builder.flowControl;
     }
 
     public QueryExecMode getExecMode() {
@@ -40,6 +43,10 @@ public String getResourcePool() {
         return this.resourcePool;
     }
 
+    public GrpcFlowControl getGrpcFlowControl() {
+        return flowControl;
+    }
+
     public static Builder newBuilder() {
         return new Builder();
     }
@@ -49,6 +56,7 @@ public static class Builder extends BaseBuilder<Builder> {
         private QueryStatsMode statsMode = QueryStatsMode.NONE;
         private boolean concurrentResultSets = false;
         private String resourcePool = null;
+        private GrpcFlowControl flowControl = null;
 
         public Builder withExecMode(QueryExecMode mode) {
             this.execMode = mode;
@@ -79,6 +87,11 @@ public Builder withResourcePool(String poolId) {
             return this;
         }
 
+        public Builder withGrpcFlowControl(GrpcFlowControl ctrl) {
+            this.flowControl = ctrl;
+            return this;
+        }
+
         @Override
         public ExecuteQuerySettings build() {
             return new ExecuteQuerySettings(this);

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

@@ -1199,9 +1199,12 @@ public GrpcReadStream<ReadTablePart> executeReadTable(String tablePath, ReadTabl
             request.addAllColumns(settings.getColumns());
         }
 
-        final GrpcRequestSettings opts = makeOptions(settings).build();
-        final GrpcReadStream<YdbTable.ReadTableResponse> origin = rpc.streamReadTable(request.build(), opts);
+        GrpcRequestSettings.Builder options = makeOptions(settings);
+        if (settings.getGrpcFlowControl() != null) {
+            options = options.withFlowControl(settings.getGrpcFlowControl());
+        }
 
+        GrpcReadStream<YdbTable.ReadTableResponse> origin = rpc.streamReadTable(request.build(), options.build());
         return new ProxyReadStream<>(origin, (response, future, observer) -> {
             StatusIds.StatusCode statusCode = response.getStatus();
             if (statusCode == StatusIds.StatusCode.SUCCESS) {
@@ -1222,18 +1225,20 @@ public GrpcReadStream<ReadTablePart> executeReadTable(String tablePath, ReadTabl
 
     @Override
     public GrpcReadStream<ResultSetReader> executeScanQuery(String query, Params params,
-                                                            ExecuteScanQuerySettings settings
-    ) {
+            ExecuteScanQuerySettings settings) {
         YdbTable.ExecuteScanQueryRequest req = YdbTable.ExecuteScanQueryRequest.newBuilder()
                 .setQuery(YdbTable.Query.newBuilder().setYqlText(query))
                 .setMode(settings.getMode().toPb())
                 .putAllParameters(params.toPb())
                 .setCollectStats(settings.getCollectStats().toPb())
                 .build();
 
-        GrpcRequestSettings opts = makeOptions(settings).build();
+        GrpcRequestSettings.Builder opts = makeOptions(settings);
+        if (settings.getGrpcFlowControl() != null) {
+            opts = opts.withFlowControl(settings.getGrpcFlowControl());
+        }
 
-        final GrpcReadStream<YdbTable.ExecuteScanQueryPartialResponse> origin = rpc.streamExecuteScanQuery(req, opts);
+        GrpcReadStream<YdbTable.ExecuteScanQueryPartialResponse> origin = rpc.streamExecuteScanQuery(req, opts.build());
         return new ProxyReadStream<>(origin, (response, future, observer) -> {
             StatusIds.StatusCode statusCode = response.getStatus();
             if (statusCode == StatusIds.StatusCode.SUCCESS) {

table/src/main/java/tech/ydb/table/settings/ExecuteScanQuerySettings.java

@@ -1,5 +1,6 @@
 package tech.ydb.table.settings;
 
+import tech.ydb.core.grpc.GrpcFlowControl;
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.proto.table.YdbTable;
 import tech.ydb.table.query.stats.QueryStatsCollectionMode;
@@ -27,11 +28,13 @@ public YdbTable.ExecuteScanQueryRequest.Mode toPb() {
 
     private final Mode mode;
     private final QueryStatsCollectionMode collectStats;
+    private final GrpcFlowControl flowControl;
 
     public ExecuteScanQuerySettings(Builder builder) {
         super(builder);
         this.mode = builder.mode;
         this.collectStats = builder.collectStats;
+        this.flowControl = builder.flowControl;
     }
 
     public static Builder newBuilder() {
@@ -41,6 +44,7 @@ public static Builder newBuilder() {
     public static final class Builder extends BaseBuilder<Builder> {
         private Mode mode = Mode.EXEC;
         private QueryStatsCollectionMode collectStats = QueryStatsCollectionMode.NONE;
+        private GrpcFlowControl flowControl = null;
 
         public Builder setMode(Mode mode) {
             this.mode = mode;
@@ -52,6 +56,11 @@ public Builder setCollectStats(QueryStatsCollectionMode collectStats) {
             return this;
         }
 
+        public Builder setGrpcFlowControl(GrpcFlowControl ctrl) {
+            this.flowControl = ctrl;
+            return this;
+        }
+
         @Override
         public ExecuteScanQuerySettings build() {
             return new ExecuteScanQuerySettings(this);
@@ -65,4 +74,8 @@ public Mode getMode() {
     public QueryStatsCollectionMode getCollectStats() {
         return collectStats;
     }
+
+    public GrpcFlowControl getGrpcFlowControl() {
+        return flowControl;
+    }
 }

table/src/main/java/tech/ydb/table/settings/ReadTableSettings.java

@@ -9,6 +9,7 @@
 
 import com.google.common.collect.ImmutableList;
 
+import tech.ydb.core.grpc.GrpcFlowControl;
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.table.values.PrimitiveValue;
 import tech.ydb.table.values.TupleValue;
@@ -31,6 +32,7 @@ public class ReadTableSettings extends BaseRequestSettings {
     private final ImmutableList<String> columns;
     private final int batchLimitBytes;
     private final int batchLimitRows;
+    private final GrpcFlowControl flowControl;
 
     private ReadTableSettings(Builder builder) {
         super(builder);
@@ -43,6 +45,7 @@ private ReadTableSettings(Builder builder) {
         this.columns = ImmutableList.copyOf(builder.columns);
         this.batchLimitBytes = builder.batchLimitBytes;
         this.batchLimitRows = builder.batchLimitRows;
+        this.flowControl = builder.flowControl;
     }
 
     public static Builder newBuilder() {
@@ -87,6 +90,10 @@ public int batchLimitRows() {
         return batchLimitRows;
     }
 
+    public GrpcFlowControl getGrpcFlowControl() {
+        return flowControl;
+    }
+
     /**
      * BUILDER
      */
@@ -100,6 +107,7 @@ public static final class Builder extends BaseBuilder<Builder> {
         private List<String> columns = Collections.emptyList();
         private int batchLimitBytes = 0;
         private int batchLimitRows = 0;
+        private GrpcFlowControl flowControl = null;
 
         public Builder orderedRead(boolean ordered) {
             this.ordered = ordered;
@@ -184,6 +192,11 @@ public Builder batchLimitRows(int batchLimitRows) {
             return this;
         }
 
+        public Builder withGrpcFlowControl(GrpcFlowControl ctrl) {
+            this.flowControl = ctrl;
+            return this;
+        }
+
         @Override
         public ReadTableSettings build() {
             return new ReadTableSettings(this);