Detect platform architecture in `ydb update` (#13669)

2 files changed, 43 insertions, 23 deletions

diff --git a/ydb/apps/ydb/CHANGELOG.md b/ydb/apps/ydb/CHANGELOG.md
index 7945569f6e..3315bd2bfe 100644
--- a/ydb/apps/ydb/CHANGELOG.md
+++ b/ydb/apps/ydb/CHANGELOG.md
@@ -1,3 +1,4 @@
+* Fixed a bug where arm64 YDB CLI binary was downloading amd64 binary to replace itself during `ydb update`. To update already installed binaries to the latest arm64 version, YDB CLI should be re-installed
 * Fixed a bug where `ydb workload tpch import generator` and `ydb workload tpcds import generator` commands were failing due to not all tables were created
 * Fixed a bug with backslashes in `ydb workload` benchmark paths on Windows
 * Added CREATE TABLE text suggestion on scheme error during `ydb import file csv`
diff --git a/ydb/public/lib/ydb_cli/common/ydb_updater.cpp b/ydb/public/lib/ydb_cli/common/ydb_updater.cpp
index be3c838d76..98a8ad55ee 100644
--- a/ydb/public/lib/ydb_cli/common/ydb_updater.cpp
+++ b/ydb/public/lib/ydb_cli/common/ydb_updater.cpp
@@ -12,29 +12,48 @@
 #include <util/system/execpath.h>
 #include <util/system/shellcommand.h>
 
+#ifndef _win32_
+#include <sys/utsname.h>
+#endif
 namespace NYdb {
 namespace NConsoleClient {
 
 const char* VersionResourceName = "version.txt";
 
+TString GetOsArchitecture() {
+#if defined(_win32_)
+    return "amd64";
+#else
+    struct utsname uts;
+    uname(&uts);
+    TString machine = uts.machine;
+    if (machine == "arm64" || machine == "aarch64") {
+        return "arm64";
+    } else {
+        return "amd64";
+    }
+#endif
+}
+
 namespace {
 #if defined(_darwin_)
-    const TString OsVersion = "darwin";
-    const TString BinaryName = "ydb";
-    const TString HomeDir = GetHomeDir();
+    const TString osVersion = "darwin";
+    const TString binaryName = "ydb";
+    const TString homeDir = GetHomeDir();
 #elif defined(_win32_)
-    const TString OsVersion = "windows";
-    const TString BinaryName = "ydb.exe";
-    const TString HomeDir = GetEnv("USERPROFILE");
+    const TString osVersion = "windows";
+    const TString binaryName = "ydb.exe";
+    const TString homeDir = GetEnv("USERPROFILE");
 #else
-    const TString OsVersion = "linux";
-    const TString BinaryName = "ydb";
-    const TString HomeDir = GetHomeDir();
+    const TString osVersion = "linux";
+    const TString binaryName = "ydb";
+    const TString homeDir = GetHomeDir();
 #endif
-    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";
+    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()
@@ -80,19 +99,19 @@ int TYdbUpdater::Update(bool forceUpdate) {
         return EXIT_FAILURE;
     }
 
-    TFsPath tmpPathToBinary(DefaultTempFile);
+    TFsPath tmpPathToBinary(defaultTempFile);
     tmpPathToBinary.Fix();
     TString corrPath = tmpPathToBinary.GetPath();
     if (!tmpPathToBinary.Parent().Exists()) {
         tmpPathToBinary.Parent().MkDirs();
     }
-    const TString DownloadUrl = TStringBuilder() << StorageUrl << LatestVersion << "/" << OsVersion << "/amd64/"
-        << BinaryName;
-    Cout << "Downloading binary from url " << DownloadUrl << Endl;
-    TShellCommand curlCmd(TStringBuilder() << "curl --max-time 60 " << DownloadUrl << " -o " << tmpPathToBinary.GetPath());
+    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());
     curlCmd.Run().Wait();
     if (curlCmd.GetExitCode() != 0) {
-        Cerr << "Failed to download from url \"" << DownloadUrl << "\". " << curlCmd.GetError() << Endl;
+        Cerr << "Failed to download from url \"" << downloadUrl << "\". " << curlCmd.GetError() << Endl;
         return EXIT_FAILURE;
     }
     Cout << "Downloaded to " << tmpPathToBinary.GetPath() << Endl;
@@ -144,7 +163,7 @@ void TYdbUpdater::SetConfigValue(const TString& name, const T& value) {
 }
 
 void TYdbUpdater::LoadConfig() {
-    TFsPath configFilePath(DefaultConfigFile);
+    TFsPath configFilePath(defaultConfigFile);
     configFilePath.Fix();
     try {
         if (configFilePath.Exists()) {
@@ -164,7 +183,7 @@ void TYdbUpdater::LoadConfig() {
 
 void TYdbUpdater::SaveConfig() {
     try {
-        TFsPath configFilePath(DefaultConfigFile);
+        TFsPath configFilePath(defaultConfigFile);
         configFilePath.Fix();
         if (!configFilePath.Parent().Exists()) {
             configFilePath.Parent().MkDirs();
@@ -202,14 +221,14 @@ bool TYdbUpdater::GetLatestVersion() {
         return true;
     }
 
-    TShellCommand curlCmd(TStringBuilder() << "curl --silent --max-time 10 " << VersionUrl);
+    TShellCommand curlCmd(TStringBuilder() << "curl --silent --max-time 10 " << versionUrl);
     curlCmd.Run().Wait();
 
     if (curlCmd.GetExitCode() == 0) {
         LatestVersion = StripString(curlCmd.GetOutput());
         return true;
     }
-    Cerr << "(!) Couldn't get latest version from url \"" << VersionUrl << "\". " << curlCmd.GetError() << Endl;
+    Cerr << "(!) Couldn't get latest version from url \"" << versionUrl << "\". " << curlCmd.GetError() << Endl;
     return false;
 }