rewrite last query info (#3321)

Author: jycor

sql/base_session.go

@@ -33,9 +33,9 @@ type BaseSession struct {
 	logger           *logrus.Entry
 	locks            map[string]bool
 	storedProcParams map[string]*StoredProcParam
-	lastQueryInfo    map[string]*atomic.Value
 	systemVars       map[string]SystemVarValue
 	statusVars       map[string]StatusVarValue
+	lastQueryInfo    *LastQueryInfo
 	idxReg           *IndexRegistry
 	viewReg          *ViewRegistry
 	transactionDb    string
@@ -497,28 +497,8 @@ func (s *BaseSession) SetViewRegistry(reg *ViewRegistry) {
 	s.viewReg = reg
 }
 
-func (s *BaseSession) SetLastQueryInfoInt(key string, value int64) {
-	s.lastQueryInfo[key].Store(value)
-}
-
-func (s *BaseSession) GetLastQueryInfoInt(key string) int64 {
-	value, ok := s.lastQueryInfo[key].Load().(int64)
-	if !ok {
-		panic(fmt.Sprintf("last query info value stored for %s is not an int64 value, but a %T", key, s.lastQueryInfo[key]))
-	}
-	return value
-}
-
-func (s *BaseSession) SetLastQueryInfoString(key string, value string) {
-	s.lastQueryInfo[key].Store(value)
-}
-
-func (s *BaseSession) GetLastQueryInfoString(key string) string {
-	value, ok := s.lastQueryInfo[key].Load().(string)
-	if !ok {
-		panic(fmt.Sprintf("last query info value stored for %s is not a string value, but a %T", key, s.lastQueryInfo[key]))
-	}
-	return value
+func (s *BaseSession) GetLastQueryInfo() *LastQueryInfo {
+	return s.lastQueryInfo
 }
 
 func (s *BaseSession) GetTransaction() Transaction {

sql/expression/auto_uuid.go

@@ -77,7 +77,7 @@ func (au *AutoUuid) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 		//       read this value too early. We should verify this isn't how MySQL behaves, and then could fix
 		//       by setting a PENDING_LAST_INSERT_UUID value in the session query info, then moving it to
 		//       LAST_INSERT_UUID in the session query info at the end of execution.
-		ctx.Session.SetLastQueryInfoString(sql.LastInsertUuid, uuidValue)
+		ctx.GetLastQueryInfo().LastInsertUUID.Store(uuidValue)
 		au.foundUuid = true
 	}
 

sql/expression/function/queryinfo.go

@@ -66,7 +66,7 @@ func (r *RowCount) IsNullable() bool {
 
 // Eval implements sql.Expression
 func (r *RowCount) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
-	return ctx.GetLastQueryInfoInt(sql.RowCount), nil
+	return ctx.GetLastQueryInfo().RowCount.Load(), nil
 }
 
 // Children implements sql.Expression
@@ -126,8 +126,8 @@ func (l *LastInsertUuid) IsNullable() bool {
 }
 
 func (l *LastInsertUuid) Eval(ctx *sql.Context, _ sql.Row) (interface{}, error) {
-	lastInsertUuid := ctx.GetLastQueryInfoString(sql.LastInsertUuid)
-	result, _, err := l.Type().Convert(ctx, lastInsertUuid)
+	lastInsertUUID := ctx.GetLastQueryInfo().LastInsertUUID.Load()
+	result, _, err := l.Type().Convert(ctx, lastInsertUUID)
 	if err != nil {
 		return nil, err
 	}
@@ -204,7 +204,7 @@ func (r *LastInsertId) IsNullable() bool {
 func (r *LastInsertId) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 	// With no arguments, just return the last insert id for this session
 	if len(r.Children()) == 0 {
-		lastInsertId := ctx.GetLastQueryInfoInt(sql.LastInsertId)
+		lastInsertId := ctx.GetLastQueryInfo().LastInsertId.Load()
 		unsigned, _, err := types.Uint64.Convert(ctx, lastInsertId)
 		if err != nil {
 			return nil, err
@@ -222,7 +222,7 @@ func (r *LastInsertId) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 		return nil, err
 	}
 
-	ctx.SetLastQueryInfoInt(sql.LastInsertId, id.(int64))
+	ctx.GetLastQueryInfo().LastInsertId.Store(id.(int64))
 	return id, nil
 }
 
@@ -296,7 +296,7 @@ func (r *FoundRows) IsNullable() bool {
 
 // Eval implements sql.Expression
 func (r *FoundRows) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
-	return ctx.GetLastQueryInfoInt(sql.FoundRows), nil
+	return ctx.GetLastQueryInfo().FoundRows.Load(), nil
 }
 
 // Children implements sql.Expression

sql/iters/rel_iters.go

@@ -67,11 +67,9 @@ func (i *topRowsIter) Next(ctx *sql.Context) (sql.Row, error) {
 
 func (i *topRowsIter) Close(ctx *sql.Context) error {
 	i.topRows = nil
-
 	if i.calcFoundRows {
-		ctx.SetLastQueryInfoInt(sql.FoundRows, i.numFoundRows)
+		ctx.GetLastQueryInfo().FoundRows.Store(i.numFoundRows)
 	}
-
 	return i.childIter.Close(ctx)
 }
 
@@ -467,9 +465,8 @@ func (li *LimitIter) Close(ctx *sql.Context) error {
 	if err != nil {
 		return err
 	}
-
 	if li.CalcFoundRows {
-		ctx.SetLastQueryInfoInt(sql.FoundRows, li.currentPos)
+		ctx.GetLastQueryInfo().FoundRows.Store(li.currentPos)
 	}
 	return nil
 }

sql/plan/process.go

@@ -354,19 +354,20 @@ func (i *TrackedRowIter) GetIter() sql.RowIter {
 }
 
 func (i *TrackedRowIter) updateSessionVars(ctx *sql.Context) {
+	// TODO: possible to just remove switch entirely?
 	switch i.QueryType {
 	case QueryTypeSelect:
-		ctx.SetLastQueryInfoInt(sql.RowCount, -1)
+		ctx.GetLastQueryInfo().RowCount.Store(-1)
 	case QueryTypeDdl:
-		ctx.SetLastQueryInfoInt(sql.RowCount, 0)
+		ctx.GetLastQueryInfo().RowCount.Store(0)
 	case QueryTypeUpdate:
 		// This is handled by RowUpdateAccumulator
 	default:
 		panic(fmt.Sprintf("Unexpected query type %v", i.QueryType))
 	}
 
 	if i.ShouldSetFoundRows {
-		ctx.SetLastQueryInfoInt(sql.FoundRows, i.numRows)
+		ctx.GetLastQueryInfo().FoundRows.Store(i.numRows)
 	}
 }
 

sql/rowexec/dml_iters.go

@@ -746,14 +746,15 @@ func (a *accumulatorIter) Next(ctx *sql.Context) (r sql.Row, err error) {
 		}
 		if err == io.EOF {
 			// TODO: The information flow here is pretty gnarly. We
-			// set some session variables based on the result, and
-			// we actually use a session variable to set
-			// InsertID. This should be improved.
+			//  set some session variables based on the result, and
+			//  we actually use a session variable to set
+			//  InsertID. This should be improved.
 
 			// UPDATE statements also set FoundRows to the number of rows that
 			// matched the WHERE clause, same as a SELECT.
+			lastQueryInfo := ctx.GetLastQueryInfo()
 			if ma, ok := a.updateRowHandler.(matchingAccumulator); ok {
-				ctx.SetLastQueryInfoInt(sql.FoundRows, ma.RowsMatched())
+				lastQueryInfo.FoundRows.Store(ma.RowsMatched())
 			}
 
 			res := a.updateRowHandler.okResult() // TODO: Should add warnings here
@@ -763,14 +764,13 @@ func (a *accumulatorIter) Next(ctx *sql.Context) (r sql.Row, err error) {
 			// to be fixed. See comment in buildRowUpdateAccumulator in rowexec/dml.go
 			switch rowHandler := a.updateRowHandler.(type) {
 			case *onDuplicateUpdateHandler, *replaceRowHandler:
-				lastInsertId := ctx.Session.GetLastQueryInfoInt(sql.LastInsertId)
+				lastInsertId := lastQueryInfo.LastInsertId.Load()
 				res.InsertID = uint64(lastInsertId)
 			case *insertRowHandler:
 				res.InsertID = rowHandler.lastInsertId
 			}
-
 			// By definition, ROW_COUNT() is equal to RowsAffected.
-			ctx.SetLastQueryInfoInt(sql.RowCount, int64(res.RowsAffected))
+			lastQueryInfo.RowCount.Store(int64(res.RowsAffected))
 
 			return sql.NewRow(res), nil
 		} else if isIg {

sql/rowexec/insert.go

@@ -355,7 +355,7 @@ func (i *insertIter) updateLastInsertId(ctx *sql.Context, row sql.Row) {
 	}
 	if i.firstGeneratedAutoIncRowIdx == 0 {
 		autoIncVal := i.getAutoIncVal(row)
-		ctx.SetLastQueryInfoInt(sql.LastInsertId, autoIncVal)
+		ctx.GetLastQueryInfo().LastInsertId.Store(autoIncVal)
 	}
 	i.firstGeneratedAutoIncRowIdx--
 }

sql/session.go

@@ -131,14 +131,8 @@ type Session interface {
 	DelLock(lockName string) error
 	// IterLocks iterates through all locks owned by this user
 	IterLocks(cb func(name string) error) error
-	// SetLastQueryInfoInt sets session-level query info for the key given, applying to the query just executed.
-	SetLastQueryInfoInt(key string, value int64)
-	// GetLastQueryInfoInt returns the session-level query info for the key given, for the query most recently executed.
-	GetLastQueryInfoInt(key string) int64
-	// SetLastQueryInfoString sets session-level query info as a string for the key given, applying to the query just executed.
-	SetLastQueryInfoString(key string, value string)
-	// GetLastQueryInfoString returns the session-level query info as a string for the key given, for the query most recently executed.
-	GetLastQueryInfoString(key string) string
+	// GetLastQueryInfo returns session-level info for the most recently executed query.
+	GetLastQueryInfo() *LastQueryInfo
 	// GetTransaction returns the active transaction, if any
 	GetTransaction() Transaction
 	// SetTransaction sets the session's transaction
@@ -253,12 +247,21 @@ type (
 	}
 )
 
-const (
-	RowCount       = "row_count"
-	FoundRows      = "found_rows"
-	LastInsertId   = "last_insert_id"
-	LastInsertUuid = "last_insert_uuid"
-)
+type LastQueryInfo struct {
+	RowCount       atomic.Int64 // Session-level Row Count for the last executed query
+	FoundRows      atomic.Int64 // Session-level Found Rows for the last executed query
+	LastInsertId   atomic.Int64 // Session-level ID for the last executed insert query
+	LastInsertUUID atomic.Value // Session-level UUID for the last executed insert query
+}
+
+func defaultLastQueryInfo() *LastQueryInfo {
+	ret := LastQueryInfo{}
+	ret.RowCount.Store(0)
+	ret.FoundRows.Store(1) // this is kind of a hack -- it handles the case of `select found_rows()` before any select statement is issued
+	ret.LastInsertId.Store(0)
+	ret.LastInsertUUID.Store("")
+	return &ret
+}
 
 // Session ID 0 used as invalid SessionID
 var autoSessionIDs uint32 = 1
@@ -703,19 +706,6 @@ func (i *spanIter) Close(ctx *Context) error {
 	return i.iter.Close(ctx)
 }
 
-func defaultLastQueryInfo() map[string]*atomic.Value {
-	ret := make(map[string]*atomic.Value)
-	ret[RowCount] = &atomic.Value{}
-	ret[RowCount].Store(int64(0))
-	ret[FoundRows] = &atomic.Value{}
-	ret[FoundRows].Store(int64(1)) // this is kind of a hack -- it handles the case of `select found_rows()` before any select statement is issue)
-	ret[LastInsertId] = &atomic.Value{}
-	ret[LastInsertId].Store(int64(0))
-	ret[LastInsertUuid] = &atomic.Value{}
-	ret[LastInsertUuid].Store("")
-	return ret
-}
-
 // cc: https://dev.mysql.com/doc/refman/8.0/en/temporary-files.html
 func GetTmpdirSessionVar() string {
 	ret := os.Getenv("TMPDIR")