Skip to content

Commit c794734

Browse files
craig[bot]yuzefovichmsbutleremnet-crlDrewKimball
committed
139799: roachtest: fix unsortedMatricesDiffWithFloatComp helper r=yuzefovich a=yuzefovich `unsortedMatricesDiffWithFloatComp` has a special logic for handling float and decimal arrays, and it figures out whether that logic is applicable based on the stringified type name. Previously, we would only match strings like `[]FLOAT4` for that, but lib/pq library returns `_FLOAT4` as the name for that array, so we would previously not apply the special logic in some cases, which would lead to spurious failures. This is now fixed. Fixes: #139727. Release note: None 140045: crosscluster: extra priv tests r=msbutler a=msbutler This patch adds further e2e privilege tests for LDR and PCR Epic: none Release note: none 140073: deps: upgrade github.com/jackc/pgx to v4.18.3 r=spilchen a=emnet-crl This PR upgrades the `jackc/pgx` library to v4.18.2. The previous version of the library was susceptible to an SQL injection vulnerability per [GO-2024-2606](1). Fixes: CRDB-46987 Release note: None [1]: https://pkg.go.dev/vuln/GO-2024-2606 140083: sql/opt: build CTEs for SQL expressions embedded in PL/pgSQL r=yuzefovich,mgartner a=DrewKimball This commit fixes an oversight from when we introduced support for CTEs within routies. Namely, SQL expressions can contain subqueries, which in turn can contain CTEs. The logic for building a SQL *statement* within a PL/pgSQL routine correctly handles CTEs by building them into the constructed `RelExpr`. The logic for SQL *expressions` was previously missing the same. To fix this issue, `plpgsqlBuilder.buildSQLExpr` now checks if the built scalar expression contained any CTEs. If it did, the scalar is wrapped in the built CTEs, which are in turn wrapped in a subquery, which is returned as the final scalar result. This prevents issues where the CTEs are built at an outer scope, which can produce invalid plans. Fixes #138273 Release note (bug fix): Fixed a bug existing only in pre-release versions of v25.1. The bug could cause creation of a PL/pgSQL routine with a CTE to fail with an error like the following: `unexpected root expression: with`. 140100: authors: add Abhinav Gupta to authors r=Abhinav1299 a=Abhinav1299 Epic: None Release note: None Co-authored-by: Yahor Yuzefovich <[email protected]> Co-authored-by: Michael Butler <[email protected]> Co-authored-by: Emnet Gossaye <[email protected]> Co-authored-by: Drew Kimball <[email protected]> Co-authored-by: Abhinav1299 <[email protected]>
6 parents 13677b4 + 261cae9 + ad8b264 + 2988626 + e8ae3ab + afbd9f3 commit c794734

File tree

14 files changed

+680
-35
lines changed

14 files changed

+680
-35
lines changed

‎AUTHORS

+1
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ Aayush Shah <[email protected]> <@cockroachlabs.com>
2626
Abbey Russell <[email protected]> Abigail Russell <[email protected]> <[email protected]>
2727
Abby Hersh <[email protected]>
2828
Abhinav Garg <[email protected]>
29+
Abhinav Gupta <[email protected]> <[email protected]>
2930
Abhishek Kedia <[email protected]>
3031
Abhishek Madan <[email protected]> <[email protected]> <[email protected]> Abhemailk [email protected] <[email protected]>
3132
Abhishek Saha <[email protected]> AbhishekSaha <[email protected]>

‎DEPS.bzl

+3-3
Original file line numberDiff line numberDiff line change
@@ -5235,10 +5235,10 @@ def go_deps():
52355235
name = "com_github_jackc_pgx_v4",
52365236
build_file_proto_mode = "disable_global",
52375237
importpath = "github.com/jackc/pgx/v4",
5238-
sha256 = "5ca92c5bf58979d9e978f6b849e02eb319d2587565375fe875a29d10d84cfadc",
5239-
strip_prefix = "github.com/jackc/pgx/[email protected].1",
5238+
sha256 = "1d5955dc65b8de8f72f9856865b33dcd9a2238f7cf9b1f2a00f1558e7d4965da",
5239+
strip_prefix = "github.com/jackc/pgx/[email protected].3",
52405240
urls = [
5241-
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v4/com_github_jackc_pgx_v4-v4.18.1.zip",
5241+
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v4/com_github_jackc_pgx_v4-v4.18.3.zip",
52425242
],
52435243
)
52445244
go_repository(

‎build/bazelutil/distdir_files.bzl

+1-1
Original file line numberDiff line numberDiff line change
@@ -684,7 +684,7 @@ DISTDIR_FILES = {
684684
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgproto3/v2/com_github_jackc_pgproto3_v2-v2.3.3.zip": "53ea236cbfe241693b439092e2d51b404c2a635ee3fe64ea7aad1527cb715189",
685685
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgservicefile/com_github_jackc_pgservicefile-v0.0.0-20221227161230-091c0ba34f0a.zip": "1f8bdf75b2a0d750e56c2a94b1d1b0b5be4b29d6df056aebd997162c29bfd8ab",
686686
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgtype/com_github_jackc_pgtype-v1.14.1.zip": "3acb69a66e7e432c010d503425810620d04c304166c45083fa8a96feca13054d",
687-
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v4/com_github_jackc_pgx_v4-v4.18.1.zip": "5ca92c5bf58979d9e978f6b849e02eb319d2587565375fe875a29d10d84cfadc",
687+
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v4/com_github_jackc_pgx_v4-v4.18.3.zip": "1d5955dc65b8de8f72f9856865b33dcd9a2238f7cf9b1f2a00f1558e7d4965da",
688688
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v5/com_github_jackc_pgx_v5-v5.5.5.zip": "221749d6187aaeeb14097a41f2510689eaf35245ef77d243a1f69dc23605d2a2",
689689
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/puddle/com_github_jackc_puddle-v1.3.0.zip": "b1eb42bb3cf9a430146af79cb183860b9dddfca51844c2d4b447dc2f43becc55",
690690
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/puddle/v2/com_github_jackc_puddle_v2-v2.2.1.zip": "6698895617fabb929fa1ac868ad5253e02a997888bf5c6004379c5b29eedee58",

‎go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ require (
8484
github.com/jackc/pgproto3/v2 v2.3.3
8585
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
8686
github.com/jackc/pgtype v1.14.1
87-
github.com/jackc/pgx/v4 v4.18.1
87+
github.com/jackc/pgx/v4 v4.18.3
8888
)
8989

9090
require (

‎go.sum

+2-6
Original file line numberDiff line numberDiff line change
@@ -1460,7 +1460,6 @@ github.com/jackc/pgconn v0.0.0-20190831204454-2fabfa3c18b7/go.mod h1:ZJKsE/KZfsU
14601460
github.com/jackc/pgconn v1.8.0/go.mod h1:1C2Pb36bGIP9QHGBYCjnyhqu7Rv3sGshaQUvmfGIB/o=
14611461
github.com/jackc/pgconn v1.9.0/go.mod h1:YctiPyvzfU11JFxoXokUOOKQXQmDMoJL9vJzHH8/2JY=
14621462
github.com/jackc/pgconn v1.9.1-0.20210724152538-d89c8390a530/go.mod h1:4z2w8XhRbP1hYxkpTuBjTS3ne3J48K83+u0zoyvg2pI=
1463-
github.com/jackc/pgconn v1.14.0/go.mod h1:9mBNlny0UvkgJdCDvdVHYSjI+8tD2rnKK69Wz8ti++E=
14641463
github.com/jackc/pgconn v1.14.3 h1:bVoTr12EGANZz66nZPkMInAV/KHD2TxH9npjXXgiB3w=
14651464
github.com/jackc/pgconn v1.14.3/go.mod h1:RZbme4uasqzybK2RK5c65VsHxoyaml09lx3tXOcO/VM=
14661465
github.com/jackc/pgio v1.0.0 h1:g12B9UwVnzGhueNavwioyEEpAmqMe1E/BN9ES+8ovkE=
@@ -1478,7 +1477,6 @@ github.com/jackc/pgproto3/v2 v2.0.0-rc3/go.mod h1:ryONWYqW6dqSg1Lw6vXNMXoBJhpzvW
14781477
github.com/jackc/pgproto3/v2 v2.0.0-rc3.0.20190831210041-4c03ce451f29/go.mod h1:ryONWYqW6dqSg1Lw6vXNMXoBJhpzvWKnT95C46ckYeM=
14791478
github.com/jackc/pgproto3/v2 v2.0.6/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwXYU5xKaEdA=
14801479
github.com/jackc/pgproto3/v2 v2.1.1/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwXYU5xKaEdA=
1481-
github.com/jackc/pgproto3/v2 v2.3.2/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwXYU5xKaEdA=
14821480
github.com/jackc/pgproto3/v2 v2.3.3 h1:1HLSx5H+tXR9pW3in3zaztoEwQYRC9SQaYUHjTSUOag=
14831481
github.com/jackc/pgproto3/v2 v2.3.3/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwXYU5xKaEdA=
14841482
github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b/go.mod h1:vsD4gTJCa9TptPL8sPkXrLZ+hDuNrZCnj29CQpr4X1E=
@@ -1488,21 +1486,19 @@ github.com/jackc/pgtype v0.0.0-20190421001408-4ed0de4755e0/go.mod h1:hdSHsc1V01C
14881486
github.com/jackc/pgtype v0.0.0-20190824184912-ab885b375b90/go.mod h1:KcahbBH1nCMSo2DXpzsoWOAfFkdEtEJpPbVLq8eE+mc=
14891487
github.com/jackc/pgtype v0.0.0-20190828014616-a8802b16cc59/go.mod h1:MWlu30kVJrUS8lot6TQqcg7mtthZ9T0EoIBFiJcmcyw=
14901488
github.com/jackc/pgtype v1.8.1-0.20210724151600-32e20a603178/go.mod h1:C516IlIV9NKqfsMCXTdChteoXmwgUceqaLfjg2e3NlM=
1491-
github.com/jackc/pgtype v1.14.0/go.mod h1:LUMuVrfsFfdKGLw+AFFVv6KtHOFMwRgDDzBt76IqCA4=
14921489
github.com/jackc/pgtype v1.14.1 h1:LyDar7M2K0tShCWqzJ/ctzF1QC3Wzc9c8a6cHE0PFdc=
14931490
github.com/jackc/pgtype v1.14.1/go.mod h1:LUMuVrfsFfdKGLw+AFFVv6KtHOFMwRgDDzBt76IqCA4=
14941491
github.com/jackc/pgx/v4 v4.0.0-20190420224344-cc3461e65d96/go.mod h1:mdxmSJJuR08CZQyj1PVQBHy9XOp5p8/SHH6a0psbY9Y=
14951492
github.com/jackc/pgx/v4 v4.0.0-20190421002000-1b8f0016e912/go.mod h1:no/Y67Jkk/9WuGR0JG/JseM9irFbnEPbuWV2EELPNuM=
14961493
github.com/jackc/pgx/v4 v4.0.0-pre1.0.20190824185557-6972a5742186/go.mod h1:X+GQnOEnf1dqHGpw7JmHqHc1NxDoalibchSk9/RWuDc=
14971494
github.com/jackc/pgx/v4 v4.12.1-0.20210724153913-640aa07df17c/go.mod h1:1QD0+tgSXP7iUjYm9C1NxKhny7lq6ee99u/z+IHFcgs=
1498-
github.com/jackc/pgx/v4 v4.18.1 h1:YP7G1KABtKpB5IHrO9vYwSrCOhs7p3uqhvhhQBptya0=
1499-
github.com/jackc/pgx/v4 v4.18.1/go.mod h1:FydWkUyadDmdNH/mHnGob881GawxeEm7TcMCzkb+qQE=
1495+
github.com/jackc/pgx/v4 v4.18.3 h1:dE2/TrEsGX3RBprb3qryqSV9Y60iZN1C6i8IrmW9/BA=
1496+
github.com/jackc/pgx/v4 v4.18.3/go.mod h1:Ey4Oru5tH5sB6tV7hDmfWFahwF15Eb7DNXlRKx2CkVw=
15001497
github.com/jackc/pgx/v5 v5.5.5 h1:amBjrZVmksIdNjxGW/IiIMzxMKZFelXbUoPNb+8sjQw=
15011498
github.com/jackc/pgx/v5 v5.5.5/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
15021499
github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
15031500
github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
15041501
github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
1505-
github.com/jackc/puddle v1.3.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
15061502
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
15071503
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
15081504
github.com/jaegertracing/jaeger v1.18.1 h1:eFqjEpTKq2FfiZ/YX53oxeCePdIZyWvDfXaTAGj0r5E=

‎pkg/cmd/roachtest/tests/query_comparison_util.go

+29-11
Original file line numberDiff line numberDiff line change
@@ -611,13 +611,30 @@ func joinAndSortRows(rowMatrix1, rowMatrix2 [][]string, sep string) (rows1, rows
611611
return rows1, rows2
612612
}
613613

614+
func isFloatArray(colType string) bool {
615+
switch colType {
616+
case "[]FLOAT4", "[]FLOAT8", "_FLOAT4", "_FLOAT8":
617+
return true
618+
default:
619+
return false
620+
}
621+
}
622+
623+
func isDecimalArray(colType string) bool {
624+
switch colType {
625+
case "[]DECIMAL", "_DECIMAL":
626+
return true
627+
default:
628+
return false
629+
}
630+
}
631+
614632
func needApproximateMatch(colType string) bool {
615633
// On s390x, check that values for both float and decimal coltypes are
616634
// approximately equal to take into account platform differences in floating
617635
// point calculations. On other architectures, check float values only.
618-
return (runtime.GOARCH == "s390x" && (colType == "DECIMAL" || colType == "[]DECIMAL")) ||
619-
colType == "FLOAT4" || colType == "[]FLOAT4" ||
620-
colType == "FLOAT8" || colType == "[]FLOAT8"
636+
return (runtime.GOARCH == "s390x" && (colType == "DECIMAL" || isDecimalArray(colType))) ||
637+
colType == "FLOAT4" || colType == "FLOAT8" || isFloatArray(colType)
621638
}
622639

623640
// sortRowsWithFloatComp is similar to joinAndSortRows, but it uses float
@@ -627,7 +644,7 @@ func sortRowsWithFloatComp(rowMatrix1, rowMatrix2 [][]string, colTypes []string)
627644
for idx := range colTypes {
628645
if needApproximateMatch(colTypes[idx]) {
629646
cmpFn := floatcmp.FloatsCmp
630-
if strings.HasPrefix(colTypes[idx], "[]") {
647+
if isFloatArray(colTypes[idx]) || isDecimalArray(colTypes[idx]) {
631648
cmpFn = floatcmp.FloatArraysCmp
632649
}
633650
res, err := cmpFn(rowMatrix[i][idx], rowMatrix[j][idx])
@@ -697,7 +714,7 @@ func unsortedMatricesDiffWithFloatComp(
697714
}
698715
var needCustomMatch bool
699716
for _, colType := range colTypes {
700-
if needApproximateMatch(colType) || colType == "DECIMAL" || colType == "[]DECIMAL" {
717+
if needApproximateMatch(colType) || colType == "DECIMAL" || isDecimalArray(colType) {
701718
needCustomMatch = true
702719
break
703720
}
@@ -712,13 +729,14 @@ func unsortedMatricesDiffWithFloatComp(
712729

713730
for j, colType := range colTypes {
714731
if needApproximateMatch(colType) {
732+
isFloatOrDecimalArray := isFloatArray(colType) || isDecimalArray(colType)
715733
cmpFn := floatcmp.FloatsMatch
716734
switch {
717-
case runtime.GOARCH == "s390x" && strings.HasPrefix(colType, "[]"):
735+
case runtime.GOARCH == "s390x" && isFloatOrDecimalArray:
718736
cmpFn = floatcmp.FloatArraysMatchApprox
719-
case runtime.GOARCH == "s390x" && !strings.HasPrefix(colType, "[]"):
737+
case runtime.GOARCH == "s390x" && !isFloatOrDecimalArray:
720738
cmpFn = floatcmp.FloatsMatchApprox
721-
case strings.HasPrefix(colType, "[]"):
739+
case isFloatOrDecimalArray:
722740
cmpFn = floatcmp.FloatArraysMatch
723741
}
724742
match, err := cmpFn(row1[j], row2[j])
@@ -729,12 +747,12 @@ func unsortedMatricesDiffWithFloatComp(
729747
return result, nil
730748
}
731749
} else {
732-
switch colType {
733-
case "DECIMAL":
750+
switch {
751+
case colType == "DECIMAL":
734752
// For decimals, remove any trailing zeroes.
735753
row1[j] = trimDecimalTrailingZeroes(row1[j])
736754
row2[j] = trimDecimalTrailingZeroes(row2[j])
737-
case "[]DECIMAL":
755+
case isDecimalArray(colType):
738756
// For decimal arrays, remove any trailing zeroes from each
739757
// decimal.
740758
row1[j] = trimDecimalsTrailingZeroesInArray(row1[j])

‎pkg/cmd/roachtest/tests/query_comparison_util_test.go

+8
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,14 @@ func TestUnsortedMatricesDiff(t *testing.T) {
148148
exactMatch: false,
149149
approxMatch: true,
150150
},
151+
{
152+
name: "multi row 0 in array matches -0 in array, lib/pq type name",
153+
colTypes: []string{"_FLOAT4"}, // this is how []FLOAT4 is named in lib/pq driver
154+
t1: [][]string{{"NULL"}, {"{-1}"}, {"{-0}"}, {"{0}"}, {"{NaN}"}},
155+
t2: [][]string{{"NULL"}, {"{-1}"}, {"{0}"}, {"{0}"}, {"{NaN}"}},
156+
exactMatch: false,
157+
approxMatch: true,
158+
},
151159
}
152160
for _, tc := range tcs {
153161
t.Run(tc.name, func(t *testing.T) {

‎pkg/crosscluster/logical/logical_replication_job_test.go

+13-6
Original file line numberDiff line numberDiff line change
@@ -1953,7 +1953,7 @@ func TestUserPrivileges(t *testing.T) {
19531953
},
19541954
}
19551955

1956-
server, s, dbA, _ := setupLogicalTestServer(t, ctx, clusterArgs, 1)
1956+
server, s, dbA, dbB := setupLogicalTestServer(t, ctx, clusterArgs, 1)
19571957
defer server.Stopper().Stop(ctx)
19581958

19591959
dbBURL := replicationtestutils.GetExternalConnectionURI(t, s, s, serverutils.DBName("b"))
@@ -1966,7 +1966,8 @@ func TestUserPrivileges(t *testing.T) {
19661966
testuser2 := sqlutils.MakeSQLRunner(s.SQLConn(t, serverutils.User(username.TestUser+"2"), serverutils.DBName("a")))
19671967

19681968
var jobAID jobspb.JobID
1969-
testuser2.QueryRow(t, "CREATE LOGICAL REPLICATION STREAM FROM TABLE tab ON $1 INTO TABLE tab", dbBURL.String()).Scan(&jobAID)
1969+
createStmt := "CREATE LOGICAL REPLICATION STREAM FROM TABLE tab ON $1 INTO TABLE tab"
1970+
testuser2.QueryRow(t, createStmt, dbBURL.String()).Scan(&jobAID)
19701971

19711972
t.Run("view-control-job", func(t *testing.T) {
19721973
showJobStmt := "select job_id from [SHOW JOBS] where job_id=$1"
@@ -2006,11 +2007,17 @@ func TestUserPrivileges(t *testing.T) {
20062007
testuser.Exec(t, fmt.Sprintf(testingUDFAcceptProposedBaseWithSchema, "testschema", "tab"))
20072008
})
20082009

2009-
t.Run("replication", func(t *testing.T) {
2010-
createWithUDFStmt := "CREATE LOGICAL REPLICATION STREAM FROM TABLE tab ON $1 INTO TABLE tab WITH DEFAULT FUNCTION = 'testschema.repl_apply'"
2011-
testuser.ExpectErr(t, "user testuser does not have REPLICATION system privilege", createWithUDFStmt, dbBURL.String())
2010+
t.Run("replication-dest", func(t *testing.T) {
2011+
testuser.ExpectErr(t, "user testuser does not have REPLICATION system privilege", createStmt, dbBURL.String())
20122012
dbA.Exec(t, fmt.Sprintf("GRANT SYSTEM REPLICATION TO %s", username.TestUser))
2013-
testuser.QueryRow(t, createWithUDFStmt, dbBURL.String()).Scan(&jobAID)
2013+
testuser.QueryRow(t, createStmt, dbBURL.String()).Scan(&jobAID)
2014+
})
2015+
t.Run("replication-src", func(t *testing.T) {
2016+
dbB.Exec(t, "CREATE USER testuser3")
2017+
dbBURL2 := replicationtestutils.GetExternalConnectionURI(t, s, s, serverutils.DBName("b"), serverutils.User(username.TestUser+"3"))
2018+
testuser.ExpectErr(t, "user testuser3 does not have REPLICATION system privilege", createStmt, dbBURL2.String())
2019+
dbB.Exec(t, fmt.Sprintf("GRANT SYSTEM REPLICATION TO %s", username.TestUser+"3"))
2020+
testuser.QueryRow(t, createStmt, dbBURL2.String()).Scan(&jobAID)
20142021
})
20152022
}
20162023

‎pkg/crosscluster/physical/replication_stream_e2e_test.go

+31
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"fmt"
1111
"net/http"
1212
"net/http/httptest"
13+
"net/url"
1314
"sync/atomic"
1415
"testing"
1516
"time"
@@ -25,6 +26,7 @@ import (
2526
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
2627
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
2728
"github.com/cockroachdb/cockroach/pkg/roachpb"
29+
"github.com/cockroachdb/cockroach/pkg/security/username"
2830
"github.com/cockroachdb/cockroach/pkg/server/telemetry"
2931
"github.com/cockroachdb/cockroach/pkg/sql"
3032
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
@@ -50,6 +52,35 @@ import (
5052
"github.com/stretchr/testify/require"
5153
)
5254

55+
func TestPCRPrivs(t *testing.T) {
56+
defer leaktest.AfterTest(t)()
57+
defer log.Scope(t).Close(t)
58+
59+
ctx := context.Background()
60+
args := replicationtestutils.DefaultTenantStreamingClustersArgs
61+
c, cleanup := replicationtestutils.CreateTenantStreamingClusters(ctx, t, args)
62+
defer cleanup()
63+
64+
c.DestSysSQL.Exec(t, fmt.Sprintf("CREATE USER %s", username.TestUser))
65+
c.SrcSysSQL.Exec(t, fmt.Sprintf("CREATE USER %s", username.TestUser+"2"))
66+
testuser := sqlutils.MakeSQLRunner(c.DestSysServer.SQLConn(t, serverutils.User(username.TestUser)))
67+
srcURL, cleanupSinkCert := sqlutils.PGUrl(t, c.SrcSysServer.AdvSQLAddr(), t.Name(), url.User(username.TestUser+"2"))
68+
defer cleanupSinkCert()
69+
70+
streamReplStmt := fmt.Sprintf("CREATE TENANT %s FROM REPLICATION OF %s ON '%s'",
71+
c.Args.DestTenantName,
72+
c.Args.SrcTenantName,
73+
srcURL.String())
74+
75+
testuser.ExpectErr(t, "user testuser does not have MANAGEVIRTUALCLUSTER system privilege", streamReplStmt)
76+
77+
c.DestSysSQL.Exec(t, fmt.Sprintf("GRANT SYSTEM MANAGEVIRTUALCLUSTER TO %s", username.TestUser))
78+
testuser.ExpectErr(t, "user testuser2 does not have REPLICATION system privilege", streamReplStmt)
79+
80+
c.SrcSysSQL.Exec(t, fmt.Sprintf("GRANT SYSTEM REPLICATION TO %s", username.TestUser+"2"))
81+
c.DestSysSQL.Exec(t, streamReplStmt)
82+
83+
}
5384
func TestTenantStreamingProducerJobTimedOut(t *testing.T) {
5485
defer leaktest.AfterTest(t)()
5586
defer log.Scope(t).Close(t)

‎pkg/sql/logictest/testdata/logic_test/udf_cte

+39
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,42 @@ query T
121121
SELECT f();
122122
----
123123
{10,20,30,40}
124+
125+
subtest regression_138273
126+
127+
statement ok
128+
CREATE SEQUENCE seq_1;
129+
130+
# Verify that the function can be successfully created with a CTE inside a
131+
# scalar expression.
132+
statement ok
133+
CREATE FUNCTION f138273_1() RETURNS TIMESTAMP LANGUAGE PLpgSQL AS $$
134+
DECLARE
135+
decl TIMESTAMP;
136+
BEGIN
137+
WHILE false LOOP
138+
IF true THEN
139+
ELSIF EXISTS (WITH cte(col) AS (SELECT * FROM (VALUES (currval('seq_1'))) AS foo) SELECT 1 FROM cte) THEN
140+
RETURN decl;
141+
END IF;
142+
END LOOP;
143+
END;
144+
$$;
145+
146+
# Test a similar function that actually executes the CTE.
147+
statement ok
148+
CREATE FUNCTION f138273_2() RETURNS INT LANGUAGE PLpgSQL AS $$
149+
BEGIN
150+
SELECT nextval('seq_1');
151+
WHILE EXISTS (WITH cte(col) AS (SELECT * FROM (VALUES (currval('seq_1'))) AS foo) SELECT 1 FROM cte) LOOP
152+
RETURN currval('seq_1');
153+
END LOOP;
154+
END;
155+
$$;
156+
157+
query I
158+
SELECT f138273_2();
159+
----
160+
1
161+
162+
subtest end

‎pkg/sql/opt/optbuilder/create_trigger.go

+9-6
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,15 @@ func (b *Builder) buildFunctionForTrigger(
179179

180180
// The trigger function can reference the NEW and OLD transition relations,
181181
// aliased in the trigger definition.
182+
if len(ct.Transitions) > 0 {
183+
// Save the existing CTEs in the builder and restore them after the function
184+
// body is built.
185+
prevCTEs := b.ctes
186+
b.ctes = nil
187+
defer func() {
188+
b.ctes = prevCTEs
189+
}()
190+
}
182191
for _, transition := range ct.Transitions {
183192
// Build a fake relational expression with a column corresponding to each
184193
// column from the table.
@@ -203,12 +212,6 @@ func (b *Builder) buildFunctionForTrigger(
203212
funcScope.ctes[string(transition.Name)] = cte
204213
b.addCTE(cte)
205214
}
206-
if len(ct.Transitions) > 0 {
207-
defer func() {
208-
// Reset the CTEs in the builder after the function body is built.
209-
b.ctes = nil
210-
}()
211-
}
212215

213216
// The trigger function takes a set of implicitly-defined parameters, two of
214217
// which are determined by the table's record type. Add them to the trigger

‎pkg/sql/opt/optbuilder/plpgsql.go

+22-1
Original file line numberDiff line numberDiff line change
@@ -2201,13 +2201,34 @@ func (b *plpgsqlBuilder) buildSQLExpr(expr ast.Expr, typ *types.T, s *scope) opt
22012201
// For lazy SQL evaluation, replace all expressions with NULL.
22022202
return memo.NullSingleton
22032203
}
2204+
// Save any outer CTEs before building the expression, which may have
2205+
// subqueries with inner CTEs.
2206+
prevCTEs := b.ob.ctes
2207+
b.ob.ctes = nil
2208+
defer func() {
2209+
b.ob.ctes = prevCTEs
2210+
}()
22042211
expr, _ = tree.WalkExpr(s, expr)
22052212
typedExpr, err := expr.TypeCheck(b.ob.ctx, b.ob.semaCtx, typ)
22062213
if err != nil {
22072214
panic(err)
22082215
}
22092216
scalar := b.ob.buildScalar(typedExpr, s, nil, nil, b.colRefs)
2210-
return b.coerceType(scalar, typ)
2217+
scalar = b.coerceType(scalar, typ)
2218+
if len(b.ob.ctes) == 0 {
2219+
return scalar
2220+
}
2221+
// There was at least one CTE within the scalar expression. It is possible to
2222+
// "hoist" them above this point, but building them eagerly here means that
2223+
// callers don't have to worry about CTE handling.
2224+
f := b.ob.factory
2225+
valuesCol := f.Metadata().AddColumn("", scalar.DataType())
2226+
valuesExpr := f.ConstructValues(
2227+
memo.ScalarListExpr{f.ConstructTuple(memo.ScalarListExpr{scalar}, scalar.DataType())},
2228+
&memo.ValuesPrivate{Cols: opt.ColList{valuesCol}, ID: f.Metadata().NextUniqueID()},
2229+
)
2230+
withExpr := b.ob.buildWiths(valuesExpr, b.ob.ctes)
2231+
return f.ConstructSubquery(withExpr, &memo.SubqueryPrivate{})
22112232
}
22122233

22132234
// buildSQLStatement type-checks and builds the given SQL statement into a

0 commit comments

Comments
 (0)