ydb-platform · GrigoriyPA · Jan 29, 2025 · Jan 29, 2025
diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf
@@ -46,13 +46,15 @@ ActorSystemConfig {
 
 ColumnShardConfig {
   DisabledOnSchemeShard: false
+  WritingInFlightRequestBytesLimit: 104857600
 }
 
 FeatureFlags {
   EnableExternalDataSources: true
   EnableScriptExecutionOperations: true
   EnableExternalSourceSchemaInference: true
   EnableTempTables: true
+  EnableReplaceIfExistsForExternalEntities: true
 }
 
 KQPConfig {

diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -88,9 +88,6 @@ struct TExecutionOptions {
 
     TRequestOptions GetSchemeQueryOptions() const {
         TString sql = SchemeQuery;
-        if (UseTemplates) {
-            ReplaceYqlTokenTemplate(sql);
-        }
 
         return {
             .Query = sql,
@@ -108,7 +105,6 @@ struct TExecutionOptions {
 
         TString sql = ScriptQueries[index];
         if (UseTemplates) {
-            ReplaceYqlTokenTemplate(sql);
             SubstGlobal(sql, "${QUERY_ID}", ToString(queryId));
         }
 
@@ -271,16 +267,6 @@ struct TExecutionOptions {
             ythrow yexception() << "Cannot format storage without real PDisks, please use --storage-path";
         }
     }
-
-private:
-    static void ReplaceYqlTokenTemplate(TString& sql) {
-        const TString variableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}";
-        if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) {
-            SubstGlobal(sql, variableName, yqlToken);
-        } else if (sql.Contains(variableName)) {
-            ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable\n";
-        }
-    }
 };
 
 
@@ -446,6 +432,7 @@ class TMain : public TMainClassArgs {
     TExecutionOptions ExecutionOptions;
     TRunnerOptions RunnerOptions;
 
+    std::unordered_map<TString, TString> Templates;
     THashMap<TString, TString> TablesMapping;
     TVector<TString> UdfsPaths;
     TString UdfsDirectory;
@@ -537,6 +524,31 @@ class TMain : public TMainClassArgs {
             .NoArgument()
             .SetFlag(&ExecutionOptions.UseTemplates);
 
+        options.AddLongOption("var-template", "Add template from environment variables or file for -s and -p queries (use variable@file for files)")
+            .RequiredArgument("variable")
+            .Handler1([this](const NLastGetopt::TOptsParser* option) {
+                TStringBuf variable;
+                TStringBuf filePath;
+                TStringBuf(option->CurVal()).Split('@', variable, filePath);
+                if (variable.empty()) {
+                    ythrow yexception() << "Variable name should not be empty";
+                }
+
+                TString value;
+                if (!filePath.empty()) {
+                    value = LoadFile(TString(filePath));
+                } else {
+                    value = GetEnv(TString(variable));
+                    if (!value) {
+                        ythrow yexception() << "Invalid env template, can not find value for variable '" << variable << "'";
+                    }
+                }
+
+                if (!Templates.emplace(variable, value).second) {
+                    ythrow yexception() << "Got duplicated template variable name '" << variable << "'";
+                }
+            });
+
         options.AddLongOption('t', "table", "File with input table (can be used by YT with -E flag), table@file")
             .RequiredArgument("table@file")
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
@@ -845,6 +857,11 @@ class TMain : public TMainClassArgs {
             .NoArgument()
             .SetFlag(&EmulateYt);
 
+        options.AddLongOption('H', "health-check", "Level of health check before start (max level 2)")
+            .RequiredArgument("uint")
+            .DefaultValue(1)
+            .StoreResult(&RunnerOptions.YdbSettings.HealthCheckLevel);
+
         options.AddLongOption("domain", "Test cluster domain name")
             .RequiredArgument("name")
             .DefaultValue(RunnerOptions.YdbSettings.DomainName)
@@ -862,9 +879,8 @@ class TMain : public TMainClassArgs {
             .RequiredArgument("path")
             .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants);
 
-        options.AddLongOption("storage-size", "Domain storage size in gigabytes")
+        options.AddLongOption("storage-size", "Domain storage size in gigabytes (32 GiB by default)")
             .RequiredArgument("uint")
-            .DefaultValue(32)
             .StoreMappedResultT<ui32>(&RunnerOptions.YdbSettings.DiskSize, [](ui32 diskSize) {
                 return static_cast<ui64>(diskSize) << 30;
             });
@@ -898,6 +914,11 @@ class TMain : public TMainClassArgs {
     int DoRun(NLastGetopt::TOptsParseResult&&) override {
         ExecutionOptions.Validate(RunnerOptions);
 
+        ReplaceTemplates(ExecutionOptions.SchemeQuery);
+        for (auto& sql : ExecutionOptions.ScriptQueries) {
+            ReplaceTemplates(sql);
+        }
+
         RunnerOptions.YdbSettings.YqlToken = YqlToken;
         RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get();
 
@@ -943,6 +964,21 @@ class TMain : public TMainClassArgs {
 
         return 0;
     }
+
+private:
+    void ReplaceTemplates(TString& sql) const {
+        for (const auto& [variable, value] : Templates) {
+            SubstGlobal(sql, TStringBuilder() << "${" << variable <<"}", value);
+        }
+        if (ExecutionOptions.UseTemplates) {
+            const TString tokenVariableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}";
+            if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) {
+                SubstGlobal(sql, tokenVariableName, yqlToken);
+            } else if (sql.Contains(tokenVariableName)) {
+                ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable";
+            }
+        }
+    }
 };
 
 
@@ -972,7 +1008,7 @@ void FloatingPointExceptionHandler(int) {
 
     Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl;
     FormatBackTrace(&Cerr);
-    Cerr << colors.Red() << "==============================================" << colors.Default() << Endl;
+    Cerr << colors.Red() << "====================================================" << colors.Default() << Endl;
 
     abort();
 }

diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp
@@ -232,40 +232,78 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaite
     static constexpr TDuration REFRESH_PERIOD = TDuration::MilliSeconds(10);
 
 public:
-    TResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount)
-        : ExpectedNodeCount_(expectedNodeCount)
+    TResourcesWaiterActor(NThreading::TPromise<void> promise, const TWaitResourcesSettings& settings)
+        : Settings_(settings)
         , Promise_(promise)
     {}
 
     void Bootstrap() {
-        Become(&TResourcesWaiterActor::StateFunc);
+        if (Settings_.HealthCheckLevel < 1) {
+            Finish();
+            return;
+        }
+
+        Become(&TResourcesWaiterActor::StateWaitNodeCont);
         CheckResourcesPublish();
     }
 
-    void Handle(NActors::TEvents::TEvWakeup::TPtr&) {
+    void HandleWaitNodeCountWakeup() {
         CheckResourcesPublish();
     }
 
     void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) {
-        if (ev->Get()->NodeCount == ExpectedNodeCount_) {
-            Promise_.SetValue();
-            PassAway();
+        const auto nodeCont = ev->Get()->NodeCount;
+        if (nodeCont == Settings_.ExpectedNodeCount) {
+            if (Settings_.HealthCheckLevel < 2) {
+                Finish();
+            } else {
+                Become(&TResourcesWaiterActor::StateWaitScript);
+                StartScriptQuery();
+            }
             return;
         }
 
+        if (Settings_.VerboseLevel >= 2) {
+            Cout << CoutColors_.Cyan() << "Retry invalid node count, got " << nodeCont << ", expected " << Settings_.ExpectedNodeCount << CoutColors_.Default() << Endl;
+        }
         Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup());
     }
 
-    STRICT_STFUNC(StateFunc,
-        hFunc(NActors::TEvents::TEvWakeup, Handle);
+    void HandleWaitScriptWakeup() {
+        StartScriptQuery();
+    }
+
+    void Handle(NKikimr::NKqp::TEvKqp::TEvScriptResponse::TPtr& ev) {
+        const auto status = ev->Get()->Status;
+        if (status == Ydb::StatusIds::SUCCESS) {
+            Finish();
+            return;
+        }
+
+        if (Settings_.VerboseLevel >= 2) {
+            Cout << CoutColors_.Cyan() << "Retry script creation fail with status " << status << ", reason:\n" << CoutColors_.Default() << ev->Get()->Issues.ToString() << Endl;
+        }
+        Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup());
+    }
+
+    STRICT_STFUNC(StateWaitNodeCont,
+        sFunc(NActors::TEvents::TEvWakeup, HandleWaitNodeCountWakeup);
         hFunc(TEvPrivate::TEvResourcesInfo, Handle);
     )
 
+    STRICT_STFUNC(StateWaitScript,
+        sFunc(NActors::TEvents::TEvWakeup, HandleWaitScriptWakeup);
+        hFunc(NKikimr::NKqp::TEvKqp::TEvScriptResponse, Handle);
+    )
+
 private:
     void CheckResourcesPublish() {
         GetResourceManager();
 
         if (!ResourceManager_) {
+            if (Settings_.VerboseLevel >= 2) {
+                Cout << CoutColors_.Cyan() << "Retry uninitialized resource manager" << CoutColors_.Default() << Endl;
+            }
             Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup());
             return;
         }
@@ -287,8 +325,27 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaite
         });
     }
 
+    void StartScriptQuery() {
+        auto event = MakeHolder<NKikimr::NKqp::TEvKqp::TEvScriptRequest>();
+        event->Record.SetUserToken(NACLib::TUserToken("", BUILTIN_ACL_ROOT, {}).SerializeAsString());
+
+        auto request = event->Record.MutableRequest();
+        request->SetQuery("SELECT 42");
+        request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT);
+        request->SetAction(NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE);
+        request->SetDatabase(Settings_.Database);
+
+        Send(NKikimr::NKqp::MakeKqpProxyID(SelfId().NodeId()), event.Release());
+    }
+
+    void Finish() {
+        Promise_.SetValue();
+        PassAway();
+    }
+
 private:
-    const i32 ExpectedNodeCount_;
+    const TWaitResourcesSettings Settings_;
+    const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout);
     NThreading::TPromise<void> Promise_;
 
     std::shared_ptr<NKikimr::NKqp::NRm::IKqpResourceManager> ResourceManager_;
@@ -415,8 +472,8 @@ NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settin
     return new TAsyncQueryRunnerActor(settings);
 }
 
-NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount) {
-    return new TResourcesWaiterActor(promise, expectedNodeCount);
+NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, const TWaitResourcesSettings& settings) {
+    return new TResourcesWaiterActor(promise, settings);
 }
 
 NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise<TString> openPromise, NThreading::TPromise<void> closePromise) {

diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h
@@ -27,6 +27,13 @@ struct TCreateSessionRequest {
     ui8 VerboseLevel;
 };
 
+struct TWaitResourcesSettings {
+    i32 ExpectedNodeCount;
+    ui8 HealthCheckLevel;
+    ui8 VerboseLevel;
+    TString Database;
+};
+
 struct TEvPrivate {
     enum EEv : ui32 {
         EvStartAsyncQuery = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
@@ -83,7 +90,7 @@ NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPr
 
 NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings);
 
-NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount);
+NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, const TWaitResourcesSettings& settings);
 
 NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise<TString> openPromise, NThreading::TPromise<void> closePromise);
 

diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h
@@ -36,12 +36,13 @@ struct TYdbSetupSettings {
     std::unordered_set<TString> SharedTenants;
     std::unordered_set<TString> ServerlessTenants;
     TDuration InitializationTimeout = TDuration::Seconds(10);
+    ui8 HealthCheckLevel = 1;
     bool SameSession = false;
 
     bool DisableDiskMock = false;
     bool FormatStorage = false;
     std::optional<TString> PDisksPath;
-    ui64 DiskSize = 32_GB;
+    std::optional<ui64> DiskSize;
 
     bool MonitoringEnabled = false;
     ui16 MonitoringPortOffset = 0;

diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto
@@ -4,4 +4,5 @@ package NKqpRun;
 
 message TStorageMeta {
     uint64 StorageGeneration = 1;
+    uint64 StorageSize = 2;
 }
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -178,7 +178,10 @@ class TYdbSetup::TImpl {
     void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const {
         TString diskPath;
         if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") {
-            diskPath = TStringBuilder() << *Settings_.PDisksPath << "/";
+            if (Settings_.PDisksPath->empty()) {
+                ythrow yexception() << "Storage directory path should not be empty";
+            }
+            diskPath = TStringBuilder() << Settings_.PDisksPath << (Settings_.PDisksPath->back() != '/' ? "/" : "");
         }
 
         bool formatDisk = true;
@@ -196,6 +199,13 @@ class TYdbSetup::TImpl {
                 formatDisk = false;
             }
 
+            if (Settings_.DiskSize && storageMeta.GetStorageSize() != *Settings_.DiskSize) {
+                if (!formatDisk) {
+                    ythrow yexception() << "Cannot change disk size without formatting storage, please use --format-storage";
+                }
+                storageMeta.SetStorageSize(*Settings_.DiskSize);
+            }
+
             TString storageMetaStr;
             google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr);
 
@@ -208,7 +218,7 @@ class TYdbSetup::TImpl {
             .UseDisk = !!Settings_.PDisksPath,
             .SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE,
             .ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE,
-            .DiskSize = Settings_.DiskSize,
+            .DiskSize = Settings_.DiskSize ? *Settings_.DiskSize : 32_GB,
             .FormatDisk = formatDisk,
             .DiskPath = diskPath
         };
@@ -351,7 +361,13 @@ class TYdbSetup::TImpl {
 
     void WaitResourcesPublishing() const {
         auto promise = NThreading::NewPromise();
-        GetRuntime()->Register(CreateResourcesWaiterActor(promise, Settings_.NodeCount), 0, GetRuntime()->GetAppData().SystemPoolId);
+        const TWaitResourcesSettings settings = {
+            .ExpectedNodeCount = static_cast<i32>(Settings_.NodeCount),
+            .HealthCheckLevel = Settings_.HealthCheckLevel,
+            .VerboseLevel = Settings_.VerboseLevel,
+            .Database = NKikimr::CanonizePath(Settings_.DomainName)
+        };
+        GetRuntime()->Register(CreateResourcesWaiterActor(promise, settings), 0, GetRuntime()->GetAppData().SystemPoolId);
 
         try {
             promise.GetFuture().GetValue(Settings_.InitializationTimeout);