fix(snapshot): reject RESTORE TABLE of referenced tables

pkg/frontend/snapshot.go

@@ -826,6 +826,12 @@ func doRestoreSnapshot(ctx context.Context, ses *Session, stmt *tree.RestoreSnap
 		return
 	}
 
+	if stmt.Level == tree.RESTORELEVELTABLE {
+		if err = validateRestoreTableTarget(ctx, ses.GetService(), bh, snapshotName, dbName, tblName, toAccountId); err != nil {
+			return stats, err
+		}
+	}
+
 	// drop foreign key related tables first
 	if err = deleteCurFkTables(ctx, ses.GetService(), bh, dbName, tblName, toAccountId); err != nil {
 		return
@@ -918,7 +924,11 @@ func doRestoreSnapshot(ctx context.Context, ses *Session, stmt *tree.RestoreSnap
 	}
 
 	if len(fkTableMap) > 0 {
-		if err = restoreTablesWithFk(ctx, ses.GetService(), bh, snapshotName, sortedFkTbls, fkTableMap, toAccountId, snapshot.ts); err != nil {
+		if err = restoreTablesWithFk(
+			ctx, ses.GetService(), bh, snapshotName, sortedFkTbls,
+			fkTableMap, toAccountId, snapshot.ts,
+			stmt.Level == tree.RESTORELEVELTABLE,
+		); err != nil {
 			return
 		}
 	}
@@ -1443,7 +1453,7 @@ func restoreToDatabaseOrTable(
 			return
 		}
 
-		if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs); err != nil {
+		if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs, restoreToTbl); err != nil {
 			return
 		}
 	}
@@ -1495,7 +1505,7 @@ func restoreSystemDatabase(
 			return
 		}
 
-		if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs); err != nil {
+		if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs, false); err != nil {
 			return
 		}
 	}
@@ -1572,7 +1582,8 @@ func restoreTablesWithFk(
 	sortedFkTbls []string,
 	fkTableMap map[string]*tableInfo,
 	toAccountId uint32,
-	snapshotTs int64) (err error) {
+	snapshotTs int64,
+	rejectMasterTable bool) (err error) {
 	getLogger(sid).Debug(fmt.Sprintf("[%s] start to drop fk related tables", snapshotName))
 
 	// recreate tables as topo order
@@ -1581,7 +1592,7 @@ func restoreTablesWithFk(
 		// e.g. t1.pk <- t2.fk, we only want to restore t2, fkTableMap[t1.key] is nil, ignore t1
 		if tblInfo := fkTableMap[key]; tblInfo != nil {
 			getLogger(sid).Debug(fmt.Sprintf("[%s] start to restore table with fk: %v, restore timestamp: %d", snapshotName, tblInfo.tblName, snapshotTs))
-			if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs); err != nil {
+			if err = recreateTable(ctx, sid, bh, snapshotName, tblInfo, toAccountId, snapshotTs, rejectMasterTable); err != nil {
 				return
 			}
 		}
@@ -1846,6 +1857,7 @@ func recreateTable(
 	tblInfo *tableInfo,
 	toAccountId uint32,
 	snapshotTs int64,
+	rejectMasterTable bool,
 ) (err error) {
 	if isExternalTable(tblInfo) {
 		return newExternalTableRestoreError(ctx, tblInfo, "snapshot")
@@ -1893,9 +1905,14 @@ func recreateTable(
 
 	ctx = defines.AttachAccountId(ctx, toAccountId)
 
-	var isMasterTable bool
-	isMasterTable, err = checkTableIsMaster(ctx, sid, bh, snapshotName, tblInfo.dbName, tblInfo.tblName)
+	isMasterTable, err := checkTableIsMaster(ctx, sid, bh, snapshotName, tblInfo.dbName, tblInfo.tblName)
+	if err != nil {
+		return err
+	}
 	if isMasterTable {
+		if rejectMasterTable {
+			return newRestoreTableForeignKeyError(ctx, tblInfo.dbName, tblInfo.tblName)
+		}
 		// skip restore the table which is master table
 		getLogger(sid).Debug(fmt.Sprintf("[%s] skip restore master table: %v.%v", snapshotName, tblInfo.dbName, tblInfo.tblName))
 		return
@@ -1986,6 +2003,30 @@ func shouldSkipRestoreTableInBulk(tblInfo *tableInfo) bool {
 	return isExternalTable(tblInfo)
 }
 
+func validateRestoreTableTarget(
+	ctx context.Context,
+	sid string,
+	bh BackgroundExec,
+	snapshotName string,
+	dbName string,
+	tblName string,
+	toAccountId uint32,
+) error {
+	toCtx := defines.AttachAccountId(ctx, toAccountId)
+	isMasterTable, err := checkTableIsMaster(toCtx, sid, bh, snapshotName, dbName, tblName)
+	if err != nil {
+		return err
+	}
+	if isMasterTable {
+		return newRestoreTableForeignKeyError(ctx, dbName, tblName)
+	}
+	return nil
+}
+
+func newRestoreTableForeignKeyError(ctx context.Context, dbName, tblName string) error {
+	return moerr.NewNotSupportedf(ctx, "can not restore table '%s.%s' referenced by some foreign key constraint", dbName, tblName)
+}
+
 func newExternalTableRestoreError(ctx context.Context, tblInfo *tableInfo, source string) error {
 	if tblInfo == nil {
 		return moerr.NewInternalError(ctx, "external table cannot be restored")

pkg/frontend/snapshot_test.go

@@ -406,7 +406,7 @@ func TestSequenceRestoreEntryPoints(t *testing.T) {
 			currentKind: catalog.SystemOrdinaryRel,
 			wantDropSQL: "drop table if exists `db1`.`seq1`",
 			run: func(bh BackgroundExec) error {
-				return recreateTable(accountCtx, "", bh, "snapshot1", sequence, 10, snapshotTS)
+				return recreateTable(accountCtx, "", bh, "snapshot1", sequence, 10, snapshotTS, false)
 			},
 		},
 		{
@@ -451,7 +451,7 @@ func TestSequenceRestoreEntryPoints(t *testing.T) {
 		bh := &backgroundExecTest{}
 		bh.init()
 
-		err := recreateTable(context.Background(), "", bh, "snapshot1", sequence, 10, snapshotTS)
+		err := recreateTable(context.Background(), "", bh, "snapshot1", sequence, 10, snapshotTS, false)
 		require.Error(t, err)
 		require.Empty(t, bh.executedSQLs)
 
@@ -959,7 +959,7 @@ func TestRestoreExternalTableDefensiveCloneGuards(t *testing.T) {
 
 		bh := &backgroundExecTest{}
 		bh.init()
-		err := recreateTable(ctx, "", bh, "sp_ext", tblInfo, uint32(sysAccountID), 100)
+		err := recreateTable(ctx, "", bh, "sp_ext", tblInfo, uint32(sysAccountID), 100, false)
 		convey.So(err, convey.ShouldNotBeNil)
 		convey.So(err.Error(), convey.ShouldContainSubstring, "external table db1.hive_ext cannot be restored from snapshot")
 		convey.So(len(bh.executedSQLs), convey.ShouldEqual, 0)
@@ -980,6 +980,196 @@ func TestRestoreExternalTableDefensiveCloneGuards(t *testing.T) {
 	})
 }
 
+func TestRecreateTableReferencedByForeignKey(t *testing.T) {
+	ctx := defines.AttachAccountId(context.Background(), uint32(sysAccountID))
+	tblInfo := &tableInfo{
+		dbName:    "db1",
+		tblName:   "parent",
+		relKind:   catalog.SystemOrdinaryRel,
+		createSql: "create table `db1`.`parent` (`id` int primary key)",
+	}
+	masterSQL := fmt.Sprintf(
+		checkTableIsMasterFormat,
+		quoteSQLStringLiteral(tblInfo.dbName),
+		quoteSQLStringLiteral(tblInfo.tblName),
+	)
+
+	for _, tc := range []struct {
+		name              string
+		rejectMasterTable bool
+		wantErr           bool
+	}{
+		{name: "explicit table restore rejects the referenced table", rejectMasterTable: true, wantErr: true},
+		{name: "bulk restore keeps skipping the referenced table", rejectMasterTable: false, wantErr: false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			bh := &backgroundExecTest{}
+			bh.init()
+			bh.sql2result[masterSQL] = newMrsForRestoreStringRows(
+				[]string{"db_name"},
+				[][]interface{}{{"db1"}},
+			)
+
+			err := recreateTable(ctx, "", bh, "snapshot", tblInfo, uint32(sysAccountID), 100, tc.rejectMasterTable)
+			if tc.wantErr {
+				require.EqualError(t, err, "not supported: can not restore table 'db1.parent' referenced by some foreign key constraint")
+			} else {
+				require.NoError(t, err)
+			}
+			require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+		})
+	}
+
+	t.Run("propagates foreign key lookup errors", func(t *testing.T) {
+		bh := &backgroundExecTest{}
+		bh.init()
+		wantErr := errors.New("check foreign key failed")
+		bh.sql2err[masterSQL] = wantErr
+
+		err := recreateTable(ctx, "", bh, "snapshot", tblInfo, uint32(sysAccountID), 100, false)
+		require.ErrorIs(t, err, wantErr)
+		require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+	})
+}
+
+func TestValidateRestoreTableTarget(t *testing.T) {
+	ctx := context.Background()
+	masterSQL := fmt.Sprintf(
+		checkTableIsMasterFormat,
+		quoteSQLStringLiteral("db1"),
+		quoteSQLStringLiteral("parent"),
+	)
+
+	t.Run("rejects referenced table", func(t *testing.T) {
+		bh := &backgroundExecTest{}
+		bh.init()
+		bh.sql2result[masterSQL] = newMrsForRestoreStringRows(
+			[]string{"db_name"},
+			[][]interface{}{{"db1"}},
+		)
+
+		err := validateRestoreTableTarget(ctx, "", bh, "snapshot", "db1", "parent", uint32(sysAccountID))
+		require.EqualError(t, err, "not supported: can not restore table 'db1.parent' referenced by some foreign key constraint")
+		require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+		require.Equal(t, []uint32{uint32(sysAccountID)}, bh.executionAccountIDs)
+	})
+
+	t.Run("allows table without foreign key dependents", func(t *testing.T) {
+		bh := &backgroundExecTest{}
+		bh.init()
+		bh.sql2result[masterSQL] = newMrsForRestoreStringRows([]string{"db_name"}, nil)
+
+		err := validateRestoreTableTarget(ctx, "", bh, "snapshot", "db1", "parent", uint32(sysAccountID))
+		require.NoError(t, err)
+		require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+		require.Equal(t, []uint32{uint32(sysAccountID)}, bh.executionAccountIDs)
+	})
+
+	t.Run("propagates foreign key lookup errors", func(t *testing.T) {
+		bh := &backgroundExecTest{}
+		bh.init()
+		wantErr := errors.New("check foreign key failed")
+		bh.sql2err[masterSQL] = wantErr
+
+		err := validateRestoreTableTarget(ctx, "", bh, "snapshot", "db1", "parent", uint32(sysAccountID))
+		require.ErrorIs(t, err, wantErr)
+		require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+		require.Equal(t, []uint32{uint32(sysAccountID)}, bh.executionAccountIDs)
+	})
+}
+
+func TestRestoreTablesWithFkRejectsReferencedTable(t *testing.T) {
+	ctx := defines.AttachAccountId(context.Background(), uint32(sysAccountID))
+	tblInfo := &tableInfo{
+		dbName:  "db1",
+		tblName: "parent",
+		relKind: catalog.SystemOrdinaryRel,
+	}
+	key := genKey(tblInfo.dbName, tblInfo.tblName)
+	masterSQL := fmt.Sprintf(
+		checkTableIsMasterFormat,
+		quoteSQLStringLiteral(tblInfo.dbName),
+		quoteSQLStringLiteral(tblInfo.tblName),
+	)
+
+	bh := &backgroundExecTest{}
+	bh.init()
+	bh.sql2result[masterSQL] = newMrsForRestoreStringRows(
+		[]string{"db_name"},
+		[][]interface{}{{tblInfo.dbName}},
+	)
+
+	err := restoreTablesWithFk(
+		ctx,
+		"",
+		bh,
+		"snapshot",
+		[]string{key},
+		map[string]*tableInfo{key: tblInfo},
+		uint32(sysAccountID),
+		100,
+		true,
+	)
+	require.EqualError(t, err, "not supported: can not restore table 'db1.parent' referenced by some foreign key constraint")
+	require.Equal(t, []string{masterSQL}, bh.executedSQLs)
+}
+
+func TestRestoreTableRejectsReferencedTableBeforeMutation(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	ses := newTestSession(t, ctrl)
+	t.Cleanup(ses.Close)
+
+	bh := &backgroundExecTest{}
+	bh.init()
+	const (
+		snapshotName = "snapshot"
+		dbName       = "db1"
+		tblName      = "parent"
+	)
+	snapshotSQL := fmt.Sprintf("%s where sname = '%s'", getSnapshotFormat, snapshotName)
+	masterSQL := fmt.Sprintf(
+		checkTableIsMasterFormat,
+		quoteSQLStringLiteral(dbName),
+		quoteSQLStringLiteral(tblName),
+	)
+	bh.sql2result[snapshotSQL] = newMrsForSnapshotRecord(
+		"snapshot-id",
+		snapshotName,
+		100,
+		tree.SNAPSHOTLEVELTABLE.String(),
+		sysAccountName,
+		dbName,
+		tblName,
+		0,
+	)
+	bh.sql2result[masterSQL] = newMrsForRestoreStringRows(
+		[]string{"db_name"},
+		[][]interface{}{{dbName}},
+	)
+
+	oldNewBackgroundExec := NewBackgroundExec
+	t.Cleanup(func() { NewBackgroundExec = oldNewBackgroundExec })
+	NewBackgroundExec = func(_ context.Context, _ FeSession, _ ...*BackgroundExecOption) BackgroundExec {
+		return bh
+	}
+
+	_, err := doRestoreSnapshot(context.Background(), ses, &tree.RestoreSnapShot{
+		Level:        tree.RESTORELEVELTABLE,
+		AccountName:  sysAccountName,
+		DatabaseName: dbName,
+		TableName:    tblName,
+		SnapShotName: snapshotName,
+	})
+	require.EqualError(t, err, "not supported: can not restore table 'db1.parent' referenced by some foreign key constraint")
+	require.Equal(t, []string{
+		"begin;",
+		catalog.ViewMetadataLifecycleGateSQL,
+		snapshotSQL,
+		masterSQL,
+		"rollback;",
+	}, bh.executedSQLs)
+}
+
 func TestBuildTableInfoListSQLEscapesLiterals(t *testing.T) {
 	for _, tableName := range []string{"tbl'name", "a_b", "a%b", `child\fk`} {
 		t.Run(tableName, func(t *testing.T) {
@@ -1217,7 +1407,7 @@ func TestRecreateUserDefinedFunctionCatalogPreservesCurrentSchema(t *testing.T)
 		bh := &backgroundExecTest{}
 		bh.init()
 		ctx := defines.AttachAccountId(t.Context(), sourceAccount)
-		require.NoError(t, recreateTable(ctx, "", bh, snapshotName, udfTable, sourceAccount, snapshotTS))
+		require.NoError(t, recreateTable(ctx, "", bh, snapshotName, udfTable, sourceAccount, snapshotTS, false))
 		require.Equal(t, []string{
 			dropTableIfExistsSQL(moCatalog, udfTable.tblName),
 			MoCatalogMoUserDefinedFunctionDDL,
@@ -1232,7 +1422,7 @@ func TestRecreateUserDefinedFunctionCatalogPreservesCurrentSchema(t *testing.T)
 		bh := &backgroundExecTest{}
 		bh.init()
 		ctx := defines.AttachAccountId(t.Context(), sourceAccount)
-		require.NoError(t, recreateTable(ctx, "", bh, snapshotName, udfTable, targetAccount, snapshotTS))
+		require.NoError(t, recreateTable(ctx, "", bh, snapshotName, udfTable, targetAccount, snapshotTS, false))
 		require.Equal(t,
 			"insert into `mo_catalog`.`mo_user_defined_function` ("+userDefinedFunctionCatalogColumns+
 				") select "+userDefinedFunctionCatalogSourceColumns+

test/distributed/cases/snapshot/cluster/restore_fk_restore_master_table.result

@@ -248,7 +248,9 @@ performancereviewid    employeeid    reviewdate    reviewerid
 4    4    2020-04-04    5
 5    5    2020-05-05    6
 6    6    2020-06-06    1
+-- @regex("not supported: can not restore table .* referenced by some foreign key constraint", true)
 restore table Company.Departments{snapshot="sys_sp"};
+not supported: can not restore table 'company.departments' referenced by some foreign key constraint
 select * from Departments;
 departmentid    name    managerid
 1    Research and Development    null

test/distributed/cases/snapshot/cluster/restore_fk_restore_master_table.sql

@@ -158,6 +158,8 @@ select * from TimeOff;
 select * from PerformanceReviews;
 
 
+-- A table referenced by foreign keys has an explicit restore limitation.
+-- @regex("not supported: can not restore table .* referenced by some foreign key constraint",true)
 restore table Company.Departments{snapshot="sys_sp"};
 select * from Departments;
 select * from Employees;

test/distributed/cases/snapshot/restore_fk_table.result

@@ -320,7 +320,9 @@ drop snapshot if exists sp04;
 create snapshot sp04 for account;
 insert into f1 values (3,20);
 insert into f1 values (4,600);
+-- @regex("not supported: can not restore table .* referenced by some foreign key constraint", true)
 restore table test04.f1{snapshot="sp04"};
+not supported: can not restore table 'test04.f1' referenced by some foreign key constraint
 use test04;
 show tables;
 Tables_in_test04
@@ -352,7 +354,9 @@ drop snapshot if exists sp05;
 create snapshot sp05 for account acc01;
 insert into f1 values (3,20);
 insert into f1 values (4,600);
+-- @regex("not supported: can not restore table .* referenced by some foreign key constraint", true)
 restore table test05.f1{snapshot="sp05"};
+not supported: can not restore table 'test05.f1' referenced by some foreign key constraint
 use test04;
 Unknown database test04
 show tables;