Merge #144685

144685: workload: add support for RLS policies in RSW workload r=Dedej-Bergin a=Dedej-Bergin

This change adds CREATE POLICY and DROP POLICY operations to the Random Schema Changer workload for testing row-level security (RLS) policy functionality in the declarative schema changer.

Informs: #137120
Epic: CRDB-11724
Release note: none

Co-authored-by: Bergin Dedej <[email protected]>

Author: craig[bot]

pkg/workload/schemachange/operation_generator.go

@@ -1655,6 +1655,14 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		return nil, err
 	}
 
+	// Check if the table has any policies
+	tableHasPolicies := false
+	if tableExists {
+		if tableHasPolicies, err = og.tableHasPolicies(ctx, tx, tableName); err != nil {
+			return nil, err
+		}
+	}
+
 	columnName, err := og.randColumn(ctx, tx, *tableName, og.pctExisting(true))
 	if err != nil {
 		return nil, err
@@ -1703,11 +1711,40 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		// It is possible the column we are dropping is in the new primary key,
 		// so a potential error is an invalid reference in this case.
 		{code: pgcode.InvalidColumnReference, condition: og.useDeclarativeSchemaChanger && hasAlterPKSchemaChange},
+		// It is possible that we cannot drop column because
+		// it is referenced in a policy expression.
+		{code: pgcode.InvalidTableDefinition, condition: tableHasPolicies},
 	})
 	stmt.sql = fmt.Sprintf(`ALTER TABLE %s DROP COLUMN %s`, tableName.String(), columnName.String())
 	return stmt, nil
 }
 
+// tableHasPolicies checks if a table has any row-level security policies defined
+func (og *operationGenerator) tableHasPolicies(
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName,
+) (bool, error) {
+	// Query to check if a table has any RLS policies
+	query := `
+	SELECT EXISTS (
+		SELECT 1
+		FROM pg_policy p
+		JOIN pg_class t ON p.polrelid = t.oid
+		JOIN pg_namespace n ON t.relnamespace = n.oid
+		WHERE t.relname = $1
+		AND n.nspname = $2
+		LIMIT 1
+	)
+	`
+
+	var hasPolicies bool
+	err := tx.QueryRow(ctx, query, tableName.Object(), tableName.Schema()).Scan(&hasPolicies)
+	if err != nil {
+		return false, err
+	}
+
+	return hasPolicies, nil
+}
+
 func (og *operationGenerator) dropColumnDefault(ctx context.Context, tx pgx.Tx) (*opStmt, error) {
 	tableName, err := og.randTable(ctx, tx, og.pctExisting(true), "")
 	if err != nil {
@@ -4905,3 +4942,232 @@ func (og *operationGenerator) alterTableRLS(ctx context.Context, tx pgx.Tx) (*op
 	opStmt.sql = sqlStatement
 	return opStmt, nil
 }
+
+// generatePolicyExpression creates a random expression suitable for use in policy USING or WITH CHECK clauses
+func (og *operationGenerator) generatePolicyExpression(
+	columns []string, preferredColumn int,
+) string {
+	if len(columns) == 0 {
+		// Fallback to simple boolean if no columns are available
+		if og.randIntn(2) == 0 {
+			return "true"
+		}
+		return "false"
+	}
+
+	// Choose a column index (or use the preferred one if valid)
+	colIdx := og.randIntn(len(columns))
+	if preferredColumn >= 0 && preferredColumn < len(columns) {
+		colIdx = preferredColumn
+	}
+
+	// Generate a basic expression (IS NULL or IS NOT NULL)
+	var expr string
+	if og.randIntn(2) == 0 {
+		// IS NULL check
+		expr = fmt.Sprintf("%s IS NULL", columns[colIdx])
+	} else {
+		// IS NOT NULL check
+		expr = fmt.Sprintf("%s IS NOT NULL", columns[colIdx])
+	}
+
+	// Sometimes add complexity to the expression
+	if og.randIntn(3) == 0 {
+		expressionType := og.randIntn(3)
+		switch expressionType {
+		case 0:
+			// Add OR TRUE/FALSE
+			if og.randIntn(2) == 0 {
+				expr = fmt.Sprintf("(%s OR TRUE)", expr)
+			} else {
+				expr = fmt.Sprintf("(%s OR FALSE)", expr)
+			}
+		case 1:
+			// Use a different column if available for a compound expression
+			if len(columns) > 1 {
+				secondColIdx := (colIdx + 1) % len(columns)
+				if og.randIntn(2) == 0 {
+					expr = fmt.Sprintf("(%s OR %s IS NOT NULL)", expr, columns[secondColIdx])
+				} else {
+					expr = fmt.Sprintf("(%s AND %s IS NULL)", expr, columns[secondColIdx])
+				}
+			}
+		case 2:
+			// Add a comparison with a literal
+			if og.randIntn(2) == 0 {
+				expr = fmt.Sprintf("((%s) OR (TRUE))", expr)
+			} else {
+				expr = fmt.Sprintf("(%s AND current_user = current_user)", expr)
+			}
+		}
+	}
+
+	return expr
+}
+
+func (og *operationGenerator) createPolicy(ctx context.Context, tx pgx.Tx) (*opStmt, error) {
+	tableName, err := og.randTable(ctx, tx, og.pctExisting(true), "")
+	if err != nil {
+		return nil, err
+	}
+
+	// Check if table exists to include appropriate expected error
+	tableExists, err := og.tableExists(ctx, tx, tableName)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get columns for the table to reference in expressions
+	var columns []string
+	if tableExists {
+		columns, err = og.tableColumnsShuffled(ctx, tx, tableName.String())
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Generate a unique policy name
+	policyName := fmt.Sprintf("policy_%s", og.newUniqueSeqNumSuffix())
+
+	// Determine which policy components to include
+	includeUsing := og.randIntn(2) == 0     // 50% chance to include a USING expression
+	includeWithCheck := og.randIntn(2) == 0 // 50% chance to include WITH CHECK
+
+	// Build the SQL statement
+	var sqlStatement strings.Builder
+	sqlStatement.WriteString(fmt.Sprintf("CREATE POLICY %s ON %s", policyName, tableName))
+
+	if includeUsing {
+		usingExpr := og.generatePolicyExpression(columns, -1) // -1 means no preferred column
+		sqlStatement.WriteString(fmt.Sprintf(" USING (%s)", usingExpr))
+	}
+
+	if includeWithCheck {
+		// Try to use a different column for WITH CHECK if possible
+		preferredColIdx := -1
+		if len(columns) > 1 && includeUsing {
+			preferredColIdx = og.randIntn(len(columns))
+		}
+
+		withCheckExpr := og.generatePolicyExpression(columns, preferredColIdx)
+		sqlStatement.WriteString(fmt.Sprintf(" WITH CHECK (%s)", withCheckExpr))
+	}
+
+	// Create the operation statement
+	opStmt := makeOpStmt(OpStmtDDL)
+	opStmt.sql = sqlStatement.String()
+
+	opStmt.expectedExecErrors.addAll(codesWithConditions{
+		{code: pgcode.FeatureNotSupported, condition: !og.useDeclarativeSchemaChanger},
+		{code: pgcode.UndefinedTable, condition: !tableExists},
+	})
+
+	return opStmt, nil
+}
+
+// policyInfo pairs a table name with a name of a policy
+type policyInfo struct {
+	table      tree.TableName
+	policyName string
+}
+
+// findExistingPolicy returns a policyInfo struct with the qualified table name and policy name.
+// It also returns a boolean indicating whether a policy was found.
+func findExistingPolicy(
+	ctx context.Context, tx pgx.Tx, og *operationGenerator,
+) (*policyInfo, bool, error) {
+	var policyWithInfo policyInfo
+	policyExists := false
+
+	// Search for tables that have policies
+	policyTableQuery := `
+		SELECT
+			t.relname as table_name,
+			n.nspname as schema_name,
+			p.polname as policy_name
+		FROM
+			pg_policy p
+			JOIN pg_class t ON p.polrelid = t.oid
+			JOIN pg_namespace n ON t.relnamespace = n.oid
+		ORDER BY random()
+		LIMIT 1
+	`
+
+	rows, err := tx.Query(ctx, policyTableQuery)
+	if err != nil {
+		return nil, policyExists, err
+	}
+	defer rows.Close()
+
+	// Check if any rows were returned
+	for rows.Next() {
+		var tableName, schemaName, policyName string
+		if err := rows.Scan(&tableName, &schemaName, &policyName); err != nil {
+			return nil, policyExists, err
+		}
+		policyWithInfo = policyInfo{
+			table: tree.MakeTableNameFromPrefix(tree.ObjectNamePrefix{
+				SchemaName:     tree.Name(schemaName),
+				ExplicitSchema: true,
+			}, tree.Name(tableName)),
+			policyName: policyName,
+		}
+		policyExists = true
+	}
+
+	return &policyWithInfo, policyExists, nil
+}
+
+func (og *operationGenerator) dropPolicy(ctx context.Context, tx pgx.Tx) (*opStmt, error) {
+	policyWithInfo, policyExists, err := findExistingPolicy(ctx, tx, og)
+	if err != nil {
+		return nil, err
+	}
+
+	tableExists := true
+
+	if !policyExists {
+		// Fall back to random table if no tables with policies were found
+		randomTable, err := og.randTable(ctx, tx, og.pctExisting(true), "")
+		if err != nil {
+			return nil, err
+		}
+		policyWithInfo.table = *randomTable
+
+		// If we didn't get a real policy name, generate a random one
+		if policyWithInfo.policyName == "" {
+			policyWithInfo.policyName = fmt.Sprintf("dummy_policy_%s", og.newUniqueSeqNumSuffix())
+		}
+
+		// Check if table exists to include appropriate expected error
+		tableExists, err = og.tableExists(ctx, tx, randomTable)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Build the SQL statement
+	var sqlStatement strings.Builder
+	sqlStatement.WriteString("DROP POLICY ")
+
+	// Randomly decide whether to include IF EXISTS (60% chance)
+	includeIfExists := og.randIntn(100) < 60
+	if includeIfExists {
+		sqlStatement.WriteString("IF EXISTS ")
+	}
+
+	sqlStatement.WriteString(fmt.Sprintf("%s ON %s", policyWithInfo.policyName, &policyWithInfo.table))
+
+	// Create the operation statement
+	opStmt := makeOpStmt(OpStmtDDL)
+	opStmt.sql = sqlStatement.String()
+
+	opStmt.expectedExecErrors.addAll(codesWithConditions{
+		// The policy might not exist
+		{code: pgcode.UndefinedObject, condition: !policyExists && !includeIfExists},
+		// Table might not exist
+		{code: pgcode.UndefinedTable, condition: !tableExists && !includeIfExists},
+	})
+
+	return opStmt, nil
+}

pkg/workload/schemachange/optype.go

@@ -122,6 +122,7 @@ const (
 	createTableAs       // CREATE TABLE <table> AS <def>
 	createView          // CREATE VIEW <view> AS <def>
 	createFunction      // CREATE FUNCTION <function> ...
+	createPolicy        // CREATE POLICY <policy> ON <table> [TO <roles>] [USING (<using_expr>)] [WITH CHECK (<check_expr>)]
 
 	// COMMENT ON ...
 
@@ -135,6 +136,7 @@ const (
 	dropSequence // DROP SEQUENCE <sequence>
 	dropTable    // DROP TABLE <table>
 	dropView     // DROP VIEW <view>
+	dropPolicy   // DROP POLICY [IF EXISTS] <policy> ON <table>
 
 	// Unimplemented operations. TODO(sql-foundations): Audit and/or implement these operations.
 	// alterDatabaseOwner
@@ -234,6 +236,7 @@ var opFuncs = []func(*operationGenerator, context.Context, pgx.Tx) (*opStmt, err
 	commentOn:                         (*operationGenerator).commentOn,
 	createFunction:                    (*operationGenerator).createFunction,
 	createIndex:                       (*operationGenerator).createIndex,
+	createPolicy:                      (*operationGenerator).createPolicy,
 	createSchema:                      (*operationGenerator).createSchema,
 	createSequence:                    (*operationGenerator).createSequence,
 	createTable:                       (*operationGenerator).createTable,
@@ -243,6 +246,7 @@ var opFuncs = []func(*operationGenerator, context.Context, pgx.Tx) (*opStmt, err
 	createView:                        (*operationGenerator).createView,
 	dropFunction:                      (*operationGenerator).dropFunction,
 	dropIndex:                         (*operationGenerator).dropIndex,
+	dropPolicy:                        (*operationGenerator).dropPolicy,
 	dropSchema:                        (*operationGenerator).dropSchema,
 	dropSequence:                      (*operationGenerator).dropSequence,
 	dropTable:                         (*operationGenerator).dropTable,
@@ -286,6 +290,7 @@ var opWeights = []int{
 	commentOn:                         1,
 	createFunction:                    1,
 	createIndex:                       1,
+	createPolicy:                      1,
 	createSchema:                      1,
 	createSequence:                    1,
 	createTable:                       10,
@@ -295,6 +300,7 @@ var opWeights = []int{
 	createView:                        1,
 	dropFunction:                      1,
 	dropIndex:                         1,
+	dropPolicy:                        1,
 	dropSchema:                        1,
 	dropSequence:                      1,
 	dropTable:                         1,
@@ -324,12 +330,14 @@ var opDeclarativeVersion = map[opType]clusterversion.Key{
 	alterTableRLS:                     clusterversion.V25_2,
 	alterTypeDropValue:                clusterversion.MinSupported,
 	commentOn:                         clusterversion.MinSupported,
-	createIndex:                       clusterversion.MinSupported,
 	createFunction:                    clusterversion.MinSupported,
+	createIndex:                       clusterversion.MinSupported,
+	createPolicy:                      clusterversion.V25_2,
 	createSchema:                      clusterversion.MinSupported,
 	createSequence:                    clusterversion.MinSupported,
-	dropIndex:                         clusterversion.MinSupported,
 	dropFunction:                      clusterversion.MinSupported,
+	dropIndex:                         clusterversion.MinSupported,
+	dropPolicy:                        clusterversion.V25_2,
 	dropSchema:                        clusterversion.MinSupported,
 	dropSequence:                      clusterversion.MinSupported,
 	dropTable:                         clusterversion.MinSupported,