From 22b8d19e18974a98c3c8f3d0a29e5a7bd9a3cece Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Tue, 17 Jan 2023 20:32:42 -0500 Subject: [PATCH] builtins: implement pg_get_indexdef as UDF Release note: None --- docs/generated/sql/functions.md | 4 ++ pkg/sql/sem/builtins/pg_builtins.go | 74 ++++++++++++----------------- 2 files changed, 34 insertions(+), 44 deletions(-) diff --git a/docs/generated/sql/functions.md b/docs/generated/sql/functions.md index 35786d7aed18..a626bc951097 100644 --- a/docs/generated/sql/functions.md +++ b/docs/generated/sql/functions.md @@ -3538,6 +3538,10 @@ table. Returns an error if validation fails.

Stable pg_get_functiondef(func_oid: oid) → string

For user-defined functions, returns the definition of the specified function. For builtin functions, returns the name of the function.

Stable +pg_get_indexdef(index_oid: oid) → string

Gets the CREATE INDEX command for index

+Stable +pg_get_indexdef(index_oid: oid, column_no: int, pretty_bool: bool) → string

Gets the CREATE INDEX command for index, or definition of just one index column when given a non-zero column number

+Stable pg_get_serial_sequence(table_name: string, column_name: string) → string

Returns the name of the sequence used by the given column_name in the table table_name.

Stable pg_get_viewdef(view_oid: oid) → string

Returns the CREATE statement for an existing view.

diff --git a/pkg/sql/sem/builtins/pg_builtins.go b/pkg/sql/sem/builtins/pg_builtins.go index 18f8cacfbb6a..7edb8a8fd117 100644 --- a/pkg/sql/sem/builtins/pg_builtins.go +++ b/pkg/sql/sem/builtins/pg_builtins.go @@ -204,48 +204,6 @@ var ( datEncodingUTF8ShortName = tree.NewDString("UTF8") ) -// Make a pg_get_indexdef function with the given arguments. -func makePGGetIndexDef(paramTypes tree.ParamTypes) tree.Overload { - return tree.Overload{ - Types: paramTypes, - ReturnType: tree.FixedReturnType(types.String), - Fn: func(ctx context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) { - colNumber := *tree.NewDInt(0) - if len(args) == 3 { - // The 1 argument and 3 argument variants are equivalent when column number 0 is passed. - colNumber = *args[1].(*tree.DInt) - } - // The 3 argument variant for column number other than 0 returns the column name. - r, err := evalCtx.Planner.QueryRowEx( - ctx, "pg_get_indexdef", - sessiondata.NoSessionDataOverride, - `SELECT CASE - WHEN $2 = 0 THEN defs.indexdef - -- If the column number does not exist in the index we return an empty string. - WHEN $2 < 0 OR $2 > array_length(i.indkey, 1) THEN '' - WHEN i.indkey[$2-1] = 0 THEN (indexprs::STRING[])[array_position(array_positions(i.indkey, 0), $2)] - ELSE a.attname - END as pg_get_indexdef - FROM pg_catalog.pg_index i - LEFT JOIN pg_attribute a ON (a.attrelid = i.indexrelid AND a.attnum = $2) - LEFT JOIN pg_indexes defs ON ($2 = 0 AND defs.crdb_oid = i.indexrelid) - WHERE i.indexrelid = $1`, args[0], colNumber) - if err != nil { - return nil, err - } - if len(r) == 0 { - return tree.DNull, nil - } - if len(r) > 1 { - return nil, errors.AssertionFailedf("pg_get_indexdef query has more than 1 result row: %+v", r) - } - return r[0], nil - }, - Info: "Gets the CREATE INDEX command for index, or definition of just one index column when given a column number", - Volatility: volatility.Stable, - } -} - // Make a pg_get_viewdef function with the given arguments. func makePGGetViewDef(paramTypes tree.ParamTypes) tree.Overload { return tree.Overload{ @@ -708,8 +666,36 @@ var pgBuiltins = map[string]builtinDefinition{ // statement. "pg_get_indexdef": makeBuiltin( tree.FunctionProperties{Category: builtinconstants.CategorySystemInfo, DistsqlBlocklist: true}, - makePGGetIndexDef(tree.ParamTypes{{Name: "index_oid", Typ: types.Oid}}), - makePGGetIndexDef(tree.ParamTypes{{Name: "index_oid", Typ: types.Oid}, {Name: "column_no", Typ: types.Int}, {Name: "pretty_bool", Typ: types.Bool}}), + tree.Overload{ + IsUDF: true, + Types: tree.ParamTypes{{Name: "index_oid", Typ: types.Oid}}, + ReturnType: tree.FixedReturnType(types.String), + Body: `SELECT indexdef FROM pg_catalog.pg_indexes WHERE crdb_oid = $1`, + Info: "Gets the CREATE INDEX command for index", + Volatility: volatility.Stable, + }, + tree.Overload{ + IsUDF: true, + Types: tree.ParamTypes{{Name: "index_oid", Typ: types.Oid}, {Name: "column_no", Typ: types.Int}, {Name: "pretty_bool", Typ: types.Bool}}, + ReturnType: tree.FixedReturnType(types.String), + Body: `SELECT CASE + WHEN $2 = 0 THEN defs.indexdef + WHEN $2 < 0 OR $2 > array_length(i.indkey, 1) THEN '' + -- array_positions(i.indkey, 0) returns the 1-based indexes of the indkey elements that are 0. + -- array_position(arr, $2) returns the 1-based index of the value $2 in arr. + -- indkey is an int2vector, which is accessed with 0-based indexes. + -- indexprs is a string[], which is accessed with 1-based indexes. + -- To put this all together, for the k-th 0 value inside of indkey, this will find the k-th indexpr. + WHEN i.indkey[$2-1] = 0 THEN (indexprs::STRING[])[array_position(array_positions(i.indkey, 0), $2)] + ELSE a.attname + END as pg_get_indexdef + FROM pg_catalog.pg_index i + LEFT JOIN pg_attribute a ON (a.attrelid = i.indexrelid AND a.attnum = $2) + LEFT JOIN pg_indexes defs ON ($2 = 0 AND defs.crdb_oid = i.indexrelid) + WHERE i.indexrelid = $1`, + Info: "Gets the CREATE INDEX command for index, or definition of just one index column when given a non-zero column number", + Volatility: volatility.Stable, + }, ), // pg_get_viewdef functions like SHOW CREATE VIEW but returns the same format as