From cc6b3561762498f7de0520a81b7f6c7596209b6a Mon Sep 17 00:00:00 2001 From: Grigoriy Pisarenko Date: Mon, 4 Mar 2024 06:53:43 +0000 Subject: [PATCH] Passed yt token --- .../external_source_factory.cpp | 2 +- .../kqp/executer_actor/kqp_data_executer.cpp | 1 + .../kqp/executer_actor/kqp_executer_impl.h | 17 +++++++ .../external_data_source/manager.cpp | 3 ++ ydb/core/kqp/gateway/kqp_metadata_loader.cpp | 15 ++++++ .../kqp/provider/yql_kikimr_datasource.cpp | 6 +++ ydb/core/kqp/provider/yql_kikimr_gateway.h | 1 + .../kqp/query_compiler/kqp_query_compiler.cpp | 50 ++++++++++++------- ydb/core/protos/flat_scheme_op.proto | 5 ++ ydb/core/protos/kqp_physical.proto | 1 + ..._operation_common_external_data_source.cpp | 2 + .../structured_token/yql_token_builder.cpp | 27 ++++++++++ .../structured_token/yql_token_builder.h | 2 + .../yt/provider/yql_yt_datasource.cpp | 5 ++ .../yt/provider/yql_yt_dq_integration.cpp | 4 +- ydb/library/yql/sql/v1/sql_translation.cpp | 6 ++- ydb/tests/tools/kqprun/kqprun.cpp | 24 +++++++-- ydb/tests/tools/kqprun/src/common.h | 2 + ydb/tests/tools/kqprun/src/kqp_runner.cpp | 8 +-- ydb/tests/tools/kqprun/src/kqp_runner.h | 2 +- ydb/tests/tools/kqprun/src/ydb_setup.cpp | 38 +++++++------- ydb/tests/tools/kqprun/src/ydb_setup.h | 2 +- 22 files changed, 170 insertions(+), 53 deletions(-) diff --git a/ydb/core/external_sources/external_source_factory.cpp b/ydb/core/external_sources/external_source_factory.cpp index 032d030968e5..98e4514f0242 100644 --- a/ydb/core/external_sources/external_source_factory.cpp +++ b/ydb/core/external_sources/external_source_factory.cpp @@ -53,7 +53,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector { protected: + void FillSecureParamsFromStage(THashMap& secureParams, const NKqpProto::TKqpPhyStage& stage) { + for (const auto& [secretName, authInfo] : stage.GetSecureParams()) { + const auto& structuredToken = NYql::CreateStructuredTokenParser(authInfo).ToBuilder().ReplaceReferences(SecureParams).ToJson(); + const auto& structuredTokenParser = NYql::CreateStructuredTokenParser(structuredToken); + YQL_ENSURE(structuredTokenParser.HasIAMToken(), "only token authentification supported for compute tasks"); + secureParams.emplace(secretName, structuredTokenParser.GetIAMToken()); + } + } + void BuildSysViewScanTasks(TStageInfo& stageInfo) { Y_DEBUG_ABORT_UNLESS(stageInfo.Meta.IsSysView()); @@ -873,6 +882,7 @@ class TKqpExecuterBase : public TActorBootstrapped { task.Meta.ReadInfo.Reverse = op.GetReadRange().GetReverse(); task.Meta.Type = TTaskMeta::TTaskType::Compute; + FillSecureParamsFromStage(task.Meta.SecureParams, stage); BuildSinks(stage, task); LOG_D("Stage " << stageInfo.Id << " create sysview scan task: " << task.Id); @@ -959,6 +969,7 @@ class TKqpExecuterBase : public TActorBootstrapped { if (structuredToken) { task.Meta.SecureParams.emplace(sourceName, structuredToken); } + FillSecureParamsFromStage(task.Meta.SecureParams, stage); if (resourceSnapshot.empty()) { task.Meta.Type = TTaskMeta::TTaskType::Compute; @@ -1051,6 +1062,7 @@ class TKqpExecuterBase : public TActorBootstrapped { YQL_ENSURE(!shardsResolved); task.Meta.ShardId = taskLocation; } + FillSecureParamsFromStage(task.Meta.SecureParams, stage); const auto& stageSource = stage.GetSources(0); auto& input = task.Inputs[stageSource.GetInputIndex()]; @@ -1306,6 +1318,7 @@ class TKqpExecuterBase : public TActorBootstrapped { auto& task = TasksGraph.AddTask(stageInfo); task.Meta.Type = TTaskMeta::TTaskType::Compute; task.Meta.ExecuterId = SelfId(); + FillSecureParamsFromStage(task.Meta.SecureParams, stage); BuildSinks(stage, task); LOG_D("Stage " << stageInfo.Id << " create compute task: " << task.Id); } @@ -1441,6 +1454,7 @@ class TKqpExecuterBase : public TActorBootstrapped { THashMap& assignedShardsCount, const bool sorted, const bool isOlapScan) { + const auto& stage = stageInfo.Meta.GetStage(stageInfo.Id); ui64 nodeId = ShardIdToNodeId.at(shardId); if (stageInfo.Meta.IsOlap() && sorted) { auto& task = TasksGraph.AddTask(stageInfo); @@ -1448,6 +1462,7 @@ class TKqpExecuterBase : public TActorBootstrapped { task.Meta.NodeId = nodeId; task.Meta.ScanTask = true; task.Meta.Type = TTaskMeta::TTaskType::Scan; + FillSecureParamsFromStage(task.Meta.SecureParams, stage); return task; } @@ -1459,6 +1474,7 @@ class TKqpExecuterBase : public TActorBootstrapped { task.Meta.NodeId = nodeId; task.Meta.ScanTask = true; task.Meta.Type = TTaskMeta::TTaskType::Scan; + FillSecureParamsFromStage(task.Meta.SecureParams, stage); tasks.push_back(task.Id); ++cnt; return task; @@ -1552,6 +1568,7 @@ class TKqpExecuterBase : public TActorBootstrapped { task.Meta.ScanTask = true; task.Meta.Type = TTaskMeta::TTaskType::Scan; task.SetMetaId(metaGlueingId); + FillSecureParamsFromStage(task.Meta.SecureParams, stage); BuildSinks(stage, task); } } diff --git a/ydb/core/kqp/gateway/behaviour/external_data_source/manager.cpp b/ydb/core/kqp/gateway/behaviour/external_data_source/manager.cpp index ad65bde5da0f..e186fcd27064 100644 --- a/ydb/core/kqp/gateway/behaviour/external_data_source/manager.cpp +++ b/ydb/core/kqp/gateway/behaviour/external_data_source/manager.cpp @@ -61,6 +61,9 @@ void FillCreateExternalDataSourceDesc(NKikimrSchemeOp::TExternalDataSourceDescri aws.SetAwsAccessKeyIdSecretName(GetOrEmpty(settings, "aws_access_key_id_secret_name")); aws.SetAwsSecretAccessKeySecretName(GetOrEmpty(settings, "aws_secret_access_key_secret_name")); aws.SetAwsRegion(GetOrEmpty(settings, "aws_region")); + } else if (authMethod == "TOKEN") { + auto& token = *externaDataSourceDesc.MutableAuth()->MutableToken(); + token.SetTokenSecretName(GetOrEmpty(settings, "token_secret_name")); } else { ythrow yexception() << "Internal error. Unknown auth method: " << authMethod; } diff --git a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp index cf0c76f3546d..a2ab6f59d1e0 100644 --- a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp +++ b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp @@ -450,6 +450,14 @@ void UpdateExternalDataSourceSecretsValue(TTableMetadataResult& externalDataSour externalDataSourceMetadata.Metadata->ExternalSource.AwsSecretAccessKey = objectDescription.SecretValues[1]; return; } + case NKikimrSchemeOp::TAuth::kToken: { + if (objectDescription.SecretValues.size() != 1) { + SetError(externalDataSourceMetadata, TStringBuilder{} << "Token auth contains invalid count of secrets: " << objectDescription.SecretValues.size() << " instead of 1"); + return; + } + externalDataSourceMetadata.Metadata->ExternalSource.Token = objectDescription.SecretValues[0]; + return; + } case NKikimrSchemeOp::TAuth::IDENTITY_NOT_SET: { SetError(externalDataSourceMetadata, "identity case is not specified in case of update external data source secrets"); return; @@ -494,6 +502,13 @@ NThreading::TFuture LoadExternalDataSo return promise.GetFuture(); } + case NKikimrSchemeOp::TAuth::kToken: { + const TString& tokenSecretId = authDescription.GetToken().GetTokenSecretName(); + auto promise = NewPromise(); + actorSystem->Register(CreateDescribeSecretsActor(userToken ? userToken->GetUserSID() : "", {tokenSecretId}, promise, maximalSecretsSnapshotWaitTime)); + return promise.GetFuture(); + } + case NKikimrSchemeOp::TAuth::IDENTITY_NOT_SET: return MakeFuture(TEvDescribeSecretsResponse::TDescription(Ydb::StatusIds::BAD_REQUEST, { NYql::TIssue("identity case is not specified") })); } diff --git a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp index af95adf40d5a..991aeb933ddc 100644 --- a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp +++ b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp @@ -105,6 +105,12 @@ TString FillAuthProperties(THashMap& properties, const TExtern properties["awsRegion"] = externalSource.DataSourceAuth.GetAws().GetAwsRegion(); return {}; + case NKikimrSchemeOp::TAuth::kToken: + properties["authMethod"] = "TOKEN"; + properties["token"] = externalSource.Token; + properties["tokenReference"] = externalSource.DataSourceAuth.GetToken().GetTokenSecretName(); + return {}; + case NKikimrSchemeOp::TAuth::IDENTITY_NOT_SET: return {"Identity case is not specified"}; } diff --git a/ydb/core/kqp/provider/yql_kikimr_gateway.h b/ydb/core/kqp/provider/yql_kikimr_gateway.h index 02fc4879f04e..357394a8d084 100644 --- a/ydb/core/kqp/provider/yql_kikimr_gateway.h +++ b/ydb/core/kqp/provider/yql_kikimr_gateway.h @@ -386,6 +386,7 @@ struct TExternalSource { TString Password; TString AwsAccessKeyId; TString AwsSecretAccessKey; + TString Token; NKikimrSchemeOp::TAuth DataSourceAuth; NKikimrSchemeOp::TExternalDataSourceProperties Properties; }; diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp index 6bcff731f56d..0d4943dca3ed 100644 --- a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp +++ b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp @@ -438,6 +438,29 @@ void FillOlapProgram(const T& node, const NKikimr::NMiniKQL::TType* miniKqlResul CompileOlapProgram(node.Process(), tableMeta, readProto, resultColNames, ctx); } +THashMap FindSecureParams(const TExprNode::TPtr& node, const TTypeAnnotationContext& typesCtx, TSet& SecretNames) { + THashMap secureParams; + NYql::NCommon::FillSecureParams(node, typesCtx, secureParams); + + for (auto& [secretName, structuredToken] : secureParams) { + const auto& tokenParser = CreateStructuredTokenParser(structuredToken); + tokenParser.ListReferences(SecretNames); + structuredToken = tokenParser.ToBuilder().RemoveSecrets().ToJson(); + } + + return secureParams; +} + +std::optional> FindOneSecureParam(const TExprNode::TPtr& node, const TTypeAnnotationContext& typesCtx, const TString& nodeName, TSet& SecretNames) { + const auto& secureParams = FindSecureParams(node, typesCtx, SecretNames); + if (secureParams.empty()) { + return std::nullopt; + } + + YQL_ENSURE(secureParams.size() == 1, "Only one SecureParams per " << nodeName << " allowed"); + return *secureParams.begin(); +} + class TKqpQueryCompiler : public IKqpQueryCompiler { public: TKqpQueryCompiler(const TString& cluster, const TIntrusivePtr tablesData, @@ -707,6 +730,9 @@ class TKqpQueryCompiler : public IKqpQueryCompiler { return true; }); + const auto& secureParams = FindSecureParams(stage.Program().Ptr(), TypesCtx, SecretNames); + stageProto.MutableSecureParams()->insert(secureParams.begin(), secureParams.end()); + auto result = stage.Program().Body(); auto resultType = result.Ref().GetTypeAnn(); ui32 outputsCount = 0; @@ -976,15 +1002,9 @@ class TKqpQueryCompiler : public IKqpQueryCompiler { externalSource.AddPartitionedTaskParams(partitionParam); } - THashMap secureParams; - NYql::NCommon::FillSecureParams(source.Ptr(), TypesCtx, secureParams); - if (!secureParams.empty()) { - YQL_ENSURE(secureParams.size() == 1, "Only one SecureParams per source allowed"); - auto it = secureParams.begin(); - externalSource.SetSourceName(it->first); - auto token = it->second; - externalSource.SetAuthInfo(CreateStructuredTokenParser(token).ToBuilder().RemoveSecrets().ToJson()); - CreateStructuredTokenParser(token).ListReferences(SecretNames); + if (const auto& secureParams = FindOneSecureParam(source.Ptr(), TypesCtx, "source", SecretNames)) { + externalSource.SetSourceName(secureParams->first); + externalSource.SetAuthInfo(secureParams->second); } google::protobuf::Any& settings = *externalSource.MutableSettings(); @@ -1062,15 +1082,9 @@ class TKqpQueryCompiler : public IKqpQueryCompiler { YQL_ENSURE(!settings.type_url().empty(), "Data sink provider \"" << dataSinkCategory << "\" did't fill dq sink settings for its dq sink node"); YQL_ENSURE(sinkType, "Data sink provider \"" << dataSinkCategory << "\" did't fill dq sink settings type for its dq sink node"); - THashMap secureParams; - NYql::NCommon::FillSecureParams(sink.Ptr(), TypesCtx, secureParams); - if (!secureParams.empty()) { - YQL_ENSURE(secureParams.size() == 1, "Only one SecureParams per sink allowed"); - auto it = secureParams.begin(); - externalSink.SetSinkName(it->first); - auto token = it->second; - externalSink.SetAuthInfo(CreateStructuredTokenParser(token).ToBuilder().RemoveSecrets().ToJson()); - CreateStructuredTokenParser(token).ListReferences(SecretNames); + if (const auto& secureParams = FindOneSecureParam(sink.Ptr(), TypesCtx, "sink", SecretNames)) { + externalSink.SetSinkName(secureParams->first); + externalSink.SetAuthInfo(secureParams->second); } } } diff --git a/ydb/core/protos/flat_scheme_op.proto b/ydb/core/protos/flat_scheme_op.proto index ee3f4460437a..2aee4a837c40 100644 --- a/ydb/core/protos/flat_scheme_op.proto +++ b/ydb/core/protos/flat_scheme_op.proto @@ -1863,6 +1863,10 @@ message TBasic { optional string PasswordSecretName = 2; } +message TToken { + optional string TokenSecretName = 2; +} + message TAuth { oneof identity { TNoneAuth None = 3; @@ -1870,6 +1874,7 @@ message TAuth { TBasic Basic = 5; TMdbBasic MdbBasic = 6; TAws Aws = 7; + TToken Token = 8; } } diff --git a/ydb/core/protos/kqp_physical.proto b/ydb/core/protos/kqp_physical.proto index 98351bfe1f1e..bc84fc8926cf 100644 --- a/ydb/core/protos/kqp_physical.proto +++ b/ydb/core/protos/kqp_physical.proto @@ -373,6 +373,7 @@ message TKqpPhyStage { repeated TKqpSource Sources = 9; bool IsSinglePartition = 10; repeated TKqpSink Sinks = 11; + map SecureParams = 12; } message TKqpPhyResult { diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_common_external_data_source.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_common_external_data_source.cpp index b9ce6755aad0..428d0feb92a4 100644 --- a/ydb/core/tx/schemeshard/schemeshard__operation_common_external_data_source.cpp +++ b/ydb/core/tx/schemeshard/schemeshard__operation_common_external_data_source.cpp @@ -74,6 +74,8 @@ bool ValidateAuth(const NKikimrSchemeOp::TAuth& auth, return CheckAuth("BASIC", availableAuthMethods, errStr); case NKikimrSchemeOp::TAuth::kAws: return CheckAuth("AWS", availableAuthMethods, errStr); + case NKikimrSchemeOp::TAuth::kToken: + return CheckAuth("TOKEN", availableAuthMethods, errStr); case NKikimrSchemeOp::TAuth::kNone: return CheckAuth("NONE", availableAuthMethods, errStr); } diff --git a/ydb/library/yql/providers/common/structured_token/yql_token_builder.cpp b/ydb/library/yql/providers/common/structured_token/yql_token_builder.cpp index 24a1bdd6cab9..8ecae8a4b5a4 100644 --- a/ydb/library/yql/providers/common/structured_token/yql_token_builder.cpp +++ b/ydb/library/yql/providers/common/structured_token/yql_token_builder.cpp @@ -35,6 +35,12 @@ TStructuredTokenBuilder& TStructuredTokenBuilder::SetBasicAuthWithSecret(const T return *this; } +TStructuredTokenBuilder& TStructuredTokenBuilder::SetTokenAuthWithSecret(const TString& tokenReference, const TString& token) { + Data.SetField("token_ref", tokenReference); + Data.SetField("token", token); + return *this; +} + TStructuredTokenBuilder& TStructuredTokenBuilder::SetIAMToken(const TString& token) { Data.SetField("token", token); return *this; @@ -56,12 +62,18 @@ TStructuredTokenBuilder& TStructuredTokenBuilder::ReplaceReferences(const std::m Data.ClearField("sa_id_signature_ref"); Data.SetField("sa_id_signature", secrets.at(reference)); } + if (Data.HasField("token_ref")) { + auto reference = Data.GetField("token_ref"); + Data.ClearField("token_ref"); + Data.SetField("token", secrets.at(reference)); + } return *this; } TStructuredTokenBuilder& TStructuredTokenBuilder::RemoveSecrets() { Data.ClearField("basic_password"); Data.ClearField("sa_id_signature"); + Data.ClearField("token"); return *this; } @@ -127,6 +139,9 @@ void TStructuredTokenParser::ListReferences(TSet& references) const { if (Data.HasField("sa_id_signature_ref")) { references.insert(Data.GetField("sa_id_signature_ref")); } + if (Data.HasField("token_ref")) { + references.insert(Data.GetField("token_ref")); + } } TStructuredTokenBuilder TStructuredTokenParser::ToBuilder() const { @@ -178,4 +193,16 @@ TString ComposeStructuredTokenJsonForBasicAuthWithSecret(const TString& login, c return result.ToJson(); } +TString ComposeStructuredTokenJsonForTokenAuthWithSecret(const TString& tokenSecretName, const TString& token) { + TStructuredTokenBuilder result; + + if (tokenSecretName && token) { + result.SetTokenAuthWithSecret(tokenSecretName, token); + return result.ToJson(); + } + + result.SetNoAuth(); + return result.ToJson(); +} + } diff --git a/ydb/library/yql/providers/common/structured_token/yql_token_builder.h b/ydb/library/yql/providers/common/structured_token/yql_token_builder.h index dd4481ea5590..2982a88b3914 100644 --- a/ydb/library/yql/providers/common/structured_token/yql_token_builder.h +++ b/ydb/library/yql/providers/common/structured_token/yql_token_builder.h @@ -16,6 +16,7 @@ class TStructuredTokenBuilder { TStructuredTokenBuilder& SetServiceAccountIdAuthWithSecret(const TString& accountId, const TString& accountIdSignatureReference, const TString& accountIdSignature); TStructuredTokenBuilder& SetBasicAuth(const TString& login, const TString& password); TStructuredTokenBuilder& SetBasicAuthWithSecret(const TString& login, const TString& passwordReference); + TStructuredTokenBuilder& SetTokenAuthWithSecret(const TString& tokenReference, const TString& token); TStructuredTokenBuilder& SetIAMToken(const TString& token); TStructuredTokenBuilder& SetNoAuth(); TStructuredTokenBuilder& ReplaceReferences(const std::map& secrets); @@ -51,4 +52,5 @@ TStructuredTokenParser CreateStructuredTokenParser(const TString& content); TString ComposeStructuredTokenJsonForServiceAccount(const TString& serviceAccountId, const TString& serviceAccountIdSignature, const TString& token); TString ComposeStructuredTokenJsonForServiceAccountWithSecret(const TString& serviceAccountId, const TString& serviceAccountIdSignatureSecretName, const TString& serviceAccountIdSignature); TString ComposeStructuredTokenJsonForBasicAuthWithSecret(const TString& login, const TString& passwordSecretName, const TString& password); +TString ComposeStructuredTokenJsonForTokenAuthWithSecret(const TString& tokenSecretName, const TString& token); } diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasource.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasource.cpp index cf6bfab84d96..0b96de14b363 100644 --- a/ydb/library/yql/providers/yt/provider/yql_yt_datasource.cpp +++ b/ydb/library/yql/providers/yt/provider/yql_yt_datasource.cpp @@ -6,6 +6,7 @@ #include "yql_yt_dq_integration.h" #include "yql_yt_dq_optimize.h" +#include #include #include #include @@ -95,11 +96,15 @@ class TYtDataSource : public TDataProviderBase { } void AddCluster(const TString& name, const THashMap& properties) override { + const TString& token = properties.Value("token", ""); + State_->Configuration->AddValidCluster(name); + State_->Configuration->Tokens[name] = ComposeStructuredTokenJsonForTokenAuthWithSecret(properties.Value("tokenReference", ""), token); TYtClusterConfig cluster; cluster.SetName(name); cluster.SetCluster(properties.Value("location", "")); + cluster.SetYTToken(token); State_->Gateway->AddCluster(cluster); } diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp index 547571031b32..49dbfed61b4e 100644 --- a/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp +++ b/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp @@ -506,8 +506,8 @@ class TYtDqIntegration: public TDqIntegrationBase { TExprNode::TPtr WrapRead(const TDqSettings&, const TExprNode::TPtr& read, TExprContext& ctx) override { if (auto maybeYtReadTable = TMaybeNode(read)) { TMaybeNode secParams; - if (State_->Configuration->Auth.Get().GetOrElse(TString())) { - const auto cluster = maybeYtReadTable.Cast().DataSource().Cluster(); + const auto cluster = maybeYtReadTable.Cast().DataSource().Cluster(); + if (State_->Configuration->Auth.Get().GetOrElse(TString()) || State_->Configuration->Tokens.contains(cluster)) { secParams = Build(ctx, read->Pos()).Name().Build(TString("cluster:default_").append(cluster)).Done(); } return Build(ctx, read->Pos()) diff --git a/ydb/library/yql/sql/v1/sql_translation.cpp b/ydb/library/yql/sql/v1/sql_translation.cpp index d67308e8f0ef..a226c8d743b9 100644 --- a/ydb/library/yql/sql/v1/sql_translation.cpp +++ b/ydb/library/yql/sql/v1/sql_translation.cpp @@ -4453,14 +4453,16 @@ bool TSqlTranslation::ValidateAuthMethod(const std::map& "password_secret_name", "aws_access_key_id_secret_name", "aws_secret_access_key_secret_name", - "aws_region" + "aws_region", + "token_secret_name" }; const static TMap> authMethodFields{ {"NONE", {}}, {"SERVICE_ACCOUNT", {"service_account_id", "service_account_secret_name"}}, {"BASIC", {"login", "password_secret_name"}}, {"AWS", {"aws_access_key_id_secret_name", "aws_secret_access_key_secret_name", "aws_region"}}, - {"MDB_BASIC", {"service_account_id", "service_account_secret_name", "login", "password_secret_name"}} + {"MDB_BASIC", {"service_account_id", "service_account_secret_name", "login", "password_secret_name"}}, + {"TOKEN", {"token_secret_name"}} }; auto authMethodIt = result.find("auth_method"); if (authMethodIt == result.end() || authMethodIt->second.GetLiteral() == nullptr) { diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index d2b620d1dae6..15ec7d3a134a 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -21,7 +21,7 @@ struct TExecutionOptions { bool ClearExecution = false; NKikimrKqp::EQueryAction ScriptQueryAction = NKikimrKqp::QUERY_ACTION_EXECUTE; - TString ScriptTraceId = "kqprun"; + TString TraceId = "kqprun"; bool HasResults() const { return ScriptQuery && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE; @@ -37,7 +37,7 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner if (executionOptions.SchemeQuery) { Cout << colors.Yellow() << "Executing scheme query..." << colors.Default() << Endl; - if (!runner.ExecuteSchemeQuery(executionOptions.SchemeQuery)) { + if (!runner.ExecuteSchemeQuery(executionOptions.SchemeQuery, executionOptions.TraceId)) { ythrow yexception() << "Scheme query execution failed"; } } @@ -45,14 +45,14 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner if (executionOptions.ScriptQuery) { Cout << colors.Yellow() << "Executing script..." << colors.Default() << Endl; if (!executionOptions.ClearExecution) { - if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) { + if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << "Script execution failed"; } if (!runner.FetchScriptResults()) { ythrow yexception() << "Fetch script results failed"; } } else { - if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) { + if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << "Query execution failed"; } } @@ -76,6 +76,15 @@ THolder SetupDefaultFileOutput(const TString& filePath, IOutputStre } +void ReplaceTemplate(const TString& variableName, const TString& variableValue, TString& query) { + TString variableTemplate = TStringBuilder() << "${" << variableName << "}"; + for (size_t position = query.find(variableTemplate); position != TString::npos; position = query.find(variableTemplate, position)) { + query.replace(position, variableTemplate.size(), variableValue); + position += variableValue.size(); + } +} + + TIntrusivePtr CreateFunctionRegistry(const TString& udfsDirectory, TVector udfsPaths) { if (!udfsDirectory.empty() || !udfsPaths.empty()) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); @@ -190,6 +199,10 @@ void RunMain(int argc, const char* argv[]) { NLastGetopt::TOptsParseResult parsedOptions(&options, argc, argv); + // Environment variables + + const TString& yqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE); + // Execution options if (!schemeQueryFile && !scriptQueryFile) { @@ -197,6 +210,7 @@ void RunMain(int argc, const char* argv[]) { } if (schemeQueryFile) { executionOptions.SchemeQuery = TFileInput(schemeQueryFile).ReadAll(); + ReplaceTemplate(NKqpRun::YQL_TOKEN_VARIABLE, yqlToken, executionOptions.SchemeQuery); } if (scriptQueryFile) { executionOptions.ScriptQuery = TFileInput(scriptQueryFile).ReadAll(); @@ -240,7 +254,7 @@ void RunMain(int argc, const char* argv[]) { std::remove(logFile.c_str()); } - runnerOptions.YdbSettings.YqlToken = GetEnv("YQL_TOKEN"); + runnerOptions.YdbSettings.YqlToken = yqlToken; runnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(udfsDirectory, udfsPaths).Get(); TString appConfigData = TFileInput(appConfigFile).ReadAll(); diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 1d57272b2dd9..e7f066cc2b16 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -9,6 +9,8 @@ namespace NKqpRun { +constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; + struct TYdbSetupSettings { TString DomainName = "Root"; diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index b4eacc1895fe..9deab0dc4209 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -20,11 +20,11 @@ class TKqpRunner::TImpl { , CoutColors_(NColorizer::AutoColors(Cout)) {} - bool ExecuteSchemeQuery(const TString& query) const { + bool ExecuteSchemeQuery(const TString& query, const TString& traceId) const { StartSchemeTraceOpt(); TSchemeMeta meta; - TRequestResult status = YdbSetup_.SchemeQueryRequest(query, meta); + TRequestResult status = YdbSetup_.SchemeQueryRequest(query, traceId, meta); TYdbSetup::StopTraceOpt(); PrintSchemeQueryAst(meta.Ast); @@ -199,8 +199,8 @@ TKqpRunner::TKqpRunner(const TRunnerOptions& options) : Impl_(new TImpl(options)) {} -bool TKqpRunner::ExecuteSchemeQuery(const TString& query) const { - return Impl_->ExecuteSchemeQuery(query); +bool TKqpRunner::ExecuteSchemeQuery(const TString& query, const TString& traceId) const { + return Impl_->ExecuteSchemeQuery(query, traceId); } bool TKqpRunner::ExecuteScript(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId) const { diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.h b/ydb/tests/tools/kqprun/src/kqp_runner.h index 273cca024d87..3171b6afa8bc 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.h +++ b/ydb/tests/tools/kqprun/src/kqp_runner.h @@ -10,7 +10,7 @@ class TKqpRunner { public: explicit TKqpRunner(const TRunnerOptions& options); - bool ExecuteSchemeQuery(const TString& query) const; + bool ExecuteSchemeQuery(const TString& query, const TString& traceId) const; bool ExecuteScript(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId) const; diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 031522a744f0..835ed9d5fa2e 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -170,9 +170,9 @@ class TYdbSetup::TImpl { InitializeServer(); } - NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr SchemeQueryRequest(const TString& query) const { + NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr SchemeQueryRequest(const TString& query, const TString& traceId) const { auto event = MakeHolder(); - FillSchemeRequest(query, *event->Record.MutableRequest()); + FillSchemeRequest(query, traceId, event->Record); return RunKqpProxyRequest(std::move(event)); } @@ -243,30 +243,30 @@ class TYdbSetup::TImpl { } private: - void FillSchemeRequest(const TString& query, NKikimrKqp::TQueryRequest& request) const { - request.SetType(NKikimrKqp::QUERY_TYPE_SQL_DDL); - request.SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE); - request.SetCollectStats(Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL); + void FillQueryRequest(const TString& query, NKikimrKqp::EQueryType type, NKikimrKqp::EQueryAction action, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { + event.SetTraceId(traceId); + event.SetUserToken(NACLib::TUserToken(Settings_.YqlToken, BUILTIN_ACL_ROOT, {}).SerializeAsString()); + + auto request = event.MutableRequest(); + request->SetQuery(query); + request->SetType(type); + request->SetAction(action); + request->SetCollectStats(Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL); + request->SetDatabase(Settings_.DomainName); + } - request.SetDatabase(Settings_.DomainName); - request.SetQuery(query); + void FillSchemeRequest(const TString& query, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { + FillQueryRequest(query, NKikimrKqp::QUERY_TYPE_SQL_DDL, NKikimrKqp::QUERY_ACTION_EXECUTE, traceId, event); } void FillScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { - event.SetTraceId(traceId); - + FillQueryRequest(script, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT, action, traceId, event); + auto request = event.MutableRequest(); if (action == NKikimrKqp::QUERY_ACTION_EXECUTE) { request->MutableTxControl()->mutable_begin_tx()->mutable_serializable_read_write(); request->MutableTxControl()->set_commit_tx(true); } - - request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT); - request->SetAction(action); - request->SetCollectStats(Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL); - - request->SetDatabase(Settings_.DomainName); - request->SetQuery(script); } private: @@ -304,8 +304,8 @@ TYdbSetup::TYdbSetup(const TYdbSetupSettings& settings) : Impl_(new TImpl(settings)) {} -TRequestResult TYdbSetup::SchemeQueryRequest(const TString& query, TSchemeMeta& meta) const { - auto schemeQueryOperationResponse = Impl_->SchemeQueryRequest(query)->Get()->Record.GetRef(); +TRequestResult TYdbSetup::SchemeQueryRequest(const TString& query, const TString& traceId, TSchemeMeta& meta) const { + auto schemeQueryOperationResponse = Impl_->SchemeQueryRequest(query, traceId)->Get()->Record.GetRef(); meta.Ast = schemeQueryOperationResponse.GetResponse().GetQueryAst(); diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h index 68a00058a3fd..b4fa05c497e5 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.h +++ b/ydb/tests/tools/kqprun/src/ydb_setup.h @@ -47,7 +47,7 @@ class TYdbSetup { public: explicit TYdbSetup(const TYdbSetupSettings& settings); - TRequestResult SchemeQueryRequest(const TString& query, TSchemeMeta& meta) const; + TRequestResult SchemeQueryRequest(const TString& query, const TString& traceId, TSchemeMeta& meta) const; TRequestResult ScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, TString& operation) const;