From 7cab9b18a28644c49332f75f41a6c6c63177cae6 Mon Sep 17 00:00:00 2001 From: David Taylor Date: Mon, 21 Feb 2022 21:05:32 +0000 Subject: [PATCH] backupccl: RESTORE TENANT x .. WITH tenant = y This adds support for restoring a tenant with a new tenant ID, allowing restoring into a cluster that already has an existing tenant (perhaps the same one) at the ID of the backed up tenant. Release note: none. --- docs/generated/sql/bnf/stmt_block.bnf | 1 + pkg/ccl/backupccl/backup_test.go | 21 ++++++++ pkg/ccl/backupccl/key_rewriter.go | 13 ++--- pkg/ccl/backupccl/restore_job.go | 18 ++++--- pkg/ccl/backupccl/restore_planning.go | 68 ++++++++++++++++++++++++-- pkg/jobs/jobspb/jobs.proto | 6 ++- pkg/sql/parser/sql.y | 5 ++ pkg/sql/parser/testdata/backup_restore | 8 +++ pkg/sql/sem/tree/backup.go | 17 ++++++- 9 files changed, 135 insertions(+), 22 deletions(-) diff --git a/docs/generated/sql/bnf/stmt_block.bnf b/docs/generated/sql/bnf/stmt_block.bnf index 288f094e518e..ba1b215c4f38 100644 --- a/docs/generated/sql/bnf/stmt_block.bnf +++ b/docs/generated/sql/bnf/stmt_block.bnf @@ -2239,6 +2239,7 @@ restore_options ::= | 'DEBUG_PAUSE_ON' '=' string_or_placeholder | 'NEW_DB_NAME' '=' string_or_placeholder | 'INCREMENTAL_LOCATION' '=' string_or_placeholder_opt_list + | 'TENANT' '=' string_or_placeholder scrub_option_list ::= ( scrub_option ) ( ( ',' scrub_option ) )* diff --git a/pkg/ccl/backupccl/backup_test.go b/pkg/ccl/backupccl/backup_test.go index 841587045313..d1a5f17e3516 100644 --- a/pkg/ccl/backupccl/backup_test.go +++ b/pkg/ccl/backupccl/backup_test.go @@ -7432,6 +7432,27 @@ func TestBackupRestoreTenant(t *testing.T) { // Check the all-tenant override. restoreTenant11.CheckQueryResults(t, `SHOW CLUSTER SETTING tenant_cost_model.kv_read_cost_per_megabyte`, [][]string{{"123"}}) + + restoreDB.Exec(t, `SELECT crdb_internal.destroy_tenant(20, true)`) + + restoreDB.Exec(t, `RESTORE TENANT 11 FROM 'nodelocal://1/clusterwide' WITH tenant = '20'`) + _, restoreConn20 := serverutils.StartTenant( + t, restoreTC.Server(0), base.TestTenantArgs{TenantID: roachpb.MakeTenantID(20), Existing: true}, + ) + defer restoreConn20.Close() + restoreTenant20 := sqlutils.MakeSQLRunner(restoreConn20) + + // Tenant 20 gets results that matched the backed up tenant 11. + restoreTenant20.CheckQueryResults(t, `select * from foo.baz`, tenant11.QueryStr(t, `select * from foo.baz`)) + // Check the all-tenant override. + restoreTenant20.CheckQueryResults(t, `SHOW CLUSTER SETTING tenant_cost_model.kv_read_cost_per_megabyte`, [][]string{{"123"}}) + + // Remove tenant 11, then confirm restoring 11 over 10 fails. + restoreDB.Exec(t, `SELECT crdb_internal.destroy_tenant(11, true)`) + restoreDB.ExpectErr(t, `exists`, `RESTORE TENANT 11 FROM 'nodelocal://1/clusterwide' WITH tenant = '10'`) + + // Verify tenant 20 is still unaffected. + restoreTenant20.CheckQueryResults(t, `select * from foo.baz`, tenant11.QueryStr(t, `select * from foo.baz`)) }) t.Run("restore-tenant10-to-ts1", func(t *testing.T) { diff --git a/pkg/ccl/backupccl/key_rewriter.go b/pkg/ccl/backupccl/key_rewriter.go index 5d9b4e18b2e7..bd187403ee5b 100644 --- a/pkg/ccl/backupccl/key_rewriter.go +++ b/pkg/ccl/backupccl/key_rewriter.go @@ -137,7 +137,7 @@ func makeKeyRewriter( ) (*KeyRewriter, error) { var prefixes prefixRewriter var tenantPrefixes prefixRewriter - tenantPrefixes.rewrites = make([]prefixRewrite, len(tenants)) + tenantPrefixes.rewrites = make([]prefixRewrite, 0, len(tenants)) seenPrefixes := make(map[string]bool) for oldID, desc := range descs { @@ -179,13 +179,10 @@ func makeKeyRewriter( fromSystemTenant = true continue } - tenantPrefixes.rewrites[i] = prefixRewrite{ - OldPrefix: keys.MakeSQLCodec(tenants[i].OldID).TenantPrefix(), - NewPrefix: keys.MakeSQLCodec(tenants[i].NewID).TenantPrefix(), - } - tenantPrefixes.rewrites[i].noop = bytes.Equal( - tenantPrefixes.rewrites[i].OldPrefix, tenantPrefixes.rewrites[i].NewPrefix, - ) + from, to := keys.MakeSQLCodec(tenants[i].OldID).TenantPrefix(), keys.MakeSQLCodec(tenants[i].NewID).TenantPrefix() + tenantPrefixes.rewrites = append(tenantPrefixes.rewrites, prefixRewrite{ + OldPrefix: from, NewPrefix: to, noop: bytes.Equal(from, to), + }) } sort.Slice(tenantPrefixes.rewrites, func(i, j int) bool { return bytes.Compare(tenantPrefixes.rewrites[i].OldPrefix, tenantPrefixes.rewrites[j].OldPrefix) < 0 diff --git a/pkg/ccl/backupccl/restore_job.go b/pkg/ccl/backupccl/restore_job.go index b8f64058a9fe..6d10119660d5 100644 --- a/pkg/ccl/backupccl/restore_job.go +++ b/pkg/ccl/backupccl/restore_job.go @@ -92,11 +92,6 @@ var restoreStatsInsertBatchSize = 10 func rewriteBackupSpanKey( codec keys.SQLCodec, kr *KeyRewriter, key roachpb.Key, ) (roachpb.Key, error) { - // TODO(dt): support rewriting tenant keys. - if bytes.HasPrefix(key, keys.TenantPrefix) { - return key, nil - } - newKey, rewritten, err := kr.RewriteKey(append([]byte(nil), key...)) if err != nil { return nil, errors.NewAssertionErrorWithWrappedErrf(err, @@ -107,6 +102,10 @@ func rewriteBackupSpanKey( return nil, errors.AssertionFailedf( "no rewrite for span start key: %s", key) } + if bytes.HasPrefix(newKey, keys.TenantPrefix) { + return newKey, nil + } + // Modify all spans that begin at the primary index to instead begin at the // start of the table. That is, change a span start key from /Table/51/1 to // /Table/51. Otherwise a permanently empty span at /Table/51-/Table/51/1 @@ -1339,9 +1338,12 @@ func (r *restoreResumer) doResume(ctx context.Context, execCtx interface{}) erro } for _, tenant := range details.Tenants { - // TODO(dt): extend the job to keep track of new ID. - id := roachpb.MakeTenantID(tenant.ID) - mainData.addTenant(id, id) + to := roachpb.MakeTenantID(tenant.ID) + from := to + if details.PreRewriteTenantId != nil { + from = *details.PreRewriteTenantId + } + mainData.addTenant(from, to) } numNodes, err := clusterNodeCount(p.ExecCfg().Gossip) diff --git a/pkg/ccl/backupccl/restore_planning.go b/pkg/ccl/backupccl/restore_planning.go index 3b59ff33c315..b080898d59bf 100644 --- a/pkg/ccl/backupccl/restore_planning.go +++ b/pkg/ccl/backupccl/restore_planning.go @@ -66,6 +66,7 @@ const ( restoreOptSkipMissingViews = "skip_missing_views" restoreOptSkipLocalitiesCheck = "skip_localities_check" restoreOptDebugPauseOn = "debug_pause_on" + restoreOptAsTenant = "tenant" // The temporary database system tables will be restored into for full // cluster backups. @@ -1106,6 +1107,36 @@ func restorePlanHook( } } + var newTenantIDFn func() (*roachpb.TenantID, error) + if restoreStmt.Options.AsTenant != nil { + if restoreStmt.DescriptorCoverage == tree.AllDescriptors || !restoreStmt.Targets.Tenant.IsSet() { + err := errors.Errorf("%q can only be used when running RESTORE TENANT for a single tenant", restoreOptAsTenant) + return nil, nil, nil, false, err + } + // TODO(dt): it'd be nice to have TypeAsInt or TypeAsTenantID and then in + // sql.y an int_or_placeholder, but right now the hook view of planner only + // has TypeAsString so we'll just atoi it. + fn, err := p.TypeAsString(ctx, restoreStmt.Options.AsTenant, "RESTORE") + if err != nil { + return nil, nil, nil, false, err + } + newTenantIDFn = func() (*roachpb.TenantID, error) { + s, err := fn() + if err != nil { + return nil, err + } + x, err := strconv.Atoi(s) + if err != nil { + return nil, err + } + if x < int(roachpb.MinTenantID.ToUint64()) { + return nil, errors.New("invalid tenant ID value") + } + id := roachpb.MakeTenantID(uint64(x)) + return &id, nil + } + } + fn := func(ctx context.Context, _ []sql.PlanNode, resultsCh chan<- tree.Datums) error { // TODO(dan): Move this span into sql. ctx, span := tracing.ChildSpan(ctx, stmt.StatementTag()) @@ -1172,6 +1203,13 @@ func restorePlanHook( return err } } + var newTenantID *roachpb.TenantID + if newTenantIDFn != nil { + newTenantID, err = newTenantIDFn() + if err != nil { + return err + } + } // incFrom will contain the directory URIs for incremental backups (i.e. // /) iff len(From)==1, regardless of the @@ -1222,7 +1260,7 @@ func restorePlanHook( } return doRestorePlan(ctx, restoreStmt, p, from, incFrom, passphrase, kms, intoDB, - newDBName, endTime, resultsCh) + newDBName, newTenantID, endTime, resultsCh) } if restoreStmt.Options.Detached { @@ -1346,6 +1384,7 @@ func doRestorePlan( kms []string, intoDB string, newDBName string, + newTenantID *roachpb.TenantID, endTime hlc.Timestamp, resultsCh chan<- tree.Datums, ) error { @@ -1540,20 +1579,40 @@ func doRestorePlan( } } + var oldTenantID *roachpb.TenantID if len(tenants) > 0 { if !p.ExecCfg().Codec.ForSystemTenant() { return pgerror.Newf(pgcode.InsufficientPrivilege, "only the system tenant can restore other tenants") } - for _, i := range tenants { + if newTenantID != nil { + if len(tenants) != 1 { + return errors.Errorf("%q option can only be used when restoring a single tenant", restoreOptAsTenant) + } res, err := p.ExecCfg().InternalExecutor.QueryRow( ctx, "restore-lookup-tenant", p.ExtendedEvalContext().Txn, - `SELECT active FROM system.tenants WHERE id = $1`, i.ID, + `SELECT active FROM system.tenants WHERE id = $1`, newTenantID.ToUint64(), ) if err != nil { return err } if res != nil { - return errors.Errorf("tenant %d already exists", i.ID) + return errors.Errorf("tenant %s already exists", newTenantID) + } + old := roachpb.MakeTenantID(tenants[0].ID) + tenants[0].ID = newTenantID.ToUint64() + oldTenantID = &old + } else { + for _, i := range tenants { + res, err := p.ExecCfg().InternalExecutor.QueryRow( + ctx, "restore-lookup-tenant", p.ExtendedEvalContext().Txn, + `SELECT active FROM system.tenants WHERE id = $1`, i.ID, + ) + if err != nil { + return err + } + if res != nil { + return errors.Errorf("tenant %d already exists", i.ID) + } } } } @@ -1713,6 +1772,7 @@ func doRestorePlan( DatabaseModifiers: databaseModifiers, DebugPauseOn: debugPauseOn, RestoreSystemUsers: restoreStmt.SystemUsers, + PreRewriteTenantId: oldTenantID, }, Progress: jobspb.RestoreProgress{}, } diff --git a/pkg/jobs/jobspb/jobs.proto b/pkg/jobs/jobspb/jobs.proto index 8c785c0be633..4b5299b21f73 100644 --- a/pkg/jobs/jobspb/jobs.proto +++ b/pkg/jobs/jobspb/jobs.proto @@ -323,7 +323,11 @@ message RestoreDetails { string debug_pause_on = 20; bool restore_system_users = 22; - // NEXT ID: 23. + // PreRewrittenTenantID is the ID of tenants[0] in the backup, aka its old ID; + // it is only valid to set this if len(tenants) == 1. + roachpb.TenantID pre_rewrite_tenant_id = 23; + + // NEXT ID: 24. } message RestoreProgress { diff --git a/pkg/sql/parser/sql.y b/pkg/sql/parser/sql.y index fcf06a302c6d..cdef75d7c5b5 100644 --- a/pkg/sql/parser/sql.y +++ b/pkg/sql/parser/sql.y @@ -3222,6 +3222,11 @@ restore_options: { $$.val = &tree.RestoreOptions{IncrementalStorage: $3.stringOrPlaceholderOptList()} } +| TENANT '=' string_or_placeholder + { + $$.val = &tree.RestoreOptions{AsTenant: $3.expr()} + } + import_format: name { diff --git a/pkg/sql/parser/testdata/backup_restore b/pkg/sql/parser/testdata/backup_restore index 5910c1b1b5ae..6b37bf61de06 100644 --- a/pkg/sql/parser/testdata/backup_restore +++ b/pkg/sql/parser/testdata/backup_restore @@ -665,6 +665,14 @@ RESTORE TENANT 36 FROM (($1), ($2)) AS OF SYSTEM TIME ('1') -- fully parenthesiz RESTORE TENANT 36 FROM ($1, $2) AS OF SYSTEM TIME '_' -- literals removed RESTORE TENANT 36 FROM ($1, $2) AS OF SYSTEM TIME '1' -- identifiers removed +parse +RESTORE TENANT 36 FROM ($1, $2) WITH tenant = '5' +---- +RESTORE TENANT 36 FROM ($1, $2) WITH tenant = '5' +RESTORE TENANT 36 FROM (($1), ($2)) WITH tenant = ('5') -- fully parenthesized +RESTORE TENANT 36 FROM ($1, $2) WITH tenant = '_' -- literals removed +RESTORE TENANT 36 FROM ($1, $2) WITH tenant = '5' -- identifiers removed + parse RESTORE TENANT 123 FROM REPLICATION STREAM FROM 'bar' ---- diff --git a/pkg/sql/sem/tree/backup.go b/pkg/sql/sem/tree/backup.go index 7760c8b9f985..e90bfa078080 100644 --- a/pkg/sql/sem/tree/backup.go +++ b/pkg/sql/sem/tree/backup.go @@ -132,6 +132,7 @@ type RestoreOptions struct { DebugPauseOn Expr NewDBName Expr IncrementalStorage StringOrPlaceholderOptList + AsTenant Expr } var _ NodeFormatter = &RestoreOptions{} @@ -391,11 +392,18 @@ func (o *RestoreOptions) Format(ctx *FmtCtx) { ctx.WriteString("new_db_name = ") ctx.FormatNode(o.NewDBName) } + if o.IncrementalStorage != nil { maybeAddSep() ctx.WriteString("incremental_location = ") ctx.FormatNode(&o.IncrementalStorage) } + + if o.AsTenant != nil { + maybeAddSep() + ctx.WriteString("tenant = ") + ctx.FormatNode(o.AsTenant) + } } // CombineWith merges other backup options into this backup options struct. @@ -485,6 +493,12 @@ func (o *RestoreOptions) CombineWith(other *RestoreOptions) error { return errors.New("incremental_location option specified multiple times") } + if o.AsTenant == nil { + o.AsTenant = other.AsTenant + } else if other.AsTenant != nil { + return errors.New("tenant option specified multiple times") + } + return nil } @@ -502,5 +516,6 @@ func (o RestoreOptions) IsDefault() bool { o.SkipLocalitiesCheck == options.SkipLocalitiesCheck && o.DebugPauseOn == options.DebugPauseOn && o.NewDBName == options.NewDBName && - cmp.Equal(o.IncrementalStorage, options.IncrementalStorage) + cmp.Equal(o.IncrementalStorage, options.IncrementalStorage) && + o.AsTenant == options.AsTenant }