workload/schemachanger: add randomized testing for TRUNCATE

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

Author: fqazi

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,
 }