Merge #154382 #154455

154382: workload/schemachanger: add randomized testing for TRUNCATE r=rafiss a=fqazi

Previously, the randomized schema changer workload had no coverage for TRUNCATE. This patch adds support for testing TRUNCATE in the randomized workload with multiple tables and with / without CASSCADE.

Note: This patch also fixes the pgcode generated by TRUNCATE when foreign key refrences exist, since previously we returned an internal error. Postgres generates a FeatureNotSupported pgcode.

Fixes: #154045

Release note: None

154455: cmd/roachtest: prevent macOS host from idle-sleeping during test r=dt a=dt

These tests run for 20mins+ so it is common to step away to make coffee or something while they are running. Unfortunately this can mean the host goes to sleep, failing the test. To avoid this we can invoke the macOS utility 'caffeinate' with '-i' to prevent only machine idle sleep (but not display sleep/lock), and with the roachtest process, so it will inhibit sleep until the roachtest process exits.

Release note: none.
Epic: none.

Co-authored-by: Faizan Qazi <faizan@cockroachlabs.com>
Co-authored-by: David Taylor <davidt@davidt.io>

Author: 3 people

pkg/cmd/roachtest/roachtestflags/flags.go

@@ -497,6 +497,14 @@ var (
 						Always collect artifacts during test teardown, even if the test did not
 						time out or fail.`,
 	})
+
+	Caffeinate bool = true
+	_               = registerRunFlag(&Caffeinate, FlagInfo{
+		Name: "caffeinate",
+		Usage: `
+						On Darwin, prevent the system from sleeping while roachtest is running
+						by invoking caffeinate -i -w <pid>. Default is true.`,
+	})
 )
 
 // The flags below override the final cluster configuration. They have no

pkg/cmd/roachtest/run.go

@@ -11,6 +11,7 @@ import (
 	"math/rand"
 	"net/http"
 	"os"
+	"os/exec"
 	"os/signal"
 	"os/user"
 	"path/filepath"
@@ -52,6 +53,21 @@ const (
 // runTests is the main function for the run and bench commands.
 // Assumes initRunFlagsBinariesAndLibraries was called.
 func runTests(register func(registry.Registry), filter *registry.TestFilter) error {
+	// On Darwin, start caffeinate to prevent the system from sleeping.
+	if runtime.GOOS == "darwin" && roachtestflags.Caffeinate {
+		pid := os.Getpid()
+		cmd := exec.Command("caffeinate", "-i", "-w", strconv.Itoa(pid))
+		if err := cmd.Start(); err != nil {
+			fmt.Fprintf(os.Stderr, "warning: failed to start caffeinate: %v\n", err)
+		} else {
+			defer func() {
+				if cmd.Process != nil {
+					_ = cmd.Process.Kill()
+				}
+			}()
+		}
+	}
+
 	globalSeed := randutil.NewPseudoSeed()
 	if globalSeedEnv := os.Getenv("ROACHTEST_GLOBAL_SEED"); globalSeedEnv != "" {
 		if parsed, err := strconv.ParseInt(globalSeedEnv, 0, 64); err == nil {

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/truncate.go

@@ -8,6 +8,8 @@ package scbuildstmt
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
@@ -51,7 +53,13 @@ func Truncate(b BuildCtx, stmt *tree.Truncate) {
 		if stmt.DropBehavior != tree.DropCascade {
 			name := b.QueryByID(tableID).FilterNamespace().MustGetOneElement()
 			refName := refElts.FilterNamespace().MustGetZeroOrOneElement()
-			panic(errors.Errorf("%q is %s table %q", name.Name, "referenced by foreign key from", refName.Name))
+			// The error code returned here matches PGSQL, even though the CASCADE
+			// operation is supported there with TRUNCATE as well.
+			panic(errors.WithHint(pgerror.Newf(pgcode.FeatureNotSupported,
+				"%q is %s table %q", name.Name,
+				"referenced by foreign key from",
+				refName.Name),
+				"truncate dependent tables at the same time or specify the CASCADE option"))
 		}
 		if !tablesToTruncate.Contains(referencingTableID) {
 			tablesToTruncate.Add(referencingTableID)

pkg/sql/truncate.go

@@ -18,6 +18,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -89,7 +91,10 @@ func (t *truncateNode) startExec(params runParams) error {
 			}
 
 			if n.DropBehavior != tree.DropCascade {
-				return errors.Errorf("%q is %s table %q", tableDesc.Name, msg, other.Name)
+				// The error code returned here matches PGSQL, even though the CASCADE
+				// operation is supported there with TRUNCATE as well.
+				return errors.WithHintf(pgerror.Newf(pgcode.FeatureNotSupported, "%q is %s table %q", tableDesc.Name, msg, other.Name),
+					"truncate dependent tables at the same time or specify the CASCADE option")
 			}
 			if err := p.CheckPrivilege(ctx, other, privilege.DROP); err != nil {
 				return err

pkg/workload/schemachange/error_screening.go

@@ -1794,3 +1794,38 @@ func (og *operationGenerator) tableHasForeignKeyMutation(
 			AND (m->'constraint'->>'foreign_key') IS NOT NULL
 		);`, tableName)
 }
+
+// getTableForeignKeyReferences returns a list of tables that reference
+// the specified table via foreign key references.
+func (og *operationGenerator) getTableForeignKeyReferences(
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName,
+) ([]tree.TableName, error) {
+	rows, err := tx.Query(ctx,
+		`WITH fk_refs AS (
+					SELECT conrelid FROM pg_constraint WHERE
+						confrelid = $1::REGCLASS AND
+						conrelid <> $1::REGCLASS
+				)
+				SELECT
+					n.nspname as schema_name,  c.relname AS object_name
+				FROM fk_refs AS f
+					JOIN pg_class AS c ON c.oid = f.conrelid
+					JOIN pg_namespace AS n ON c.relnamespace = n.oid
+`,
+		tableName.String())
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var result []tree.TableName
+	for rows.Next() {
+		var table, schema string
+		err = rows.Scan(&schema, &table)
+		if err != nil {
+			return nil, err
+		}
+		name := tree.MakeTableNameWithSchema(tableName.CatalogName, tree.Name(schema), tree.Name(table))
+		result = append(result, name)
+	}
+	return result, rows.Err()
+}

pkg/workload/schemachange/operation_generator.go

@@ -5580,3 +5580,67 @@ func (og *operationGenerator) findExistingTrigger(
 
 	return &triggerWithInfo, nil
 }
+
+// truncateTable generates a TRUNCATE TABLE statement.
+func (og *operationGenerator) truncateTable(ctx context.Context, tx pgx.Tx) (*opStmt, error) {
+	tbls, err := Collect(ctx, og, tx, pgx.RowTo[string],
+		`SELECT quote_ident(schema_name) || '.' || quote_ident(table_name) FROM [SHOW TABLES] WHERE type = 'table'`)
+	if err != nil {
+		return nil, err
+	}
+
+	const MaxTruncateTables = 8
+	stmt, expectedCode, err := Generate[*tree.Truncate](og.params.rng, og.produceError(), []GenerationCase{
+		{pgcode.SuccessfulCompletion, `TRUNCATE TABLE {MultipleTableNames}`},
+		{pgcode.SuccessfulCompletion, `TRUNCATE TABLE {MultipleTableNames} CASCADE`},
+		{pgcode.UndefinedTable, `TRUNCATE TABLE {TableNameDoesNotExist}`},
+		{pgcode.UndefinedTable, `TRUNCATE TABLE {TableNameDoesNotExist} CASCADE`},
+	}, template.FuncMap{
+		"MultipleTableNames": func() (string, error) {
+			selectedTables, err := PickBetween(og.params.rng, 1, MaxTruncateTables, tbls)
+			if err != nil {
+				return "", err
+			}
+			return strings.Join(selectedTables, ","), nil
+		},
+		"TableNameDoesNotExist": func() (string, error) {
+			return "TableThatDoesNotExist", nil
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	op := newOpStmt(stmt, codesWithConditions{
+		{expectedCode, true},
+	})
+
+	// If the the TRUNCATE is not cascaded, then the operation can fail
+	// if foreign key references exist.
+	if expectedCode == pgcode.SuccessfulCompletion &&
+		stmt.DropBehavior != tree.DropCascade {
+		tableSet := map[string]struct{}{}
+		fkSet := map[string]struct{}{}
+		// Gather the set of tables handled by this statement, and
+		// any foreign keys that reference these tables.
+		for _, table := range stmt.Tables {
+			tableSet[table.FQString()] = struct{}{}
+			fkReferenceTables, err := og.getTableForeignKeyReferences(ctx, tx, &table)
+			if err != nil {
+				return nil, err
+			}
+			for _, fk := range fkReferenceTables {
+				fkSet[fk.FQString()] = struct{}{}
+			}
+		}
+		// Check if any of the foreign keys that reference these
+		// tables are not truncated. If they are, then the TRUNCATE
+		// will fail with an error.
+		for fk := range fkSet {
+			if _, hasTable := tableSet[fk]; !hasTable {
+				op.expectedExecErrors.add(pgcode.FeatureNotSupported)
+				break
+			}
+		}
+	}
+	return op, nil
+}

pkg/workload/schemachange/optype.go

@@ -143,6 +143,7 @@ const (
 	dropTable    // DROP TABLE <table>
 	dropTrigger  // DROP TRIGGER <trigger> ON <table>
 	dropView     // DROP VIEW <view>
+	truncateTable
 
 	// Unimplemented operations. TODO(sql-foundations): Audit and/or implement these operations.
 	// alterDatabaseOwner
@@ -264,6 +265,7 @@ var opFuncs = []func(*operationGenerator, context.Context, pgx.Tx) (*opStmt, err
 	renameSequence:                    (*operationGenerator).renameSequence,
 	renameTable:                       (*operationGenerator).renameTable,
 	renameView:                        (*operationGenerator).renameView,
+	truncateTable:                     (*operationGenerator).truncateTable,
 }
 
 var opWeights = []int{
@@ -321,6 +323,7 @@ var opWeights = []int{
 	renameSequence:                    1,
 	renameTable:                       1,
 	renameView:                        1,
+	truncateTable:                     1,
 }
 
 // This workload will maintain its own list of minimal supported versions for
@@ -358,4 +361,5 @@ var opDeclarativeVersion = map[opType]clusterversion.Key{
 	dropTable:                         clusterversion.MinSupported,
 	dropTrigger:                       clusterversion.MinSupported,
 	dropView:                          clusterversion.MinSupported,
+	truncateTable:                     clusterversion.V25_4,
 }