diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp index 11da07090688..7e6ec315071d 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp @@ -646,8 +646,6 @@ class TKqpProxyService : public TActorBootstrapped { return; } - Cerr << "------------------------------- DatabaseId: " << ev->Get()->GetDatabaseId() << "\n"; - const TString& database = ev->Get()->GetDatabase(); const TString& traceId = ev->Get()->GetTraceId(); const auto queryType = ev->Get()->GetType(); diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp index 883ec7d9198e..d61d643e0588 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp @@ -2,7 +2,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -65,6 +67,63 @@ TString CreateSession(TTestActorRuntime* runtime, const TActorId& kqpProxy, cons return sessionId; } +class TDatabaseCacheTestActor : public TActorBootstrapped { +public: + TDatabaseCacheTestActor(const TString& database, const TString& expectedDatabaseId, bool fromCache, TDatabasesCache& cache, NThreading::TPromise promise) + : Database(database) + , ExpectedDatabaseId(expectedDatabaseId) + , Cache(cache) + , Promise(promise) + , FromCache(fromCache) + {} + + void Bootstrap() { + Become(&TDatabaseCacheTestActor::StateFunc); + + auto event = MakeHolder(); + event->Record.MutableRequest()->SetDatabase(Database); + Send(SelfId(), event.Release()); + } + + void Handle(TEvKqp::TEvUpdateDatabaseInfo::TPtr& ev) { + Cache.UpdateDatabaseInfo(ev, ActorContext()); + } + + void Handle(TEvKqp::TEvQueryRequest::TPtr& ev) { + auto success = Cache.SetDatabaseIdOrDeffer(ev, [this](Ydb::StatusIds::StatusCode status, NYql::TIssues issues){ + UNIT_ASSERT_C(false, TStringBuilder() << "Unexpected fail, " << GetErrorString() << ", status: " << status << ", reason: " << issues.ToOneLineString()); + }, ActorContext()); + + if (FromCache) { + UNIT_ASSERT_C(success, TStringBuilder() << "Expected database id from cache, " << GetErrorString()); + UNIT_ASSERT_STRING_CONTAINS_C(ev->Get()->GetDatabaseId(), ExpectedDatabaseId, GetErrorString()); + Promise.SetValue(); + PassAway(); + } else { + UNIT_ASSERT_C(!success, TStringBuilder() << "Unexpected database id from cache, " << GetErrorString()); + FromCache = true; + } + } + + STRICT_STFUNC(StateFunc, + hFunc(TEvKqp::TEvUpdateDatabaseInfo, Handle); + hFunc(TEvKqp::TEvQueryRequest, Handle); + ) + +private: + TString GetErrorString() const { + return TStringBuilder() << "database: " << Database << ", from cache: " << FromCache << "\n"; + } + +private: + const TString Database; + const TString ExpectedDatabaseId; + TDatabasesCache& Cache; + NThreading::TPromise Promise; + + bool FromCache = false; +}; + } Y_UNIT_TEST_SUITE(KqpProxy) { @@ -542,5 +601,32 @@ Y_UNIT_TEST_SUITE(KqpProxy) { UNIT_ASSERT(allDoneOk); } + + Y_UNIT_TEST(DatabasesCacheForServerless) { + auto ydb = NWorkload::TYdbSetupSettings() + .CreateSampleTenants(true) + .Create(); + + auto& runtime = *ydb->GetRuntime(); + TDatabasesCache cache; + + auto checkCache = [&](const TString& database, const TString& expectedDatabaseId, bool fromCache) { + auto promise = NThreading::NewPromise(); + runtime.Register(new TDatabaseCacheTestActor(database, expectedDatabaseId, fromCache, cache, promise)); + promise.GetFuture().GetValueSync(); + }; + + const auto& dedicatedTennant = ydb->GetSettings().GetDedicatedTenantName(); + checkCache(dedicatedTennant, dedicatedTennant, false); + checkCache(dedicatedTennant, dedicatedTennant, true); + + const auto& sharedTennant = ydb->GetSettings().GetSharedTenantName(); + checkCache(sharedTennant, sharedTennant, false); + checkCache(sharedTennant, sharedTennant, true); + + const auto& serverlessTennant = ydb->GetSettings().GetServerlessTenantName(); + checkCache(serverlessTennant, TStringBuilder() << ":4:" << serverlessTennant, false); + checkCache(serverlessTennant, TStringBuilder() << ":4:" << serverlessTennant, true); + } } // namspace NKqp } // namespace NKikimr