diff --git a/store/opentracinglayer/opentracinglayer.go b/store/opentracinglayer/opentracinglayer.go index 87b443640a..129ef5254b 100644 --- a/store/opentracinglayer/opentracinglayer.go +++ b/store/opentracinglayer/opentracinglayer.go @@ -3197,7 +3197,7 @@ func (s *OpenTracingLayerComplianceStore) GetAll(offset int, limit int) (model.C return result, err } -func (s *OpenTracingLayerComplianceStore) MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { +func (s *OpenTracingLayerComplianceStore) MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { origCtx := s.Root.Store.Context() span, newCtx := tracing.StartSpanWithParentByContext(s.Root.Store.Context(), "ComplianceStore.MessageExport") s.Root.Store.SetContext(newCtx) @@ -3206,7 +3206,7 @@ func (s *OpenTracingLayerComplianceStore) MessageExport(cursor model.MessageExpo }() defer span.Finish() - result, resultVar1, err := s.ComplianceStore.MessageExport(cursor, limit) + result, resultVar1, err := s.ComplianceStore.MessageExport(ctx, cursor, limit) if err != nil { span.LogFields(spanlog.Error(err)) ext.Error.Set(span, true) diff --git a/store/retrylayer/retrylayer.go b/store/retrylayer/retrylayer.go index 38a3f52dbf..065a1400a9 100644 --- a/store/retrylayer/retrylayer.go +++ b/store/retrylayer/retrylayer.go @@ -3567,11 +3567,11 @@ func (s *RetryLayerComplianceStore) GetAll(offset int, limit int) (model.Complia } -func (s *RetryLayerComplianceStore) MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { +func (s *RetryLayerComplianceStore) MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { tries := 0 for { - result, resultVar1, err := s.ComplianceStore.MessageExport(cursor, limit) + result, resultVar1, err := s.ComplianceStore.MessageExport(ctx, cursor, limit) if err == nil { return result, resultVar1, nil } diff --git a/store/sqlstore/compliance_store.go b/store/sqlstore/compliance_store.go index 6e2ae38ac1..9c76cea3d6 100644 --- a/store/sqlstore/compliance_store.go +++ b/store/sqlstore/compliance_store.go @@ -4,6 +4,7 @@ package sqlstore import ( + "context" "database/sql" "fmt" "strings" @@ -270,7 +271,7 @@ func (s SqlComplianceStore) ComplianceExport(job *model.Compliance, cursor model return append(channelPosts, directMessagePosts...), cursor, nil } -func (s SqlComplianceStore) MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { +func (s SqlComplianceStore) MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { var args []any args = append(args, model.ChannelTypeDirect, model.ChannelTypeGroup, cursor.LastPostUpdateAt, cursor.LastPostUpdateAt, cursor.LastPostId, limit) query := @@ -317,7 +318,7 @@ func (s SqlComplianceStore) MessageExport(cursor model.MessageExportCursor, limi LIMIT ?` cposts := []*model.MessageExport{} - if err := s.GetReplicaX().Select(&cposts, query, args...); err != nil { + if err := s.GetReplicaX().SelectCtx(ctx, &cposts, query, args...); err != nil { return nil, cursor, errors.Wrap(err, "unable to export messages") } if len(cposts) > 0 { diff --git a/store/sqlstore/sqlx_wrapper.go b/store/sqlstore/sqlx_wrapper.go index 95b274aae9..3d215ff8a5 100644 --- a/store/sqlstore/sqlx_wrapper.go +++ b/store/sqlstore/sqlx_wrapper.go @@ -224,8 +224,12 @@ func (w *sqlxDBWrapper) QueryX(query string, args ...any) (*sqlx.Rows, error) { } func (w *sqlxDBWrapper) Select(dest any, query string, args ...any) error { + return w.SelectCtx(context.Background(), dest, query, args...) +} + +func (w *sqlxDBWrapper) SelectCtx(ctx context.Context, dest any, query string, args ...any) error { query = w.DB.Rebind(query) - ctx, cancel := context.WithTimeout(context.Background(), w.queryTimeout) + ctx, cancel := context.WithTimeout(ctx, w.queryTimeout) defer cancel() if w.trace { diff --git a/store/store.go b/store/store.go index 49490afdbb..8a1c69d91a 100644 --- a/store/store.go +++ b/store/store.go @@ -549,7 +549,7 @@ type ComplianceStore interface { Get(id string) (*model.Compliance, error) GetAll(offset, limit int) (model.Compliances, error) ComplianceExport(compliance *model.Compliance, cursor model.ComplianceExportCursor, limit int) ([]*model.CompliancePost, model.ComplianceExportCursor, error) - MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) + MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) } type OAuthStore interface { diff --git a/store/storetest/compliance_store.go b/store/storetest/compliance_store.go index f39b1c43ef..7c2d106325 100644 --- a/store/storetest/compliance_store.go +++ b/store/storetest/compliance_store.go @@ -4,6 +4,7 @@ package storetest import ( + "context" "encoding/json" "testing" "time" @@ -399,7 +400,7 @@ func testMessageExportPublicChannel(t *testing.T, ss store.Store) { // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -469,7 +470,7 @@ func testMessageExportPublicChannel(t *testing.T, ss store.Store) { // fetch the message exports for both posts that user1 sent messageExportMap := map[string]model.MessageExport{} - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 2, len(messages)) @@ -503,7 +504,7 @@ func testMessageExportPrivateChannel(t *testing.T, ss store.Store) { // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -573,7 +574,7 @@ func testMessageExportPrivateChannel(t *testing.T, ss store.Store) { // fetch the message exports for both posts that user1 sent messageExportMap := map[string]model.MessageExport{} - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 2, len(messages)) @@ -609,7 +610,7 @@ func testMessageExportDirectMessageChannel(t *testing.T, ss store.Store) { // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -664,7 +665,7 @@ func testMessageExportDirectMessageChannel(t *testing.T, ss store.Store) { // fetch the message export for the post that user1 sent messageExportMap := map[string]model.MessageExport{} - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) @@ -690,7 +691,7 @@ func testMessageExportGroupMessageChannel(t *testing.T, ss store.Store) { // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -762,7 +763,7 @@ func testMessageExportGroupMessageChannel(t *testing.T, ss store.Store) { // fetch the message export for the post that user1 sent messageExportMap := map[string]model.MessageExport{} - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 10}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) @@ -787,7 +788,7 @@ func testEditExportMessage(t *testing.T, ss store.Store) { defer cleanupStoreState(t, ss) // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -842,7 +843,7 @@ func testEditExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports from the start - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 2, len(messages)) @@ -879,7 +880,7 @@ func testEditAfterExportMessage(t *testing.T, ss store.Store) { defer cleanupStoreState(t, ss) // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -927,7 +928,7 @@ func testEditAfterExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports from the start - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) @@ -953,7 +954,7 @@ func testEditAfterExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports after edit - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: postEditTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: postEditTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 2, len(messages)) @@ -990,7 +991,7 @@ func testDeleteExportMessage(t *testing.T, ss store.Store) { defer cleanupStoreState(t, ss) // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -1043,7 +1044,7 @@ func testDeleteExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports from the start - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) @@ -1075,7 +1076,7 @@ func testDeleteAfterExportMessage(t *testing.T, ss store.Store) { defer cleanupStoreState(t, ss) // get the starting number of message export entries startTime := model.GetMillis() - messages, _, err := ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err := ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 0, len(messages)) @@ -1123,7 +1124,7 @@ func testDeleteAfterExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports from the start - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: startTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) @@ -1146,7 +1147,7 @@ func testDeleteAfterExportMessage(t *testing.T, ss store.Store) { require.NoError(t, err) // fetch the message exports after delete - messages, _, err = ss.Compliance().MessageExport(model.MessageExportCursor{LastPostUpdateAt: postDeleteTime - 1}, 10) + messages, _, err = ss.Compliance().MessageExport(context.Background(), model.MessageExportCursor{LastPostUpdateAt: postDeleteTime - 1}, 10) require.NoError(t, err) assert.Equal(t, 1, len(messages)) diff --git a/store/storetest/mocks/ComplianceStore.go b/store/storetest/mocks/ComplianceStore.go index 657670b023..a1041e3bae 100644 --- a/store/storetest/mocks/ComplianceStore.go +++ b/store/storetest/mocks/ComplianceStore.go @@ -5,6 +5,8 @@ package mocks import ( + context "context" + model "github.com/mattermost/mattermost-server/v6/model" mock "github.com/stretchr/testify/mock" ) @@ -90,13 +92,13 @@ func (_m *ComplianceStore) GetAll(offset int, limit int) (model.Compliances, err return r0, r1 } -// MessageExport provides a mock function with given fields: cursor, limit -func (_m *ComplianceStore) MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { - ret := _m.Called(cursor, limit) +// MessageExport provides a mock function with given fields: ctx, cursor, limit +func (_m *ComplianceStore) MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { + ret := _m.Called(ctx, cursor, limit) var r0 []*model.MessageExport - if rf, ok := ret.Get(0).(func(model.MessageExportCursor, int) []*model.MessageExport); ok { - r0 = rf(cursor, limit) + if rf, ok := ret.Get(0).(func(context.Context, model.MessageExportCursor, int) []*model.MessageExport); ok { + r0 = rf(ctx, cursor, limit) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]*model.MessageExport) @@ -104,15 +106,15 @@ func (_m *ComplianceStore) MessageExport(cursor model.MessageExportCursor, limit } var r1 model.MessageExportCursor - if rf, ok := ret.Get(1).(func(model.MessageExportCursor, int) model.MessageExportCursor); ok { - r1 = rf(cursor, limit) + if rf, ok := ret.Get(1).(func(context.Context, model.MessageExportCursor, int) model.MessageExportCursor); ok { + r1 = rf(ctx, cursor, limit) } else { r1 = ret.Get(1).(model.MessageExportCursor) } var r2 error - if rf, ok := ret.Get(2).(func(model.MessageExportCursor, int) error); ok { - r2 = rf(cursor, limit) + if rf, ok := ret.Get(2).(func(context.Context, model.MessageExportCursor, int) error); ok { + r2 = rf(ctx, cursor, limit) } else { r2 = ret.Error(2) } diff --git a/store/timerlayer/timerlayer.go b/store/timerlayer/timerlayer.go index 02be7b110d..19ccaf656d 100644 --- a/store/timerlayer/timerlayer.go +++ b/store/timerlayer/timerlayer.go @@ -2932,10 +2932,10 @@ func (s *TimerLayerComplianceStore) GetAll(offset int, limit int) (model.Complia return result, err } -func (s *TimerLayerComplianceStore) MessageExport(cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { +func (s *TimerLayerComplianceStore) MessageExport(ctx context.Context, cursor model.MessageExportCursor, limit int) ([]*model.MessageExport, model.MessageExportCursor, error) { start := time.Now() - result, resultVar1, err := s.ComplianceStore.MessageExport(cursor, limit) + result, resultVar1, err := s.ComplianceStore.MessageExport(ctx, cursor, limit) elapsed := float64(time.Since(start)) / float64(time.Second) if s.Root.Metrics != nil {