Merge #69553

69553: sessiondata: fix ListSessions with SET LOCAL / SET ROLE r=rafiss a=otan

Previously, two incorrect phenomena was incorrect due to recent changes:
* SET ROLE would report the user the current session was impersonated
  as, rather than the user who was logged in. The latter seems more
  correct to report here.
* There is a race between a BEGIN/COMMIT/ROLLBACK and listing all
  sessions using ListSessions as they both try to access the stack.

This is both fixed by:
* Changing the returned user of SET ROLE to be the SessionUser.
* Introducing a `base` method to SessionDataStack, which returns the
  base of the stack. This is *always* set, so allow accordingly.

Release justification: fix to new functionality

Resolves: #69508

Release note: None

Co-authored-by: Oliver Tan <[email protected]>

pkg/sql/conn_executor.go

@@ -2809,13 +2809,17 @@ func (ex *connExecutor) serialize() serverpb.Session {
 		lastActiveQueryNoConstants = truncateSQL(formatStatementHideConstants(ex.mu.LastActiveQuery))
 	}
 
+	// We always use base here as the fields from the SessionData should always
+	// be that of the root session.
+	sd := ex.sessionDataStack.Base()
+
 	remoteStr := "<admin>"
-	if ex.sessionData().RemoteAddr != nil {
-		remoteStr = ex.sessionData().RemoteAddr.String()
+	if sd.RemoteAddr != nil {
+		remoteStr = sd.RemoteAddr.String()
 	}
 
 	return serverpb.Session{
-		Username:                   ex.sessionData().User().Normalized(),
+		Username:                   sd.SessionUser().Normalized(),
 		ClientAddress:              remoteStr,
 		ApplicationName:            ex.applicationName.Load().(string),
 		Start:                      ex.phaseTimes.GetSessionPhaseTime(sessionphase.SessionInit).UTC(),

pkg/sql/logictest/testdata/logic_test/set_role

@@ -325,11 +325,13 @@ testuser  testuser
 
 
 # SET ROLE is specially cased in SET LOCAL as it uses SetWithPlanner,
-# so test it behaves as appropriate.
+# so test it behaves as appropriate. Also ensure that the node_sessions
+# is correctly attributed to root instead of testuser.
 
 user root
 
 statement ok
+GRANT ADMIN TO testuser;
 BEGIN;
 SET LOCAL ROLE testuser
 
@@ -338,10 +340,28 @@ SELECT current_user, session_user()
 ----
 testuser  root
 
+query T
+SELECT user_name FROM crdb_internal.node_sessions
+WHERE active_queries LIKE 'SELECT user_name%'
+----
+root
+
 statement ok
 ROLLBACK
 
 query TT
 SELECT current_user, session_user()
 ----
 root  root
+
+statement ok
+SET ROLE testuser
+
+query T
+SELECT user_name FROM crdb_internal.node_sessions
+WHERE active_queries LIKE 'SELECT user_name%'
+----
+root
+
+statement ok
+RESET ROLE

pkg/sql/sessiondata/session_data.go

@@ -203,11 +203,15 @@ func (s *SessionData) GetTemporarySchemaIDForDb(dbID uint32) (uint32, bool) {
 type Stack struct {
 	// Use an internal variable to prevent abstraction leakage.
 	stack []*SessionData
+	// base is a pointer to the first element of the stack.
+	// This avoids a race with stack being reassigned, as the first element
+	// is *always* set.
+	base *SessionData
 }
 
 // NewStack creates a new tack.
 func NewStack(firstElem *SessionData) *Stack {
-	return &Stack{stack: []*SessionData{firstElem}}
+	return &Stack{stack: []*SessionData{firstElem}, base: firstElem}
 }
 
 // Clone clones the current stack.
@@ -229,6 +233,12 @@ func (s *Stack) Top() *SessionData {
 	return s.stack[len(s.stack)-1]
 }
 
+// Base returns the bottom element of the stack.
+// This is a non-racy structure, as the bottom element is always constant.
+func (s *Stack) Base() *SessionData {
+	return s.base
+}
+
 // Push pushes a SessionData element to the stack.
 func (s *Stack) Push(elem *SessionData) {
 	s.stack = append(s.stack, elem)

pkg/sql/sessiondata/session_data_test.go

@@ -67,4 +67,6 @@ func TestStack(t *testing.T) {
 	require.Error(t, s.PopN(3), "there must always be at least one element in the SessionData stack")
 	require.NoError(t, s.PopN(2))
 	require.Equal(t, s.Elems(), []*SessionData{initialElem})
+
+	require.Equal(t, s.Base(), initialElem)
 }