Skip to content

Commit

Permalink
workload/schemachange: support UDFs invoking other UDFs
Browse files Browse the repository at this point in the history 
Previously, the randomized schema changer workload only created routines
that referred to types or tables. Since we added support for UDFs
calling other UDFs we need to add support. To address this, this patch
adds support for creating UDFs that can refer/invoke to other UDFs.

Fixes: #120671

Release note: None
  • Loading branch information
@fqazi
fqazi committed Mar 26, 2024
1 parent f03ac0c commit 6def369
Show file tree
Hide file tree
Showing 3 changed files with 143 additions and 9 deletions.
18 changes: 18 additions & 0 deletions pkg/workload/schemachange/error_screening.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1301,6 +1301,24 @@ SELECT EXISTS(
`, tableName.String(), columnName)
}

func (og *operationGenerator) schemaContainsHasReferredFunctions(
ctx context.Context, tx pgx.Tx, schemaName string,
) (bool, error) {
ctes := []CTE{
{"descriptors", descJSONQuery},
{"functions", functionDescsQuery},
{"pg_depends_from_diff_schema", `
SELECT refobjid FROM pg_depend as d, functions as src_function, functions as dst_function
WHERE src_function.schema_id <> $1::REGNAMESPACE::INT8 AND dst_function.schema_id=$1::REGNAMESPACE::INT8 AND
d.refobjid=(src_function.id+100000) AND d.objid=(dst_function.id) AND
d.classid = 'pg_catalog.pg_proc'::REGCLASS::INT8 AND d.refclassid = 'pg_catalog.pg_proc'::REGCLASS::INT8`},
}

result, err := Collect(ctx, og, tx, pgx.RowToMap, With(ctes,
"SELECT * FROM pg_depends_from_diff_schema;"), schemaName)
return len(result) > 0, err
}

func (og *operationGenerator) schemaContainsTypesWithCrossSchemaReferences(
ctx context.Context, tx pgx.Tx, schemaName string,
) (bool, error) {
Expand Down
126 changes: 117 additions & 9 deletions pkg/workload/schemachange/operation_generator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3875,17 +3875,24 @@ func (og *operationGenerator) dropSchema(ctx context.Context, tx pgx.Tx) (*opStm
return nil, err
}
crossReferences := false
crossSchemaFunctionReferences := false
if schemaExists {
crossReferences, err = og.schemaContainsTypesWithCrossSchemaReferences(ctx, tx, schemaName)
if err != nil {
return nil, err
}
crossSchemaFunctionReferences, err = og.schemaContainsHasReferredFunctions(ctx, tx, schemaName)
if err != nil {
return nil, err
}
}

stmt := makeOpStmt(OpStmtDDL)
stmt.expectedExecErrors.addAll(codesWithConditions{
{pgcode.UndefinedSchema, !schemaExists},
{pgcode.InvalidSchemaName, schemaName == catconstants.PublicSchemaName},
{pgcode.FeatureNotSupported, crossReferences && !og.useDeclarativeSchemaChanger},
{pgcode.DependentObjectsStillExist, crossSchemaFunctionReferences && !og.useDeclarativeSchemaChanger},
})

stmt.sql = fmt.Sprintf(`DROP SCHEMA "%s" CASCADE`, schemaName)
Expand Down Expand Up @@ -3923,6 +3930,16 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o
{"descriptors", descJSONQuery},
}, `SELECT quote_ident(name) FROM descriptors WHERE descriptor ? 'schema'`)

functionsQuery := With([]CTE{
{"descriptors", descJSONQuery},
{"functions", functionDescsQuery},
}, `SELECT
quote_ident(schema_id::REGNAMESPACE::STRING) AS schema,
quote_ident(name) AS name,
array_to_string(proargnames, ',') AS args
FROM
functions
INNER JOIN pg_catalog.pg_proc ON oid = (id + 100000);`)
enums, err := Collect(ctx, og, tx, pgx.RowToMap, enumQuery)
if err != nil {
return nil, err
Expand All @@ -3935,6 +3952,10 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o
if err != nil {
return nil, err
}
functions, err := Collect(ctx, og, tx, pgx.RowToMap, functionsQuery)
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.
Expand Down Expand Up @@ -3971,6 +3992,25 @@ func (og *operationGenerator) createFunction(ctx context.Context, tx pgx.Tx) (*o
possibleBodyReferences = append(possibleBodyReferences, fmt.Sprintf(`((SELECT count(*) FROM %s LIMIT 0) = 0)`, table))
}

// For each function generate a possible invocation passing in null arguments.
for _, function := range functions {
args := ""
if function["args"] != nil {
args = function["args"].(string)
argIn := strings.Builder{}
// TODO(fqazi): Longer term we should populate actual arguments and
// not just NULLs.
for range strings.Split(args, ",") {
if argIn.Len() > 0 {
argIn.WriteString(",")
}
argIn.WriteString("NULL")
}
args = argIn.String()
}
possibleBodyReferences = append(possibleBodyReferences, fmt.Sprintf("(SELECT %s.%s(%s) IS NOT NULL)", function["schema"].(string), function["name"].(string), args))
}

placeholderMap := template.FuncMap{
"UniqueName": func() *tree.Name {
name := tree.Name(fmt.Sprintf("udf_%s", og.newUniqueSeqNumSuffix()))
Expand Down Expand Up @@ -4057,7 +4097,8 @@ func (og *operationGenerator) dropFunction(ctx context.Context, tx pgx.Tx) (*opS
{"descriptors", descJSONQuery},
{"functions", functionDescsQuery},
}, `SELECT
quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) || '(' || array_to_string(funcargs, ', ') || ')'
quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) || '(' || array_to_string(funcargs, ', ') || ')' as name,
(id + 100000) as func_oid
FROM functions
JOIN LATERAL (
SELECT
Expand All @@ -4070,18 +4111,55 @@ func (og *operationGenerator) dropFunction(ctx context.Context, tx pgx.Tx) (*opS
`,
)

functions, err := Collect(ctx, og, tx, pgx.RowTo[string], q)
functions, err := Collect(ctx, og, tx, pgx.RowToMap, q)
if err != nil {
return nil, err
}

functionDeps, err := Collect(ctx, og, tx, pgx.RowToMap,
With([]CTE{
{"function_deps", functionDepsQuery},
},
`SELECT DISTINCT to_oid::INT8 FROM function_deps;`,
))
if err != nil {
return nil, err
}

functionDepsMap := make(map[int64]struct{})
for _, f := range functionDeps {
functionDepsMap[f["to_oid"].(int64)] = struct{}{}
}

functionWithDeps := make([]map[string]any, 0, len(functions))
functionWithoutDeps := make([]map[string]any, 0, len(functions))
for _, f := range functions {
if _, ok := functionDepsMap[f["func_oid"].(int64)]; ok {
functionWithDeps = append(functionWithDeps, f)
} else {
functionWithoutDeps = append(functionWithoutDeps, f)
}
}

stmt, expectedCode, err := Generate[*tree.DropRoutine](og.params.rng, og.produceError(), []GenerationCase{
{pgcode.UndefinedFunction, `DROP FUNCTION "NoSuchFunction"`},
{pgcode.SuccessfulCompletion, `DROP FUNCTION IF EXISTS "NoSuchFunction"`},
{pgcode.SuccessfulCompletion, `DROP FUNCTION { Function }`},
{pgcode.SuccessfulCompletion, `DROP FUNCTION { FunctionWithoutDeps }`},
{pgcode.DependentObjectsStillExist, `DROP FUNCTION { FunctionWithDeps }`},
}, template.FuncMap{
"Function": func() (string, error) {
return PickOne(og.params.rng, functions)
"FunctionWithoutDeps": func() (string, error) {
one, err := PickOne(og.params.rng, functionWithoutDeps)
if err != nil {
return "", err
}
return one["name"].(string), nil
},
"FunctionWithDeps": func() (string, error) {
one, err := PickOne(og.params.rng, functionWithDeps)
if err != nil {
return "", err
}
return one["name"].(string), nil
},
})

Expand All @@ -4100,7 +4178,8 @@ func (og *operationGenerator) alterFunctionRename(ctx context.Context, tx pgx.Tx
}, `SELECT
quote_ident(schema_id::REGNAMESPACE::TEXT) AS schema,
quote_ident(name) AS name,
quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) || '(' || array_to_string(funcargs, ', ') || ')' AS qualified_name
quote_ident(schema_id::REGNAMESPACE::TEXT) || '.' || quote_ident(name) || '(' || array_to_string(funcargs, ', ') || ')' AS qualified_name,
(id + 100000) as func_oid
FROM functions
JOIN LATERAL (
SELECT
Expand All @@ -4117,6 +4196,31 @@ func (og *operationGenerator) alterFunctionRename(ctx context.Context, tx pgx.Tx
return nil, err
}

functionDeps, err := Collect(ctx, og, tx, pgx.RowToMap,
With([]CTE{
{"function_deps", functionDepsQuery},
},
`SELECT DISTINCT to_oid::INT8 FROM function_deps;`,
))
if err != nil {
return nil, err
}

functionDepsMap := make(map[int64]struct{})
for _, f := range functionDeps {
functionDepsMap[f["to_oid"].(int64)] = struct{}{}
}

functionWithDeps := make([]map[string]any, 0, len(functions))
functionWithoutDeps := make([]map[string]any, 0, len(functions))
for _, f := range functions {
if _, ok := functionDepsMap[f["func_oid"].(int64)]; ok {
functionWithDeps = append(functionWithDeps, f)
} else {
functionWithoutDeps = append(functionWithoutDeps, f)
}
}

stmt, expectedCode, err := Generate[*tree.AlterRoutineRename](og.params.rng, og.produceError(), []GenerationCase{
{pgcode.UndefinedFunction, `ALTER FUNCTION "NoSuchFunction" RENAME TO "IrrelevantFunctionName"`},
// TODO(chrisseto): Neither of these seem to work as expected. Renaming a
Expand All @@ -4125,14 +4229,18 @@ func (og *operationGenerator) alterFunctionRename(ctx context.Context, tx pgx.Tx
// something to do with search paths and/or function overloads.
// {pgcode.DuplicateFunction, `{ with ExistingFunction } ALTER FUNCTION { .qualified_name } RENAME TO { ConflictingName . } { end }`},
// {pgcode.DuplicateFunction, `{ with ExistingFunction } ALTER FUNCTION { .qualified_name } RENAME TO { .name } { end }`},
{pgcode.SuccessfulCompletion, `ALTER FUNCTION { (ExistingFunction).qualified_name } RENAME TO { UniqueName }`},
{pgcode.SuccessfulCompletion, `ALTER FUNCTION { (ExistingFunctionWithoutDeps).qualified_name } RENAME TO { UniqueName }`},
{pgcode.FeatureNotSupported, `ALTER FUNCTION { (ExistingFunctionWithDeps).qualified_name } RENAME TO { UniqueName }`},
}, template.FuncMap{
"UniqueName": func() *tree.Name {
name := tree.Name(fmt.Sprintf("udf_%s", og.newUniqueSeqNumSuffix()))
return &name
},
"ExistingFunction": func() (map[string]any, error) {
return PickOne(og.params.rng, functions)
"ExistingFunctionWithoutDeps": func() (map[string]any, error) {
return PickOne(og.params.rng, functionWithoutDeps)
},
"ExistingFunctionWithDeps": func() (map[string]any, error) {
return PickOne(og.params.rng, functionWithDeps)
},
"ConflictingName": func(existing map[string]any) (string, error) {
selected, err := PickOne(og.params.rng, util.Filter(functions, func(other map[string]any) bool {
Expand Down
8 changes: 8 additions & 0 deletions pkg/workload/schemachange/query_util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,14 @@ const (
functionDescsQuery = `SELECT id, schema_id, name, descriptor->'function' AS descriptor FROM descriptors WHERE descriptor ? 'function'`

regionsFromClusterQuery = `SELECT * FROM [SHOW REGIONS FROM CLUSTER]`

functionDepsQuery = `SELECT
objid AS from_oid, refobjid AS to_oid
FROM
pg_depend AS d
WHERE
d.classid = 'pg_catalog.pg_proc'::REGCLASS::INT8
AND d.refclassid = 'pg_catalog.pg_proc'::REGCLASS::INT8`
)

func regionsFromDatabaseQuery(database string) string {
Expand Down

0 comments on commit 6def369

Please sign in to comment.