sessiondata: introduce sessiondata.Stack

This stack keeps a track of SessionData objects for use in
transaction-scoped variables.

Release justification: low risk, high pri change

Release note: None

pkg/sql/sessiondata/BUILD.bazel

@@ -21,13 +21,21 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/timeutil/pgdate",
+        "@com_github_cockroachdb_errors//:errors",
     ],
 )
 
 go_test(
     name = "sessiondata_test",
     size = "small",
-    srcs = ["search_path_test.go"],
+    srcs = [
+        "search_path_test.go",
+        "session_data_test.go",
+    ],
     embed = [":sessiondata"],
-    deps = ["@com_github_stretchr_testify//assert"],
+    deps = [
+        "//pkg/sql/sessiondatapb",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/sql/sessiondata/session_data.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
+	"github.com/cockroachdb/errors"
 )
 
 // SessionData contains session parameters. They are all user-configurable.
@@ -184,3 +185,39 @@ func (s *SessionData) GetTemporarySchemaIDForDb(dbID uint32) (uint32, bool) {
 	schemaID, found := s.DatabaseIDToTempSchemaID[dbID]
 	return schemaID, found
 }
+
+// Stack represents a stack of SessionData objects.
+// This is used to support transaction-scoped variables, where SET LOCAL only
+// affects the top of the stack.
+// There is always guaranteed to be one element in the stack.
+type Stack struct {
+	// Use an internal variable to prevent abstraction leakage.
+	stack []*SessionData
+}
+
+// NewStack creates a new tack.
+func NewStack(firstElem *SessionData) *Stack {
+	return &Stack{stack: []*SessionData{firstElem}}
+}
+
+// Top returns the top element of the stack.
+func (s *Stack) Top() *SessionData {
+	if len(s.stack) == 0 {
+		return nil
+	}
+	return s.stack[len(s.stack)-1]
+}
+
+// Push pushes a SessionData element to the stack.
+func (s *Stack) Push(elem *SessionData) {
+	s.stack = append(s.stack, elem)
+}
+
+// Pop removes the top SessionData element from the stack.
+func (s *Stack) Pop() error {
+	if len(s.stack) <= 1 {
+		return errors.AssertionFailedf("there must always be at least one element in the SessionData stack")
+	}
+	s.stack = s.stack[:len(s.stack)-1]
+	return nil
+}

pkg/sql/sessiondata/session_data_test.go

@@ -0,0 +1,50 @@
+// Copyright 2021 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 sessiondata
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStack(t *testing.T) {
+	initialElem := &SessionData{
+		SessionData: sessiondatapb.SessionData{
+			ApplicationName: "bob",
+		},
+	}
+	secondElem := &SessionData{
+		SessionData: sessiondatapb.SessionData{
+			ApplicationName: "jane",
+		},
+	}
+	thirdElem := &SessionData{
+		SessionData: sessiondatapb.SessionData{
+			ApplicationName: "t-marts",
+		},
+	}
+	s := NewStack(initialElem)
+	require.Equal(t, s.Top(), initialElem)
+	require.EqualError(t, s.Pop(), "there must always be at least one element in the SessionData stack")
+
+	s.Push(secondElem)
+	require.Equal(t, s.Top(), secondElem)
+	s.Push(thirdElem)
+	require.Equal(t, s.Top(), thirdElem)
+
+	require.NoError(t, s.Pop())
+	require.Equal(t, s.Top(), secondElem)
+	require.NoError(t, s.Pop())
+	require.Equal(t, s.Top(), initialElem)
+	require.EqualError(t, s.Pop(), "there must always be at least one element in the SessionData stack")
+}