diff --git a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp index 5aebbc47a632..67fbc7105da8 100644 --- a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp +++ b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp @@ -753,7 +753,7 @@ NThreading::TFuture TKqpTableMetadataLoader::LoadTableMeta switch (entry.Kind) { case EKind::KindExternalDataSource: { auto externalDataSourceMetadata = GetLoadTableMetadataResult(entry, cluster, mainCluster, table); - if (!externalDataSourceMetadata.Success()) { + if (!externalDataSourceMetadata.Success() || !settings.RequestAuthInfo_) { promise.SetValue(externalDataSourceMetadata); return; } diff --git a/ydb/core/kqp/provider/yql_kikimr_datasink.cpp b/ydb/core/kqp/provider/yql_kikimr_datasink.cpp index 5c17b5891341..69544ba145b2 100644 --- a/ydb/core/kqp/provider/yql_kikimr_datasink.cpp +++ b/ydb/core/kqp/provider/yql_kikimr_datasink.cpp @@ -181,7 +181,7 @@ class TKiSinkIntentDeterminationTransformer: public TKiSinkVisitorTransformer { auto tableType = settings.TableType.IsValid() ? GetTableTypeFromString(settings.TableType.Cast()) : ETableType::Table; // v0, pg support - ctx->Tables().GetOrAddTable(TString(cluster), ctx->GetDatabase(), key.GetTablePath(), tableType); + ctx->Tables().GetOrAddTable(TString(cluster), ctx->GetDatabase(), key.GetTablePath(), tableType).DisableAuthInfo(); } TStatus HandleWrite(TExprBase node, TExprContext& ctx) override { @@ -665,18 +665,28 @@ class TKikimrDataSink : public TDataProviderBase return true; } - if (tableDesc.Metadata->ExternalSource.SourceType == ESourceType::ExternalDataSource && tableDesc.Metadata->TableType == NYql::ETableType::Unknown) { - ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Attempt to write to external data source \"" << key.GetTablePath() << "\" without table. Please specify table to write to")); - return false; - } - if (tableDesc.Metadata->ExternalSource.SourceType != ESourceType::ExternalDataSource && tableDesc.Metadata->ExternalSource.SourceType != ESourceType::ExternalTable) { YQL_CVLOG(NLog::ELevel::ERROR, NLog::EComponent::ProviderKikimr) << "Skip RewriteIO for external entity: unknown entity type: " << (int)tableDesc.Metadata->ExternalSource.SourceType; return true; } if (mode != "insert_abort") { - ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Write mode '" << static_cast(mode) << "' is not supported for external entities")); + if (mode == "drop" || mode == "drop_if_exists") { + TString dropHint; + if (tableDesc.Metadata->ExternalSource.SourceType == ESourceType::ExternalDataSource) { + dropHint = "DROP EXTERNAL DATA SOURCE"; + } else if (tableDesc.Metadata->ExternalSource.SourceType == ESourceType::ExternalTable) { + dropHint = "DROP EXTERNAL TABLE"; + } + ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Cannot drop external entity by using DROP TABLE" << (dropHint ? ". Please use " : "") << dropHint)); + } else { + ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Write mode '" << static_cast(mode) << "' is not supported for external entities")); + } + return false; + } + + if (tableDesc.Metadata->ExternalSource.SourceType == ESourceType::ExternalDataSource && tableDesc.Metadata->TableType == NYql::ETableType::Unknown) { + ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Attempt to write to external data source \"" << key.GetTablePath() << "\" without table. Please specify table to write to")); return false; } diff --git a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp index af95adf40d5a..3ba7d447d957 100644 --- a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp +++ b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp @@ -238,6 +238,7 @@ class TKiSourceLoadTableMetadataTransformer : public TGraphTransformerBase { .WithTableStats(table.GetNeedsStats()) .WithPrivateTables(IsInternalCall) .WithExternalDatasources(SessionCtx->Config().FeatureFlags.GetEnableExternalDataSources()) + .WithAuthInfo(table.GetNeedAuthInfo()) ); futures.push_back(future.Apply([result, queryType] diff --git a/ydb/core/kqp/provider/yql_kikimr_gateway.h b/ydb/core/kqp/provider/yql_kikimr_gateway.h index 3c9b541eee70..30b55a1058bf 100644 --- a/ydb/core/kqp/provider/yql_kikimr_gateway.h +++ b/ydb/core/kqp/provider/yql_kikimr_gateway.h @@ -765,9 +765,15 @@ class IKikimrGateway : public TThrRefBase { return *this; } + TLoadTableMetadataSettings& WithAuthInfo(bool enable) { + RequestAuthInfo_ = enable; + return *this; + } + bool RequestStats_ = false; bool WithPrivateTables_ = false; bool WithExternalDatasources_ = false; + bool RequestAuthInfo_ = true; }; class IKqpTableMetadataLoader : public std::enable_shared_from_this { diff --git a/ydb/core/kqp/provider/yql_kikimr_provider.h b/ydb/core/kqp/provider/yql_kikimr_provider.h index abcabe0c416e..af1f0d07af2f 100644 --- a/ydb/core/kqp/provider/yql_kikimr_provider.h +++ b/ydb/core/kqp/provider/yql_kikimr_provider.h @@ -167,12 +167,15 @@ class TKikimrTableDescription { void RequireStats() { NeedsStats = true; } bool GetNeedsStats() const { return NeedsStats; } + void DisableAuthInfo() { NeedAuthInfo = false; } + bool GetNeedAuthInfo() const { return NeedAuthInfo; } ETableType GetTableType() const { return TableType; } void SetTableType(ETableType tableType) { TableType = tableType; } private: THashMap ColumnTypes; bool NeedsStats = false; + bool NeedAuthInfo = true; ETableType TableType; }; diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_scheme_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_scheme_ut.cpp index 45506581cff1..cfe4dcb4c9b6 100644 --- a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_scheme_ut.cpp +++ b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_scheme_ut.cpp @@ -145,6 +145,77 @@ Y_UNIT_TEST_SUITE(KqpFederatedSchemeTest) { checkTableExists(false, 1); checkDrop(true, EEx::IfExists, 1); } + + void TestInvalidDropForExternalTableWithAuth(std::function(const TString&)> queryExecuter, TString tableSuffix) { + const TString externalDataSourceName = "test_data_source_" + tableSuffix; + const TString externalTableName = "test_table_" + tableSuffix; + + // Create external table + { + const TString sql = TStringBuilder() << R"( + UPSERT OBJECT mysasignature (TYPE SECRET) WITH (value = "mysasignaturevalue"); + CREATE EXTERNAL DATA SOURCE `)" << externalDataSourceName << R"(` WITH ( + SOURCE_TYPE="ObjectStorage", + LOCATION="my-bucket", + AUTH_METHOD="SERVICE_ACCOUNT", + SERVICE_ACCOUNT_ID="mysa", + SERVICE_ACCOUNT_SECRET_NAME="mysasignature" + ); + CREATE EXTERNAL TABLE `)" << externalTableName << R"(` ( + Key Uint64 + ) WITH ( + DATA_SOURCE=")" << externalDataSourceName << R"(", + LOCATION="/", + FORMAT="json_each_row" + );)"; + const auto& [success, issues] = queryExecuter(sql); + UNIT_ASSERT_C(success, issues); + } + + // Drop secret object + { + const TString sql = "DROP OBJECT mysasignature (TYPE SECRET)"; + const auto& [success, issues] = queryExecuter(sql); + UNIT_ASSERT_C(success, issues); + } + + // Drop external table + { + const TString sql = TStringBuilder() << "DROP TABLE `" << externalTableName << "`"; + const auto& [success, issues] = queryExecuter(sql); + UNIT_ASSERT(!success); + UNIT_ASSERT_STRING_CONTAINS(issues, "Cannot drop external entity by using DROP TABLE. Please use DROP EXTERNAL TABLE"); + } + + // Drop external data source + { + const TString sql = TStringBuilder() << "DROP TABLE `" << externalDataSourceName << "`"; + const auto& [success, issues] = queryExecuter(sql); + UNIT_ASSERT(!success); + UNIT_ASSERT_STRING_CONTAINS(issues, "Cannot drop external entity by using DROP TABLE. Please use DROP EXTERNAL DATA SOURCE"); + } + } + + Y_UNIT_TEST(InvalidDropForExternalTableWithAuth) { + auto kikimr = MakeKikimrRunner(NYql::IHTTPGateway::Make()); + + auto driver = kikimr->GetDriver(); + NScripting::TScriptingClient yqlScriptClient(driver); + auto yqlScriptClientExecutor = [&](const TString& sql) { + Cerr << "Execute sql by yql script client:\n" << sql << Endl; + auto result = yqlScriptClient.ExecuteYqlScript(sql).GetValueSync(); + return std::make_pair(result.IsSuccess(), result.GetIssues().ToString()); + }; + TestInvalidDropForExternalTableWithAuth(yqlScriptClientExecutor, "yql_script"); + + auto queryClient = kikimr->GetQueryClient(); + auto queryClientExecutor = [&](const TString& sql) { + Cerr << "Execute sql by query client:\n" << sql << Endl; + auto result = queryClient.ExecuteQuery(sql, TTxControl::NoTx()).GetValueSync(); + return std::make_pair(result.IsSuccess(), result.GetIssues().ToString()); + }; + TestInvalidDropForExternalTableWithAuth(queryClientExecutor, "generic_query"); + } } } // namespace NKikimr::NKqp diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 52f476ff220d..0b3f6677d846 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -11,18 +11,25 @@ #include #include +#include #include struct TExecutionOptions { + enum class EClearExecutionCase { + Disabled, + GenericQuery, + YqlScript + }; + TString ScriptQuery; TString SchemeQuery; - bool ClearExecution = false; bool ForgetExecution = false; + EClearExecutionCase ClearExecution = EClearExecutionCase::Disabled; NKikimrKqp::EQueryAction ScriptQueryAction = NKikimrKqp::QUERY_ACTION_EXECUTE; - TString ScriptTraceId = "kqprun"; + TString TraceId = "kqprun"; bool HasResults() const { return ScriptQuery && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE; @@ -38,15 +45,16 @@ 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"; } } if (executionOptions.ScriptQuery) { Cout << colors.Yellow() << "Executing script..." << colors.Default() << Endl; - if (!executionOptions.ClearExecution) { - if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) { + switch (executionOptions.ClearExecution) { + case TExecutionOptions::EClearExecutionCase::Disabled: + if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << "Script execution failed"; } Cout << colors.Yellow() << "Fetching script results..." << colors.Default() << Endl; @@ -59,10 +67,19 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner ythrow yexception() << "Forget script execution operation failed"; } } - } else { - if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) { + break; + + case TExecutionOptions::EClearExecutionCase::GenericQuery: + if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << "Query execution failed"; } + break; + + case TExecutionOptions::EClearExecutionCase::YqlScript: + if (!runner.ExecuteYqlScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { + ythrow yexception() << "Yql script execution failed"; + } + break; } } @@ -86,7 +103,16 @@ THolder SetupDefaultFileOutput(const TString& filePath, IOutputStre } -TIntrusivePtr CreateFunctionRegistry(const TString& udfsDirectory, TVector udfsPaths) { +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, bool excludeLinkedUdfs) { if (!udfsDirectory.empty() || !udfsPaths.empty()) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); Cout << colors.Yellow() << "Fetching udfs..." << colors.Default() << Endl; @@ -94,7 +120,17 @@ TIntrusivePtr CreateFunctionRegistr NKikimr::NMiniKQL::FindUdfsInDir(udfsDirectory, &udfsPaths); auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&NYql::NBacktrace::KikimrBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, udfsPaths)->Clone(); - NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); + + if (excludeLinkedUdfs) { + for (const auto& wrapper : NYql::NUdf::GetStaticUdfModuleWrapperList()) { + auto [name, ptr] = wrapper(); + if (!functionRegistry->IsLoadedUdfModule(name)) { + functionRegistry->AddModule(TString(NKikimr::NMiniKQL::StaticModulePrefix) + name, name, std::move(ptr)); + } + } + } else { + NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); + } return functionRegistry; } @@ -113,6 +149,7 @@ void RunMain(int argc, const char* argv[]) { TString logFile = "-"; TString appConfigFile = "./configuration/app_config.conf"; + TString clearExecutionType = "disabled"; TString traceOptType = "disabled"; TString scriptQueryAction = "execute"; TString planOutputFormat = "pretty"; @@ -121,6 +158,7 @@ void RunMain(int argc, const char* argv[]) { TVector udfsPaths; TString udfsDirectory; + bool excludeLinkedUdfs = false; NLastGetopt::TOpts options = NLastGetopt::TOpts::Default(); options.AddLongOption('p', "script-query", "Script query to execute") @@ -158,11 +196,11 @@ void RunMain(int argc, const char* argv[]) { .RequiredArgument("FILE") .StoreResult(&scriptQueryPlanFile); - options.AddLongOption('C', "clear-execution", "Execute script query without RunScriptActor in one query request") + options.AddLongOption('C', "clear-execution", "Execute script query without creating additional tables, one of { query | yql-script }") .Optional() - .NoArgument() - .DefaultValue(executionOptions.ClearExecution) - .SetFlag(&executionOptions.ClearExecution); + .RequiredArgument("STR") + .DefaultValue(clearExecutionType) + .StoreResult(&clearExecutionType); options.AddLongOption('F', "forget", "Forget script execution operation after fetching results, cannot be used with -C") .Optional() .NoArgument() @@ -202,9 +240,18 @@ void RunMain(int argc, const char* argv[]) { .Optional() .RequiredArgument("PATH") .StoreResult(&udfsDirectory); + options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") + .Optional() + .NoArgument() + .DefaultValue(excludeLinkedUdfs) + .SetFlag(&excludeLinkedUdfs); NLastGetopt::TOptsParseResult parsedOptions(&options, argc, argv); + // Environment variables + + const TString& yqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE); + // Execution options if (!schemeQueryFile && !scriptQueryFile) { @@ -212,11 +259,18 @@ 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(); } + executionOptions.ClearExecution = + (clearExecutionType == TStringBuf("query")) ? TExecutionOptions::EClearExecutionCase::GenericQuery + : (clearExecutionType == TStringBuf("yql-script")) ? TExecutionOptions::EClearExecutionCase::YqlScript + : (clearExecutionType == TStringBuf("disabled")) ? TExecutionOptions::EClearExecutionCase::Disabled + : TExecutionOptions::EClearExecutionCase::Disabled; + executionOptions.ScriptQueryAction = (scriptQueryAction == TStringBuf("execute")) ? NKikimrKqp::QUERY_ACTION_EXECUTE : (scriptQueryAction == TStringBuf("explain")) ? NKikimrKqp::QUERY_ACTION_EXPLAIN @@ -255,8 +309,8 @@ void RunMain(int argc, const char* argv[]) { std::remove(logFile.c_str()); } - runnerOptions.YdbSettings.YqlToken = GetEnv("YQL_TOKEN"); - runnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(udfsDirectory, udfsPaths).Get(); + runnerOptions.YdbSettings.YqlToken = yqlToken; + runnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(udfsDirectory, udfsPaths, excludeLinkedUdfs).Get(); TString appConfigData = TFileInput(appConfigFile).ReadAll(); if (!google::protobuf::TextFormat::ParseFromString(appConfigData, &runnerOptions.YdbSettings.AppConfig)) { diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 1d57272b2dd9..e69331ab9aed 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"; @@ -16,7 +18,7 @@ struct TYdbSetupSettings { TMaybe LogOutputFile; TString YqlToken; - TIntrusivePtr FunctionRegistry = nullptr; + TIntrusivePtr FunctionRegistry; NKikimrConfig::TAppConfig AppConfig; }; diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index d39d0d7406e5..4bf0971935d8 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -13,6 +13,11 @@ namespace NKqpRun { class TKqpRunner::TImpl { public: + enum class EQueryType { + ScriptQuery, + YqlScriptQuery + }; + explicit TImpl(const TRunnerOptions& options) : Options_(options) , YdbSetup_(options.YdbSettings) @@ -20,11 +25,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); @@ -50,11 +55,21 @@ class TKqpRunner::TImpl { return WaitScriptExecutionOperation(); } - bool ExecuteQuery(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) { + bool ExecuteQuery(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, EQueryType queryType) { StartScriptTraceOpt(); TQueryMeta meta; - TRequestResult status = YdbSetup_.QueryRequest(query, action, traceId, meta, ResultSets_); + TRequestResult status; + switch (queryType) { + case EQueryType::ScriptQuery: + status = YdbSetup_.QueryRequest(query, action, traceId, meta, ResultSets_); + break; + + case EQueryType::YqlScriptQuery: + status = YdbSetup_.YqlScriptRequest(query, action, traceId, meta, ResultSets_); + break; + } + TYdbSetup::StopTraceOpt(); PrintScriptAst(meta.Ast); @@ -64,6 +79,10 @@ class TKqpRunner::TImpl { return false; } + if (!status.Issues.Empty()) { + Cerr << CerrColors_.Red() << "Request finished with issues:" << CerrColors_.Default() << Endl << status.Issues.ToString() << Endl; + } + PrintScriptPlan(meta.Plan); return true; @@ -72,6 +91,7 @@ class TKqpRunner::TImpl { bool FetchScriptResults() { TYdbSetup::StopTraceOpt(); + ResultSets_.clear(); ResultSets_.resize(ExecutionMeta_.ResultSetsCount); for (i32 resultSetId = 0; resultSetId < ExecutionMeta_.ResultSetsCount; ++resultSetId) { TRequestResult status = YdbSetup_.FetchScriptExecutionResultsRequest(ExecutionOperation_, resultSetId, ResultSets_[resultSetId]); @@ -212,8 +232,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 { @@ -221,7 +241,11 @@ bool TKqpRunner::ExecuteScript(const TString& script, NKikimrKqp::EQueryAction a } bool TKqpRunner::ExecuteQuery(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const { - return Impl_->ExecuteQuery(query, action, traceId); + return Impl_->ExecuteQuery(query, action, traceId, TImpl::EQueryType::ScriptQuery); +} + +bool TKqpRunner::ExecuteYqlScript(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const { + return Impl_->ExecuteQuery(query, action, traceId, TImpl::EQueryType::YqlScriptQuery); } bool TKqpRunner::FetchScriptResults() { diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.h b/ydb/tests/tools/kqprun/src/kqp_runner.h index 8585bca6a3ac..b263bbedbf55 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.h +++ b/ydb/tests/tools/kqprun/src/kqp_runner.h @@ -2,6 +2,7 @@ #include "common.h" +#include namespace NKqpRun { @@ -9,12 +10,14 @@ 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; bool ExecuteQuery(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const; + bool ExecuteYqlScript(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const; + bool FetchScriptResults(); bool ForgetExecutionOperation(); diff --git a/ydb/tests/tools/kqprun/src/ya.make b/ydb/tests/tools/kqprun/src/ya.make index a6481facb877..8b74f61204d0 100644 --- a/ydb/tests/tools/kqprun/src/ya.make +++ b/ydb/tests/tools/kqprun/src/ya.make @@ -8,9 +8,6 @@ SRCS( PEERDIR( ydb/core/testlib - - ydb/library/yql/parser/pg_wrapper - ydb/library/yql/sql/pg ) YQL_LAST_ABI_VERSION() diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index d1f5e5760239..f83c6f0e124c 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)); } @@ -196,6 +196,13 @@ class TYdbSetup::TImpl { return promise.GetFuture().GetValueSync(); } + NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr YqlScriptRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const { + auto event = MakeHolder(); + FillYqlScriptRequest(query, action, traceId, event->Record); + + return RunKqpProxyRequest(std::move(event)); + } + NKikimr::NKqp::TEvGetScriptExecutionOperationResponse::TPtr GetScriptExecutionOperationRequest(const TString& operation) const { NKikimr::NOperationId::TOperationId operationId(operation); auto event = MakeHolder(Settings_.DomainName, operationId); @@ -250,30 +257,34 @@ 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); + void FillYqlScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { + FillQueryRequest(script, NKikimrKqp::QUERY_TYPE_SQL_SCRIPT, action, traceId, event); } private: @@ -296,6 +307,12 @@ TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TI , Issues(issues) {} +TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues) + : Status(status) +{ + NYql::IssuesFromMessage(issues, Issues); +} + bool TRequestResult::IsSuccess() const { return Status == Ydb::StatusIds::SUCCESS; } @@ -311,15 +328,13 @@ 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(); + const auto& responseRecord = schemeQueryOperationResponse.GetResponse(); - meta.Ast = schemeQueryOperationResponse.GetResponse().GetQueryAst(); + meta.Ast = responseRecord.GetQueryAst(); - NYql::TIssues issues; - NYql::IssuesFromMessage(schemeQueryOperationResponse.GetResponse().GetQueryIssues(), issues); - - return TRequestResult(schemeQueryOperationResponse.GetYdbStatus(), issues); + return TRequestResult(schemeQueryOperationResponse.GetYdbStatus(), responseRecord.GetQueryIssues()); } TRequestResult TYdbSetup::ScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, TString& operation) const { @@ -331,15 +346,33 @@ TRequestResult TYdbSetup::ScriptRequest(const TString& script, NKikimrKqp::EQuer } TRequestResult TYdbSetup::QueryRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, TQueryMeta& meta, std::vector& resultSets) const { + resultSets.clear(); + auto queryOperationResponse = Impl_->QueryRequest(query, action, traceId, resultSets)->Get()->Record.GetRef(); + const auto& responseRecord = queryOperationResponse.GetResponse(); - meta.Ast = queryOperationResponse.GetResponse().GetQueryAst(); - meta.Plan = queryOperationResponse.GetResponse().GetQueryPlan(); + meta.Ast = responseRecord.GetQueryAst(); + meta.Plan = responseRecord.GetQueryPlan(); - NYql::TIssues issues; - NYql::IssuesFromMessage(queryOperationResponse.GetResponse().GetQueryIssues(), issues); + return TRequestResult(queryOperationResponse.GetYdbStatus(), responseRecord.GetQueryIssues()); +} + +TRequestResult TYdbSetup::YqlScriptRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, TQueryMeta& meta, std::vector& resultSets) const { + resultSets.clear(); + + auto yqlQueryOperationResponse = Impl_->YqlScriptRequest(query, action, traceId)->Get()->Record.GetRef(); + const auto& responseRecord = yqlQueryOperationResponse.GetResponse(); + + meta.Ast = responseRecord.GetQueryAst(); + meta.Plan = responseRecord.GetQueryPlan(); - return TRequestResult(queryOperationResponse.GetYdbStatus(), issues); + resultSets.reserve(responseRecord.results_size()); + for (const auto& result : responseRecord.results()) { + resultSets.emplace_back(); + NKikimr::NKqp::ConvertKqpQueryResultToDbResult(result, &resultSets.back()); + } + + return TRequestResult(yqlQueryOperationResponse.GetYdbStatus(), responseRecord.GetQueryIssues()); } TRequestResult TYdbSetup::GetScriptExecutionOperationRequest(const TString& operation, TExecutionMeta& meta) const { @@ -366,10 +399,7 @@ TRequestResult TYdbSetup::FetchScriptExecutionResultsRequest(const TString& oper resultSet = scriptExecutionResults.GetResultSet(); - NYql::TIssues issues; - NYql::IssuesFromMessage(scriptExecutionResults.GetIssues(), issues); - - return TRequestResult(scriptExecutionResults.GetStatus(), issues); + return TRequestResult(scriptExecutionResults.GetStatus(), scriptExecutionResults.GetIssues()); } TRequestResult TYdbSetup::ForgetScriptExecutionOperationRequest(const TString& operation) const { diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h index eeb3aa34edcc..f31a3b249891 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.h +++ b/ydb/tests/tools/kqprun/src/ydb_setup.h @@ -38,6 +38,8 @@ struct TRequestResult { TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues); + TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues); + bool IsSuccess() const; TString ToString() const; @@ -48,12 +50,14 @@ 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; TRequestResult QueryRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, TQueryMeta& meta, std::vector& resultSets) const; + TRequestResult YqlScriptRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, TQueryMeta& meta, std::vector& resultSets) const; + TRequestResult GetScriptExecutionOperationRequest(const TString& operation, TExecutionMeta& meta) const; TRequestResult FetchScriptExecutionResultsRequest(const TString& operation, i32 resultSetId, Ydb::ResultSet& resultSet) const; diff --git a/ydb/tests/tools/kqprun/ya.make b/ydb/tests/tools/kqprun/ya.make index 0aefcbe444e8..7ba9c73011e5 100644 --- a/ydb/tests/tools/kqprun/ya.make +++ b/ydb/tests/tools/kqprun/ya.make @@ -7,6 +7,9 @@ SRCS( PEERDIR( library/cpp/getopt + ydb/library/yql/parser/pg_wrapper + ydb/library/yql/sql/pg + ydb/tests/tools/kqprun/src )