feat(storage): add request batch assignments

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",

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
+}

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",

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",

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",

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
+}

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())
+		})
+	}
+}