Merge cockroachdb#124866

124866: backupccl: add keyRewriter.RewriteSpan r=dt a=dt

Span bounds should always be rewritten, even when they fall within the tables for which individual keys are elided during a restore, such as the liveness table. This elision is typically controlled by the key rewriter returning false for keys in these tables, indicating that they should not be restored. However if it returns false when one of these keys appears as a span boundary, we still want to rewrite it. Previously when the rewrite of one of these keys incorrectly returned false, we would return false from rewriteSpan. Unfortunately, this early return was done quietly, leaving the span half rewritten.

This pulls span rewriting into the key rewriter, parameterizing various methods to know when they are rewriting span bounds versus keys, so that they only check for and return false for keys in the elided tables when rewriting a key and not a span.

Release note: none.
Epic: none.

Co-authored-by: David Taylor <[email protected]>

pkg/ccl/backupccl/key_rewriter.go

@@ -221,8 +221,12 @@ func MakeKeyRewriterPrefixIgnoringInterleaved(tableID descpb.ID, indexID descpb.
 
 // RewriteTenant rewrites a tenant key.
 func (kr *KeyRewriter) RewriteTenant(key []byte) ([]byte, bool, error) {
+	return kr.rewriteTenant(key, false)
+}
+
+func (kr *KeyRewriter) rewriteTenant(key []byte, forSpan bool) ([]byte, bool, error) {
 	k, ok := kr.tenants.rewriteKey(key)
-	if ok {
+	if ok && !forSpan {
 		// Skip keys from ephemeral cluster status tables so that the restored
 		// cluster does not observe stale leases/liveness until it expires.
 		noTenantPrefix, _, err := keys.DecodeTenantPrefix(key)
@@ -247,14 +251,20 @@ func (kr *KeyRewriter) RewriteTenant(key []byte) ([]byte, bool, error) {
 // more details.
 func (kr *KeyRewriter) RewriteKey(
 	key []byte, walltimeForImportElision int64,
+) ([]byte, bool, error) {
+	return kr.rewriteKey(key, walltimeForImportElision, false)
+}
+
+func (kr *KeyRewriter) rewriteKey(
+	key []byte, walltimeForImportElision int64, forSpan bool,
 ) ([]byte, bool, error) {
 	// If we are reading a system tenant backup and this is a tenant key then it
 	// is part of a backup *of* that tenant, so we only restore it if we have a
 	// tenant rekey for it, i.e. we're restoring that tenant.
 	// We also enable rekeying if we are restoring a tenant from a replication stream
 	// in which case we are restoring as a system tenant.
 	if kr.fromSystemTenant && bytes.HasPrefix(key, keys.TenantPrefix) {
-		return kr.RewriteTenant(key)
+		return kr.rewriteTenant(key, forSpan)
 	}
 
 	// At this point we know we're not restoring a tenant, however the keys we're
@@ -266,7 +276,7 @@ func (kr *KeyRewriter) RewriteKey(
 		return nil, false, err
 	}
 
-	rekeyed, ok, err := kr.checkAndRewriteTableKey(noTenantPrefix, walltimeForImportElision)
+	rekeyed, ok, err := kr.checkAndRewriteTableKey(noTenantPrefix, walltimeForImportElision, forSpan)
 	if err != nil || !ok {
 		return nil, false, err
 	}
@@ -318,17 +328,19 @@ func (kr *KeyRewriter) RewriteKey(
 // filtering occurs. Filtering is necessary during restore because the restoring
 // cluster should not contain keys from an in-progress import.
 func (kr *KeyRewriter) checkAndRewriteTableKey(
-	key []byte, walltimeForImportElision int64,
+	key []byte, walltimeForImportElision int64, forSpan bool,
 ) ([]byte, bool, error) {
 	// Fetch the original table ID for descriptor lookup. Ignore errors because
 	// they will be caught later on if tableID isn't in descs or kr doesn't
 	// perform a rewrite.
 	_, tableID, _ := keys.SystemSQLCodec.DecodeTablePrefix(key)
 
-	// Skip keys from ephemeral cluster status tables so that the restored cluster
-	// does not observe stale leases/liveness until it expires.
-	if tableID == keys.SQLInstancesTableID || tableID == keys.SqllivenessID || tableID == keys.LeaseTableID {
-		return nil, false, nil
+	if !forSpan {
+		// Skip keys from ephemeral cluster status tables so that the restored cluster
+		// does not observe stale leases/liveness until it expires.
+		if tableID == keys.SQLInstancesTableID || tableID == keys.SqllivenessID || tableID == keys.LeaseTableID {
+			return nil, false, nil
+		}
 	}
 
 	desc := kr.descs[descpb.ID(tableID)]
@@ -352,3 +364,26 @@ func (kr *KeyRewriter) checkAndRewriteTableKey(
 	}
 	return key, true, nil
 }
+
+func (kr *KeyRewriter) RewriteSpan(span roachpb.Span) (roachpb.Span, error) {
+	var (
+		ok  bool
+		err error
+	)
+	span.Key, ok, err = kr.rewriteKey(span.Key, 0, true)
+	if err != nil {
+		return roachpb.Span{}, errors.Wrapf(err, "span start key %s was not rewritten", span.Key)
+	}
+	if !ok {
+		return roachpb.Span{}, errors.Newf("rewriting span start key %s failed", span.Key)
+	}
+
+	span.EndKey, ok, err = kr.rewriteKey(span.EndKey, 0, true)
+	if err != nil {
+		return roachpb.Span{}, errors.Wrapf(err, "rewriting span end key %s failed", span.EndKey)
+	}
+	if !ok {
+		return roachpb.Span{}, errors.Newf("span end key %s was not rewritten", span.EndKey)
+	}
+	return span, nil
+}

pkg/ccl/backupccl/restore_online.go

@@ -226,22 +226,10 @@ func assertCommonPrefix(span roachpb.Span, elidedPrefixType execinfrapb.ElidePre
 func rewriteSpan(
 	kr *KeyRewriter, span roachpb.Span, elidedPrefixType execinfrapb.ElidePrefix,
 ) (roachpb.Span, error) {
-	var (
-		ok  bool
-		err error
-	)
-	if err = assertCommonPrefix(span, elidedPrefixType); err != nil {
-		return span, err
-	}
-	span.Key, ok, err = kr.RewriteKey(span.Key, 0)
-	if !ok || err != nil {
-		return span, errors.Wrapf(err, "span start key %s was not rewritten", span.Key)
-	}
-	span.EndKey, ok, err = kr.RewriteKey(span.EndKey, 0)
-	if !ok || err != nil {
-		return span, errors.Wrapf(err, "span end key %s was not rewritten ", span.Key)
+	if err := assertCommonPrefix(span, elidedPrefixType); err != nil {
+		return roachpb.Span{}, err
 	}
-	return span, nil
+	return kr.RewriteSpan(span)
 }
 
 // linkExternalFiles runs through all entries produced by genSpans and links in