Merge #124437

124437: pgwire: support results_buffer_size in connection string options r=fqazi a=rafiss informs #124360 Release note (bug fix): The results_buffer_size session variable previously could not be configured by using the "options" query parameter in the connection string; it could only be configured as a top-level query parameter. Now, it can be configured in either part of the connection string. (This variable still cannot be changed with the SET command after the session begins.) Co-authored-by: Rafi Shamim <[email protected]>
cockroachdb · May 28, 2024 · 4c1e969 · 4c1e969
2 parents 2cd70a3 + 219aebd
commit 4c1e969
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/pkg/sql/pgwire/conn_test.go b/pkg/sql/pgwire/conn_test.go
@@ -1491,11 +1491,12 @@ func TestParseClientProvidedSessionParameters(t *testing.T) {
 			},
 		},
 		{
-			desc:  "results_buffer_size is not configurable from options",
-			query: "user=root&options=-c%20results_buffer_size=42",
+			desc:  "results_buffer_size is configurable from options",
+			query: "user=root&options=-c%20results_buffer_size=512kb",
 			assert: func(t *testing.T, args sql.SessionArgs, err error) {
-				require.Error(t, err)
-				require.Regexp(t, "options: parameter \"results_buffer_size\" cannot be changed", err)
+				require.NoError(t, err)
+				require.Equal(t, "root", args.User.Normalized())
+				require.EqualValues(t, 512000, args.ConnResultsBufferSize)
 			},
 		},
 		{

diff --git a/pkg/sql/pgwire/pre_serve_options.go b/pkg/sql/pgwire/pre_serve_options.go
@@ -194,6 +194,18 @@ func parseClientProvidedSessionParameters(
 					args.tenantName = parts[0]
 					hasTenantSelectOption = true
 					continue
+				case "results_buffer_size":
+					if args.ConnResultsBufferSize, err = humanizeutil.ParseBytes(optvalue); err != nil {
+						return args, errors.WithSecondaryError(
+							pgerror.Newf(pgcode.ProtocolViolation,
+								"error parsing results_buffer_size option value '%s' as bytes", optvalue), err)
+					}
+					if args.ConnResultsBufferSize < 0 {
+						return args, pgerror.Newf(pgcode.ProtocolViolation,
+							"results_buffer_size option value '%s' cannot be negative", value)
+					}
+					args.foundBufferSize = true
+					continue
 				}
 				err = loadParameter(ctx, opt, optvalue, &args.SessionArgs)
 				if err != nil {