From 8fb47cd0a2349eec486ff44aa727ded9d33299fe Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Fri, 9 Feb 2024 16:00:12 -0500 Subject: [PATCH] workload/schemachange: fix and re-enable ALTER FUNCTION A few fixes were needed: - Include UDF params when generating ALTER FUNCTION statement. - Changed an errors.Is call to errors.HasType. - Avoid only using the public schema when creating functions. Release note: None --- pkg/workload/schemachange/generate.go | 4 +++ .../schemachange/operation_generator.go | 30 ++++++++++++++++--- pkg/workload/schemachange/optype.go | 8 ++--- 3 files changed, 34 insertions(+), 8 deletions(-) diff --git a/pkg/workload/schemachange/generate.go b/pkg/workload/schemachange/generate.go index d147c3c176ec..1f7c2eb93b5c 100644 --- a/pkg/workload/schemachange/generate.go +++ b/pkg/workload/schemachange/generate.go @@ -187,6 +187,10 @@ func Generate[T tree.Statement]( aIsSuccess := a.Code == pgcode.SuccessfulCompletion bIsSuccess := b.Code == pgcode.SuccessfulCompletion + if aIsSuccess && bIsSuccess { + // If both are valid, choose randomly. + return rng.Float64() < 0.5 + } return aIsSuccess && !bIsSuccess }) diff --git a/pkg/workload/schemachange/operation_generator.go b/pkg/workload/schemachange/operation_generator.go index c2987257578c..dfc26b22a6a1 100644 --- a/pkg/workload/schemachange/operation_generator.go +++ b/pkg/workload/schemachange/operation_generator.go @@ -198,7 +198,7 @@ func (og *operationGenerator) randOp( // We can only ignore this error, if no other PgErrors // were set in the clean up process. if errors.Is(err, pgx.ErrNoRows) && - !errors.Is(err, &pgconn.PgError{}) { + !errors.HasType(err, &pgconn.PgError{}) { continue } // Table select had a primary key swap, so no statement @@ -3878,11 +3878,18 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o COALESCE(member->>'direction' = 'REMOVE', false) AS dropping FROM enum_members `) + schemasQuery := With([]CTE{ + {"descriptors", descJSONQuery}, + }, `SELECT quote_ident(name) FROM descriptors WHERE descriptor ? 'schema'`) enums, err := Collect(ctx, og, tx, pgx.RowToMap, enumQuery) if err != nil { return nil, err } + schemas, err := Collect(ctx, og, tx, pgx.RowTo[string], schemasQuery) + if err != nil { + return nil, err + } // Roll some variables to ensure we have variance in the types of references // that we aside from being bound by what we could make references to. @@ -3916,6 +3923,9 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o name := tree.Name(fmt.Sprintf("udf_%s", og.newUniqueSeqNumSuffix())) return &name }, + "Schema": func() (string, error) { + return PickOne(og.params.rng, schemas) + }, "DroppingEnum": func() (string, error) { return PickOne(og.params.rng, droppingEnums) }, @@ -3947,9 +3957,9 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o // to the schema workload but may become a nice to have. stmt, expectedCode, err := Generate[*tree.CreateRoutine](og.params.rng, og.produceError(), []GenerationCase{ // 1. Nothing special, fully self contained function. - {pgcode.SuccessfulCompletion, `CREATE FUNCTION { UniqueName } (i int, j int) RETURNS VOID LANGUAGE SQL AS $$ SELECT NULL $$`}, + {pgcode.SuccessfulCompletion, `CREATE FUNCTION { Schema } . { UniqueName } (i int, j int) RETURNS VOID LANGUAGE SQL AS $$ SELECT NULL $$`}, // 2. 1 or more table or type references spread across parameters, return types, or the function body. - {pgcode.SuccessfulCompletion, `CREATE FUNCTION { UniqueName } ({ ParamRefs }) RETURNS { ReturnRefs } LANGUAGE SQL AS $$ SELECT NULL WHERE { BodyRefs } $$`}, + {pgcode.SuccessfulCompletion, `CREATE FUNCTION { Schema } . { UniqueName } ({ ParamRefs }) RETURNS { ReturnRefs } LANGUAGE SQL AS $$ SELECT NULL WHERE { BodyRefs } $$`}, // 3. Reference a table that does not exist. {pgcode.UndefinedTable, `CREATE FUNCTION { UniqueName } () RETURNS VOID LANGUAGE SQL AS $$ SELECT * FROM "ThisTableDoesNotExist" $$`}, // 4. Reference a UDT that does not exist. @@ -4092,7 +4102,19 @@ func (og *operationGenerator) alterFunctionSetSchema( functionsQuery := With([]CTE{ {"descriptors", descJSONQuery}, {"functions", functionDescsQuery}, - }, `SELECT quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) FROM functions`) + }, `SELECT + quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) || '(' || array_to_string(funcargs, ', ') || ')' + FROM functions + JOIN LATERAL ( + SELECT + COALESCE(array_agg(quote_ident(typnamespace::REGNAMESPACE::TEXT) || '.' || quote_ident(typname)), '{}') AS funcargs + FROM pg_catalog.pg_type + JOIN LATERAL ( + SELECT unnest(proargtypes) AS oid FROM pg_catalog.pg_proc WHERE oid = (id + 100000) + ) args ON args.oid = pg_type.oid + ) funcargs ON TRUE + `, + ) schemasQuery := With([]CTE{ {"descriptors", descJSONQuery}, diff --git a/pkg/workload/schemachange/optype.go b/pkg/workload/schemachange/optype.go index 6085b95a382d..164b67b5a170 100644 --- a/pkg/workload/schemachange/optype.go +++ b/pkg/workload/schemachange/optype.go @@ -262,10 +262,10 @@ var opWeights = []int{ alterDatabaseDropSuperRegion: 0, // Disabled and tracked with #111299 alterDatabasePrimaryRegion: 0, // Disabled and tracked with #83831 alterDatabaseSurvivalGoal: 0, // Disabled and tracked with #83831 - alterFunctionRename: 0, // Disabled and tracked with #116794. - alterFunctionSetSchema: 0, // Disabled and tracked with #116794. + alterFunctionRename: 1, + alterFunctionSetSchema: 1, alterTableAddColumn: 1, - alterTableAddConstraintForeignKey: 1, // Tentatively re-enabled, see #91195. + alterTableAddConstraintForeignKey: 1, alterTableAddConstraintUnique: 0, alterTableAlterColumnType: 0, // Disabled and tracked with #66662. alterTableAlterPrimaryKey: 1, @@ -288,7 +288,7 @@ var opWeights = []int{ createTableAs: 1, createTypeEnum: 1, createView: 1, - dropFunction: 0, // Disabled and tracked with #116794. + dropFunction: 1, dropIndex: 1, dropSchema: 0, // Disabled and tracked with 116792. dropSequence: 1,