Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

release-23.2: backupccl: replace UDT IDs within routine bodies and views #116841

Merged
merged 1 commit into from Jan 3, 2024
Merged

release-23.2: backupccl: replace UDT IDs within routine bodies and views #116841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
18 changes: 18 additions & 0 deletions pkg/ccl/backupccl/testdata/backup-restore/plpgsql_procedures
View file
Open in desktop
Expand Up @@ -10,8 +10,11 @@ CREATE TABLE sc1.tbl1(a INT PRIMARY KEY);
CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE PROCEDURE sc1.p1(a sc1.enum1) LANGUAGE PLpgSQL AS $$
DECLARE
foobar sc1.enum1;
BEGIN
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
END
$$;
Expand Down Expand Up @@ -66,8 +69,11 @@ SELECT create_statement FROM [SHOW CREATE PROCEDURE sc1.p1]
CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE plpgsql
AS $$
DECLARE
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
END

Expand Down Expand Up @@ -107,8 +113,11 @@ SELECT create_statement FROM [SHOW CREATE PROCEDURE sc1.p1]
CREATE PROCEDURE sc1.p1(IN a db1_new.sc1.enum1)
LANGUAGE plpgsql
AS $$
DECLARE
foobar db1_new.sc1.enum1;
BEGIN
SELECT a FROM db1_new.sc1.tbl1;
SELECT 'Good':::db1_new.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
END

Expand Down Expand Up @@ -183,8 +192,11 @@ CREATE TABLE sc1.tbl1(a INT PRIMARY KEY);
CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE PROCEDURE sc1.p1(a sc1.enum1) LANGUAGE PLpgSQL AS $$
DECLARE
foobar sc1.enum1;
BEGIN
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
END
$$;
Expand Down Expand Up @@ -244,8 +256,11 @@ SELECT create_statement FROM [SHOW CREATE PROCEDURE sc1.p1]
CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE plpgsql
AS $$
DECLARE
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
END

Expand Down Expand Up @@ -286,8 +301,11 @@ SELECT create_statement FROM [SHOW CREATE PROCEDURE sc1.p1]
CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE plpgsql
AS $$
DECLARE
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
END

Expand Down
12 changes: 12 additions & 0 deletions pkg/ccl/backupccl/testdata/backup-restore/plpgsql_user_defined_functions
View file
Open in desktop
Expand Up @@ -12,8 +12,10 @@ CREATE SEQUENCE sc1.sq1;
CREATE FUNCTION sc1.f1(a sc1.enum1) RETURNS INT LANGUAGE PLpgSQL AS $$
DECLARE
x INT := 0;
foobar sc1.enum1;
BEGIN
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
RETURN nextval('sc1.sq1');
END
$$;
Expand Down Expand Up @@ -78,8 +80,10 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
AS $$
DECLARE
x INT8 := 0;
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
RETURN nextval('sc1.sq1'::REGCLASS);
END

Expand Down Expand Up @@ -121,8 +125,10 @@ CREATE FUNCTION sc1.f1(IN a db1_new.sc1.enum1)
AS $$
DECLARE
x INT8 := 0;
foobar db1_new.sc1.enum1;
BEGIN
SELECT a FROM db1_new.sc1.tbl1;
SELECT 'Good':::db1_new.sc1.enum1;
RETURN nextval('sc1.sq1'::REGCLASS);
END

Expand Down Expand Up @@ -188,8 +194,10 @@ CREATE SEQUENCE sc1.sq1;
CREATE FUNCTION sc1.f1(a sc1.enum1) RETURNS INT LANGUAGE PLpgSQL AS $$
DECLARE
x INT;
foobar sc1.enum1;
BEGIN
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1') INTO x;
RETURN x;
END
Expand Down Expand Up @@ -257,8 +265,10 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
AS $$
DECLARE
x INT8;
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS) INTO x;
RETURN x;
END
Expand Down Expand Up @@ -302,8 +312,10 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
AS $$
DECLARE
x INT8;
foobar db1.sc1.enum1;
BEGIN
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS) INTO x;
RETURN x;
END
Expand Down
6 changes: 6 additions & 0 deletions pkg/ccl/backupccl/testdata/backup-restore/procedures
View file
Open in desktop
Expand Up @@ -11,6 +11,7 @@ CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE PROCEDURE sc1.p1(a sc1.enum1) LANGUAGE SQL AS $$
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
$$;
CREATE SCHEMA sc2;
Expand Down Expand Up @@ -62,6 +63,7 @@ CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -97,6 +99,7 @@ CREATE PROCEDURE sc1.p1(IN a db1_new.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1_new.sc1.tbl1;
SELECT 'Good':::db1_new.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -168,6 +171,7 @@ CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE PROCEDURE sc1.p1(a sc1.enum1) LANGUAGE SQL AS $$
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
$$;
CREATE SCHEMA sc2;
Expand Down Expand Up @@ -224,6 +228,7 @@ CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -260,6 +265,7 @@ CREATE PROCEDURE sc1.p1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down
6 changes: 6 additions & 0 deletions pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions
View file
Open in desktop
Expand Up @@ -11,6 +11,7 @@ CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE FUNCTION sc1.f1(a sc1.enum1) RETURNS INT LANGUAGE SQL AS $$
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
$$;
CREATE SCHEMA sc2;
Expand Down Expand Up @@ -65,6 +66,7 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -100,6 +102,7 @@ CREATE FUNCTION sc1.f1(IN a db1_new.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1_new.sc1.tbl1;
SELECT 'Good':::db1_new.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -160,6 +163,7 @@ CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE SEQUENCE sc1.sq1;
CREATE FUNCTION sc1.f1(a sc1.enum1) RETURNS INT LANGUAGE SQL AS $$
SELECT a FROM sc1.tbl1;
SELECT 'Good'::sc1.enum1;
SELECT nextval('sc1.sq1');
$$;
CREATE SCHEMA sc2;
Expand Down Expand Up @@ -219,6 +223,7 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down Expand Up @@ -255,6 +260,7 @@ CREATE FUNCTION sc1.f1(IN a db1.sc1.enum1)
LANGUAGE SQL
AS $$
SELECT a FROM db1.sc1.tbl1;
SELECT 'Good':::db1.sc1.enum1;
SELECT nextval('sc1.sq1'::REGCLASS);
$$

Expand Down
23 changes: 17 additions & 6 deletions pkg/ccl/backupccl/testdata/backup-restore/views
View file
Open in desktop
Expand Up @@ -6,8 +6,9 @@ exec-sql
CREATE DATABASE db1;
USE db1;
CREATE SCHEMA sc1;
CREATE TYPE sc1.enum1 AS ENUM('Good');
CREATE TABLE sc1.tbl1(a INT PRIMARY KEY);
CREATE VIEW sc1.v1 AS SELECT a FROM sc1.tbl1;
CREATE VIEW sc1.v1 AS SELECT a, 'Good'::sc1.enum1 FROM sc1.tbl1;
----

exec-sql
Expand All @@ -17,26 +18,36 @@ INSERT INTO sc1.tbl1 VALUES (123);
query-sql
SELECT * FROM sc1.v1;
----
123
123 Good

query-sql
SHOW CREATE VIEW sc1.v1;
----
sc1.v1 CREATE VIEW sc1.v1 (
a
) AS SELECT a FROM db1.sc1.tbl1
a,
enum1
) AS SELECT a, 'Good':::db1.sc1.enum1 FROM db1.sc1.tbl1

exec-sql
BACKUP DATABASE db1 INTO 'nodelocal://1/test/'
----

exec-sql
DROP DATABASE db1
----

exec-sql
RESTORE DATABASE db1 FROM LATEST IN 'nodelocal://1/test/' WITH new_db_name = db1_new
----

exec-sql
USE db1_new
----

query-sql
SHOW CREATE VIEW db1_new.sc1.v1;
----
db1_new.sc1.v1 CREATE VIEW sc1.v1 (
a
) AS SELECT a FROM db1_new.sc1.tbl1
a,
enum1
) AS SELECT a, 'Good':::db1_new.sc1.enum1 FROM db1_new.sc1.tbl1
66 changes: 63 additions & 3 deletions pkg/sql/catalog/rewrite/rewrite.go
View file
Open in desktop
Expand Up @@ -117,6 +117,10 @@ func TableDescs(
if err != nil {
return err
}
viewQuery, err = rewriteTypesInView(viewQuery, descriptorRewrites)
if err != nil {
return err
}
table.ViewQuery = viewQuery
}

Expand Down Expand Up @@ -375,8 +379,62 @@ func rewriteTypesInExpr(expr string, rewrites jobspb.DescRewriteMap) (string, er
if err != nil {
return "", err
}
ctx := makeTypeReplaceFmtCtx(rewrites)
ctx.FormatNode(parsed)
return ctx.CloseAndGetString(), nil
}

ctx := tree.NewFmtCtx(
// rewriteTypesInView rewrites all explicit ID type references in the input view
// query string according to rewrites.
func rewriteTypesInView(viewQuery string, rewrites jobspb.DescRewriteMap) (string, error) {
stmt, err := parser.ParseOne(viewQuery)
if err != nil {
return "", err
}
ctx := makeTypeReplaceFmtCtx(rewrites)
ctx.FormatNode(stmt.AST)
return ctx.CloseAndGetString(), nil
}

func rewriteTypesInRoutine(
fnBody string, rewrites jobspb.DescRewriteMap, lang catpb.Function_Language,
) (string, error) {
switch lang {
case catpb.Function_SQL:
fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
stmts, err := parser.Parse(fnBody)
if err != nil {
return "", err
}
for i, stmt := range stmts {
if i > 0 {
fmtCtx.WriteString("\n")
}
typeReplaceCtx := makeTypeReplaceFmtCtx(rewrites)
typeReplaceCtx.FormatNode(stmt.AST)
fmtCtx.WriteString(typeReplaceCtx.CloseAndGetString())
fmtCtx.WriteString(";")
}
return fmtCtx.CloseAndGetString(), nil

case catpb.Function_PLPGSQL:
stmt, err := plpgsqlparser.Parse(fnBody)
if err != nil {
return "", err
}
typeReplaceCtx := makeTypeReplaceFmtCtx(rewrites)
typeReplaceCtx.FormatNode(stmt.AST)
return typeReplaceCtx.CloseAndGetString(), nil

default:
return "", errors.AssertionFailedf("unexpected function language: %v", lang)
}
}

// makeTypeReplaceFmtCtx returns a FmtCtx which rewrites explicit ID references
// according to the rewrites map.
func makeTypeReplaceFmtCtx(rewrites jobspb.DescRewriteMap) *tree.FmtCtx {
return tree.NewFmtCtx(
tree.FmtSerializable,
tree.FmtIndexedTypeFormat(func(ctx *tree.FmtCtx, ref *tree.OIDTypeReference) {
newRef := ref
Expand All @@ -387,8 +445,6 @@ func rewriteTypesInExpr(expr string, rewrites jobspb.DescRewriteMap) (string, er
ctx.WriteString(newRef.SQLString())
}),
)
ctx.FormatNode(parsed)
return ctx.CloseAndGetString(), nil
}

// rewriteSequencesInExpr rewrites all sequence IDs in the input expression
Expand Down Expand Up @@ -962,6 +1018,10 @@ func FunctionDescs(
if err != nil {
return err
}
fnBody, err = rewriteTypesInRoutine(fnBody, descriptorRewrites, fnDesc.Lang)
if err != nil {
return err
}
fnDesc.FunctionBody = fnBody

// Rewrite type IDs.
Expand Down