From e9e822105c6427aadc97d9d5f4a4afaee3ce529f Mon Sep 17 00:00:00 2001 From: abettigole Date: Tue, 28 Jul 2026 22:33:39 +0000 Subject: [PATCH] feat(storage): add request batch assignments Add an immutable request-to-batch mapping store and wire its MySQL implementation into the storage extension. Jira Issues: CODEM-304 --- submitqueue/entity/BUILD.bazel | 1 + submitqueue/entity/request_batch.go | 25 +++ submitqueue/extension/storage/BUILD.bazel | 1 + .../extension/storage/mock/BUILD.bazel | 1 + .../storage/mock/request_batch_store_mock.go | 71 +++++++++ .../extension/storage/mock/storage_mock.go | 14 ++ .../extension/storage/mysql/BUILD.bazel | 2 + .../storage/mysql/request_batch_store.go | 75 +++++++++ .../storage/mysql/request_batch_store_test.go | 147 ++++++++++++++++++ .../extension/storage/mysql/storage.go | 7 + .../extension/storage/mysql/storage_test.go | 1 + .../extension/storage/request_batch_store.go | 32 ++++ submitqueue/extension/storage/storage.go | 3 + 13 files changed, 380 insertions(+) create mode 100644 submitqueue/entity/request_batch.go create mode 100644 submitqueue/extension/storage/mock/request_batch_store_mock.go create mode 100644 submitqueue/extension/storage/mysql/request_batch_store.go create mode 100644 submitqueue/extension/storage/mysql/request_batch_store_test.go create mode 100644 submitqueue/extension/storage/request_batch_store.go diff --git a/submitqueue/entity/BUILD.bazel b/submitqueue/entity/BUILD.bazel index 2ffad326..dd72b2b1 100644 --- a/submitqueue/entity/BUILD.bazel +++ b/submitqueue/entity/BUILD.bazel @@ -17,6 +17,7 @@ go_library( "push_result.go", "queue_config.go", "request.go", + "request_batch.go", "request_history.go", "request_log.go", "request_summary.go", diff --git a/submitqueue/entity/request_batch.go b/submitqueue/entity/request_batch.go new file mode 100644 index 00000000..f88a373a --- /dev/null +++ b/submitqueue/entity/request_batch.go @@ -0,0 +1,25 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package entity + +// RequestBatch is the immutable assignment of a request to its owning batch. +type RequestBatch struct { + // RequestID is the globally unique request identifier. + RequestID string + // BatchID is the globally unique identifier of the batch containing the request. + BatchID string + // Version is the version of the assignment. Immutable assignments start at version 1. + Version int32 +} diff --git a/submitqueue/extension/storage/BUILD.bazel b/submitqueue/extension/storage/BUILD.bazel index 8fa93965..2c1b3bb9 100644 --- a/submitqueue/extension/storage/BUILD.bazel +++ b/submitqueue/extension/storage/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", "request_store.go", diff --git a/submitqueue/extension/storage/mock/BUILD.bazel b/submitqueue/extension/storage/mock/BUILD.bazel index 546c0848..8f7c1f8b 100644 --- a/submitqueue/extension/storage/mock/BUILD.bazel +++ b/submitqueue/extension/storage/mock/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store_mock.go", "build_store_mock.go", "change_store_mock.go", + "request_batch_store_mock.go", "request_log_store_mock.go", "request_queue_summary_store_mock.go", "request_store_mock.go", diff --git a/submitqueue/extension/storage/mock/request_batch_store_mock.go b/submitqueue/extension/storage/mock/request_batch_store_mock.go new file mode 100644 index 00000000..220cba33 --- /dev/null +++ b/submitqueue/extension/storage/mock/request_batch_store_mock.go @@ -0,0 +1,71 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: submitqueue/extension/storage/request_batch_store.go +// +// Generated by this command: +// +// mockgen -source=submitqueue/extension/storage/request_batch_store.go -destination=submitqueue/extension/storage/mock/request_batch_store_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + entity "github.com/uber/submitqueue/submitqueue/entity" + gomock "go.uber.org/mock/gomock" +) + +// MockRequestBatchStore is a mock of RequestBatchStore interface. +type MockRequestBatchStore struct { + ctrl *gomock.Controller + recorder *MockRequestBatchStoreMockRecorder + isgomock struct{} +} + +// MockRequestBatchStoreMockRecorder is the mock recorder for MockRequestBatchStore. +type MockRequestBatchStoreMockRecorder struct { + mock *MockRequestBatchStore +} + +// NewMockRequestBatchStore creates a new mock instance. +func NewMockRequestBatchStore(ctrl *gomock.Controller) *MockRequestBatchStore { + mock := &MockRequestBatchStore{ctrl: ctrl} + mock.recorder = &MockRequestBatchStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRequestBatchStore) EXPECT() *MockRequestBatchStoreMockRecorder { + return m.recorder +} + +// Create mocks base method. +func (m *MockRequestBatchStore) Create(ctx context.Context, assignment entity.RequestBatch) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Create", ctx, assignment) + ret0, _ := ret[0].(error) + return ret0 +} + +// Create indicates an expected call of Create. +func (mr *MockRequestBatchStoreMockRecorder) Create(ctx, assignment any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockRequestBatchStore)(nil).Create), ctx, assignment) +} + +// Get mocks base method. +func (m *MockRequestBatchStore) Get(ctx context.Context, requestID string) (entity.RequestBatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Get", ctx, requestID) + ret0, _ := ret[0].(entity.RequestBatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Get indicates an expected call of Get. +func (mr *MockRequestBatchStoreMockRecorder) Get(ctx, requestID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockRequestBatchStore)(nil).Get), ctx, requestID) +} diff --git a/submitqueue/extension/storage/mock/storage_mock.go b/submitqueue/extension/storage/mock/storage_mock.go index 3bd7789c..d32eb70a 100644 --- a/submitqueue/extension/storage/mock/storage_mock.go +++ b/submitqueue/extension/storage/mock/storage_mock.go @@ -110,6 +110,20 @@ func (mr *MockStorageMockRecorder) GetChangeStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChangeStore", reflect.TypeOf((*MockStorage)(nil).GetChangeStore)) } +// GetRequestBatchStore mocks base method. +func (m *MockStorage) GetRequestBatchStore() storage.RequestBatchStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetRequestBatchStore") + ret0, _ := ret[0].(storage.RequestBatchStore) + return ret0 +} + +// GetRequestBatchStore indicates an expected call of GetRequestBatchStore. +func (mr *MockStorageMockRecorder) GetRequestBatchStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestBatchStore", reflect.TypeOf((*MockStorage)(nil).GetRequestBatchStore)) +} + // GetRequestLogStore mocks base method. func (m *MockStorage) GetRequestLogStore() storage.RequestLogStore { m.ctrl.T.Helper() diff --git a/submitqueue/extension/storage/mysql/BUILD.bazel b/submitqueue/extension/storage/mysql/BUILD.bazel index aac0822b..af189d37 100644 --- a/submitqueue/extension/storage/mysql/BUILD.bazel +++ b/submitqueue/extension/storage/mysql/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", "request_store.go", @@ -32,6 +33,7 @@ go_test( "batch_store_test.go", "build_store_test.go", "change_store_test.go", + "request_batch_store_test.go", "request_log_store_test.go", "request_queue_summary_store_test.go", "request_store_test.go", diff --git a/submitqueue/extension/storage/mysql/request_batch_store.go b/submitqueue/extension/storage/mysql/request_batch_store.go new file mode 100644 index 00000000..142e6aaf --- /dev/null +++ b/submitqueue/extension/storage/mysql/request_batch_store.go @@ -0,0 +1,75 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mysql + +import ( + "context" + "database/sql" + "errors" + "fmt" + + "github.com/go-sql-driver/mysql" + "github.com/uber-go/tally" + + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +type requestBatchStore struct { + db *sql.DB + scope tally.Scope +} + +// NewRequestBatchStore creates a MySQL-backed RequestBatchStore. +func NewRequestBatchStore(db *sql.DB, scope tally.Scope) storage.RequestBatchStore { + return &requestBatchStore{db: db, scope: scope} +} + +func (s *requestBatchStore) Get(ctx context.Context, requestID string) (ret entity.RequestBatch, retErr error) { + op := metrics.Begin(s.scope, "get", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + var assignment entity.RequestBatch + err := s.db.QueryRowContext(ctx, + "SELECT request_id, batch_id, version FROM request_batch WHERE request_id = ?", + requestID, + ).Scan(&assignment.RequestID, &assignment.BatchID, &assignment.Version) + if errors.Is(err, sql.ErrNoRows) { + return entity.RequestBatch{}, storage.WrapNotFound(err) + } + if err != nil { + return entity.RequestBatch{}, fmt.Errorf("failed to get request batch assignment requestID=%s: %w", requestID, err) + } + return assignment, nil +} + +func (s *requestBatchStore) Create(ctx context.Context, assignment entity.RequestBatch) (retErr error) { + op := metrics.Begin(s.scope, "create", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + _, err := s.db.ExecContext(ctx, + "INSERT INTO request_batch (request_id, batch_id, version) VALUES (?, ?, ?)", + assignment.RequestID, assignment.BatchID, assignment.Version, + ) + if err != nil { + var mysqlErr *mysql.MySQLError + if errors.As(err, &mysqlErr) && mysqlErr.Number == mysqlErrDuplicateEntry { + return fmt.Errorf("request batch assignment requestID=%s: %w", assignment.RequestID, storage.ErrAlreadyExists) + } + return fmt.Errorf("failed to create request batch assignment requestID=%s batchID=%s: %w", assignment.RequestID, assignment.BatchID, err) + } + return nil +} diff --git a/submitqueue/extension/storage/mysql/request_batch_store_test.go b/submitqueue/extension/storage/mysql/request_batch_store_test.go new file mode 100644 index 00000000..57e4e0f0 --- /dev/null +++ b/submitqueue/extension/storage/mysql/request_batch_store_test.go @@ -0,0 +1,147 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mysql + +import ( + "context" + "database/sql" + "fmt" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/go-sql-driver/mysql" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +func setupRequestBatchStoreTest(t *testing.T) (*sql.DB, sqlmock.Sqlmock, storage.RequestBatchStore) { + t.Helper() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + return db, mock, NewRequestBatchStore(db, testMetrics()) +} + +func TestRequestBatchStore_Get(t *testing.T) { + assignment := entity.RequestBatch{ + RequestID: "monorepo/1", + BatchID: "monorepo/batch/2", + Version: 1, + } + tests := map[string]struct { + setup func(sqlmock.Sqlmock) + want entity.RequestBatch + errMsg string + }{ + "not found": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(assignment.RequestID). + WillReturnError(sql.ErrNoRows) + }, + errMsg: storage.ErrNotFound.Error(), + }, + "query fails": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(assignment.RequestID). + WillReturnError(fmt.Errorf("connection reset")) + }, + errMsg: "connection reset", + }, + "success": { + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"request_id", "batch_id", "version"}). + AddRow(assignment.RequestID, assignment.BatchID, assignment.Version) + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(assignment.RequestID). + WillReturnRows(rows) + }, + want: assignment, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + db, mock, store := setupRequestBatchStoreTest(t) + defer db.Close() + tt.setup(mock) + + got, err := store.Get(context.Background(), assignment.RequestID) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.want, got) + } + assert.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestRequestBatchStore_Create(t *testing.T) { + assignment := entity.RequestBatch{ + RequestID: "monorepo/1", + BatchID: "monorepo/batch/2", + Version: 1, + } + tests := map[string]struct { + setup func(sqlmock.Sqlmock) + errMsg string + }{ + "duplicate request returns already exists": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(assignment.RequestID, assignment.BatchID, assignment.Version). + WillReturnError(&mysql.MySQLError{Number: mysqlErrDuplicateEntry}) + }, + errMsg: storage.ErrAlreadyExists.Error(), + }, + "insert fails": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(assignment.RequestID, assignment.BatchID, assignment.Version). + WillReturnError(fmt.Errorf("connection reset")) + }, + errMsg: "connection reset", + }, + "success": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(assignment.RequestID, assignment.BatchID, assignment.Version). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + db, mock, store := setupRequestBatchStoreTest(t) + defer db.Close() + tt.setup(mock) + + err := store.Create(context.Background(), assignment) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + } else { + assert.NoError(t, err) + } + assert.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/submitqueue/extension/storage/mysql/storage.go b/submitqueue/extension/storage/mysql/storage.go index f810f827..053c4686 100644 --- a/submitqueue/extension/storage/mysql/storage.go +++ b/submitqueue/extension/storage/mysql/storage.go @@ -30,6 +30,7 @@ const mysqlErrDuplicateEntry = 1062 type mysqlStorage struct { db *sql.DB requestStore storage.RequestStore + requestBatchStore storage.RequestBatchStore changeStore storage.ChangeStore batchStore storage.BatchStore batchDependentStore storage.BatchDependentStore @@ -45,6 +46,7 @@ func NewStorage(db *sql.DB, scope tally.Scope) (storage.Storage, error) { return &mysqlStorage{ db: db, requestStore: NewRequestStore(db, scope.SubScope("request_store")), + requestBatchStore: NewRequestBatchStore(db, scope.SubScope("request_batch_store")), changeStore: NewChangeStore(db, scope.SubScope("change_store")), batchStore: NewBatchStore(db, scope.SubScope("batch_store")), batchDependentStore: NewBatchDependentStore(db, scope.SubScope("batch_dependent_store")), @@ -61,6 +63,11 @@ func (f *mysqlStorage) GetRequestStore() storage.RequestStore { return f.requestStore } +// GetRequestBatchStore returns the MySQL-backed RequestBatchStore. +func (f *mysqlStorage) GetRequestBatchStore() storage.RequestBatchStore { + return f.requestBatchStore +} + // GetChangeStore returns the MySQL-backed ChangeStore. func (f *mysqlStorage) GetChangeStore() storage.ChangeStore { return f.changeStore diff --git a/submitqueue/extension/storage/mysql/storage_test.go b/submitqueue/extension/storage/mysql/storage_test.go index 6055ad5d..5dc14a0b 100644 --- a/submitqueue/extension/storage/mysql/storage_test.go +++ b/submitqueue/extension/storage/mysql/storage_test.go @@ -37,6 +37,7 @@ func TestNewStorage(t *testing.T) { require.NoError(t, err) assert.NotNil(t, s.GetRequestStore()) + assert.NotNil(t, s.GetRequestBatchStore()) assert.NotNil(t, s.GetChangeStore()) assert.NotNil(t, s.GetBatchStore()) assert.NotNil(t, s.GetBatchDependentStore()) diff --git a/submitqueue/extension/storage/request_batch_store.go b/submitqueue/extension/storage/request_batch_store.go new file mode 100644 index 00000000..87b361b6 --- /dev/null +++ b/submitqueue/extension/storage/request_batch_store.go @@ -0,0 +1,32 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package storage + +//go:generate mockgen -source=request_batch_store.go -destination=mock/request_batch_store_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// RequestBatchStore persists the immutable assignment from a request ID to its owning batch ID. +type RequestBatchStore interface { + // Get retrieves the assignment by request ID. Returns ErrNotFound if no assignment exists. + Get(ctx context.Context, requestID string) (entity.RequestBatch, error) + + // Create inserts an immutable assignment. Returns ErrAlreadyExists if the request is already assigned. + Create(ctx context.Context, assignment entity.RequestBatch) error +} diff --git a/submitqueue/extension/storage/storage.go b/submitqueue/extension/storage/storage.go index c1fd91f8..c2233298 100644 --- a/submitqueue/extension/storage/storage.go +++ b/submitqueue/extension/storage/storage.go @@ -49,6 +49,9 @@ type Storage interface { // GetRequestStore returns the RequestStore instance. GetRequestStore() RequestStore + // GetRequestBatchStore returns the RequestBatchStore instance. + GetRequestBatchStore() RequestBatchStore + // GetChangeStore returns the ChangeStore instance. GetChangeStore() ChangeStore