From ae2c3d658990bb2066c5a9ab47ed0a82dcdb1793 Mon Sep 17 00:00:00 2001
From: Michael Butler <butler@cockroachlabs.com>
Date: Sat, 6 Aug 2022 12:45:35 -0400
Subject: [PATCH] backupccl: add import cmd and regex arg to backup data driven
 driver

Release note: none
---
 pkg/ccl/backupccl/datadriven_test.go | 40 +++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/pkg/ccl/backupccl/datadriven_test.go b/pkg/ccl/backupccl/datadriven_test.go
index 4b3823fe9ac7..b913ba1f226e 100644
--- a/pkg/ccl/backupccl/datadriven_test.go
+++ b/pkg/ccl/backupccl/datadriven_test.go
@@ -13,6 +13,7 @@ import (
 	gosql "database/sql"
 	"fmt"
 	"net/url"
+	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -240,9 +241,12 @@ func (d *datadrivenTestState) getSQLDB(t *testing.T, server string, user string)
 //   + ignore-notice: does not print out the notice that is buffered during
 //   query execution.
 //
-// - "query-sql [server=<name>] [user=<name>]"
+// - "query-sql [server=<name>] [user=<name>] [regex=<regex pattern>]"
 //   Executes the input SQL query and print the results.
 //
+//   + regex: return true if the query result matches the regex pattern and
+//   false otherwise.
+//
 // - "reset"
 //    Clear all state associated with the test.
 //
@@ -472,6 +476,16 @@ func TestDataDriven(t *testing.T) {
 				}
 				output, err := sqlutils.RowsToDataDrivenOutput(rows)
 				require.NoError(t, err)
+				if d.HasArg("regex") {
+					var pattern string
+					d.ScanArgs(t, "regex", &pattern)
+					matched, err := regexp.MatchString(pattern, output)
+					require.NoError(t, err)
+					if matched {
+						return "true"
+					}
+					return "false"
+				}
 				return output
 
 			case "let":
@@ -525,6 +539,30 @@ func TestDataDriven(t *testing.T) {
 				require.NoError(t, err)
 				return ""
 
+			case "import":
+				server := lastCreatedServer
+				user := "root"
+				jobType := "IMPORT"
+
+				// First, run the backup.
+				_, err := ds.getSQLDB(t, server, user).Exec(d.Input)
+
+				// Tag the job.
+				if d.HasArg("tag") {
+					tagJob(t, server, user, jobType, ds, d)
+				}
+
+				// Check if we expect a pausepoint error.
+				if d.HasArg("expect-pausepoint") {
+					expectPausepoint(t, err, jobType, server, user, ds)
+					ret := append(ds.noticeBuffer, "job paused at pausepoint")
+					return strings.Join(ret, "\n")
+				}
+
+				// All other errors are bad.
+				require.NoError(t, err)
+				return ""
+
 			case "restore":
 				server := lastCreatedServer
 				user := "root"