Fixed bug with not null columns in ydb tools restore

Fixed bug with not null columns in ydb tools restore

Pull Request resolved: #307

3 files changed, 106 insertions, 69 deletions

diff --git a/ydb/library/backup/query_builder.cpp b/ydb/library/backup/query_builder.cpp
index dd931c82bb..f2aabc1802 100644
--- a/ydb/library/backup/query_builder.cpp
+++ b/ydb/library/backup/query_builder.cpp
@@ -39,11 +39,7 @@ TString TQueryBuilder::BuildQuery(const TString &path) {
 }
 
 template<typename T>
-TMaybe<T> TryParse(const TStringBuf& buf) {
-    if (buf == "null") {
-        return {};
-    }
-
+T TryParse(const TStringBuf& buf) {
     T tmp;
     TMemoryInput stream(buf);
     stream >> tmp;
@@ -51,11 +47,7 @@ TMaybe<T> TryParse(const TStringBuf& buf) {
 }
 
 template<>
-TMaybe<TString> TryParse(const TStringBuf& buf) {
-    if (buf == "null") {
-        return {};
-    }
-
+TString TryParse(const TStringBuf& buf) {
     Y_ENSURE(buf.Size() >= 1 && buf.front() == '"' && buf.back() == '"',
             "Source string neither surrounded by quotes nor equals to null, string# " << TString{buf}.Quote());
     TString tmp;
@@ -66,126 +58,113 @@ TMaybe<TString> TryParse(const TStringBuf& buf) {
 }
 
 template<>
-TMaybe<TInstant> TryParse(const TStringBuf& buf) {
-    if (buf == "null") {
-        return {};
-    }
-
+TInstant TryParse(const TStringBuf& buf) {
     return TInstant::ParseIso8601(buf);
 }
 
 template<>
-TMaybe<bool> TryParse(const TStringBuf& buf) {
-    auto tmp = TryParse<ui32>(buf);
-    if (tmp) {
-        return *tmp ? true : false;
-    } else {
-        return {};
-    }
+bool TryParse(const TStringBuf& buf) {
+    return TryParse<ui32>(buf) ? true : false;
 }
 
 void TQueryBuilder::AddPrimitiveMember(EPrimitiveType type, TStringBuf buf) {
     switch (type) {
 
     case EPrimitiveType::Bool:
-        Value.OptionalBool(TryParse<bool>(buf));
+        Value.Bool(TryParse<bool>(buf));
         break;
 
     case EPrimitiveType::Int8:
-        Value.OptionalInt8(TryParse<i32>(buf));
+        Value.Int8(TryParse<i32>(buf));
         break;
 
     case EPrimitiveType::Uint8:
-        Value.OptionalUint8(TryParse<ui32>(buf));
+        Value.Uint8(TryParse<ui32>(buf));
         break;
 
     case EPrimitiveType::Int16:
-        Value.OptionalInt16(TryParse<i32>(buf));
+        Value.Int16(TryParse<i32>(buf));
         break;
 
     case EPrimitiveType::Uint16:
-        Value.OptionalUint16(TryParse<ui32>(buf));
+        Value.Uint16(TryParse<ui32>(buf));
         break;
 
     case EPrimitiveType::Int32:
-        Value.OptionalInt32(TryParse<i32>(buf));
+        Value.Int32(TryParse<i32>(buf));
         break;
 
     case EPrimitiveType::Uint32:
-        Value.OptionalUint32(TryParse<ui32>(buf));
+        Value.Uint32(TryParse<ui32>(buf));
         break;
 
     case EPrimitiveType::Int64:
-        Value.OptionalInt64(TryParse<i64>(buf));
+        Value.Int64(TryParse<i64>(buf));
         break;
 
     case EPrimitiveType::Uint64:
-        Value.OptionalUint64(TryParse<ui64>(buf));
+        Value.Uint64(TryParse<ui64>(buf));
         break;
 
     case EPrimitiveType::Float:
-        Value.OptionalFloat(TryParse<float>(buf));
+        Value.Float(TryParse<float>(buf));
         break;
 
     case EPrimitiveType::Double:
-        Value.OptionalDouble(TryParse<double>(buf));
+        Value.Double(TryParse<double>(buf));
         break;
 
     case EPrimitiveType::Date:
-        Value.OptionalDate(TryParse<TInstant>(buf));
+        Value.Date(TryParse<TInstant>(buf));
         break;
 
     case EPrimitiveType::Datetime:
-        Value.OptionalDatetime(TryParse<TInstant>(buf));
+        Value.Datetime(TryParse<TInstant>(buf));
         break;
 
     case EPrimitiveType::Timestamp:
-        Value.OptionalTimestamp(TryParse<TInstant>(buf));
+        Value.Timestamp(TryParse<TInstant>(buf));
         break;
 
     case EPrimitiveType::Interval:
-        Value.OptionalInterval(TryParse<i64>(buf));
+        Value.Interval(TryParse<i64>(buf));
         break;
 
     case EPrimitiveType::TzDate:
-        Value.OptionalTzDate(TryParse<TString>(buf));
+        Value.TzDate(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::TzDatetime:
-        Value.OptionalTzDatetime(TryParse<TString>(buf));
+        Value.TzDatetime(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::TzTimestamp:
-        Value.OptionalTzTimestamp(TryParse<TString>(buf));
+        Value.TzTimestamp(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::String:
-        Value.OptionalString(TryParse<TString>(buf));
+        Value.String(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::Utf8:
-        Value.OptionalUtf8(TryParse<TString>(buf));
+        Value.Utf8(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::Yson:
-        Value.OptionalYson(TryParse<TString>(buf));
+        Value.Yson(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::Json:
-        Value.OptionalJson(TryParse<TString>(buf));
+        Value.Json(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::JsonDocument:
-        Value.OptionalJsonDocument(TryParse<TString>(buf));
+        Value.JsonDocument(TryParse<TString>(buf));
         break;
 
     case EPrimitiveType::DyNumber:
-        if (buf == "null") {
-            Value.OptionalDyNumber(Nothing());
-        } else {
-            Y_ENSURE(NKikimr::NDyNumber::IsValidDyNumberString(buf));
-            Value.OptionalDyNumber(TString(buf));
-        }
+        Y_ENSURE(NKikimr::NDyNumber::IsValidDyNumberString(buf));
+        Value.DyNumber(TString(buf));
         break;
 
     case EPrimitiveType::Uuid:
@@ -195,27 +174,61 @@ void TQueryBuilder::AddPrimitiveMember(EPrimitiveType type, TStringBuf buf) {
     }
 }
 
-void TQueryBuilder::AddMemberFromString(const TColumn &col, TStringBuf buf) {
-    TTypeParser type(col.Type);
-    Y_ENSURE(type.GetKind() == TTypeParser::ETypeKind::Optional);
-    type.OpenOptional();
+void TQueryBuilder::BuildType(TTypeParser& typeParser, TTypeBuilder& typeBuilder, const TString& name) {
+    switch (typeParser.GetKind()) {
+    case TTypeParser::ETypeKind::Primitive:
+        typeBuilder.Primitive(typeParser.GetPrimitive());
+        break;
+    case TTypeParser::ETypeKind::Decimal:
+        typeBuilder.Decimal(typeParser.GetDecimal());
+        break;
+    case TTypeParser::ETypeKind::Optional:
+        typeParser.OpenOptional();
+        typeBuilder.BeginOptional();
+        BuildType(typeParser, typeBuilder, name);
+        typeBuilder.EndOptional();
+        typeParser.CloseOptional();
+        break;
+    default:
+        throw yexception() << "Unsupported type kind \"" << typeParser.GetKind() << "\" for column: " << name;
+    }
+}
+
+TType TQueryBuilder::GetType(TTypeParser& typeParser, const TString& name) {
+    TTypeBuilder typeBuilder;
+    BuildType(typeParser, typeBuilder, name);
+    return typeBuilder.Build();
+}
 
-    Value.AddMember(col.Name);
+void TQueryBuilder::CheckNull(const TString& name, TStringBuf buf) {
+    if (buf == "null") {
+        throw yexception() << "Wrong value \"null\" for non-optional column: " << name;
+    }
+}
+
+void TQueryBuilder::AddMemberFromString(TTypeParser& type, const TString& name, TStringBuf buf) {
     switch (type.GetKind()) {
         case TTypeParser::ETypeKind::Primitive:
+            CheckNull(name, buf);
             AddPrimitiveMember(type.GetPrimitive(), buf);
             break;
-        case TTypeParser::ETypeKind::Decimal:
+        case TTypeParser::ETypeKind::Optional:
+            type.OpenOptional();
             if (buf == "null") {
-                Value.EmptyOptional();
+                Value.EmptyOptional(GetType(type, name));
             } else {
                 Value.BeginOptional();
-                Value.Decimal(TDecimalValue(TString(buf), type.GetDecimal().Precision, type.GetDecimal().Scale));
+                AddMemberFromString(type, name, buf);
                 Value.EndOptional();
             }
+            type.CloseOptional();
+            break;
+        case TTypeParser::ETypeKind::Decimal:
+            CheckNull(name, buf);
+            Value.Decimal(TDecimalValue(TString(buf), type.GetDecimal().Precision, type.GetDecimal().Scale));
             break;
         default:
-            Y_FAIL("");
+            throw yexception() << "Unsupported type kind \"" << type.GetKind() << "\" for column: " << name;
     }
 }
 
@@ -224,12 +237,15 @@ void TQueryBuilder::Begin() {
 }
 
 void TQueryBuilder::AddLine(TStringBuf line) {
+    Cerr << line << Endl;
     Value.AddListItem();
     Value.BeginStruct();
     for (const auto& col : Columns) {
         TStringBuf tok = line.NextTok(',');
         Y_ENSURE(tok, "Empty token on line");
-        AddMemberFromString(col, tok);
+        TTypeParser type(col.Type);
+        Value.AddMember(col.Name);
+        AddMemberFromString(type, col.Name, tok);
     }
     Value.EndStruct();
 }
diff --git a/ydb/library/backup/query_builder.h b/ydb/library/backup/query_builder.h
index 559f66ed99..05f57d1d3f 100644
--- a/ydb/library/backup/query_builder.h
+++ b/ydb/library/backup/query_builder.h
@@ -18,8 +18,11 @@ class TQueryBuilder {
     TValueBuilder Value;
 
     TString BuildQuery(const TString& path);
-    void AddMemberFromString(const TColumn& col, TStringBuf ss);
+    void AddMemberFromString(TTypeParser& type, const TString& name, TStringBuf ss);
     void AddPrimitiveMember(EPrimitiveType type, TStringBuf buf);
+    static void CheckNull(const TString& name, TStringBuf buf);
+    static void BuildType(TTypeParser& typeParser, TTypeBuilder& typeBuilder, const TString& name);
+    static TType GetType(TTypeParser& typeParser, const TString& name);
 
 public:
     TQueryBuilder(const TString& path, TVector<TColumn> columns)
@@ -29,7 +32,6 @@ public:
 
     void Begin();
     void AddLine(TStringBuf line);
-    TValue GetLinesAsValue();
     TValue EndAndGetResultingValue();
     TParams EndAndGetResultingParams();
     TString GetQueryString() const;
diff --git a/ydb/tests/functional/ydb_cli/test_ydb_backup.py b/ydb/tests/functional/ydb_cli/test_ydb_backup.py
index 0c016253a2..3601d9c38e 100644
--- a/ydb/tests/functional/ydb_cli/test_ydb_backup.py
+++ b/ydb/tests/functional/ydb_cli/test_ydb_backup.py
@@ -54,12 +54,12 @@ def columns_to_string(columns):
     return list_to_string(columns, lambda col: col.name + ":" + str(col.type.item).strip())
 
 
-def create_table_with_data(session, path):
+def create_table_with_data(session, path, not_null=False):
     path = "/Root/" + path
     session.create_table(
         path,
         ydb.TableDescription()
-        .with_column(ydb.Column("id", ydb.OptionalType(ydb.PrimitiveType.Uint32)))
+        .with_column(ydb.Column("id", ydb.PrimitiveType.Uint32 if not_null else ydb.OptionalType(ydb.PrimitiveType.Uint32)))
         .with_column(ydb.Column("number", ydb.OptionalType(ydb.PrimitiveType.Uint64)))
         .with_column(ydb.Column("string", ydb.OptionalType(ydb.PrimitiveType.String)))
         .with_column(ydb.Column("fixed_point", ydb.OptionalType(ydb.DecimalType())))
@@ -217,6 +217,22 @@ class TestBackupSingle(BaseTestBackupInFiles):
         )
 
 
+class TestBackupSingleNotNull(BaseTestBackupInFiles):
+    def test_single_table_backup(self):
+        session = self.driver.table_client.session().create()
+        # Create table
+        path = "table"
+        create_table_with_data(session, path, True)
+
+        # Backup table
+        self.create_backup(path, [path], False)
+
+        assert_that(
+            [child.name for child in self.driver.scheme_client.list_directory("/Root").children],
+            is_(["table", ".sys"])
+        )
+
+
 class TestBaseSingleFromDifPlaces(BaseTestBackupInFiles):
     def test_single_table_backup_from_different_places(self):
         session = self.driver.table_client.session().create()
@@ -375,20 +391,23 @@ class TestRecursiveConsistent(BaseTestBackupInFiles):
 
 class TestSingleBackupRestore(BaseTestBackupInFiles):
     def test_single_table_with_data_backup_restore(self):
-        self.test_single_table_with_data_backup_restore_impl(False)
-        self.test_single_table_with_data_backup_restore_impl(True)
+        self.test_single_table_with_data_backup_restore_impl(False, False)
+        self.test_single_table_with_data_backup_restore_impl(False, True)
+        self.test_single_table_with_data_backup_restore_impl(True, False)
+        self.test_single_table_with_data_backup_restore_impl(True, True)
 
     @classmethod
-    def test_single_table_with_data_backup_restore_impl(self, use_bulk_upsert):
+    def test_single_table_with_data_backup_restore_impl(self, use_bulk_upsert, not_null):
         self.driver.scheme_client.make_directory(
             '/Root/folder'
         )
         postfix = '_bulk_upsert' if use_bulk_upsert else ''
+        postfix += '_not_null' if not_null else ''
 
         session = self.driver.table_client.session().create()
 
         # Create table and fill with data
-        create_table_with_data(session, "folder/table")
+        create_table_with_data(session, "folder/table", not_null)
 
         # Backup table
         backup_files_dir = output_path(self.test_name, 'test_single_table_with_data_backup_restore' + postfix, "backup_files_dir")