feat(entities) Adds build store interface and mysql implementation

Author: manjari25

extension/storage/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "batch_dependent_store.go",
         "batch_store.go",
+        "build_store.go",
         "change_provider_store.go",
         "request_store.go",
         "storage.go",

extension/storage/build_store.go

@@ -0,0 +1,20 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/uber/submitqueue/entity"
+)
+
+// BuildStore is an interface that defines methods for managing builds in the database.
+type BuildStore interface {
+	// Get retrieves a build by ID. Returns ErrNotFound if the build is not found.
+	Get(ctx context.Context, id string) (entity.Build, error)
+
+	// Create creates a new build. The build must have a unique ID already assigned.
+	// Returns ErrAlreadyExists if a build with the same ID already exists.
+	Create(ctx context.Context, build entity.Build) error
+
+	// UpdateStatus updates the status of a build.
+	UpdateStatus(ctx context.Context, id string, newStatus entity.BuildStatus) error
+}

extension/storage/mysql/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "batch_dependent_store.go",
         "batch_store.go",
+        "build_store.go",
         "change_provider_store.go",
         "request_store.go",
         "storage.go",

extension/storage/mysql/build_store.go

@@ -0,0 +1,91 @@
+package mysql
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+	"fmt"
+
+	"github.com/go-sql-driver/mysql"
+
+	"github.com/uber/submitqueue/entity"
+	"github.com/uber/submitqueue/extension/storage"
+)
+
+type buildStore struct {
+	db *sql.DB
+}
+
+// NewBuildStore creates a new MySQL-backed BuildStore.
+func NewBuildStore(db *sql.DB) storage.BuildStore {
+	return &buildStore{db: db}
+}
+
+// Get retrieves a build by ID. Returns ErrNotFound if the build is not found.
+func (s *buildStore) Get(ctx context.Context, id string) (entity.Build, error) {
+	var build entity.Build
+	var speculationPathJSON []byte
+
+	err := s.db.QueryRowContext(ctx,
+		"SELECT id, batch_id, speculation_path, score, status FROM build WHERE id = ?",
+		id,
+	).Scan(&build.ID, &build.BatchID, &speculationPathJSON, &build.Score, &build.Status)
+
+	if errors.Is(err, sql.ErrNoRows) {
+		return entity.Build{}, storage.WrapNotFound(err)
+	}
+	if err != nil {
+		return entity.Build{}, fmt.Errorf("failed to get build entity id=%s from the database: %w", id, err)
+	}
+
+	if err := json.Unmarshal(speculationPathJSON, &build.SpeculationPath); err != nil {
+		return entity.Build{}, fmt.Errorf("failed to unmarshal speculation_path for build entity id=%s from the database: %w", id, err)
+	}
+
+	return build, nil
+}
+
+// Create creates a new build. The build must have a unique ID already assigned. Returns ErrAlreadyExists if the build ID already exists.
+func (s *buildStore) Create(ctx context.Context, build entity.Build) error {
+	speculationPathJSON, err := json.Marshal(build.SpeculationPath)
+	if err != nil {
+		return fmt.Errorf("failed to marshal speculation_path id=%s for Create build entity: %w", build.ID, err)
+	}
+
+	_, err = s.db.ExecContext(ctx,
+		"INSERT INTO build (id, batch_id, speculation_path, score, status) VALUES (?, ?, ?, ?, ?)",
+		build.ID, build.BatchID, speculationPathJSON, build.Score, build.Status,
+	)
+	if err != nil {
+		var mysqlErr *mysql.MySQLError
+		if errors.As(err, &mysqlErr) && mysqlErr.Number == 1062 {
+			return fmt.Errorf("build entity id=%s: %w", build.ID, storage.ErrAlreadyExists)
+		}
+		return fmt.Errorf("failed to insert build entity id=%s: %w", build.ID, err)
+	}
+
+	return nil
+}
+
+// UpdateStatus updates the status of a build. Returns ErrNotFound if the build is not found.
+func (s *buildStore) UpdateStatus(ctx context.Context, id string, newStatus entity.BuildStatus) error {
+	result, err := s.db.ExecContext(ctx,
+		"UPDATE build SET status = ? WHERE id = ?",
+		newStatus, id,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to update build status for id=%q newStatus=%v: %w", id, newStatus, err)
+	}
+
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		return fmt.Errorf("failed to get rows affected from update for id=%q newStatus=%v: %w", id, newStatus, err)
+	}
+
+	if rowsAffected != 1 {
+		return storage.WrapNotFound(fmt.Errorf("build entity id=%s", id))
+	}
+
+	return nil
+}

extension/storage/mysql/storage.go

@@ -14,6 +14,7 @@ type mysqlStorage struct {
 	changeProviderStore   storage.ChangeProviderStore
 	batchStore            storage.BatchStore
 	batchDependentStore   storage.BatchDependentStore
+	buildStore            storage.BuildStore
 }
 
 // NewStorage creates a new MySQL storage.
@@ -24,6 +25,7 @@ func NewStorage(db *sql.DB) (storage.Storage, error) {
 		changeProviderStore:   NewChangeProviderStore(db),
 		batchStore:            NewBatchStore(db),
 		batchDependentStore:   NewBatchDependentStore(db),
+		buildStore:            NewBuildStore(db),
 	}, nil
 }
 
@@ -47,6 +49,11 @@ func (f *mysqlStorage) GetBatchDependentStore() storage.BatchDependentStore {
 	return f.batchDependentStore
 }
 
+// GetBuildStore returns the MySQL-backed BuildStore.
+func (f *mysqlStorage) GetBuildStore() storage.BuildStore {
+	return f.buildStore
+}
+
 // Close closes the underlying database connection.
 func (f *mysqlStorage) Close() error {
 	return f.db.Close()

extension/storage/storage.go

@@ -38,6 +38,9 @@ type Storage interface {
 	// GetBatchDependentStore returns the BatchDependentStore instance.
 	GetBatchDependentStore() BatchDependentStore
 
+	// GetBuildStore returns the BuildStore instance.
+	GetBuildStore() BuildStore
+
 	// Close closes the storage and all underlying connections. Should only be called once at the end of the program.
 	Close() error
 }

gateway/controller/land_test.go

@@ -97,11 +97,39 @@ func (m *mockBatchDependentStore) Create(ctx context.Context, batchDependent ent
 	return nil
 }
 
+type mockBuildStore struct {
+	createFunc       func(ctx context.Context, build entity.Build) error
+	getFunc          func(ctx context.Context, id string) (entity.Build, error)
+	updateStatusFunc func(ctx context.Context, id string, newStatus entity.BuildStatus) error
+}
+
+func (m *mockBuildStore) Get(ctx context.Context, id string) (entity.Build, error) {
+	if m.getFunc != nil {
+		return m.getFunc(ctx, id)
+	}
+	return entity.Build{}, nil
+}
+
+func (m *mockBuildStore) Create(ctx context.Context, build entity.Build) error {
+	if m.createFunc != nil {
+		return m.createFunc(ctx, build)
+	}
+	return nil
+}
+
+func (m *mockBuildStore) UpdateStatus(ctx context.Context, id string, newStatus entity.BuildStatus) error {
+	if m.updateStatusFunc != nil {
+		return m.updateStatusFunc(ctx, id, newStatus)
+	}
+	return nil
+}
+
 type mockStorage struct {
 	requestStore          storage.RequestStore
 	changeProviderStore   storage.ChangeProviderStore
 	batchStore            storage.BatchStore
 	batchDependentStore   storage.BatchDependentStore
+	buildStore            storage.BuildStore
 }
 
 func (m *mockStorage) GetRequestStore() storage.RequestStore {
@@ -120,6 +148,10 @@ func (m *mockStorage) GetBatchDependentStore() storage.BatchDependentStore {
 	return m.batchDependentStore
 }
 
+func (m *mockStorage) GetBuildStore() storage.BuildStore {
+	return m.buildStore
+}
+
 func (m *mockStorage) Close() error {
 	return nil
 }