YDB CLI. Use CLI settings parameter instead of compile flag to disable updates (#14451)

pnv1 · web-flow · commit 108ab9a4612c · 2025-02-11T16:21:08.000Z
diff --git a/ydb/apps/ydb/commands/ya.make b/ydb/apps/ydb/commands/ya.make
@@ -1,20 +1,11 @@
 LIBRARY(commands)
 
 SRCS(
-    ydb_cloud_root.cpp
+    ydb_root.cpp
+    ydb_update.cpp
     ydb_version.cpp
 )
 
-IF (YDB_CERTIFIED)
-    CFLAGS(
-        -DDISABLE_UPDATE
-    )
-ELSE()
-    SRCS(
-        ydb_update.cpp
-    )
-ENDIF ()
-
 PEERDIR(
     ydb/public/sdk/cpp/src/client/iam
     ydb/public/lib/ydb_cli/commands
diff --git a/ydb/apps/ydb/commands/ydb_root.cpp b/ydb/apps/ydb/commands/ydb_root.cpp
@@ -1,13 +1,11 @@
-#include "ydb_cloud_root.h"
+#include "ydb_root.h"
+#include "ydb_update.h"
 #include "ydb_version.h"
 
 #include <ydb-cpp-sdk/client/iam/iam.h>
 #include <ydb-cpp-sdk/client/types/credentials/oauth2_token_exchange/from_file.h>
 
-#ifndef DISABLE_UPDATE
-#include "ydb_update.h"
 #include <ydb/public/lib/ydb_cli/common/ydb_updater.h>
-#endif
 
 #include <filesystem>
 
@@ -56,12 +54,12 @@ void TClientCommandRoot::SetCredentialsGetter(TConfig& config) {
     };
 }
 
-TYCloudClientCommandRoot::TYCloudClientCommandRoot(const TString& name, const TClientSettings& settings)
+TYdbClientCommandRoot::TYdbClientCommandRoot(const TString& name, const TClientSettings& settings)
     : TClientCommandRoot(name, settings)
 {
-#ifndef DISABLE_UPDATE
-    AddCommand(std::make_unique<TCommandUpdate>());
-#endif
+    if (settings.StorageUrl.has_value()) {
+        AddCommand(std::make_unique<TCommandUpdate>());
+    }
     AddCommand(std::make_unique<TCommandVersion>());
 }
 
@@ -76,17 +74,16 @@ namespace {
     }
 }
 
-void TYCloudClientCommandRoot::Config(TConfig& config) {
+void TYdbClientCommandRoot::Config(TConfig& config) {
     TClientCommandRoot::Config(config);
 
     NLastGetopt::TOpts& opts = *config.Opts;
     RemoveOption(opts, "svnrevision");
 }
 
-int TYCloudClientCommandRoot::Run(TConfig& config) {
-#ifndef DISABLE_UPDATE
-    if (config.NeedToCheckForUpdate) {
-        TYdbUpdater updater;
+int TYdbClientCommandRoot::Run(TConfig& config) {
+    if (config.StorageUrl.has_value() && config.NeedToCheckForUpdate) {
+        TYdbUpdater updater(config.StorageUrl.value());
         if (config.ForceVersionCheck) {
             Cout << "Force checking if there is a newer version..." << Endl;
         }
@@ -101,12 +98,11 @@ int TYCloudClientCommandRoot::Run(TConfig& config) {
                 << colors.OldColor() << Endl;
         }
     }
-#endif
 
     return TClientCommandRoot::Run(config);
 }
 
-int NewYCloudClient(int argc, char** argv) {
+int NewYdbClient(int argc, char** argv) {
     NYdb::NConsoleClient::TClientSettings settings;
     settings.EnableSsl = true;
     settings.UseAccessToken = true;
@@ -116,9 +112,10 @@ int NewYCloudClient(int argc, char** argv) {
     settings.UseOauth2TokenExchange = true;
     settings.UseExportToYt = false;
     settings.MentionUserAccount = false;
+    settings.StorageUrl = "https://storage.yandexcloud.net/yandexcloud-ydb/release";
     settings.YdbDir = "ydb";
 
-    auto commandsRoot = MakeHolder<TYCloudClientCommandRoot>(std::filesystem::path(argv[0]).stem().string(), settings);
+    auto commandsRoot = MakeHolder<TYdbClientCommandRoot>(std::filesystem::path(argv[0]).stem().string(), settings);
     commandsRoot->Opts.SetTitle("YDB client");
     TClientCommand::TConfig config(argc, argv);
     return commandsRoot->Process(config);
diff --git a/ydb/apps/ydb/commands/ydb_root.h b/ydb/apps/ydb/commands/ydb_root.h
@@ -14,14 +14,14 @@ class TClientCommandRoot : public TClientCommandRootCommon {
     void SetCredentialsGetter(TConfig& config) override;
 };
 
-class TYCloudClientCommandRoot : public TClientCommandRoot {
+class TYdbClientCommandRoot : public TClientCommandRoot {
 public:
-    TYCloudClientCommandRoot(const TString& name, const TClientSettings& settings);
+    TYdbClientCommandRoot(const TString& name, const TClientSettings& settings);
     void Config(TConfig& config) override;
     int Run(TConfig& config) override;
 };
 
-int NewYCloudClient(int argc, char** argv);
+int NewYdbClient(int argc, char** argv);
 
 }
 }
diff --git a/ydb/apps/ydb/commands/ydb_update.cpp b/ydb/apps/ydb/commands/ydb_update.cpp
@@ -20,10 +20,7 @@ void TCommandUpdate::Config(TConfig& config) {
 }
 
 int TCommandUpdate::Run(TConfig& config) {
-    Y_UNUSED(config);
-
-
-    TYdbUpdater updater;
+    TYdbUpdater updater(config.StorageUrl.value());
     return updater.Update(ForceUpdate);
 }
 
diff --git a/ydb/apps/ydb/commands/ydb_version.cpp b/ydb/apps/ydb/commands/ydb_version.cpp
@@ -1,8 +1,6 @@
 #include "ydb_version.h"
 
-#ifndef DISABLE_UPDATE
-    #include <ydb/public/lib/ydb_cli/common/ydb_updater.h>
-#endif
+#include <ydb/public/lib/ydb_cli/common/ydb_updater.h>
 
 #include <library/cpp/resource/resource.h>
 #include <util/string/strip.h>
@@ -24,8 +22,11 @@ void TCommandVersion::Config(TConfig& config) {
 
     config.Opts->AddLongOption("semantic", "Print semantic version only")
         .StoreTrue(&Semantic);
+    
+    if (!config.StorageUrl.has_value()) {
+        return;
+    }
 
-#ifndef DISABLE_UPDATE
     config.Opts->AddLongOption("check", "Force to check latest version available")
         .StoreTrue(&config.ForceVersionCheck);
     config.Opts->AddLongOption("disable-checks", "Disable version checks. CLI will not check whether there is a newer version available")
@@ -39,8 +40,6 @@ void TCommandVersion::Config(TConfig& config) {
     config.Opts->MutuallyExclusive("enable-checks", "semantic");
     config.Opts->MutuallyExclusive("enable-checks", "check");
     config.Opts->MutuallyExclusive("check", "semantic");
-#endif
-
 }
 
 void TCommandVersion::Parse(TConfig& config) {
@@ -52,20 +51,20 @@ void TCommandVersion::Parse(TConfig& config) {
 int TCommandVersion::Run(TConfig& config) {
     Y_UNUSED(config);
 
-#ifndef DISABLE_UPDATE
-    if (EnableChecks) {
-        TYdbUpdater updater;
-        updater.SetCheckVersion(true);
-        Cout << "Latest version checks enabled" << Endl;
-        return EXIT_SUCCESS;
-    }
-    if (DisableChecks) {
-        TYdbUpdater updater;
-        updater.SetCheckVersion(false);
-        Cout << "Latest version checks disabled" << Endl;
-        return EXIT_SUCCESS;
+    if (config.StorageUrl.has_value()) {
+        if (EnableChecks) {
+            TYdbUpdater updater(config.StorageUrl.value());
+            updater.SetCheckVersion(true);
+            Cout << "Latest version checks enabled" << Endl;
+            return EXIT_SUCCESS;
+        }
+        if (DisableChecks) {
+            TYdbUpdater updater(config.StorageUrl.value());
+            updater.SetCheckVersion(false);
+            Cout << "Latest version checks disabled" << Endl;
+            return EXIT_SUCCESS;
+        }
     }
-#endif
 
     if (!Semantic) {
         Cout << "YDB CLI ";
diff --git a/ydb/apps/ydb/main.cpp b/ydb/apps/ydb/main.cpp
@@ -1,4 +1,4 @@
-#include <ydb/apps/ydb/commands/ydb_cloud_root.h>
+#include <ydb/apps/ydb/commands/ydb_root.h>
 #include <ydb/public/lib/ydb_cli/commands/ydb_service_topic.h>
 
 TVector<NYdb::NTopic::ECodec> NYdb::NConsoleClient::InitAllowedCodecs() {
@@ -11,7 +11,7 @@ TVector<NYdb::NTopic::ECodec> NYdb::NConsoleClient::InitAllowedCodecs() {
 
 int main(int argc, char **argv) {
     try {
-        return NYdb::NConsoleClient::NewYCloudClient(argc, argv);
+        return NYdb::NConsoleClient::NewYdbClient(argc, argv);
     }
     catch (const NYdb::NConsoleClient::TMisuseWithHelpException& e) {
         // command help is already printed. Just exit(1)
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_root_common.cpp b/ydb/public/lib/ydb_cli/commands/ydb_root_common.cpp
@@ -88,6 +88,7 @@ void TClientCommandRootCommon::FillConfig(TConfig& config) {
     config.UseStaticCredentials = Settings.UseStaticCredentials.GetRef();
     config.UseOauth2TokenExchange = Settings.UseOauth2TokenExchange.GetRef();
     config.UseExportToYt = Settings.UseExportToYt.GetRef();
+    config.StorageUrl = Settings.StorageUrl;
     SetCredentialsGetter(config);
 }
 
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_root_common.h b/ydb/public/lib/ydb_cli/commands/ydb_root_common.h
@@ -23,6 +23,9 @@ struct TClientSettings {
     TMaybe<bool> UseExportToYt;
     // Whether to mention user account in --help command or not
     TMaybe<bool> MentionUserAccount;
+    // A storage url to get latest YDB CLI version and to update from
+    // If not set, than no updates nor latest version checks will be available
+    std::optional<std::string> StorageUrl = std::nullopt;
     // Name of a directory in user home directory to save profile config
     TString YdbDir;
 };
diff --git a/ydb/public/lib/ydb_cli/common/command.h b/ydb/public/lib/ydb_cli/common/command.h
@@ -148,6 +148,7 @@ class TClientCommand {
         bool AllowEmptyAddress = false;
         bool OnlyExplicitProfile = false;
         bool AssumeYes = false;
+        std::optional<std::string> StorageUrl = std::nullopt;
 
         TCredentialsGetter CredentialsGetter;
         std::shared_ptr<ICredentialsProviderFactory> SingletonCredentialsProviderFactory = nullptr;
diff --git a/ydb/public/lib/ydb_cli/common/ya.make b/ydb/public/lib/ydb_cli/common/ya.make
@@ -30,19 +30,10 @@ SRCS(
     sys.cpp
     tabbed_table.cpp
     waiting_bar.cpp
+    ydb_updater.cpp
     yt.cpp
 )
 
-IF (YDB_CERTIFIED)
-    CFLAGS(
-        -DDISABLE_UPDATE
-    )
-ELSE()
-    SRCS(
-        ydb_updater.cpp
-    )
-ENDIF ()
-
 PEERDIR(
     contrib/libs/aws-sdk-cpp/aws-cpp-sdk-s3
     library/cpp/config
diff --git a/ydb/public/lib/ydb_cli/common/ydb_updater.cpp b/ydb/public/lib/ydb_cli/common/ydb_updater.cpp
@@ -50,12 +50,11 @@ namespace {
     const TString osArch = GetOsArchitecture();
     const TString defaultConfigFile = TStringBuilder() << homeDir << "/ydb/bin/config.json";
     const TString defaultTempFile = TStringBuilder() << homeDir << "/ydb/install/" << binaryName;
-    const TString storageUrl = "https://storage.yandexcloud.net/yandexcloud-ydb/release/";
-    const TString versionUrl = TStringBuilder() << storageUrl << "stable";
 }
 
-TYdbUpdater::TYdbUpdater()
+TYdbUpdater::TYdbUpdater(std::string storageUrl)
     : MyVersion(StripString(NResource::Find(TStringBuf(VersionResourceName))))
+    , StorageUrl(storageUrl)
 {
     LoadConfig();
 }
@@ -103,7 +102,7 @@ int TYdbUpdater::Update(bool forceUpdate) {
     if (!tmpPathToBinary.Parent().Exists()) {
         tmpPathToBinary.Parent().MkDirs();
     }
-    const TString downloadUrl = TStringBuilder() << storageUrl << LatestVersion << '/' << osVersion
+    const TString downloadUrl = TStringBuilder() << StorageUrl << '/' << LatestVersion << '/' << osVersion
         << '/' << osArch << '/' << binaryName;
     Cout << "Downloading binary from url " << downloadUrl << Endl;
     TShellCommand curlCmd(TStringBuilder() << "curl --max-time 60 " << downloadUrl << " -o " << tmpPathToBinary.GetPath());
@@ -218,7 +217,7 @@ bool TYdbUpdater::GetLatestVersion() {
     if (LatestVersion) {
         return true;
     }
-
+    std::string versionUrl = StorageUrl + "/stable";
     TShellCommand curlCmd(TStringBuilder() << "curl --silent --max-time 10 " << versionUrl);
     curlCmd.Run().Wait();
 
diff --git a/ydb/public/lib/ydb_cli/common/ydb_updater.h b/ydb/public/lib/ydb_cli/common/ydb_updater.h
@@ -8,7 +8,7 @@ extern const char* VersionResourceName;
 
 class TYdbUpdater {
 public:
-    TYdbUpdater();
+    TYdbUpdater(std::string storageUrl);
     ~TYdbUpdater();
 
     bool CheckIfUpdateNeeded(bool forceRequest = false);
@@ -29,6 +29,7 @@ class TYdbUpdater {
     NJson::TJsonValue Config;
     TString MyVersion;
     TString LatestVersion;
+    std::string StorageUrl;
     bool ConfigChanged = false;
 };
 

Original file line number	Diff line number	Diff line change
`@@ -20,10 +20,7 @@ void TCommandUpdate::Config(TConfig& config) {`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`int TCommandUpdate::Run(TConfig& config) {`
`23`		`- Y_UNUSED(config);`
`24`		`-`
`25`		`-`
`26`		`- TYdbUpdater updater;`
	`23`	`+ TYdbUpdater updater(config.StorageUrl.value());`
`27`	`24`	`return updater.Update(ForceUpdate);`
`28`	`25`	`}`
`29`	`26`
Original file line number	Diff line number	Diff line change
`@@ -88,6 +88,7 @@ void TClientCommandRootCommon::FillConfig(TConfig& config) {`
`88`	`88`	`config.UseStaticCredentials = Settings.UseStaticCredentials.GetRef();`
`89`	`89`	`config.UseOauth2TokenExchange = Settings.UseOauth2TokenExchange.GetRef();`
`90`	`90`	`config.UseExportToYt = Settings.UseExportToYt.GetRef();`
	`91`	`+ config.StorageUrl = Settings.StorageUrl;`
`91`	`92`	`SetCredentialsGetter(config);`
`92`	`93`	`}`
`93`	`94`