diff --git a/extension/storage/BUILD.bazel b/extension/storage/BUILD.bazel index dfc26dcd..0b3e3d06 100644 --- a/extension/storage/BUILD.bazel +++ b/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", diff --git a/extension/storage/build_store.go b/extension/storage/build_store.go new file mode 100644 index 00000000..0f31795a --- /dev/null +++ b/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 +} diff --git a/extension/storage/mysql/BUILD.bazel b/extension/storage/mysql/BUILD.bazel index f33770ab..e14b8769 100644 --- a/extension/storage/mysql/BUILD.bazel +++ b/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", diff --git a/extension/storage/mysql/build_store.go b/extension/storage/mysql/build_store.go new file mode 100644 index 00000000..1db9e06f --- /dev/null +++ b/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 +} diff --git a/extension/storage/mysql/storage.go b/extension/storage/mysql/storage.go index 310b5e86..8b42ddb1 100644 --- a/extension/storage/mysql/storage.go +++ b/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() diff --git a/extension/storage/storage.go b/extension/storage/storage.go index 4ed8e809..3f42c461 100644 --- a/extension/storage/storage.go +++ b/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 } diff --git a/gateway/controller/land_test.go b/gateway/controller/land_test.go index a6b657fd..1bfa40e6 100644 --- a/gateway/controller/land_test.go +++ b/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 }