proto: make sessiondatapb.SessionData compatible with protoreflect.Me…

…ssageToJSON

Fixes #107823

This PR makes `VectorizeExecMode`'s protobuf and JSON string representations
the same to prevent errors like this when `SessionData.VectorizeMode` has a
non-zero value:
```
unknown value "on" for enum cockroach.sql.sessiondatapb.VectorizeExecMode
```

Release note: None

pkg/BUILD.bazel

@@ -549,6 +549,7 @@ ALL_TESTS = [
     "//pkg/sql/sem/tree:tree_disallowed_imports_test",
     "//pkg/sql/sem/tree:tree_test",
     "//pkg/sql/sessiondata:sessiondata_test",
+    "//pkg/sql/sessiondatapb:sessiondatapb_test",
     "//pkg/sql/sessioninit:sessioninit_test",
     "//pkg/sql/span:span_test",
     "//pkg/sql/sqlinstance/instancestorage:instancestorage_test",
@@ -2050,6 +2051,7 @@ GO_TARGETS = [
     "//pkg/sql/sessiondata:sessiondata",
     "//pkg/sql/sessiondata:sessiondata_test",
     "//pkg/sql/sessiondatapb:sessiondatapb",
+    "//pkg/sql/sessiondatapb:sessiondatapb_test",
     "//pkg/sql/sessioninit:sessioninit",
     "//pkg/sql/sessioninit:sessioninit_test",
     "//pkg/sql/sessionphase:sessionphase",

pkg/sql/colexec/colbuilder/execplan.go

@@ -257,7 +257,7 @@ var (
 	errExporterWrap                   = errors.New("core.Exporter is not supported (not an execinfra.RowSource)")
 	errSamplerWrap                    = errors.New("core.Sampler is not supported (not an execinfra.RowSource)")
 	errSampleAggregatorWrap           = errors.New("core.SampleAggregator is not supported (not an execinfra.RowSource)")
-	errExperimentalWrappingProhibited = errors.New("wrapping for non-JoinReader and non-LocalPlanNode cores is prohibited in vectorize=experimental_always")
+	errExperimentalWrappingProhibited = errors.Newf("wrapping for non-JoinReader and non-LocalPlanNode cores is prohibited in vectorize=%s", sessiondatapb.VectorizeExperimentalAlways)
 	errWrappedCast                    = errors.New("mismatched types in NewColOperator and unsupported casts")
 	errLookupJoinUnsupported          = errors.New("lookup join reader is unsupported in vectorized")
 	errFilteringAggregation           = errors.New("filtering aggregation not supported")

pkg/sql/exec_util.go

@@ -588,12 +588,13 @@ var VectorizeClusterMode = settings.RegisterEnumSetting(
 	VectorizeClusterSettingName,
 	"default vectorize mode",
 	"on",
-	map[int64]string{
-		int64(sessiondatapb.VectorizeUnset):              "on",
-		int64(sessiondatapb.VectorizeOn):                 "on",
-		int64(sessiondatapb.VectorizeExperimentalAlways): "experimental_always",
-		int64(sessiondatapb.VectorizeOff):                "off",
-	},
+	func() map[int64]string {
+		m := make(map[int64]string, len(sessiondatapb.VectorizeExecMode_name))
+		for k := range sessiondatapb.VectorizeExecMode_name {
+			m[int64(k)] = sessiondatapb.VectorizeExecMode(k).String()
+		}
+		return m
+	}(),
 ).WithPublic()
 
 // DistSQLClusterExecMode controls the cluster default for when DistSQL is used.

pkg/sql/logictest/testdata/logic_test/set

@@ -243,7 +243,7 @@ SET vectorize = experimental_always
 statement ok
 SET vectorize = off
 
-statement error invalid value for parameter "vectorize": "bogus"
+statement error invalid value for parameter "vectorize": "bogus"\nHINT: Available values: off,on,experimental_always
 SET vectorize = bogus
 
 statement ok

pkg/sql/sessiondatapb/BUILD.bazel

@@ -1,6 +1,6 @@
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "sessiondatapb",
@@ -54,3 +54,14 @@ go_proto_library(
         "@com_github_gogo_protobuf//gogoproto",
     ],
 )
+
+go_test(
+    name = "sessiondatapb_test",
+    srcs = ["session_data_test.go"],
+    args = ["-test.timeout=295s"],
+    embed = [":sessiondatapb"],
+    deps = [
+        "//pkg/sql/protoreflect",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/sql/sessiondatapb/session_data.go

@@ -55,30 +55,26 @@ func (c DataConversionConfig) GetFloatPrec(typ *types.T) int {
 }
 
 func (m VectorizeExecMode) String() string {
-	switch m {
-	case VectorizeOn, VectorizeUnset:
-		return "on"
-	case VectorizeExperimentalAlways:
-		return "experimental_always"
-	case VectorizeOff:
-		return "off"
-	default:
+	if m == VectorizeUnset {
+		m = VectorizeOn
+	}
+	name, ok := VectorizeExecMode_name[int32(m)]
+	if !ok {
 		return fmt.Sprintf("invalid (%d)", m)
 	}
+	return name
 }
 
 // VectorizeExecModeFromString converts a string into a VectorizeExecMode.
 // False is returned if the conversion was unsuccessful.
 func VectorizeExecModeFromString(val string) (VectorizeExecMode, bool) {
-	var m VectorizeExecMode
-	switch strings.ToUpper(val) {
-	case "ON":
-		m = VectorizeOn
-	case "EXPERIMENTAL_ALWAYS":
-		m = VectorizeExperimentalAlways
-	case "OFF":
-		m = VectorizeOff
-	default:
+	lowerVal := strings.ToLower(val)
+	mInt, ok := VectorizeExecMode_value[lowerVal]
+	if !ok {
+		return 0, false
+	}
+	m := VectorizeExecMode(mInt)
+	if m == VectorizeUnset {
 		return 0, false
 	}
 	return m, true

pkg/sql/sessiondatapb/session_data.proto

@@ -139,16 +139,16 @@ enum VectorizeExecMode {
   // VectorizeUnset means the VectorizeExecMode wasn't explicitly set. Having
   // the first enum value as zero is required by proto3. This is mapped to
   // VectorizeOn.
-  VectorizeUnset = 0;
+  unset = 0 [(gogoproto.enumvalue_customname) = "VectorizeUnset"];
   reserved 1;
   // VectorizeOn means that any supported queries will be run using the
   // columnar execution.
-  VectorizeOn = 2;
+  on = 2 [(gogoproto.enumvalue_customname) = "VectorizeOn"];
   // VectorizeExperimentalAlways means that we attempt to vectorize all
   // queries; unsupported queries will fail. Mostly used for testing.
-  VectorizeExperimentalAlways = 3;
+  experimental_always = 3 [(gogoproto.enumvalue_customname) = "VectorizeExperimentalAlways"];
   // VectorizeOff means that columnar execution is disabled.
-  VectorizeOff = 4;
+  off = 4 [(gogoproto.enumvalue_customname) = "VectorizeOff"];
 }
 
 // SequenceState is used to marshall the sessiondata.SequenceState struct.

pkg/sql/sessiondatapb/session_data_test.go

@@ -0,0 +1,30 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package sessiondatapb
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSessionDataJsonCompat(t *testing.T) {
+	expectedSessionData := SessionData{
+		VectorizeMode: VectorizeOn,
+	}
+	json, err := protoreflect.MessageToJSON(&expectedSessionData, protoreflect.FmtFlags{})
+	require.NoError(t, err)
+	actualSessionData := SessionData{}
+	_, err = protoreflect.JSONBMarshalToMessage(json, &actualSessionData)
+	require.NoError(t, err)
+	require.Equal(t, expectedSessionData, actualSessionData)
+}

pkg/sql/vars.go

@@ -688,8 +688,13 @@ var varGen = map[string]sessionVar{
 		Set: func(_ context.Context, m sessionDataMutator, s string) error {
 			mode, ok := sessiondatapb.VectorizeExecModeFromString(s)
 			if !ok {
-				return newVarValueError(`vectorize`, s,
-					"off", "on", "experimental_always")
+				return newVarValueError(
+					`vectorize`,
+					s,
+					sessiondatapb.VectorizeOff.String(),
+					sessiondatapb.VectorizeOn.String(),
+					sessiondatapb.VectorizeExperimentalAlways.String(),
+				)
 			}
 			m.SetVectorize(mode)
 			return nil