Merge 82d5b85 into 7c5a9d4

shnikd · web-flow · commit f90c9a32f674 · 2024-03-25T12:43:00.000+03:00
diff --git a/ydb/core/protos/flat_scheme_op.proto b/ydb/core/protos/flat_scheme_op.proto
@@ -1293,6 +1293,11 @@ message TMoveIndex {
     optional bool AllowOverwrite = 4;
 }
 
+message TSetVal {
+    optional sint64 NextValue = 1;
+    optional bool NextUsed = 2;
+}
+
 message TSequenceDescription {
     optional string Name = 1; // mandatory
     optional NKikimrProto.TPathID PathId = 2; // sequence path id, assigned by schemeshard
@@ -1305,6 +1310,7 @@ message TSequenceDescription {
     optional uint64 Cache = 8; // number of items to cache, defaults to 1
     optional sint64 Increment = 9; // increment at each call, defaults to 1
     optional bool Cycle = 10; // true when cycle on overflow is allowed
+    optional TSetVal SetVal = 11; // SetVal(NextValue, NextUsed) is executed atomically when creating
 }
 
 message TSequenceSharding {
diff --git a/ydb/core/protos/tx_sequenceshard.proto b/ydb/core/protos/tx_sequenceshard.proto
@@ -17,6 +17,11 @@ message TEvMarkSchemeShardPipe {
     uint64 Round = 3;
 }
 
+message TSetVal {
+    sint64 NextValue = 1;
+    bool NextUsed = 2;
+}
+
 message TEvCreateSequence {
     NKikimrProto.TPathID PathId = 1;
     uint64 TxId = 2;
@@ -40,6 +45,9 @@ message TEvCreateSequence {
         bool Cycle = 9;
     }
     bool Frozen = 10; // defaults to false
+    oneof OptionalSetVal {
+        TSetVal SetVal = 11;
+    }
 }
 
 message TEvCreateSequenceResult {
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_consistent_copy_tables.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_consistent_copy_tables.cpp
@@ -137,10 +137,6 @@ TVector<ISubOperation::TPtr> CreateConsistentCopyTables(TOperationId nextId, con
         result.push_back(CreateCopyTable(NextPartId(nextId, result),
             CopyTableTask(srcPath, dstPath, descr.GetOmitFollowers(), descr.GetIsBackup()), sequences));
 
-        if (descr.GetOmitIndexes()) {
-            continue;
-        }
-
         TVector<NKikimrSchemeOp::TSequenceDescription> sequenceDescriptions;
         for (const auto& child: srcPath.Base()->GetChildren()) {
             const auto& name = child.first;
@@ -160,6 +156,10 @@ TVector<ISubOperation::TPtr> CreateConsistentCopyTables(TOperationId nextId, con
                 continue;
             }
 
+            if (descr.GetOmitIndexes()) {
+                continue;
+            }
+
             if (!srcIndexPath.IsTableIndex()) {
                 continue;
             }
@@ -185,9 +185,11 @@ TVector<ISubOperation::TPtr> CreateConsistentCopyTables(TOperationId nextId, con
                 NKikimrSchemeOp::EOperationType::ESchemeOpCreateSequence);
             scheme.SetFailOnExist(true);
 
+            auto* copySequence = scheme.MutableCopySequence();
+            copySequence->SetCopyFrom(srcPath.PathString() + "/" + sequenceDescription.GetName());
             *scheme.MutableSequence() = std::move(sequenceDescription);
 
-            result.push_back(CreateNewSequence(NextPartId(nextId, result), scheme));
+            result.push_back(CreateCopySequence(NextPartId(nextId, result), scheme));
         }
     }
 
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_copy_sequence.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_copy_sequence.cpp
@@ -115,7 +115,7 @@ class TConfigureParts : public TSubOperationState {
             event->Record.SetFrozen(true);
 
             LOG_DEBUG_S(context.Ctx, NKikimrServices::FLAT_TX_SCHEMESHARD,
-                        "TCoptSequence TConfigureParts ProgressState"
+                        "TCopySequence TConfigureParts ProgressState"
                         << " sending TEvCreateSequence to tablet " << tabletId
                         << " operationId# " << OperationId
                         << " at tablet " << ssId);
@@ -274,6 +274,18 @@ class TProposedCopySequence : public TSubOperationState {
                 << " operationId#" << OperationId;
     }
 
+    void UpdateSequenceDescription(NKikimrSchemeOp::TSequenceDescription& descr) {
+        descr.SetStartValue(GetSequenceResult.GetStartValue());
+        descr.SetMinValue(GetSequenceResult.GetMinValue());
+        descr.SetMaxValue(GetSequenceResult.GetMaxValue());
+        descr.SetCache(GetSequenceResult.GetCache());
+        descr.SetIncrement(GetSequenceResult.GetIncrement());
+        descr.SetCycle(GetSequenceResult.GetCycle());
+        auto* setValMsg = descr.MutableSetVal();
+        setValMsg->SetNextValue(GetSequenceResult.GetNextValue());
+        setValMsg->SetNextUsed(GetSequenceResult.GetNextUsed());
+    }
+
 public:
     TProposedCopySequence(TOperationId id)
         : OperationId(id)
@@ -333,7 +345,15 @@ class TProposedCopySequence : public TSubOperationState {
             return false;
         }
 
+        TPathId pathId = txState->TargetPathId;
+
         NIceDb::TNiceDb db(context.GetDB());
+
+        auto sequenceInfo = context.SS->Sequences.at(pathId);
+        UpdateSequenceDescription(sequenceInfo->Description);
+
+        context.SS->PersistSequence(db, pathId, *sequenceInfo);
+
         context.SS->ChangeTxState(db, OperationId, TTxState::Done);
         context.OnComplete.ActivateTx(OperationId);
         return true;
@@ -387,7 +407,7 @@ class TProposedCopySequence : public TSubOperationState {
             return false;
         }
 
-        auto getSequenceResult = ev->Get()->Record;
+        GetSequenceResult = ev->Get()->Record;
 
         Y_ABORT_UNLESS(txState->Shards.size() == 1);
         for (auto shard : txState->Shards) {
@@ -397,7 +417,8 @@ class TProposedCopySequence : public TSubOperationState {
             Y_ABORT_UNLESS(currentTabletId != InvalidTabletId);
 
             auto event = MakeHolder<NSequenceShard::TEvSequenceShard::TEvRestoreSequence>(
-                txState->TargetPathId, getSequenceResult);
+                txState->TargetPathId, GetSequenceResult);
+
             event->Record.SetTxId(ui64(OperationId.GetTxId()));
             event->Record.SetTxPartId(OperationId.GetSubTxId());
 
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_create_sequence.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_create_sequence.cpp
@@ -142,6 +142,10 @@ class TConfigureParts : public TSubOperationState {
             if (alterData->Description.HasCycle()) {
                 event->Record.SetCycle(alterData->Description.GetCycle());
             }
+            if (alterData->Description.HasSetVal()) {
+                event->Record.MutableSetVal()->SetNextValue(alterData->Description.GetSetVal().GetNextValue());
+                event->Record.MutableSetVal()->SetNextUsed(alterData->Description.GetSetVal().GetNextUsed());
+            }
 
             LOG_DEBUG_S(context.Ctx, NKikimrServices::FLAT_TX_SCHEMESHARD,
                         "TCreateSequence TConfigureParts ProgressState"
diff --git a/ydb/core/tx/schemeshard/schemeshard_export_flow_proposals.cpp b/ydb/core/tx/schemeshard/schemeshard_export_flow_proposals.cpp
@@ -106,8 +106,11 @@ THolder<TEvSchemeShard::TEvModifySchemeTransaction> BackupPropose(
     task.SetNeedToBill(!exportInfo->UserSID || !ss->SystemBackupSIDs.contains(*exportInfo->UserSID));
 
     const TPath sourcePath = TPath::Init(exportInfo->Items[itemIdx].SourcePathId, ss);
-    if (sourcePath.IsResolved()) {
-        task.MutableTable()->CopyFrom(GetTableDescription(ss, sourcePath.Base()->PathId));
+    const TPath exportPathItem = exportPath.Child(ToString(itemIdx));
+    if (sourcePath.IsResolved() && exportPathItem.IsResolved()) {
+        auto exportDescription = GetTableDescription(ss, exportPathItem.Base()->PathId);
+        exportDescription.MutableTable()->SetName(sourcePath.LeafName());
+        task.MutableTable()->CopyFrom(exportDescription);
     }
 
     task.SetSnapshotStep(exportInfo->SnapshotStep);
diff --git a/ydb/core/tx/schemeshard/schemeshard_import_flow_proposals.cpp b/ydb/core/tx/schemeshard/schemeshard_import_flow_proposals.cpp
@@ -72,6 +72,11 @@ THolder<TEvSchemeShard::TEvModifySchemeTransaction> CreateTablePropose(
                 if (fromSequence.has_cycle()) {
                     seqDesc->SetCycle(fromSequence.cycle());
                 }
+                if (fromSequence.has_set_val()) {
+                    auto* setVal = seqDesc->MutableSetVal();
+                    setVal->SetNextUsed(fromSequence.set_val().next_used());
+                    setVal->SetNextValue(fromSequence.set_val().next_value());
+                }
 
                 break;
             }
diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp b/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp
@@ -2362,16 +2362,17 @@ namespace NSchemeShardUT_Private {
         runtime.Send(new IEventHandle(NSequenceProxy::MakeSequenceProxyServiceID(), sender, request.Release()));
     }
 
-    i64 WaitNextValResult(TTestActorRuntime& runtime, const TActorId& sender) {
+    i64 WaitNextValResult(
+            TTestActorRuntime& runtime, const TActorId& sender, Ydb::StatusIds::StatusCode expectedStatus) {
         auto ev = runtime.GrabEdgeEventRethrow<NSequenceProxy::TEvSequenceProxy::TEvNextValResult>(sender);
         auto* msg = ev->Get();
-        UNIT_ASSERT_VALUES_EQUAL(msg->Status, Ydb::StatusIds::SUCCESS);
+        UNIT_ASSERT_VALUES_EQUAL(msg->Status, expectedStatus);
         return msg->Value;
     }
 
-    i64 DoNextVal(TTestActorRuntime& runtime, const TString& path) {
+    i64 DoNextVal(TTestActorRuntime& runtime, const TString& path, Ydb::StatusIds::StatusCode expectedStatus) {
         auto sender = runtime.AllocateEdgeActor(0);
         SendNextValRequest(runtime, sender, path);
-        return WaitNextValResult(runtime, sender);
+        return WaitNextValResult(runtime, sender, expectedStatus);
     }
 }
diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.h b/ydb/core/tx/schemeshard/ut_helpers/helpers.h
@@ -553,7 +553,11 @@ namespace NSchemeShardUT_Private {
     void WriteRow(TTestActorRuntime& runtime, const ui64 txId, const TString& tablePath, int partitionIdx, const ui32 key, const TString& value, bool successIsExpected = true);
 
     void SendNextValRequest(TTestActorRuntime& runtime, const TActorId& sender, const TString& path);
-    i64 WaitNextValResult(TTestActorRuntime& runtime, const TActorId& sender);
-    i64 DoNextVal(TTestActorRuntime& runtime, const TString& path);
+    i64 WaitNextValResult(
+        TTestActorRuntime& runtime, const TActorId& sender,
+        Ydb::StatusIds::StatusCode expectedStatus = Ydb::StatusIds::SUCCESS);
+    i64 DoNextVal(
+        TTestActorRuntime& runtime, const TString& path,
+        Ydb::StatusIds::StatusCode expectedStatus = Ydb::StatusIds::SUCCESS);
 
 } //NSchemeShardUT_Private
diff --git a/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp b/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp
@@ -955,6 +955,153 @@ value {
         UNIT_ASSERT_C(CheckDefaultFromSequence(table), "Invalid default value");
     }
 
+    Y_UNIT_TEST(ShouldRestoreSequence) {
+        TPortManager portManager;
+        const ui16 port = portManager.GetPort();
+
+        TS3Mock s3Mock({}, TS3Mock::TSettings(port));
+        UNIT_ASSERT(s3Mock.Start());
+
+        TTestBasicRuntime runtime;
+        TTestEnv env(runtime);
+
+        ui64 txId = 100;
+
+        runtime.SetLogPriority(NKikimrServices::DATASHARD_BACKUP, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::DATASHARD_RESTORE, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::EXPORT, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::IMPORT, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::SEQUENCEPROXY, NActors::NLog::PRI_TRACE);
+
+        TestCreateIndexedTable(runtime, ++txId, "/MyRoot", R"(
+            TableDescription {
+                Name: "Original"
+                Columns { Name: "key" Type: "Uint64" DefaultFromSequence: "myseq" }
+                Columns { Name: "value" Type: "Uint64" }
+                KeyColumnNames: ["key"]
+            }
+            SequenceDescription {
+                Name: "myseq"
+            }
+        )");
+        env.TestWaitNotification(runtime, txId);
+
+        i64 value = DoNextVal(runtime, "/MyRoot/Original/myseq");
+        UNIT_ASSERT_VALUES_EQUAL(value, 1);
+
+        TestExport(runtime, ++txId, "/MyRoot", Sprintf(R"(
+            ExportToS3Settings {
+              endpoint: "localhost:%d"
+              scheme: HTTP
+              items {
+                source_path: "/MyRoot/Original"
+                destination_prefix: ""
+              }
+            }
+        )", port));
+        env.TestWaitNotification(runtime, txId);
+        TestGetExport(runtime, txId, "/MyRoot");
+
+        TestImport(runtime, ++txId, "/MyRoot", Sprintf(R"(
+            ImportFromS3Settings {
+              endpoint: "localhost:%d"
+              scheme: HTTP
+              items {
+                source_prefix: ""
+                destination_path: "/MyRoot/Restored"
+              }
+            }
+        )", port));
+        env.TestWaitNotification(runtime, txId);
+        TestGetImport(runtime, txId, "/MyRoot");
+
+        const auto desc = DescribePath(runtime, "/MyRoot/Restored", true, true);
+        UNIT_ASSERT_VALUES_EQUAL(desc.GetStatus(), NKikimrScheme::StatusSuccess);
+
+        const auto& table = desc.GetPathDescription().GetTable();
+
+        value = DoNextVal(runtime, "/MyRoot/Restored/myseq");
+        UNIT_ASSERT_VALUES_EQUAL(value, 2);
+
+        UNIT_ASSERT_C(CheckDefaultFromSequence(table), "Invalid default value");
+    }
+
+    Y_UNIT_TEST(ShouldRestoreSequenceWithOverflow) {
+        TPortManager portManager;
+        const ui16 port = portManager.GetPort();
+
+        TS3Mock s3Mock({}, TS3Mock::TSettings(port));
+        UNIT_ASSERT(s3Mock.Start());
+
+        TTestBasicRuntime runtime;
+        TTestEnv env(runtime);
+
+        ui64 txId = 100;
+
+        runtime.SetLogPriority(NKikimrServices::DATASHARD_BACKUP, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::DATASHARD_RESTORE, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::EXPORT, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::IMPORT, NActors::NLog::PRI_TRACE);
+        runtime.SetLogPriority(NKikimrServices::SEQUENCEPROXY, NActors::NLog::PRI_TRACE);
+
+        TestCreateIndexedTable(runtime, ++txId, "/MyRoot", R"(
+            TableDescription {
+                Name: "Original"
+                Columns { Name: "key" Type: "Uint64" DefaultFromSequence: "myseq" }
+                Columns { Name: "value" Type: "Uint64" }
+                KeyColumnNames: ["key"]
+            }
+            SequenceDescription {
+                Name: "myseq"
+                MinValue: 1
+                MaxValue: 2
+            }
+        )");
+        env.TestWaitNotification(runtime, txId);
+
+        i64 value = DoNextVal(runtime, "/MyRoot/Original/myseq");
+        UNIT_ASSERT_VALUES_EQUAL(value, 1);
+
+        value = DoNextVal(runtime, "/MyRoot/Original/myseq");
+        UNIT_ASSERT_VALUES_EQUAL(value, 2);
+
+        TestExport(runtime, ++txId, "/MyRoot", Sprintf(R"(
+            ExportToS3Settings {
+              endpoint: "localhost:%d"
+              scheme: HTTP
+              items {
+                source_path: "/MyRoot/Original"
+                destination_prefix: ""
+              }
+            }
+        )", port));
+        env.TestWaitNotification(runtime, txId);
+        TestGetExport(runtime, txId, "/MyRoot");
+
+        TestImport(runtime, ++txId, "/MyRoot", Sprintf(R"(
+            ImportFromS3Settings {
+              endpoint: "localhost:%d"
+              scheme: HTTP
+              items {
+                source_prefix: ""
+                destination_path: "/MyRoot/Restored"
+              }
+            }
+        )", port));
+        env.TestWaitNotification(runtime, txId);
+        TestGetImport(runtime, txId, "/MyRoot");
+
+        const auto desc = DescribePath(runtime, "/MyRoot/Restored", true, true);
+        UNIT_ASSERT_VALUES_EQUAL(desc.GetStatus(), NKikimrScheme::StatusSuccess);
+
+        const auto& table = desc.GetPathDescription().GetTable();
+
+        value = DoNextVal(runtime, "/MyRoot/Restored/myseq", Ydb::StatusIds::SCHEME_ERROR);
+        UNIT_ASSERT_VALUES_EQUAL(value, 2);
+
+        UNIT_ASSERT_C(CheckDefaultFromSequence(table), "Invalid default value");
+    }
+
     Y_UNIT_TEST(ExportImportPg) {
         TTestBasicRuntime runtime;
         TTestEnv env(runtime, TTestEnvOptions().EnableTablePgTypes(true));
diff --git a/ydb/core/tx/schemeshard/ut_restore/ya.make b/ydb/core/tx/schemeshard/ut_restore/ya.make
@@ -21,7 +21,7 @@ PEERDIR(
     ydb/core/wrappers/ut_helpers
     ydb/core/ydb_convert
     ydb/library/yql/sql/pg
-    ydb/library/yql/parser/pg_wrapper    
+    ydb/library/yql/parser/pg_wrapper
 )
 
 SRCS(
diff --git a/ydb/core/tx/sequenceshard/tx_create_sequence.cpp b/ydb/core/tx/sequenceshard/tx_create_sequence.cpp
@@ -74,8 +74,15 @@ namespace NSequenceShard {
                 sequence.Cycle = msg->Record.GetCycle();
             }
 
-            sequence.NextValue = sequence.StartValue;
-            sequence.NextUsed = false;
+
+            if (msg->Record.OptionalSetVal_case() == NKikimrTxSequenceShard::TEvCreateSequence::kSetVal) {
+                sequence.NextValue = msg->Record.GetSetVal().GetNextValue();
+                sequence.NextUsed = msg->Record.GetSetVal().GetNextUsed();
+            } else {
+                sequence.NextUsed = false;
+                sequence.NextValue = sequence.StartValue;
+            }
+
             if (msg->Record.OptionalCache_case() == NKikimrTxSequenceShard::TEvCreateSequence::kCache) {
                 sequence.Cache = msg->Record.GetCache();
                 if (sequence.Cache < 1) {
diff --git a/ydb/core/ydb_convert/table_description.cpp b/ydb/core/ydb_convert/table_description.cpp
diff --git a/ydb/public/api/protos/ydb_table.proto b/ydb/public/api/protos/ydb_table.proto

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,11 @@ message TEvMarkSchemeShardPipe {`
`17`	`17`	`uint64 Round = 3;`
`18`	`18`	`}`
`19`	`19`
	`20`	`+message TSetVal {`
	`21`	`+ sint64 NextValue = 1;`
	`22`	`+ bool NextUsed = 2;`
	`23`	`+}`
	`24`	`+`
`20`	`25`	`message TEvCreateSequence {`
`21`	`26`	`NKikimrProto.TPathID PathId = 1;`
`22`	`27`	`uint64 TxId = 2;`
`@@ -40,6 +45,9 @@ message TEvCreateSequence {`
`40`	`45`	`bool Cycle = 9;`
`41`	`46`	`}`
`42`	`47`	`bool Frozen = 10; // defaults to false`
	`48`	`+ oneof OptionalSetVal {`
	`49`	`+ TSetVal SetVal = 11;`
	`50`	`+ }`
`43`	`51`	`}`
`44`	`52`
`45`	`53`	`message TEvCreateSequenceResult {`
Original file line number	Diff line number	Diff line change
`@@ -72,6 +72,11 @@ THolder<TEvSchemeShard::TEvModifySchemeTransaction> CreateTablePropose(`
`72`	`72`	`if (fromSequence.has_cycle()) {`
`73`	`73`	`seqDesc->SetCycle(fromSequence.cycle());`
`74`	`74`	`}`
	`75`	`+ if (fromSequence.has_set_val()) {`
	`76`	`+ auto* setVal = seqDesc->MutableSetVal();`
	`77`	`+ setVal->SetNextUsed(fromSequence.set_val().next_used());`
	`78`	`+ setVal->SetNextValue(fromSequence.set_val().next_value());`
	`79`	`+ }`
`75`	`80`
`76`	`81`	`break;`
`77`	`82`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2362,16 +2362,17 @@ namespace NSchemeShardUT_Private {`
`2362`	`2362`	`runtime.Send(new IEventHandle(NSequenceProxy::MakeSequenceProxyServiceID(), sender, request.Release()));`
`2363`	`2363`	`}`
`2364`	`2364`
`2365`		`- i64 WaitNextValResult(TTestActorRuntime& runtime, const TActorId& sender) {`
	`2365`	`+ i64 WaitNextValResult(`
	`2366`	`+ TTestActorRuntime& runtime, const TActorId& sender, Ydb::StatusIds::StatusCode expectedStatus) {`
`2366`	`2367`	`auto ev = runtime.GrabEdgeEventRethrow<NSequenceProxy::TEvSequenceProxy::TEvNextValResult>(sender);`
`2367`	`2368`	`auto* msg = ev->Get();`
`2368`		`- UNIT_ASSERT_VALUES_EQUAL(msg->Status, Ydb::StatusIds::SUCCESS);`
	`2369`	`+ UNIT_ASSERT_VALUES_EQUAL(msg->Status, expectedStatus);`
`2369`	`2370`	`return msg->Value;`
`2370`	`2371`	`}`
`2371`	`2372`
`2372`		`- i64 DoNextVal(TTestActorRuntime& runtime, const TString& path) {`
	`2373`	`+ i64 DoNextVal(TTestActorRuntime& runtime, const TString& path, Ydb::StatusIds::StatusCode expectedStatus) {`
`2373`	`2374`	`auto sender = runtime.AllocateEdgeActor(0);`
`2374`	`2375`	`SendNextValRequest(runtime, sender, path);`
`2375`		`- return WaitNextValResult(runtime, sender);`
	`2376`	`+ return WaitNextValResult(runtime, sender, expectedStatus);`
`2376`	`2377`	`}`
`2377`	`2378`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ PEERDIR(`
`21`	`21`	`ydb/core/wrappers/ut_helpers`
`22`	`22`	`ydb/core/ydb_convert`
`23`	`23`	`ydb/library/yql/sql/pg`
`24`		`- ydb/library/yql/parser/pg_wrapper`
	`24`	`+ ydb/library/yql/parser/pg_wrapper`
`25`	`25`	`)`
`26`	`26`
`27`	`27`	`SRCS(`