Added database id into create session request

ydb-platform · Sep 24, 2024 · 9cae004 · 9cae004
1 parent 5ae39b9
commit 9cae004
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/ydb/core/kqp/common/events/events.h b/ydb/core/kqp/common/events/events.h
@@ -26,7 +26,20 @@ struct TEvKqp {
         NKikimrKqp::TEvCloseSessionRequest, TKqpEvents::EvCloseSessionRequest> {};
 
     struct TEvCreateSessionRequest : public TEventPB<TEvCreateSessionRequest,
-        NKikimrKqp::TEvCreateSessionRequest, TKqpEvents::EvCreateSessionRequest> {};
+        NKikimrKqp::TEvCreateSessionRequest, TKqpEvents::EvCreateSessionRequest> {
+
+        const TString& GetDatabase() const {
+            return Record.GetRequest().GetDatabase();
+        }
+
+        const TString& GetDatabaseId() const {
+            return Record.GetRequest().GetDatabaseId();
+        }
+
+        void SetDatabaseId(const TString& databaseId) {
+            Record.MutableRequest()->SetDatabaseId(databaseId);
+        }
+    };
 
     struct TEvPingSessionRequest : public TEventPB<TEvPingSessionRequest,
         NKikimrKqp::TEvPingSessionRequest, TKqpEvents::EvPingSessionRequest> {};

diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
@@ -590,6 +590,16 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
     }
 
     void Handle(TEvKqp::TEvCreateSessionRequest::TPtr& ev) {
+        const auto errorHandler = [this, sender = ev->Sender](Ydb::StatusIds::StatusCode status, NYql::TIssues issues){
+            auto response = std::make_unique<TEvKqp::TEvCreateSessionResponse>();
+            response->Record.SetYdbStatus(status);
+            response->Record.SetError(issues.ToOneLineString());
+            Send(sender, std::move(response));
+        };
+        if (!DatabasesCache.SetDatabaseIdOrDeffer(ev, errorHandler, ActorContext())) {
+            return;
+        }
+
         auto& event = ev->Get()->Record;
         auto& request = event.GetRequest();
         TKqpRequestInfo requestInfo(event.GetTraceId());

diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto
@@ -346,6 +346,7 @@ message TKqpDefaultSettings {
 
 message TCreateSessionRequest {
     optional string Database = 1;
+    optional string DatabaseId = 2;
 }
 
 message TEvCreateSessionRequest {