Merge cockroachdb#69394

69394: sql: refactor sessionDataMutator as sessionDataMutatorIterator r=rafiss a=otan

Part of a body of work for SET LOCAL (refs cockroachdb#32562)

See individual commits for details.

Release justification: high pri work

Co-authored-by: Oliver Tan <otan@cockroachlabs.com>

pkg/ccl/importccl/import_stmt.go

@@ -1791,8 +1791,7 @@ func (r *importResumer) parseBundleSchemaIfNeeded(ctx context.Context, phs inter
 
 	owner := r.job.Payload().UsernameProto.Decode()
 
-	sessionMutator := p.ExtendedEvalContext().SessionMutator
-	sessionMutator.SetDefaultIntSize(details.DefaultIntSize)
+	p.ExtendedEvalContext().SessionMutatorIterator.SetSessionDefaultIntSize(details.DefaultIntSize)
 
 	if details.ParseBundleSchema {
 		var span *tracing.Span

pkg/sql/conn_executor.go

@@ -582,11 +582,13 @@ func (s *Server) SetupConn(
 	memMetrics MemoryMetrics,
 ) (ConnectionHandler, error) {
 	sd := s.newSessionData(args)
-
+	sds := sessiondata.NewStack(sd)
 	// Set the SessionData from args.SessionDefaults. This also validates the
 	// respective values.
-	sdMut := s.makeSessionDataMutator(sd, args.SessionDefaults)
-	if err := resetSessionVars(ctx, &sdMut); err != nil {
+	sdMutIterator := s.makeSessionDataMutatorIterator(sds, args.SessionDefaults)
+	if err := sdMutIterator.forEachMutatorError(func(m *sessionDataMutator) error {
+		return resetSessionVars(ctx, m)
+	}); err != nil {
 		log.Errorf(ctx, "error setting up client session: %s", err)
 		return ConnectionHandler{}, err
 	}
@@ -627,7 +629,7 @@ func (h ConnectionHandler) GetParamStatus(ctx context.Context, varName string) s
 		log.Fatalf(ctx, "programming error: status param %q must be defined session var", varName)
 		return ""
 	}
-	hasDefault, defVal := getSessionVarDefaultString(name, v, h.ex.dataMutator)
+	hasDefault, defVal := getSessionVarDefaultString(name, v, h.ex.dataMutatorIterator.sessionDataMutatorBase)
 	if !hasDefault {
 		log.Fatalf(ctx, "programming error: status param %q must have a default value", varName)
 		return ""
@@ -672,14 +674,16 @@ func (s *Server) newSessionData(args SessionArgs) *sessiondata.SessionData {
 	return sd
 }
 
-func (s *Server) makeSessionDataMutator(
-	sd *sessiondata.SessionData, defaults SessionDefaults,
-) sessionDataMutator {
-	return sessionDataMutator{
-		data:               sd,
-		defaults:           defaults,
-		settings:           s.cfg.Settings,
-		paramStatusUpdater: &noopParamStatusUpdater{},
+func (s *Server) makeSessionDataMutatorIterator(
+	sds *sessiondata.Stack, defaults SessionDefaults,
+) *sessionDataMutatorIterator {
+	return &sessionDataMutatorIterator{
+		sds: sds,
+		sessionDataMutatorBase: sessionDataMutatorBase{
+			defaults: defaults,
+			settings: s.cfg.Settings,
+		},
+		sessionDataMutatorCallbacks: sessionDataMutatorCallbacks{},
 	}
 }
 
@@ -739,8 +743,6 @@ func (s *Server) newConnExecutor(
 	)
 
 	nodeIDOrZero, _ := s.cfg.NodeID.OptionalNodeID()
-	sdMutator := new(sessionDataMutator)
-	*sdMutator = s.makeSessionDataMutator(sd, sdDefaults)
 	ex := &connExecutor{
 		server:           s,
 		metrics:          srvMetrics,
@@ -749,7 +751,6 @@ func (s *Server) newConnExecutor(
 		mon:              sessionRootMon,
 		sessionMon:       sessionMon,
 		sessionDataStack: sessiondata.NewStack(sd),
-		dataMutator:      sdMutator,
 		state: txnState{
 			mon:                          txnMon,
 			connCtx:                      ctx,
@@ -781,20 +782,23 @@ func (s *Server) newConnExecutor(
 
 	ex.state.txnAbortCount = ex.metrics.EngineMetrics.TxnAbortCount
 
+	ex.dataMutatorIterator = s.makeSessionDataMutatorIterator(
+		ex.sessionDataStack,
+		sdDefaults,
+	)
 	// The transaction_read_only variable is special; its updates need to be
 	// hooked-up to the executor.
-	sdMutator.setCurTxnReadOnly = func(val bool) {
+	ex.dataMutatorIterator.setCurTxnReadOnly = func(val bool) {
 		ex.state.readOnly = val
 	}
-
-	sdMutator.onTempSchemaCreation = func() {
+	ex.dataMutatorIterator.onTempSchemaCreation = func() {
 		ex.hasCreatedTemporarySchema = true
 	}
 
 	ex.applicationName.Store(ex.sessionData().ApplicationName)
 	ex.statsWriter = statsWriter
 	ex.statsCollector = sslocal.NewStatsCollector(statsWriter, ex.phaseTimes)
-	sdMutator.RegisterOnSessionDataChange("application_name", func(newName string) {
+	ex.dataMutatorIterator.RegisterOnSessionDataChange("application_name", func(newName string) {
 		ex.applicationName.Store(newName)
 		ex.statsWriter = ex.server.sqlStats.GetWriterForApplication(newName)
 	})
@@ -852,7 +856,9 @@ func (s *Server) newConnExecutorWithTxn(
 	if txn.Type() == kv.LeafTxn {
 		// If the txn is a leaf txn it is not allowed to perform mutations. For
 		// sanity, set read only on the session.
-		ex.dataMutator.SetReadOnly(true)
+		ex.dataMutatorIterator.forEachMutator(func(m *sessionDataMutator) {
+			m.SetReadOnly(true)
+		})
 	}
 
 	// The new transaction stuff below requires active monitors and traces, so
@@ -1226,9 +1232,10 @@ type connExecutor struct {
 
 	// sessionDataStack contains the user-configurable connection variables.
 	sessionDataStack *sessiondata.Stack
-	// dataMutator is nil for session-bound internal executors; we shouldn't issue
-	// statements that manipulate session state to an internal executor.
-	dataMutator *sessionDataMutator
+	// dataMutatorIterator is nil for session-bound internal executors; we
+	// shouldn't issue statements that manipulate session state to an internal
+	// executor.
+	dataMutatorIterator *sessionDataMutatorIterator
 
 	// statsWriter is a writer interface for recording per-application SQL usage
 	// statistics. It is maintained to represent statistics for the application
@@ -1477,7 +1484,9 @@ func (ex *connExecutor) Ctx() context.Context {
 	return ctx
 }
 
-// sessionData returns the top SessionData on the executor.
+// sessionData returns the top SessionData in the executor's sessionDataStack.
+// This should be how callers should reference SessionData objects, as it
+// will always contain the "latest" SessionData object in the transaction.
 func (ex *connExecutor) sessionData() *sessiondata.SessionData {
 	if ex.sessionDataStack == nil {
 		return nil
@@ -2380,7 +2389,7 @@ func (ex *connExecutor) initEvalCtx(ctx context.Context, evalCtx *extendedEvalCo
 			SQLStatsController:     ex.server.sqlStatsController,
 			CompactEngineSpan:      ex.server.cfg.CompactEngineSpanFunc,
 		},
-		SessionMutator:         ex.dataMutator,
+		SessionMutatorIterator: ex.dataMutatorIterator,
 		VirtualSchemas:         ex.server.cfg.VirtualSchemas,
 		Tracing:                &ex.sessionTracing,
 		NodesStatusServer:      ex.server.cfg.NodesStatusServer,
@@ -2462,7 +2471,7 @@ func (ex *connExecutor) initPlanner(ctx context.Context, p *planner) {
 
 	ex.initEvalCtx(ctx, &p.extendedEvalCtx, p)
 
-	p.sessionDataMutator = ex.dataMutator
+	p.sessionDataMutatorIterator = ex.dataMutatorIterator
 	p.noticeSender = nil
 	p.preparedStatements = ex.getPrepStmtsAccessor()
 

pkg/sql/conn_executor_exec.go

@@ -363,7 +363,7 @@ func (ex *connExecutor) execStmtInOpenState(
 	stmtTS := ex.server.cfg.Clock.PhysicalTime()
 	ex.statsCollector.Reset(ex.statsWriter, ex.phaseTimes)
 	ex.resetPlanner(ctx, p, ex.state.mu.txn, stmtTS)
-	p.sessionDataMutator.paramStatusUpdater = res
+	p.sessionDataMutatorIterator.paramStatusUpdater = res
 	p.noticeSender = res
 	ih := &p.instrumentation
 

pkg/sql/discard.go

@@ -30,7 +30,11 @@ func (p *planner) Discard(ctx context.Context, s *tree.Discard) (planNode, error
 		}
 
 		// RESET ALL
-		if err := resetSessionVars(ctx, p.sessionDataMutator); err != nil {
+		if err := p.sessionDataMutatorIterator.forEachMutatorError(
+			func(m *sessionDataMutator) error {
+				return resetSessionVars(ctx, m)
+			},
+		); err != nil {
 			return nil, err
 		}
 
@@ -46,7 +50,7 @@ func resetSessionVars(ctx context.Context, m *sessionDataMutator) error {
 	for _, varName := range varNames {
 		v := varGen[varName]
 		if v.Set != nil {
-			hasDefault, defVal := getSessionVarDefaultString(varName, v, m)
+			hasDefault, defVal := getSessionVarDefaultString(varName, v, m.sessionDataMutatorBase)
 			if hasDefault {
 				if err := v.Set(ctx, m, defVal); err != nil {
 					return err

pkg/sql/exec_util.go

@@ -2391,38 +2391,108 @@ type paramStatusUpdater interface {
 	BufferParamStatusUpdate(string, string)
 }
 
-// noopParamStatusUpdater implements paramStatusUpdater by performing a no-op.
-type noopParamStatusUpdater struct{}
-
-var _ paramStatusUpdater = (*noopParamStatusUpdater)(nil)
-
-func (noopParamStatusUpdater) BufferParamStatusUpdate(string, string) {}
+// sessionDataMutatorBase contains elements in a sessionDataMutator
+// which is the same across all SessionData elements in the sessiondata.Stack.
+type sessionDataMutatorBase struct {
+	defaults SessionDefaults
+	settings *cluster.Settings
+}
 
-// sessionDataMutator is the interface used by sessionVars to change the session
-// state. It mostly mutates the Session's SessionData, but not exclusively (e.g.
-// see curTxnReadOnly).
-type sessionDataMutator struct {
-	data               *sessiondata.SessionData
-	defaults           SessionDefaults
-	settings           *cluster.Settings
+// RegisterOnSessionDataChange adds a listener to execute when a change on the
+// given key is made using the mutator object.
+func (it *sessionDataMutatorIterator) RegisterOnSessionDataChange(key string, f func(val string)) {
+	if it.onSessionDataChangeListeners == nil {
+		it.onSessionDataChangeListeners = make(map[string][]func(val string))
+	}
+	it.onSessionDataChangeListeners[key] = append(it.onSessionDataChangeListeners[key], f)
+}
+
+// sessionDataMutatorCallbacks contains elements in a sessionDataMutator
+// which are only populated when mutating the "top" sessionData element.
+// It is intended for functions which should only be called once per SET
+// (e.g. param status updates, which only should be sent once within
+// a transaction where there may be two or more SessionData elements in
+// the stack)
+type sessionDataMutatorCallbacks struct {
+	// paramStatusUpdater is called when there is a ParamStatusUpdate.
+	// It can be nil, in which case nothing triggers on execution.
 	paramStatusUpdater paramStatusUpdater
 	// setCurTxnReadOnly is called when we execute SET transaction_read_only = ...
+	// It can be nil, in which case nothing triggers on execution.
 	setCurTxnReadOnly func(val bool)
 	// onTempSchemaCreation is called when the temporary schema is set
 	// on the search path (the first and only time).
+	// It can be nil, in which case nothing triggers on execution.
 	onTempSchemaCreation func()
 	// onSessionDataChangeListeners stores all the observers to execute when
 	// session data is modified, keyed by the value to change on.
 	onSessionDataChangeListeners map[string][]func(val string)
 }
 
-// RegisterOnSessionDataChange adds a listener to execute when a change on the
-// given key is made using the mutator object.
-func (m *sessionDataMutator) RegisterOnSessionDataChange(key string, f func(val string)) {
-	if m.onSessionDataChangeListeners == nil {
-		m.onSessionDataChangeListeners = make(map[string][]func(val string))
+// sessionDataMutatorIterator generates sessionDataMutators which allow
+// the changing of SessionData on some element inside the sessiondata Stack.
+type sessionDataMutatorIterator struct {
+	sessionDataMutatorBase
+	sds *sessiondata.Stack
+	sessionDataMutatorCallbacks
+}
+
+// mutator returns a mutator for the given sessionData.
+func (it *sessionDataMutatorIterator) mutator(
+	applyCallbacks bool, sd *sessiondata.SessionData,
+) *sessionDataMutator {
+	ret := &sessionDataMutator{
+		data:                   sd,
+		sessionDataMutatorBase: it.sessionDataMutatorBase,
+	}
+	// We usually apply callbacks on the first element in the stack, as the txn
+	// rollback will always reset to the first element we touch in the stack,
+	// in which case it should be up-to-date by default.
+	if applyCallbacks {
+		ret.sessionDataMutatorCallbacks = it.sessionDataMutatorCallbacks
 	}
-	m.onSessionDataChangeListeners[key] = append(m.onSessionDataChangeListeners[key], f)
+	return ret
+}
+
+// SetSessionDefaultIntSize sets the default int size for the session.
+// It is exported for use in import which is a CCL package.
+func (it *sessionDataMutatorIterator) SetSessionDefaultIntSize(size int32) {
+	it.forEachMutator(func(m *sessionDataMutator) {
+		m.SetDefaultIntSize(size)
+	})
+}
+
+// forEachMutator iterates over each mutator over all SessionData elements
+// in the stack and applies the given function to them.
+// It is the equivalent of SET SESSION x = y.
+func (it *sessionDataMutatorIterator) forEachMutator(applyFunc func(m *sessionDataMutator)) {
+	elems := it.sds.Elems()
+	for i, sd := range elems {
+		applyFunc(it.mutator(i == 0, sd))
+	}
+}
+
+// forEachMutatorError is the same as forEachMutator, but takes in a function
+// that can return an error, erroring if any of applications error.
+func (it *sessionDataMutatorIterator) forEachMutatorError(
+	applyFunc func(m *sessionDataMutator) error,
+) error {
+	elems := it.sds.Elems()
+	for i, sd := range elems {
+		if err := applyFunc(it.mutator(i == 0, sd)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// sessionDataMutator is the object used by sessionVars to change the session
+// state. It mostly mutates the session's SessionData, but not exclusively (e.g.
+// see curTxnReadOnly).
+type sessionDataMutator struct {
+	data *sessiondata.SessionData
+	sessionDataMutatorBase
+	sessionDataMutatorCallbacks
 }
 
 func (m *sessionDataMutator) notifyOnDataChangeListeners(key string, val string) {
@@ -2431,11 +2501,17 @@ func (m *sessionDataMutator) notifyOnDataChangeListeners(key string, val string)
 	}
 }
 
+func (m *sessionDataMutator) bufferParamStatusUpdate(param string, status string) {
+	if m.paramStatusUpdater != nil {
+		m.paramStatusUpdater.BufferParamStatusUpdate(param, status)
+	}
+}
+
 // SetApplicationName sets the application name.
 func (m *sessionDataMutator) SetApplicationName(appName string) {
 	m.data.ApplicationName = appName
 	m.notifyOnDataChangeListeners("application_name", appName)
-	m.paramStatusUpdater.BufferParamStatusUpdate("application_name", appName)
+	m.bufferParamStatusUpdate("application_name", appName)
 }
 
 func (m *sessionDataMutator) SetBytesEncodeFormat(val sessiondatapb.BytesEncodeFormat) {
@@ -2451,7 +2527,9 @@ func (m *sessionDataMutator) SetDatabase(dbName string) {
 }
 
 func (m *sessionDataMutator) SetTemporarySchemaName(scName string) {
-	m.onTempSchemaCreation()
+	if m.onTempSchemaCreation != nil {
+		m.onTempSchemaCreation()
+	}
 	m.data.SearchPath = m.data.SearchPath.WithTemporarySchemaName(scName)
 }
 
@@ -2570,7 +2648,7 @@ func (m *sessionDataMutator) UpdateSearchPath(paths []string) {
 
 func (m *sessionDataMutator) SetLocation(loc *time.Location) {
 	m.data.Location = loc
-	m.paramStatusUpdater.BufferParamStatusUpdate("TimeZone", sessionDataTimeZoneFormat(loc))
+	m.bufferParamStatusUpdate("TimeZone", sessionDataTimeZoneFormat(loc))
 }
 
 func (m *sessionDataMutator) SetReadOnly(val bool) {
@@ -2674,13 +2752,13 @@ func (m *sessionDataMutator) initSequenceCache() {
 // SetIntervalStyle sets the IntervalStyle for the given session.
 func (m *sessionDataMutator) SetIntervalStyle(style duration.IntervalStyle) {
 	m.data.DataConversionConfig.IntervalStyle = style
-	m.paramStatusUpdater.BufferParamStatusUpdate("IntervalStyle", strings.ToLower(style.String()))
+	m.bufferParamStatusUpdate("IntervalStyle", strings.ToLower(style.String()))
 }
 
 // SetDateStyle sets the DateStyle for the given session.
 func (m *sessionDataMutator) SetDateStyle(style pgdate.DateStyle) {
 	m.data.DataConversionConfig.DateStyle = style
-	m.paramStatusUpdater.BufferParamStatusUpdate("DateStyle", style.SQLString())
+	m.bufferParamStatusUpdate("DateStyle", style.SQLString())
 }
 
 // SetIntervalStyleEnabled sets the IntervalStyleEnabled for the given session.

pkg/sql/pg_catalog.go

@@ -2398,7 +2398,11 @@ https://www.postgresql.org/docs/9.5/catalog-pg-settings.html`,
 					globalDefVal := gen.GlobalDefault(&p.EvalContext().Settings.SV)
 					bootDatum = tree.NewDString(globalDefVal)
 				}
-				if hasDefault, defVal := getSessionVarDefaultString(vName, gen, p.sessionDataMutator); hasDefault {
+				if hasDefault, defVal := getSessionVarDefaultString(
+					vName,
+					gen,
+					p.sessionDataMutatorIterator.sessionDataMutatorBase,
+				); hasDefault {
 					resetDatum = tree.NewDString(defVal)
 				}
 			}