Report nodeId with create session response. (#562)

4 files changed, 13 insertions, 6 deletions

diff --git a/ydb/core/grpc_services/rpc_common/rpc_common_kqp_session.cpp b/ydb/core/grpc_services/rpc_common/rpc_common_kqp_session.cpp
index 102f5a7c1eb..041b1c9e1f9 100644
--- a/ydb/core/grpc_services/rpc_common/rpc_common_kqp_session.cpp
+++ b/ydb/core/grpc_services/rpc_common/rpc_common_kqp_session.cpp
@@ -124,7 +124,7 @@ private:
                 // We already lost the client, so the client should not see this status
                 Reply(Ydb::StatusIds::INTERNAL_ERROR);
             } else {
-                SendSessionResult(kqpResponse.GetSessionId());
+                SendSessionResult(kqpResponse);
                 PassAway();
                 return;
             }
@@ -134,7 +134,7 @@ private:
     }
 
 private:
-    virtual void SendSessionResult(const TString& id) = 0;
+    virtual void SendSessionResult(const NKikimrKqp::TCreateSessionResponse& kqpResponse) = 0;
 
     template<typename TResp>
     void ReplyResponseError(const TResp& kqpResponse) {
@@ -171,9 +171,9 @@ public:
     }
 
 private:
-    void SendSessionResult(const TString& id) override {
+    void SendSessionResult(const NKikimrKqp::TCreateSessionResponse& kqpResponse) override {
         Ydb::Table::CreateSessionResult result;
-        result.set_session_id(id);
+        result.set_session_id(kqpResponse.GetSessionId());
         static_cast<TCtx*>(Request.get())->SendResult(result, Ydb::StatusIds::SUCCESS);
     };
 };
@@ -186,11 +186,15 @@ public:
     }
 
 private:
-    void SendSessionResult(const TString& id) override {
+    void SendSessionResult(const NKikimrKqp::TCreateSessionResponse& kqpResponse) override {
         using TRes = Ydb::Query::CreateSessionResponse;
         auto res = google::protobuf::Arena::CreateMessage<TRes>(Request->GetArena());;
         res->set_status(Ydb::StatusIds::SUCCESS);
-        res->set_session_id(id);
+        res->set_session_id(kqpResponse.GetSessionId());
+
+        if (kqpResponse.HasNodeId())
+            res->set_node_id(kqpResponse.GetNodeId());
+
         Request->Reply(res, Ydb::StatusIds::SUCCESS);
     };
 };
diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
index 537aef1ac23..b68763142a1 100644
--- a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
+++ b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
@@ -570,6 +570,7 @@ public:
         {
             auto& response = *responseEv->Record.MutableResponse();
             response.SetSessionId(result.Value->SessionId);
+            response.SetNodeId(SelfId().NodeId());
             dbCounters = result.Value->DbCounters;
         } else {
             dbCounters = Counters->GetDbCounters(request.GetDatabase());
diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto
index 30e6c740aa0..eb64dabcb1f 100644
--- a/ydb/core/protos/kqp.proto
+++ b/ydb/core/protos/kqp.proto
@@ -332,6 +332,7 @@ message TEvCreateSessionRequest {
 
 message TCreateSessionResponse {
     optional bytes SessionId = 1;
+    optional int64 NodeId = 2;
 }
 
 message TEvCreateSessionResponse {
diff --git a/ydb/public/lib/ut_helpers/ut_helpers_query.cpp b/ydb/public/lib/ut_helpers/ut_helpers_query.cpp
index e23bbadecb4..1d2193fbac2 100644
--- a/ydb/public/lib/ut_helpers/ut_helpers_query.cpp
+++ b/ydb/public/lib/ut_helpers/ut_helpers_query.cpp
@@ -26,6 +26,7 @@ TString CreateQuerySession(const TGRpcClientConfig& clientConfig) {
             UNIT_ASSERT_C(grpcStatus.GRpcStatusCode == 0, grpcStatus.Msg + " " + grpcStatus.Details);
             UNIT_ASSERT_VALUES_EQUAL(response.status(), Ydb::StatusIds::SUCCESS);
             UNIT_ASSERT(response.session_id() != "");
+            UNIT_ASSERT(response.node_id() != 0);
             sessionId = response.session_id();
     };