From 5bea8fe57e7314a55ad80ff9e227164ec1412ec6 Mon Sep 17 00:00:00 2001
From: Paul Bardea <pbardea@gmail.com>
Date: Mon, 12 Jul 2021 12:53:29 -0400
Subject: [PATCH] backupccl: stop splitting at invalid split keys

Restore could perform invalid splits during the split and scatter phase.
This commit ensures that splits are always performed at safe keys during
restore.

Release note (bug fix): Fix a bug where restores of data with multiple
column families could be split illegally (within a single SQL row). This
could result in temporary data unavailability until the ranges on either
side of the invalid split are merged.
---
 .../backupccl/split_and_scatter_processor.go  | 10 +++++++
 .../testdata/backup-restore/column-families   | 29 +++++++++++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 pkg/ccl/backupccl/testdata/backup-restore/column-families

diff --git a/pkg/ccl/backupccl/split_and_scatter_processor.go b/pkg/ccl/backupccl/split_and_scatter_processor.go
index 110c05026fd1..3143c6ecdb8f 100644
--- a/pkg/ccl/backupccl/split_and_scatter_processor.go
+++ b/pkg/ccl/backupccl/split_and_scatter_processor.go
@@ -84,6 +84,11 @@ func (s dbSplitAndScatterer) split(
 	if err != nil {
 		return err
 	}
+	if splitAt, err := keys.EnsureSafeSplitKey(newSplitKey); err != nil {
+		// Ignore the error, not all keys are table keys.
+	} else if len(splitAt) != 0 {
+		newSplitKey = splitAt
+	}
 	log.VEventf(ctx, 1, "presplitting new key %+v", newSplitKey)
 	if err := s.db.AdminSplit(ctx, newSplitKey, expirationTime); err != nil {
 		return errors.Wrapf(err, "splitting key %s", newSplitKey)
@@ -107,6 +112,11 @@ func (s dbSplitAndScatterer) scatter(
 	if err != nil {
 		return 0, err
 	}
+	if scatterAt, err := keys.EnsureSafeSplitKey(newScatterKey); err != nil {
+		// Ignore the error, not all keys are table keys.
+	} else if len(scatterAt) != 0 {
+		newScatterKey = scatterAt
+	}
 
 	log.VEventf(ctx, 1, "scattering new key %+v", newScatterKey)
 	req := &roachpb.AdminScatterRequest{
diff --git a/pkg/ccl/backupccl/testdata/backup-restore/column-families b/pkg/ccl/backupccl/testdata/backup-restore/column-families
new file mode 100644
index 000000000000..09ad3f69b35d
--- /dev/null
+++ b/pkg/ccl/backupccl/testdata/backup-restore/column-families
@@ -0,0 +1,29 @@
+new-server name=s1 localities=us-east-1,us-west-1,us-west-2,eu-central-1
+----
+
+exec-sql
+CREATE DATABASE orig;
+USE orig;
+CREATE TABLE cfs (a INT PRIMARY KEY, b STRING, c STRING, d STRING, FAMILY (b), FAMILY (c));
+INSERT INTO cfs SELECT x, repeat('abc', 100), repeat('abc', 100) FROM generate_series(0, 3) AS x;
+ALTER TABLE cfs SPLIT AT SELECT a FROM cfs;
+-- Split the output files very small to catch output SSTs mid-row.
+SET CLUSTER SETTING kv.bulk_sst.target_size = '1';
+BACKUP cfs TO 'nodelocal://1/foo';
+CREATE DATABASE r1;
+RESTORE cfs FROM 'nodelocal://1/foo' WITH into_db='r1';
+----
+
+query-sql
+SELECT max(length(start_key)) FROM [SHOW RANGES FROM TABLE orig.cfs];
+----
+2
+
+# Regression test for #67488.
+# This can return up to 6 if RESTORE improperly splits the ranges in the middle
+# of a SQL row since column family keys are longer. E.g. Keys are expected to be
+# of the form '/1' (row with PK 1), and not '/3/1/1' (row with PK 3, and CF 1).
+query-sql
+SELECT max(length(start_key)) FROM [SHOW RANGES FROM TABLE r1.cfs];
+----
+2