backupccl: add keyRewriter.RewriteSpan

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.

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