fix restore sub database(#2.0-dev) (#20099)

fix restore sub database Approved by: @aressu1985, @daviszhen, @sukki37
matrixorigin · Nov 15, 2024 · 1cb3265 · 1cb3265
1 parent 0079c99
commit 1cb3265
Show file tree

Hide file tree

Showing 5 changed files with 1,009 additions and 2 deletions.
diff --git a/pkg/frontend/pitr.go b/pkg/frontend/pitr.go
@@ -1304,6 +1304,7 @@ func restoreToDatabaseOrTableWithPitr(
 	var (
 		createDbSql string
 		tableInfos  []*tableInfo
+		isSubDb     bool
 	)
 	createDbSql, err = getCreateDatabaseSqlInPitr(ctx, sid, bh, pitrName, dbName, curAccount, ts)
 	if err != nil {
@@ -1313,7 +1314,10 @@ func restoreToDatabaseOrTableWithPitr(
 	restoreToTbl := tblName != ""
 
 	// if restore to table, check if the db is sub db
-	isSubDb := strings.Contains(createDbSql, "from") && strings.Contains(createDbSql, "publication")
+	isSubDb, err = checkDbIsSubDb(ctx, createDbSql)
+	if err != nil {
+		return
+	}
 	if isSubDb && restoreToTbl {
 		return moerr.NewInternalError(ctx, "can't restore to table for sub db")
 	}

diff --git a/pkg/frontend/pitr_test.go b/pkg/frontend/pitr_test.go
@@ -2890,3 +2890,46 @@ func Test_RestorePitrFaultTolerance(t *testing.T) {
 		assert.Error(t, err)
 	})
 }
+
+func TestCheckDbIsSubDb(t *testing.T) {
+	ctx := context.Background()
+
+	tests := []struct {
+		name        string
+		createDbsql string
+		want        bool
+		wantErr     bool
+	}{
+		{
+			name:        "SubscriptionOption exists",
+			createDbsql: "create database sub01 from acc01 publication pub01;",
+			want:        true,
+			wantErr:     false,
+		},
+		{
+			name:        "SubscriptionOption does not exist",
+			createDbsql: "CREATE DATABASE test",
+			want:        false,
+			wantErr:     false,
+		},
+		{
+			name:        "Invalid SQL",
+			createDbsql: "INVALID SQL",
+			want:        false,
+			wantErr:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := checkDbIsSubDb(ctx, tt.createDbsql)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("checkDbIsSubDb() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("checkDbIsSubDb() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/pkg/frontend/snapshot.go b/pkg/frontend/snapshot.go
@@ -731,6 +731,7 @@ func restoreToDatabaseOrTable(
 	}
 
 	var createDbSql string
+	var isSubDb bool
 	createDbSql, err = getCreateDatabaseSql(ctx, sid, bh, snapshotName, dbName, restoreAccount)
 	if err != nil {
 		return
@@ -740,7 +741,10 @@ func restoreToDatabaseOrTable(
 	restoreToTbl := tblName != ""
 
 	// if restore to table, check if the db is sub db
-	isSubDb := strings.Contains(createDbSql, "from") && strings.Contains(createDbSql, "publication")
+	isSubDb, err = checkDbIsSubDb(toCtx, createDbSql)
+	if err != nil {
+		return
+	}
 	if isSubDb && restoreToTbl {
 		return moerr.NewInternalError(ctx, "can't restore to table for sub db")
 	}
@@ -2018,3 +2022,21 @@ func checkSubscriptionExist(
 
 	return true, nil
 }
+
+func checkDbIsSubDb(ctx context.Context, createDbsql string) (bool, error) {
+	var (
+		err error
+		ast []tree.Statement
+	)
+	ast, err = mysql.Parse(ctx, createDbsql, 1)
+	if err != nil {
+		return false, err
+	}
+
+	if createDb, ok := ast[0].(*tree.CreateDatabase); ok {
+		if createDb.SubscriptionOption != nil {
+			return true, nil
+		}
+	}
+	return false, nil
+}