[MM-48946] Fix read after write issue when uploading data (#21868)

* Fix read after write issue when uploading data

* Prefer request.CTX interface

store/opentracinglayer/opentracinglayer.go

@@ -10612,7 +10612,7 @@ func (s *OpenTracingLayerUploadSessionStore) Delete(id string) error {
 	return err
 }
 
-func (s *OpenTracingLayerUploadSessionStore) Get(id string) (*model.UploadSession, error) {
+func (s *OpenTracingLayerUploadSessionStore) Get(ctx context.Context, id string) (*model.UploadSession, error) {
 	origCtx := s.Root.Store.Context()
 	span, newCtx := tracing.StartSpanWithParentByContext(s.Root.Store.Context(), "UploadSessionStore.Get")
 	s.Root.Store.SetContext(newCtx)
@@ -10621,7 +10621,7 @@ func (s *OpenTracingLayerUploadSessionStore) Get(id string) (*model.UploadSessio
 	}()
 
 	defer span.Finish()
-	result, err := s.UploadSessionStore.Get(id)
+	result, err := s.UploadSessionStore.Get(ctx, id)
 	if err != nil {
 		span.LogFields(spanlog.Error(err))
 		ext.Error.Set(span, true)

store/retrylayer/retrylayer.go

@@ -12126,11 +12126,11 @@ func (s *RetryLayerUploadSessionStore) Delete(id string) error {
 
 }
 
-func (s *RetryLayerUploadSessionStore) Get(id string) (*model.UploadSession, error) {
+func (s *RetryLayerUploadSessionStore) Get(ctx context.Context, id string) (*model.UploadSession, error) {
 
 	tries := 0
 	for {
-		result, err := s.UploadSessionStore.Get(id)
+		result, err := s.UploadSessionStore.Get(ctx, id)
 		if err == nil {
 			return result, nil
 		}

store/sqlstore/upload_session_store.go

@@ -4,6 +4,7 @@
 package sqlstore
 
 import (
+	"context"
 	"database/sql"
 
 	sq "github.com/mattermost/squirrel"
@@ -78,7 +79,7 @@ func (us SqlUploadSessionStore) Update(session *model.UploadSession) error {
 	return nil
 }
 
-func (us SqlUploadSessionStore) Get(id string) (*model.UploadSession, error) {
+func (us SqlUploadSessionStore) Get(ctx context.Context, id string) (*model.UploadSession, error) {
 	if !model.IsValidId(id) {
 		return nil, errors.New("SqlUploadSessionStore.Get: id is not valid")
 	}
@@ -91,7 +92,7 @@ func (us SqlUploadSessionStore) Get(id string) (*model.UploadSession, error) {
 		return nil, errors.Wrap(err, "SqlUploadSessionStore.Get: failed to build query")
 	}
 	var session model.UploadSession
-	if err := us.GetReplicaX().Get(&session, query, args...); err != nil {
+	if err := us.DBXFromContext(ctx).Get(&session, query, args...); err != nil {
 		if err == sql.ErrNoRows {
 			return nil, store.NewErrNotFound("UploadSession", id)
 		}

store/store.go

@@ -714,7 +714,7 @@ type FileInfoStore interface {
 type UploadSessionStore interface {
 	Save(session *model.UploadSession) (*model.UploadSession, error)
 	Update(session *model.UploadSession) error
-	Get(id string) (*model.UploadSession, error)
+	Get(ctx context.Context, id string) (*model.UploadSession, error)
 	GetForUser(userID string) ([]*model.UploadSession, error)
 	Delete(id string) error
 }

store/storetest/mocks/UploadSessionStore.go

@@ -5,6 +5,8 @@
 package mocks
 
 import (
+	context "context"
+
 	model "github.com/mattermost/mattermost-server/v6/model"
 	mock "github.com/stretchr/testify/mock"
 )
@@ -28,13 +30,13 @@ func (_m *UploadSessionStore) Delete(id string) error {
 	return r0
 }
 
-// Get provides a mock function with given fields: id
-func (_m *UploadSessionStore) Get(id string) (*model.UploadSession, error) {
-	ret := _m.Called(id)
+// Get provides a mock function with given fields: ctx, id
+func (_m *UploadSessionStore) Get(ctx context.Context, id string) (*model.UploadSession, error) {
+	ret := _m.Called(ctx, id)
 
 	var r0 *model.UploadSession
-	if rf, ok := ret.Get(0).(func(string) *model.UploadSession); ok {
-		r0 = rf(id)
+	if rf, ok := ret.Get(0).(func(context.Context, string) *model.UploadSession); ok {
+		r0 = rf(ctx, id)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*model.UploadSession)
@@ -42,8 +44,8 @@ func (_m *UploadSessionStore) Get(id string) (*model.UploadSession, error) {
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(id)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, id)
 	} else {
 		r1 = ret.Error(1)
 	}

store/storetest/upload_session_store.go

@@ -4,6 +4,7 @@
 package storetest
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -52,13 +53,13 @@ func testUploadSessionStoreSaveGet(t *testing.T, ss store.Store) {
 	})
 
 	t.Run("getting non-existing session should fail", func(t *testing.T) {
-		us, err := ss.UploadSession().Get("fake")
+		us, err := ss.UploadSession().Get(context.Background(), "fake")
 		require.Error(t, err)
 		require.Nil(t, us)
 	})
 
 	t.Run("getting existing session should succeed", func(t *testing.T) {
-		us, err := ss.UploadSession().Get(session.Id)
+		us, err := ss.UploadSession().Get(context.Background(), session.Id)
 		require.NoError(t, err)
 		require.NotNil(t, us)
 		require.Equal(t, session, us)
@@ -100,7 +101,7 @@ func testUploadSessionStoreUpdate(t *testing.T, ss store.Store) {
 		err = ss.UploadSession().Update(us)
 		require.NoError(t, err)
 
-		updated, err := ss.UploadSession().Get(us.Id)
+		updated, err := ss.UploadSession().Get(context.Background(), us.Id)
 		require.NoError(t, err)
 		require.NotNil(t, us)
 		require.Equal(t, us, updated)
@@ -199,7 +200,7 @@ func testUploadSessionStoreDelete(t *testing.T, ss store.Store) {
 		err = ss.UploadSession().Delete(session.Id)
 		require.NoError(t, err)
 
-		us, err = ss.UploadSession().Get(us.Id)
+		us, err = ss.UploadSession().Get(context.Background(), us.Id)
 		require.Error(t, err)
 		require.Nil(t, us)
 		require.IsType(t, &store.ErrNotFound{}, err)

store/timerlayer/timerlayer.go

@@ -9549,10 +9549,10 @@ func (s *TimerLayerUploadSessionStore) Delete(id string) error {
 	return err
 }
 
-func (s *TimerLayerUploadSessionStore) Get(id string) (*model.UploadSession, error) {
+func (s *TimerLayerUploadSessionStore) Get(ctx context.Context, id string) (*model.UploadSession, error) {
 	start := time.Now()
 
-	result, err := s.UploadSessionStore.Get(id)
+	result, err := s.UploadSessionStore.Get(ctx, id)
 
 	elapsed := float64(time.Since(start)) / float64(time.Second)
 	if s.Root.Metrics != nil {