Skip to content

Commit 4a35432

Browse files
authored
Add create context to all requests (#16319)
1 parent 5f9dbe8 commit 4a35432

File tree

1 file changed

+20
-5
lines changed
  • ydb/library/yql/providers/generic/connector/libcpp

1 file changed

+20
-5
lines changed

ydb/library/yql/providers/generic/connector/libcpp/client.cpp

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -252,16 +252,19 @@ namespace NYql::NConnector {
252252
promise.SetValue({std::move(status), std::move(NApi::TDescribeTableResponse())});
253253
return promise.GetFuture();
254254
}
255+
256+
auto context = CreateClientContext();
255257

256-
auto callback = [promise](NYdbGrpc::TGrpcStatus&& status, NApi::TDescribeTableResponse&& resp) mutable {
258+
auto callback = [context, promise](NYdbGrpc::TGrpcStatus&& status, NApi::TDescribeTableResponse&& resp) mutable {
257259
promise.SetValue({std::move(status), std::move(resp)});
258260
};
259261

260262
GetConnection(kind)->DoRequest<NApi::TDescribeTableRequest, NApi::TDescribeTableResponse>(
261263
std::move(request),
262264
std::move(callback),
263265
&NApi::Connector::Stub::AsyncDescribeTable,
264-
{ .Timeout = timeout }
266+
{ .Timeout = timeout },
267+
context.get()
265268
);
266269

267270
return promise.GetFuture();
@@ -307,6 +310,16 @@ namespace NYql::NConnector {
307310
}
308311

309312
private:
313+
NYdbGrpc::IQueueClientContextPtr CreateClientContext() {
314+
auto context = GrpcClient_->CreateContext();
315+
316+
if (!context) {
317+
throw yexception() << "Client is being shut down";
318+
}
319+
320+
return context;
321+
}
322+
310323
void Init(const TGenericGatewayConfig& config) {
311324
// TODO: place in a config file ?
312325
GrpcClient_ = std::make_shared<NYdbGrpc::TGRpcClientLow>(DEFAULT_CONNECTION_MANAGER_NUM_THREADS);
@@ -378,10 +391,11 @@ namespace NYql::NConnector {
378391
typename TRpcCallback = typename NYdbGrpc::TStreamRequestReadProcessor<NApi::Connector::Stub, TRequest, TResponse>::TAsyncRequest
379392
>
380393
TIteratorAsyncResult<IStreamIterator<TResponse>> ServerSideStreamingCall(
381-
const NYql::EGenericDataSourceKind& kind, const TRequest& request, TRpcCallback rpc, TDuration timeout = {}) const {
394+
const NYql::EGenericDataSourceKind& kind, const TRequest& request, TRpcCallback rpc, TDuration timeout = {}) {
382395
auto promise = NThreading::NewPromise<TIteratorResult<IStreamIterator<TResponse>>>();
396+
auto context = CreateClientContext();
383397

384-
auto callback = [promise](NYdbGrpc::TGrpcStatus&& status, NYdbGrpc::IStreamRequestReadProcessor<TResponse>::TPtr streamProcessor) mutable {
398+
auto callback = [context, promise](NYdbGrpc::TGrpcStatus&& status, NYdbGrpc::IStreamRequestReadProcessor<TResponse>::TPtr streamProcessor) mutable {
385399
if (!streamProcessor) {
386400
promise.SetValue({std::move(status), nullptr});
387401
return;
@@ -396,7 +410,8 @@ namespace NYql::NConnector {
396410
std::move(request),
397411
std::move(callback),
398412
rpc,
399-
{ .Timeout = timeout }
413+
{ .Timeout = timeout },
414+
context.get()
400415
);
401416

402417
return promise.GetFuture();

0 commit comments

Comments
 (0)