Skip to content

Commit 2ca6e52

Browse files
committed
proto: make sessiondatapb.SessionData compatible with protoreflect.MessageToJSON
Fixes cockroachdb#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
1 parent 08d18aa commit 2ca6e52

File tree

9 files changed

+77
-32
lines changed

9 files changed

+77
-32
lines changed

pkg/BUILD.bazel

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -549,6 +549,7 @@ ALL_TESTS = [
549549
"//pkg/sql/sem/tree:tree_disallowed_imports_test",
550550
"//pkg/sql/sem/tree:tree_test",
551551
"//pkg/sql/sessiondata:sessiondata_test",
552+
"//pkg/sql/sessiondatapb:sessiondatapb_test",
552553
"//pkg/sql/sessioninit:sessioninit_test",
553554
"//pkg/sql/span:span_test",
554555
"//pkg/sql/sqlinstance/instancestorage:instancestorage_test",
@@ -2050,6 +2051,7 @@ GO_TARGETS = [
20502051
"//pkg/sql/sessiondata:sessiondata",
20512052
"//pkg/sql/sessiondata:sessiondata_test",
20522053
"//pkg/sql/sessiondatapb:sessiondatapb",
2054+
"//pkg/sql/sessiondatapb:sessiondatapb_test",
20532055
"//pkg/sql/sessioninit:sessioninit",
20542056
"//pkg/sql/sessioninit:sessioninit_test",
20552057
"//pkg/sql/sessionphase:sessionphase",

pkg/sql/colexec/colbuilder/execplan.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ var (
257257
errExporterWrap = errors.New("core.Exporter is not supported (not an execinfra.RowSource)")
258258
errSamplerWrap = errors.New("core.Sampler is not supported (not an execinfra.RowSource)")
259259
errSampleAggregatorWrap = errors.New("core.SampleAggregator is not supported (not an execinfra.RowSource)")
260-
errExperimentalWrappingProhibited = errors.New("wrapping for non-JoinReader and non-LocalPlanNode cores is prohibited in vectorize=experimental_always")
260+
errExperimentalWrappingProhibited = errors.Newf("wrapping for non-JoinReader and non-LocalPlanNode cores is prohibited in vectorize=%s", sessiondatapb.VectorizeExperimentalAlways)
261261
errWrappedCast = errors.New("mismatched types in NewColOperator and unsupported casts")
262262
errLookupJoinUnsupported = errors.New("lookup join reader is unsupported in vectorized")
263263
errFilteringAggregation = errors.New("filtering aggregation not supported")

pkg/sql/exec_util.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -588,12 +588,13 @@ var VectorizeClusterMode = settings.RegisterEnumSetting(
588588
VectorizeClusterSettingName,
589589
"default vectorize mode",
590590
"on",
591-
map[int64]string{
592-
int64(sessiondatapb.VectorizeUnset): "on",
593-
int64(sessiondatapb.VectorizeOn): "on",
594-
int64(sessiondatapb.VectorizeExperimentalAlways): "experimental_always",
595-
int64(sessiondatapb.VectorizeOff): "off",
596-
},
591+
func() map[int64]string {
592+
m := make(map[int64]string, len(sessiondatapb.VectorizeExecMode_name))
593+
for k := range sessiondatapb.VectorizeExecMode_name {
594+
m[int64(k)] = sessiondatapb.VectorizeExecMode(k).String()
595+
}
596+
return m
597+
}(),
597598
).WithPublic()
598599

599600
// DistSQLClusterExecMode controls the cluster default for when DistSQL is used.

pkg/sql/logictest/testdata/logic_test/set

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ SET vectorize = experimental_always
243243
statement ok
244244
SET vectorize = off
245245

246-
statement error invalid value for parameter "vectorize": "bogus"
246+
statement error invalid value for parameter "vectorize": "bogus"\nHINT: Available values: off,on,experimental_always
247247
SET vectorize = bogus
248248

249249
statement ok

pkg/sql/sessiondatapb/BUILD.bazel

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
load("@rules_proto//proto:defs.bzl", "proto_library")
22
load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
3-
load("@io_bazel_rules_go//go:def.bzl", "go_library")
3+
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
44

55
go_library(
66
name = "sessiondatapb",
@@ -54,3 +54,14 @@ go_proto_library(
5454
"@com_github_gogo_protobuf//gogoproto",
5555
],
5656
)
57+
58+
go_test(
59+
name = "sessiondatapb_test",
60+
srcs = ["session_data_test.go"],
61+
args = ["-test.timeout=295s"],
62+
embed = [":sessiondatapb"],
63+
deps = [
64+
"//pkg/sql/protoreflect",
65+
"@com_github_stretchr_testify//require",
66+
],
67+
)

pkg/sql/sessiondatapb/session_data.go

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -55,30 +55,26 @@ func (c DataConversionConfig) GetFloatPrec(typ *types.T) int {
5555
}
5656

5757
func (m VectorizeExecMode) String() string {
58-
switch m {
59-
case VectorizeOn, VectorizeUnset:
60-
return "on"
61-
case VectorizeExperimentalAlways:
62-
return "experimental_always"
63-
case VectorizeOff:
64-
return "off"
65-
default:
58+
if m == VectorizeUnset {
59+
m = VectorizeOn
60+
}
61+
name, ok := VectorizeExecMode_name[int32(m)]
62+
if !ok {
6663
return fmt.Sprintf("invalid (%d)", m)
6764
}
65+
return name
6866
}
6967

7068
// VectorizeExecModeFromString converts a string into a VectorizeExecMode.
7169
// False is returned if the conversion was unsuccessful.
7270
func VectorizeExecModeFromString(val string) (VectorizeExecMode, bool) {
73-
var m VectorizeExecMode
74-
switch strings.ToUpper(val) {
75-
case "ON":
76-
m = VectorizeOn
77-
case "EXPERIMENTAL_ALWAYS":
78-
m = VectorizeExperimentalAlways
79-
case "OFF":
80-
m = VectorizeOff
81-
default:
71+
lowerVal := strings.ToLower(val)
72+
mInt, ok := VectorizeExecMode_value[lowerVal]
73+
if !ok {
74+
return 0, false
75+
}
76+
m := VectorizeExecMode(mInt)
77+
if m == VectorizeUnset {
8278
return 0, false
8379
}
8480
return m, true

pkg/sql/sessiondatapb/session_data.proto

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -139,16 +139,16 @@ enum VectorizeExecMode {
139139
// VectorizeUnset means the VectorizeExecMode wasn't explicitly set. Having
140140
// the first enum value as zero is required by proto3. This is mapped to
141141
// VectorizeOn.
142-
VectorizeUnset = 0;
142+
unset = 0 [(gogoproto.enumvalue_customname) = "VectorizeUnset"];
143143
reserved 1;
144144
// VectorizeOn means that any supported queries will be run using the
145145
// columnar execution.
146-
VectorizeOn = 2;
146+
on = 2 [(gogoproto.enumvalue_customname) = "VectorizeOn"];
147147
// VectorizeExperimentalAlways means that we attempt to vectorize all
148148
// queries; unsupported queries will fail. Mostly used for testing.
149-
VectorizeExperimentalAlways = 3;
149+
experimental_always = 3 [(gogoproto.enumvalue_customname) = "VectorizeExperimentalAlways"];
150150
// VectorizeOff means that columnar execution is disabled.
151-
VectorizeOff = 4;
151+
off = 4 [(gogoproto.enumvalue_customname) = "VectorizeOff"];
152152
}
153153

154154
// SequenceState is used to marshall the sessiondata.SequenceState struct.
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// Copyright 2023 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the Business Source License
4+
// included in the file licenses/BSL.txt.
5+
//
6+
// As of the Change Date specified in that file, in accordance with
7+
// the Business Source License, use of this software will be governed
8+
// by the Apache License, Version 2.0, included in the file
9+
// licenses/APL.txt.
10+
11+
package sessiondatapb
12+
13+
import (
14+
"testing"
15+
16+
"github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
17+
"github.com/stretchr/testify/require"
18+
)
19+
20+
func TestSessionDataJsonCompat(t *testing.T) {
21+
expectedSessionData := SessionData{
22+
VectorizeMode: VectorizeOn,
23+
}
24+
json, err := protoreflect.MessageToJSON(&expectedSessionData, protoreflect.FmtFlags{})
25+
require.NoError(t, err)
26+
actualSessionData := SessionData{}
27+
_, err = protoreflect.JSONBMarshalToMessage(json, &actualSessionData)
28+
require.NoError(t, err)
29+
require.Equal(t, expectedSessionData, actualSessionData)
30+
}

pkg/sql/vars.go

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -688,8 +688,13 @@ var varGen = map[string]sessionVar{
688688
Set: func(_ context.Context, m sessionDataMutator, s string) error {
689689
mode, ok := sessiondatapb.VectorizeExecModeFromString(s)
690690
if !ok {
691-
return newVarValueError(`vectorize`, s,
692-
"off", "on", "experimental_always")
691+
return newVarValueError(
692+
`vectorize`,
693+
s,
694+
sessiondatapb.VectorizeOff.String(),
695+
sessiondatapb.VectorizeOn.String(),
696+
sessiondatapb.VectorizeExperimentalAlways.String(),
697+
)
693698
}
694699
m.SetVectorize(mode)
695700
return nil

0 commit comments

Comments
 (0)