Skip to content

Commit 987c9c3

Browse files
committed
Refactoring of grpc settings creation
1 parent e88a254 commit 987c9c3

File tree

2 files changed

+64
-117
lines changed

2 files changed

+64
-117
lines changed

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

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -119,8 +119,7 @@ public CompletableFuture<Result<QueryTransaction>> beginTransaction(TxMode tx, B
119119
.setTxSettings(TxControl.txSettings(tx))
120120
.build();
121121

122-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
123-
return rpc.beginTransaction(request, grpcSettings).thenApply(result -> {
122+
return rpc.beginTransaction(request, makeOptions(settings).build()).thenApply(result -> {
124123
updateSessionState(result.getStatus());
125124
return result.map(resp -> updateTransaction(new TransactionImpl(tx, resp.getTxMeta().getId())));
126125
});
@@ -138,7 +137,7 @@ GrpcReadStream<Status> attach(AttachSessionSettings settings) {
138137
YdbQuery.AttachSessionRequest request = YdbQuery.AttachSessionRequest.newBuilder()
139138
.setSessionId(sessionId)
140139
.build();
141-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
140+
GrpcRequestSettings grpcSettings = makeOptions(settings).build();
142141
return new ProxyReadStream<>(rpc.attachSession(request, grpcSettings), (message, promise, observer) -> {
143142
logger.trace("session '{}' got attach stream message {}", sessionId, TextFormat.shortDebugString(message));
144143
Status status = Status.of(StatusCode.fromProto(message.getStatus()), Issue.fromPb(message.getIssuesList()));
@@ -147,13 +146,12 @@ GrpcReadStream<Status> attach(AttachSessionSettings settings) {
147146
});
148147
}
149148

150-
private GrpcRequestSettings makeGrpcRequestSettings(BaseRequestSettings settings) {
149+
private GrpcRequestSettings.Builder makeOptions(BaseRequestSettings settings) {
151150
String traceId = settings.getTraceId() == null ? UUID.randomUUID().toString() : settings.getTraceId();
152151
return GrpcRequestSettings.newBuilder()
153152
.withDeadline(settings.getRequestTimeout())
154153
.withPreferredNodeID((int) nodeID)
155-
.withTraceId(traceId)
156-
.build();
154+
.withTraceId(traceId);
157155
}
158156

159157
private static YdbQuery.ExecMode mapExecMode(QueryExecMode mode) {
@@ -207,8 +205,7 @@ GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
207205
}
208206

209207
YdbQuery.ExecuteQueryRequest request = requestBuilder.build();
210-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
211-
return rpc.executeQuery(request, grpcSettings);
208+
return rpc.executeQuery(request, makeOptions(settings).build());
212209
}
213210

214211
@Override
@@ -229,8 +226,7 @@ public CompletableFuture<Result<YdbQuery.DeleteSessionResponse>> delete(DeleteSe
229226
.setSessionId(sessionId)
230227
.build();
231228

232-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
233-
return rpc.deleteSession(request, grpcSettings).thenApply(DELETE_SESSION);
229+
return rpc.deleteSession(request, makeOptions(settings).build()).thenApply(DELETE_SESSION);
234230
}
235231

236232
static CompletableFuture<Result<YdbQuery.CreateSessionResponse>> createSession(
@@ -405,8 +401,7 @@ public CompletableFuture<Result<QueryInfo>> commit(CommitTransactionSettings set
405401
.setSessionId(sessionId)
406402
.setTxId(transactionId)
407403
.build();
408-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
409-
return rpc.commitTransaction(request, grpcSettings)
404+
return rpc.commitTransaction(request, makeOptions(settings).build())
410405
.thenApply(res -> {
411406
Status status = res.getStatus();
412407
currentStatusFuture.complete(status);
@@ -438,8 +433,7 @@ public CompletableFuture<Status> rollback(RollbackTransactionSettings settings)
438433
.setSessionId(sessionId)
439434
.setTxId(transactionId)
440435
.build();
441-
GrpcRequestSettings grpcSettings = makeGrpcRequestSettings(settings);
442-
return rpc.rollbackTransaction(request, grpcSettings)
436+
return rpc.rollbackTransaction(request, makeOptions(settings).build())
443437
.thenApply(result -> {
444438
updateSessionState(result.getStatus());
445439
if (!txId.compareAndSet(transactionId, null)) {

0 commit comments

Comments
 (0)