Fix UUID handling in backup/restore in YDB CLI (#17198)

Fix YQL, BulkUpsert and ImportData backup and restore modes so that they can handle UUID fields

Co-authored-by: jepett0 <111313089+jepett0@users.noreply.github.com>

6 files changed, 112 insertions, 4 deletions

diff --git a/ydb/library/backup/backup.cpp b/ydb/library/backup/backup.cpp
index b2baf248a22..b70d687252a 100644
--- a/ydb/library/backup/backup.cpp
+++ b/ydb/library/backup/backup.cpp
@@ -152,6 +152,7 @@ void PrintPrimitive(IOutputStream& out, const TValueParser& parser) {
         CASE_PRINT_PRIMITIVE_TYPE(out, Datetime64);
         CASE_PRINT_PRIMITIVE_TYPE(out, Timestamp64);
         CASE_PRINT_PRIMITIVE_TYPE(out, Interval64);
+        CASE_PRINT_PRIMITIVE_TYPE(out, Uuid);
         CASE_PRINT_PRIMITIVE_STRING_TYPE(out, TzDate);
         CASE_PRINT_PRIMITIVE_STRING_TYPE(out, TzDatetime);
         CASE_PRINT_PRIMITIVE_STRING_TYPE(out, TzTimestamp);
diff --git a/ydb/library/backup/query_builder.cpp b/ydb/library/backup/query_builder.cpp
index 503da2457b9..865b590652f 100644
--- a/ydb/library/backup/query_builder.cpp
+++ b/ydb/library/backup/query_builder.cpp
@@ -3,6 +3,7 @@
 #include "backup.h"
 
 #include <yql/essentials/types/dynumber/dynumber.h>
+#include <yql/essentials/types/uuid/uuid.h>
 #include <ydb/public/api/protos/ydb_value.pb.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/proto/accessor.h>
 
@@ -136,11 +137,11 @@ void TQueryBuilder::AddPrimitiveMember(EPrimitiveType type, TStringBuf buf) {
 
     case EPrimitiveType::Datetime64:
         Value.Datetime64(TryParse<i64>(buf));
-        break;        
+        break;
 
     case EPrimitiveType::Timestamp64:
         Value.Timestamp64(TryParse<i64>(buf));
-        break;        
+        break;
 
     case EPrimitiveType::Interval64:
         Value.Interval64(TryParse<i64>(buf));
@@ -184,7 +185,8 @@ void TQueryBuilder::AddPrimitiveMember(EPrimitiveType type, TStringBuf buf) {
         break;
 
     case EPrimitiveType::Uuid:
-        Y_ENSURE(false, TStringBuilder() << "Unexpected Primitive kind while parsing line: " << type);
+        Y_ENSURE(NKikimr::NUuid::IsValidUuid(buf));
+        Value.Uuid(TUuidValue(std::string(buf.begin(), buf.end())));
         break;
 
     }
diff --git a/ydb/public/lib/ydb_cli/dump/restore_import_data.cpp b/ydb/public/lib/ydb_cli/dump/restore_import_data.cpp
index 64c22297b6f..0c592bd8ecf 100644
--- a/ydb/public/lib/ydb_cli/dump/restore_import_data.cpp
+++ b/ydb/public/lib/ydb_cli/dump/restore_import_data.cpp
@@ -22,6 +22,20 @@
 #include <util/system/mutex.h>
 #include <util/thread/pool.h>
 
+namespace NYdb {
+
+bool operator<(const TUuidValue& lhs, const TUuidValue& rhs) {
+    // Lexicographical comparison of UUIDs for TValue comparison.
+    // It works just like TCell::CompareCellsAsByteString.
+    // We need it since RPC Import Data expects keys to be sorted.
+    const char* pa = lhs.Buf_.Bytes;
+    const char* pb = rhs.Buf_.Bytes;
+    int cmp = memcmp(pa, pb, 16);
+    return cmp < 0;
+}
+
+}
+
 namespace NYdb::NDump {
 
 using namespace NImport;
@@ -44,6 +58,7 @@ class TValue {
         Null,
         String,
         Pod,
+        Uuid
     };
 
     inline EType GetType() const {
@@ -54,6 +69,8 @@ class TValue {
             return EType::Null;
         case 2:
             return EType::String;
+        case 9:
+            return EType::Uuid;
         default:
             return EType::Pod;
         }
@@ -110,7 +127,8 @@ private:
         i32,
         ui32,
         i64,
-        ui64
+        ui64,
+        TUuidValue
     > Value;
 
 }; // TValue
@@ -149,6 +167,8 @@ class TValueConverter {
             return TValue(Parser.GetString());
         case EPrimitiveType::Utf8:
             return TValue(Parser.GetUtf8());
+        case EPrimitiveType::Uuid:
+            return TValue(Parser.GetUuid());
         default:
             Y_ENSURE(false, "Unexpected primitive type: " << type);
         }
@@ -290,6 +310,10 @@ public:
         return CheckedUnescape();
     }
 
+    TUuidValue GetUuid() const {
+        return TUuidValue(std::string(Value));
+    }
+
     bool IsNull() const {
         return Value == "null";
     }
diff --git a/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/value/value.h b/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/value/value.h
index a87a4ca9e35..61eb9ce2c0f 100644
--- a/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/value/value.h
+++ b/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/value/value.h
@@ -539,3 +539,6 @@ public:
 };
 
 } // namespace NYdb
+
+template<>
+void Out<NYdb::TUuidValue>(IOutputStream& o, const NYdb::TUuidValue& value);
diff --git a/ydb/public/sdk/cpp/src/client/value/value.cpp b/ydb/public/sdk/cpp/src/client/value/value.cpp
index 1f48b7f8158..2516f3a4226 100644
--- a/ydb/public/sdk/cpp/src/client/value/value.cpp
+++ b/ydb/public/sdk/cpp/src/client/value/value.cpp
@@ -3374,3 +3374,8 @@ TValue TValueBuilder::Build() {
 }
 
 } // namespace NYdb
+
+template<>
+void Out<NYdb::TUuidValue>(IOutputStream& o, const NYdb::TUuidValue& value) {
+    o << value.ToString();
+}
diff --git a/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp b/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
index 93444678320..cfcb607b4b2 100644
--- a/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
+++ b/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
@@ -1372,6 +1372,79 @@ Y_UNIT_TEST_SUITE(BackupRestore) {
             [](const TYdbErrorException& e) { return e.GetStatus().GetStatus() == EStatus::SCHEME_ERROR; });
     }
 
+    Y_UNIT_TEST(BackupUuid) {
+        TKikimrWithGrpcAndRootSchema server;
+        auto driver = TDriver(TDriverConfig().SetEndpoint(Sprintf("localhost:%u", server.GetPort())));
+        TTableClient tableClient(driver);
+        auto session = tableClient.GetSession().ExtractValueSync().GetSession();
+        TTempDir tempDir;
+        const auto& pathToBackup = tempDir.Path();
+
+        constexpr const char* dbPath = "/Root";
+        constexpr const char* table = "/Root/table";
+
+        ExecuteDataDefinitionQuery(session, Sprintf(R"(
+                CREATE TABLE `%s` (
+                    Key Uuid,
+                    Value Utf8,
+                    PRIMARY KEY (Key)
+                );
+            )",
+            table
+        ));
+
+        std::vector<std::string> uuids = {
+            "5b99a330-04ef-4f1a-9b64-ba6d5f44eafe",
+            "706cca52-b00a-4cbd-a21e-6538de188271",
+            "81b1e345-f2ae-4c9e-8d1a-75447be314f2",
+            "be2765f2-9f4c-4a22-8d2c-a1b77d84f4fb",
+            "d3f9e0a2-5871-4afe-a23a-8db160b449cd",
+            "d3f9e0a2-0000-0000-0000-8db160b449cd"
+        };
+
+        ExecuteDataModificationQuery(session, Sprintf(R"(
+                UPSERT INTO `%s` (Key, Value)
+                VALUES 
+                    (Uuid("%s"), "one"),
+                    (Uuid("%s"), "two"),
+                    (Uuid("%s"), "three"),
+                    (Uuid("%s"), "four"),
+                    (Uuid("%s"), "five"),
+                    (Uuid("%s"), "six");
+            )", table, uuids[0].c_str(), uuids[1].c_str(), uuids[2].c_str(), uuids[3].c_str(), uuids[4].c_str(), uuids[5].c_str()
+        ));
+
+        const auto originalContent = GetTableContent(session, table);
+
+        NDump::TClient backupClient(driver);
+        {
+            const auto result = backupClient.Dump(dbPath, pathToBackup, NDump::TDumpSettings().Database(dbPath));
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        // Check that backup file contains all uuids as strings, making sure we stringify UUIDs correctly in backups
+        TString backupFileContent = TFileInput(pathToBackup.GetPath() + "/table/data_00.csv").ReadAll();
+        for (const auto& uuid : uuids) {
+            UNIT_ASSERT_C(backupFileContent.find(uuid) != TString::npos, "UUID not found in backup file");
+        }
+
+        for (auto backupMode : {NDump::TRestoreSettings::EMode::BulkUpsert, NDump::TRestoreSettings::EMode::ImportData, NDump::TRestoreSettings::EMode::Yql}) {
+            auto opts = NDump::TRestoreSettings().Mode(backupMode);
+
+            ExecuteDataDefinitionQuery(session, Sprintf(R"(
+                    DROP TABLE `%s`;
+                )", table
+            ));
+
+            auto result = backupClient.Restore(pathToBackup, dbPath, opts);
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+            const auto newContent = GetTableContent(session, table);
+
+            CompareResults(newContent, originalContent);
+        }
+    }
+
     // TO DO: test index impl table split boundaries restoration from a backup
 
     Y_UNIT_TEST(RestoreViewQueryText) {