Merge #145219

145219: workload/schemachanger: add testing for schema_locked tables r=fqazi a=fqazi

Previously, the schema changer workload lacked randomized testing for schema_locked tables. This resulted in a lack of coverage for schema changes involving these tables. This patch adds support for creating and using
schema_locked tables within the declarative schema changer tests.

This patch also adds support for setting creating tables with schema_locked with in a transaction to simplify testing.

Fixes: #144591
Release note: None

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/sql/alter_table.go

@@ -731,7 +731,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 			}
 
 		case *tree.AlterTableSetStorageParams:
-			setter := tablestorageparam.NewSetter(n.tableDesc)
+			setter := tablestorageparam.NewSetter(n.tableDesc, false /* isNewObject */)
 			if err := storageparam.Set(
 				params.ctx,
 				params.p.SemaCtx(),
@@ -761,7 +761,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 			}
 
 		case *tree.AlterTableResetStorageParams:
-			setter := tablestorageparam.NewSetter(n.tableDesc)
+			setter := tablestorageparam.NewSetter(n.tableDesc, false /* isNewObject */)
 			if err := storageparam.Reset(
 				params.ctx,
 				params.EvalContext(),

pkg/sql/create_table.go

@@ -1420,8 +1420,7 @@ func NewTableDesc(
 	desc := tabledesc.InitTableDescriptor(
 		id, dbID, sc.GetID(), n.Table.Table(), creationTime, privileges, persistence,
 	)
-
-	setter := tablestorageparam.NewSetter(&desc)
+	setter := tablestorageparam.NewSetter(&desc, true /* isNewObject */)
 	if err := storageparam.Set(
 		ctx,
 		semaCtx,

pkg/sql/storageparam/indexstorageparam/index_storage_param.go

@@ -223,3 +223,9 @@ func (po *Setter) RunPostChecks() error {
 
 	return nil
 }
+
+// IsNewTableObject implements the Setter interface.
+func (po *Setter) IsNewTableObject() bool {
+	//Not applicable to indexes.
+	panic(errors.AssertionFailedf("not-implemented for indexes"))
+}

pkg/sql/storageparam/storage_param.go

@@ -34,6 +34,9 @@ type Setter interface {
 	// This allows checking whether multiple storage parameters together
 	// form a valid configuration.
 	RunPostChecks() error
+	// IsNewTableObject returns true if the storage parameter is being set on a new
+	// table.
+	IsNewTableObject() bool
 }
 
 // Set sets the given storage parameters using the
@@ -45,7 +48,7 @@ func Set(
 	params tree.StorageParams,
 	setter Setter,
 ) error {
-	if err := storageParamPreChecks(ctx, evalCtx, params, nil /* resetParams */); err != nil {
+	if err := storageParamPreChecks(ctx, evalCtx, setter, params, nil /* resetParams */); err != nil {
 		return err
 	}
 	for _, sp := range params {
@@ -92,7 +95,7 @@ func Set(
 func Reset(
 	ctx context.Context, evalCtx *eval.Context, params []string, paramObserver Setter,
 ) error {
-	if err := storageParamPreChecks(ctx, evalCtx, nil /* setParam */, params); err != nil {
+	if err := storageParamPreChecks(ctx, evalCtx, paramObserver, nil /* setParam */, params); err != nil {
 		return err
 	}
 	for _, p := range params {
@@ -126,7 +129,11 @@ func SetFillFactor(ctx context.Context, evalCtx *eval.Context, key string, datum
 // storageParamPreChecks is where we specify pre-conditions for setting/resetting
 // storage parameters `param`.
 func storageParamPreChecks(
-	ctx context.Context, evalCtx *eval.Context, setParams tree.StorageParams, resetParams []string,
+	ctx context.Context,
+	evalCtx *eval.Context,
+	setter Setter,
+	setParams tree.StorageParams,
+	resetParams []string,
 ) error {
 	if setParams != nil && resetParams != nil {
 		return errors.AssertionFailedf("only one of setParams and resetParams should be non-nil.")
@@ -152,7 +159,11 @@ func storageParamPreChecks(
 			// change we make to the descriptor in the transaction, so we can uphold
 			// the "one-version invariant" as discussed further in RFC
 			// https://github.com/ajwerner/cockroach/blob/ajwerner/low-latency-rfc-take-3/docs/RFCS/20230328_low_latency_changefeeds.md
-			if len(keys) > 1 || !evalCtx.TxnImplicit || !evalCtx.TxnIsSingleStmt {
+			// For newly created tables we will allow schema_locked to be set upon creation,
+			// since later operations cannot unset schema_locked (i.e. only implicit single
+			// statement transactions are allowed to manipulate schema_locked, see
+			// checkSchemaChangeIsAllowed).
+			if !setter.IsNewTableObject() && (len(keys) > 1 || !evalCtx.TxnImplicit || !evalCtx.TxnIsSingleStmt) {
 				return pgerror.Newf(pgcode.InvalidParameterValue, "%q can only be set/reset on "+
 					"its own without other parameters in a single-statement implicit transaction.", key)
 			}

pkg/sql/storageparam/tablestorageparam/table_storage_param.go

@@ -36,18 +36,22 @@ type Setter struct {
 	// UpdatedRowLevelTTL is kept separate from the RowLevelTTL in TableDesc
 	// in case changes need to be made in schema changer.
 	UpdatedRowLevelTTL *catpb.RowLevelTTL
+
+	// NewObject bool tracks if this is a newly created object.
+	NewObject bool
 }
 
 var _ storageparam.Setter = (*Setter)(nil)
 
 // NewSetter returns a new Setter.
-func NewSetter(tableDesc *tabledesc.Mutable) *Setter {
+func NewSetter(tableDesc *tabledesc.Mutable, isNewObject bool) *Setter {
 	var updatedRowLevelTTL *catpb.RowLevelTTL
 	if tableDesc.HasRowLevelTTL() {
 		updatedRowLevelTTL = protoutil.Clone(tableDesc.GetRowLevelTTL()).(*catpb.RowLevelTTL)
 	}
 	return &Setter{
 		TableDesc:          tableDesc,
+		NewObject:          isNewObject,
 		UpdatedRowLevelTTL: updatedRowLevelTTL,
 	}
 }
@@ -60,6 +64,11 @@ func (po *Setter) RunPostChecks() error {
 	return nil
 }
 
+// IsNewTableObject implements the Setter interface.
+func (po *Setter) IsNewTableObject() bool {
+	return po.NewObject
+}
+
 func boolFromDatum(
 	ctx context.Context, evalCtx *eval.Context, key string, datum tree.Datum,
 ) (bool, error) {

pkg/workload/schemachange/operation_generator.go

@@ -1306,6 +1306,14 @@ func (og *operationGenerator) createTable(ctx context.Context, tx pgx.Tx) (*opSt
 		return false
 	}()
 
+	// Randomly create as schema locked table.
+	if og.params.rng.Intn(2) == 0 {
+		stmt.StorageParams = append(stmt.StorageParams, tree.StorageParam{
+			Key:   "schema_locked",
+			Value: tree.DBoolTrue,
+		})
+	}
+
 	tableExists, err := og.tableExists(ctx, tx, tableName)
 	if err != nil {
 		return nil, err
@@ -3821,6 +3829,14 @@ func (og *operationGenerator) randTypeName(
 func (og *operationGenerator) randTable(
 	ctx context.Context, tx pgx.Tx, pctExisting int, desiredSchema string,
 ) (*tree.TableName, error) {
+	// Because the declarative schema change can automatically set / unset
+	// schema_locked on tables, we will allow random table selection include
+	// schema_locked tables. When working with the legacy schema changer, we
+	// will intentionally only select non-schema locked tables.
+	excludeSchemaLocked := "  AND create_statement NOT LIKE '%schema_locked%' "
+	if og.useDeclarativeSchemaChanger {
+		excludeSchemaLocked = ""
+	}
 	if err := og.setSeedInDB(ctx, tx); err != nil {
 		return nil, err
 	}
@@ -3833,13 +3849,15 @@ func (og *operationGenerator) randTable(
 			return &treeTableName, nil
 		}
 		q := fmt.Sprintf(`
-		  SELECT table_name
-		    FROM [SHOW TABLES]
-		   WHERE table_name SIMILAR TO 'table_w[0-9]_+%%'
+		  SELECT descriptor_name 
+		    FROM crdb_internal.create_statements
+		   WHERE descriptor_name SIMILAR TO 'table_w[0-9]_+%%'
 				 AND schema_name = '%s'
+		     AND descriptor_type='table'
+		 	   %s 
 		ORDER BY random()
 		   LIMIT 1;
-		`, desiredSchema)
+		`, desiredSchema, excludeSchemaLocked)
 
 		var tableName string
 		if err := tx.QueryRow(ctx, q).Scan(&tableName); err != nil {
@@ -3870,13 +3888,17 @@ func (og *operationGenerator) randTable(
 		return &treeTableName, nil
 	}
 
-	const q = `
-  SELECT schema_name, table_name
-    FROM [SHOW TABLES]
-   WHERE table_name SIMILAR TO 'table_w[0-9]_+%'
-ORDER BY random()
-   LIMIT 1;
-`
+	q := fmt.Sprintf(`
+SELECT schema_name, descriptor_name 
+		    FROM crdb_internal.create_statements
+		   WHERE descriptor_name SIMILAR TO 'table_w[0-9]_+%%'
+		     AND descriptor_type='table'
+		 	   %s 
+		ORDER BY random()
+		   LIMIT 1;
+`,
+		excludeSchemaLocked)
+
 	var schemaName string
 	var tableName string
 	if err := tx.QueryRow(ctx, q).Scan(&schemaName, &tableName); err != nil {