Add batch get repository size api

Author: pulltheflower

_mocks/opencsg.com/csghub-server/builder/store/database/mock_RepositoryStatisticsStore.go

@@ -2046,6 +2046,47 @@ func (h *RepoHandler) GetRepoSizeByBranch(ctx *gin.Context) {
 	httpbase.OK(ctx, size)
 }
 
+// BatchGetRepoExtra godoc
+// @Security     ApiKey
+// @Summary      Batch get extra information for multiple repositories
+// @Tags         Repository
+// @Accept       json
+// @Produce      json
+// @Param        current_user query string false "current user name"
+// @Param        body body types.BatchRepoExtraReq true "request body"
+// @Success      200  {object}  types.Response{data=[]types.RepoExtraItem} "OK"
+// @Failure      400  {object}  types.APIBadRequest "Bad request"
+// @Failure      500  {object}  types.APIInternalServerError "Internal server error"
+// @Router       /repos/extra [post]
+func (h *RepoHandler) BatchGetRepoExtra(ctx *gin.Context) {
+	var req types.BatchRepoExtraReq
+	if err := ctx.ShouldBindJSON(&req); err != nil {
+		slog.ErrorContext(ctx.Request.Context(), "Bad request format", "error", err)
+		httpbase.BadRequestWithExt(ctx, errorx.ReqBodyFormat(err, nil))
+		return
+	}
+
+	if len(req.RepoIDs) > h.config.MaxRepoBatchNum {
+		httpbase.BadRequestWithExt(ctx, errorx.ReqParamInvalid(
+			fmt.Errorf("too many repository ids, max %d, got %d", h.config.MaxRepoBatchNum, len(req.RepoIDs)),
+			errorx.Ctx().Set("max", h.config.MaxRepoBatchNum).Set("got", len(req.RepoIDs)),
+		))
+		return
+	}
+
+	currentUser := httpbase.GetCurrentUser(ctx)
+
+	extras, err := h.c.BatchGetRepoExtra(ctx.Request.Context(), req.RepoIDs, currentUser)
+	if err != nil {
+		slog.ErrorContext(ctx.Request.Context(), "Failed to batch get repo extras", slog.Any("error", err))
+		httpbase.ServerError(ctx, err)
+		return
+	}
+
+	slog.Debug("Batch get repo extras succeed", slog.Int("count", len(extras)))
+	httpbase.OK(ctx, extras)
+}
+
 // DownloadCodeZip godoc
 // @Summary      Download code repository as zip archive
 // @Description  Download code repository as zip archive

_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go

@@ -56,7 +56,9 @@ func NewRepoTester(t *testing.T) *RepoTester {
 		m:        tester.mocks.model,
 		d:        tester.mocks.dataset,
 		temporal: tester.mocks.workflow,
-		config:   &config.Config{},
+		config: &config.Config{
+			MaxRepoBatchNum: 500,
+		},
 	}
 	tester.WithParam("name", "r")
 	tester.WithParam("namespace", "u")
@@ -1975,3 +1977,60 @@ func TestRepoHandler_GetRepos(t *testing.T) {
 		t, 200, tester.OKText, []string{},
 	)
 }
+
+func TestRepoHandler_BatchGetRepoExtra(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		tester := NewRepoTester(t).WithHandleFunc(func(rp *RepoHandler) gin.HandlerFunc {
+			return rp.BatchGetRepoExtra
+		})
+		tester.WithBody(t, types.BatchRepoExtraReq{RepoIDs: []int64{1, 2}}).WithUser()
+
+		tester.mocks.repo.EXPECT().BatchGetRepoExtra(tester.Ctx(), []int64{1, 2}, "u").Return([]types.RepoExtraItem{{RepoID: 1, Size: 100}, {RepoID: 2, Size: 200}}, nil)
+
+		tester.Execute()
+
+		tester.ResponseEq(t, http.StatusOK, tester.OKText, []types.RepoExtraItem{
+			{RepoID: 1, Size: 100},
+			{RepoID: 2, Size: 200},
+		})
+	})
+
+	t.Run("too many IDs", func(t *testing.T) {
+		ids := make([]int64, 501)
+		for i := range ids {
+			ids[i] = int64(i + 1)
+		}
+		tester := NewRepoTester(t).WithHandleFunc(func(rp *RepoHandler) gin.HandlerFunc {
+			return rp.BatchGetRepoExtra
+		})
+		tester.WithBody(t, types.BatchRepoExtraReq{RepoIDs: ids}).WithUser()
+
+		tester.Execute()
+
+		tester.ResponseEqCode(t, http.StatusBadRequest)
+	})
+
+	t.Run("invalid JSON body", func(t *testing.T) {
+		tester := NewRepoTester(t).WithHandleFunc(func(rp *RepoHandler) gin.HandlerFunc {
+			return rp.BatchGetRepoExtra
+		})
+		tester.WithBody(t, "not a valid json").WithUser()
+
+		tester.Execute()
+
+		tester.ResponseEqCode(t, http.StatusBadRequest)
+	})
+
+	t.Run("server error", func(t *testing.T) {
+		tester := NewRepoTester(t).WithHandleFunc(func(rp *RepoHandler) gin.HandlerFunc {
+			return rp.BatchGetRepoExtra
+		})
+		tester.WithBody(t, types.BatchRepoExtraReq{RepoIDs: []int64{1}}).WithUser()
+
+		tester.mocks.repo.EXPECT().BatchGetRepoExtra(tester.Ctx(), []int64{1}, "u").Return(nil, errors.New("db error"))
+
+		tester.Execute()
+
+		tester.ResponseEqCode(t, http.StatusInternalServerError)
+	})
+}

api/handler/repo.go

@@ -421,6 +421,9 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
 	// Admin user get repo path list
 	adminGroup.GET("/repos", repoCommonHandler.GetRepos)
 
+	// Batch get repository extra information for multiple repositories
+	apiGroup.POST("/repos/extra", repoCommonHandler.BatchGetRepoExtra)
+
 	// routes for broadcast
 	broadcastHandler, err := handler.NewBroadcastHandler()
 	if err != nil {

api/handler/repo_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/uptrace/bun"
 	"opencsg.com/csghub-server/common/errorx"
 )
 
@@ -32,6 +33,9 @@ type RepositoryStatisticsStore interface {
 	// FindByRepositoryIDAndBranch finds repository statistics by repository ID and branch
 	FindByRepositoryIDAndBranch(ctx context.Context, repoID int64, branch string) (*RepositoryStatistics, error)
 
+	// FindByRepositoryIDs finds repository statistics for multiple repository IDs
+	FindByRepositoryIDs(ctx context.Context, repoIDs []int64) ([]*RepositoryStatistics, error)
+
 	// Update updates repository statistics
 	Update(ctx context.Context, stats *RepositoryStatistics) error
 
@@ -87,6 +91,19 @@ func (s *RepositoryStatisticsStoreImpl) FindByRepositoryIDAndBranch(ctx context.
 	return &stats, nil
 }
 
+// FindByRepositoryIDs finds repository statistics for multiple repository IDs
+func (s *RepositoryStatisticsStoreImpl) FindByRepositoryIDs(ctx context.Context, repoIDs []int64) ([]*RepositoryStatistics, error) {
+	if len(repoIDs) == 0 {
+		return nil, nil
+	}
+	var stats []*RepositoryStatistics
+	_, err := s.db.Operator.Core.NewSelect().Model(&stats).Where("repository_id IN (?)", bun.In(repoIDs)).Exec(ctx, &stats)
+	if err != nil {
+		return nil, errorx.HandleDBError(err, nil)
+	}
+	return stats, nil
+}
+
 // Update updates repository statistics
 func (s *RepositoryStatisticsStoreImpl) Update(ctx context.Context, stats *RepositoryStatistics) error {
 	_, err := s.db.Operator.Core.NewUpdate().Model(stats).Where("id = ?", stats.ID).Exec(ctx)

api/router/api.go

@@ -240,3 +240,78 @@ func TestRepositoryStatisticsStore_FindByRepositoryIDAndBranch(t *testing.T) {
 	_, err = store.FindByRepositoryIDAndBranch(ctx, expectedStats.RepositoryID, "non-existent-branch")
 	assert.Error(t, err)
 }
+
+func TestRepositoryStatisticsStore_FindByRepositoryIDs(t *testing.T) {
+	ctx := context.Background()
+	db := tests.InitTestDB()
+	defer db.Close()
+
+	store := database.NewRepositoryStatisticsStoreWithDB(db)
+
+	// Create test statistics for multiple repositories
+	stats1 := &database.RepositoryStatistics{
+		RepositoryID: 8,
+		Branch:       "main",
+		TotalSize:    1024,
+		NonLfsSize:   512,
+		LfsSize:      512,
+		CreatedAt:    time.Now(),
+		UpdatedAt:    time.Now(),
+	}
+	stats2 := &database.RepositoryStatistics{
+		RepositoryID: 9,
+		Branch:       "main",
+		TotalSize:    2048,
+		NonLfsSize:   1024,
+		LfsSize:      1024,
+		CreatedAt:    time.Now(),
+		UpdatedAt:    time.Now(),
+	}
+	stats3 := &database.RepositoryStatistics{
+		RepositoryID: 10,
+		Branch:       "dev",
+		TotalSize:    512,
+		NonLfsSize:   256,
+		LfsSize:      256,
+		CreatedAt:    time.Now(),
+		UpdatedAt:    time.Now(),
+	}
+
+	err := store.Create(ctx, stats1)
+	assert.NoError(t, err)
+	err = store.Create(ctx, stats2)
+	assert.NoError(t, err)
+	err = store.Create(ctx, stats3)
+	assert.NoError(t, err)
+
+	t.Run("find by multiple existing IDs", func(t *testing.T) {
+		result, err := store.FindByRepositoryIDs(ctx, []int64{8, 9})
+		assert.NoError(t, err)
+		assert.Len(t, result, 2)
+	})
+
+	t.Run("find by single ID", func(t *testing.T) {
+		result, err := store.FindByRepositoryIDs(ctx, []int64{8})
+		assert.NoError(t, err)
+		assert.Len(t, result, 1)
+		assert.Equal(t, int64(1024), result[0].TotalSize)
+	})
+
+	t.Run("find by non-existent ID returns empty", func(t *testing.T) {
+		result, err := store.FindByRepositoryIDs(ctx, []int64{999})
+		assert.NoError(t, err)
+		assert.Len(t, result, 0)
+	})
+
+	t.Run("find by empty IDs returns nil", func(t *testing.T) {
+		result, err := store.FindByRepositoryIDs(ctx, []int64{})
+		assert.NoError(t, err)
+		assert.Nil(t, result)
+	})
+
+	t.Run("find by mixed existing and non-existing IDs", func(t *testing.T) {
+		result, err := store.FindByRepositoryIDs(ctx, []int64{8, 999, 10})
+		assert.NoError(t, err)
+		assert.Len(t, result, 2)
+	})
+}

builder/store/database/repository_statistics.go

@@ -728,6 +728,8 @@ type Config struct {
 	ActivityLog struct {
 		Enabled bool `env:"STARHUB_SERVER_ACTIVITY_LOG_ENABLE" default:"true"`
 	}
+
+	MaxRepoBatchNum int `env:"STARHUB_SERVER_MAX_REPO_SIZE_BATCH_NUM" default:"500"`
 }
 
 type MemoryConfig struct {