db: use `context` and go-mockgen for `LFSStore` (#7038)

1 files changed, 19 insertions, 18 deletions

diff --git a/internal/db/lfs.go b/internal/db/lfs.go
index 43515d8c..1a34b802 100644
--- a/internal/db/lfs.go
+++ b/internal/db/lfs.go
@@ -5,6 +5,7 @@
 package db
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -19,24 +20,24 @@ import (
 // NOTE: All methods are sorted in alphabetical order.
 type LFSStore interface {
 	// CreateObject creates a LFS object record in database.
-	CreateObject(repoID int64, oid lfsutil.OID, size int64, storage lfsutil.Storage) error
-	// GetObjectByOID returns the LFS object with given OID. It returns ErrLFSObjectNotExist
-	// when not found.
-	GetObjectByOID(repoID int64, oid lfsutil.OID) (*LFSObject, error)
-	// GetObjectsByOIDs returns LFS objects found within "oids". The returned list could have
-	// less elements if some oids were not found.
-	GetObjectsByOIDs(repoID int64, oids ...lfsutil.OID) ([]*LFSObject, error)
+	CreateObject(ctx context.Context, repoID int64, oid lfsutil.OID, size int64, storage lfsutil.Storage) error
+	// GetObjectByOID returns the LFS object with given OID. It returns
+	// ErrLFSObjectNotExist when not found.
+	GetObjectByOID(ctx context.Context, repoID int64, oid lfsutil.OID) (*LFSObject, error)
+	// GetObjectsByOIDs returns LFS objects found within "oids". The returned list
+	// could have less elements if some oids were not found.
+	GetObjectsByOIDs(ctx context.Context, repoID int64, oids ...lfsutil.OID) ([]*LFSObject, error)
 }
 
 var LFS LFSStore
 
 // LFSObject is the relation between an LFS object and a repository.
 type LFSObject struct {
-	RepoID    int64           `gorm:"PRIMARY_KEY;AUTO_INCREMENT:false"`
-	OID       lfsutil.OID     `gorm:"PRIMARY_KEY;COLUMN:oid"`
-	Size      int64           `gorm:"NOT NULL"`
-	Storage   lfsutil.Storage `gorm:"NOT NULL"`
-	CreatedAt time.Time       `gorm:"NOT NULL"`
+	RepoID    int64           `gorm:"primary_key;auto_increment:false"`
+	OID       lfsutil.OID     `gorm:"primary_key;column:oid"`
+	Size      int64           `gorm:"not null"`
+	Storage   lfsutil.Storage `gorm:"not null"`
+	CreatedAt time.Time       `gorm:"not null"`
 }
 
 var _ LFSStore = (*lfs)(nil)
@@ -45,14 +46,14 @@ type lfs struct {
 	*gorm.DB
 }
 
-func (db *lfs) CreateObject(repoID int64, oid lfsutil.OID, size int64, storage lfsutil.Storage) error {
+func (db *lfs) CreateObject(ctx context.Context, repoID int64, oid lfsutil.OID, size int64, storage lfsutil.Storage) error {
 	object := &LFSObject{
 		RepoID:  repoID,
 		OID:     oid,
 		Size:    size,
 		Storage: storage,
 	}
-	return db.DB.Create(object).Error
+	return db.WithContext(ctx).Create(object).Error
 }
 
 type ErrLFSObjectNotExist struct {
@@ -72,9 +73,9 @@ func (ErrLFSObjectNotExist) NotFound() bool {
 	return true
 }
 
-func (db *lfs) GetObjectByOID(repoID int64, oid lfsutil.OID) (*LFSObject, error) {
+func (db *lfs) GetObjectByOID(ctx context.Context, repoID int64, oid lfsutil.OID) (*LFSObject, error) {
 	object := new(LFSObject)
-	err := db.Where("repo_id = ? AND oid = ?", repoID, oid).First(object).Error
+	err := db.WithContext(ctx).Where("repo_id = ? AND oid = ?", repoID, oid).First(object).Error
 	if err != nil {
 		if err == gorm.ErrRecordNotFound {
 			return nil, ErrLFSObjectNotExist{args: errutil.Args{"repoID": repoID, "oid": oid}}
@@ -84,13 +85,13 @@ func (db *lfs) GetObjectByOID(repoID int64, oid lfsutil.OID) (*LFSObject, error)
 	return object, err
 }
 
-func (db *lfs) GetObjectsByOIDs(repoID int64, oids ...lfsutil.OID) ([]*LFSObject, error) {
+func (db *lfs) GetObjectsByOIDs(ctx context.Context, repoID int64, oids ...lfsutil.OID) ([]*LFSObject, error) {
 	if len(oids) == 0 {
 		return []*LFSObject{}, nil
 	}
 
 	objects := make([]*LFSObject, 0, len(oids))
-	err := db.Where("repo_id = ? AND oid IN (?)", repoID, oids).Find(&objects).Error
+	err := db.WithContext(ctx).Where("repo_id = ? AND oid IN (?)", repoID, oids).Find(&objects).Error
 	if err != nil && err != gorm.ErrRecordNotFound {
 		return nil, err
 	}