diff --git a/_mocks/opencsg.com/csghub-server/builder/store/database/mock_RepositoryStatisticsStore.go b/_mocks/opencsg.com/csghub-server/builder/store/database/mock_RepositoryStatisticsStore.go index c66d5d1a4..46cd746ba 100644 --- a/_mocks/opencsg.com/csghub-server/builder/store/database/mock_RepositoryStatisticsStore.go +++ b/_mocks/opencsg.com/csghub-server/builder/store/database/mock_RepositoryStatisticsStore.go @@ -282,6 +282,65 @@ func (_c *MockRepositoryStatisticsStore_FindByRepositoryIDAndBranch_Call) RunAnd return _c } +// FindByRepositoryIDs provides a mock function with given fields: ctx, repoIDs +func (_m *MockRepositoryStatisticsStore) FindByRepositoryIDs(ctx context.Context, repoIDs []int64) ([]*database.RepositoryStatistics, error) { + ret := _m.Called(ctx, repoIDs) + + if len(ret) == 0 { + panic("no return value specified for FindByRepositoryIDs") + } + + var r0 []*database.RepositoryStatistics + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, []int64) ([]*database.RepositoryStatistics, error)); ok { + return rf(ctx, repoIDs) + } + if rf, ok := ret.Get(0).(func(context.Context, []int64) []*database.RepositoryStatistics); ok { + r0 = rf(ctx, repoIDs) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*database.RepositoryStatistics) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, []int64) error); ok { + r1 = rf(ctx, repoIDs) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockRepositoryStatisticsStore_FindByRepositoryIDs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'FindByRepositoryIDs' +type MockRepositoryStatisticsStore_FindByRepositoryIDs_Call struct { + *mock.Call +} + +// FindByRepositoryIDs is a helper method to define mock.On call +// - ctx context.Context +// - repoIDs []int64 +func (_e *MockRepositoryStatisticsStore_Expecter) FindByRepositoryIDs(ctx interface{}, repoIDs interface{}) *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call { + return &MockRepositoryStatisticsStore_FindByRepositoryIDs_Call{Call: _e.mock.On("FindByRepositoryIDs", ctx, repoIDs)} +} + +func (_c *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call) Run(run func(ctx context.Context, repoIDs []int64)) *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].([]int64)) + }) + return _c +} + +func (_c *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call) Return(_a0 []*database.RepositoryStatistics, _a1 error) *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call) RunAndReturn(run func(context.Context, []int64) ([]*database.RepositoryStatistics, error)) *MockRepositoryStatisticsStore_FindByRepositoryIDs_Call { + _c.Call.Return(run) + return _c +} + // Update provides a mock function with given fields: ctx, stats func (_m *MockRepositoryStatisticsStore) Update(ctx context.Context, stats *database.RepositoryStatistics) error { ret := _m.Called(ctx, stats) diff --git a/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go b/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go index 1d810d3c7..ec93808d4 100644 --- a/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go +++ b/_mocks/opencsg.com/csghub-server/component/mock_RepoComponent.go @@ -567,6 +567,66 @@ func (_c *MockRepoComponent_AllowWriteAccess_Call) RunAndReturn(run func(context return _c } +// BatchGetRepoExtra provides a mock function with given fields: ctx, repoIDs, currentUser +func (_m *MockRepoComponent) BatchGetRepoExtra(ctx context.Context, repoIDs []int64, currentUser string) ([]types.RepoExtraItem, error) { + ret := _m.Called(ctx, repoIDs, currentUser) + + if len(ret) == 0 { + panic("no return value specified for BatchGetRepoExtra") + } + + var r0 []types.RepoExtraItem + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, []int64, string) ([]types.RepoExtraItem, error)); ok { + return rf(ctx, repoIDs, currentUser) + } + if rf, ok := ret.Get(0).(func(context.Context, []int64, string) []types.RepoExtraItem); ok { + r0 = rf(ctx, repoIDs, currentUser) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]types.RepoExtraItem) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, []int64, string) error); ok { + r1 = rf(ctx, repoIDs, currentUser) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockRepoComponent_BatchGetRepoExtra_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BatchGetRepoExtra' +type MockRepoComponent_BatchGetRepoExtra_Call struct { + *mock.Call +} + +// BatchGetRepoExtra is a helper method to define mock.On call +// - ctx context.Context +// - repoIDs []int64 +// - currentUser string +func (_e *MockRepoComponent_Expecter) BatchGetRepoExtra(ctx interface{}, repoIDs interface{}, currentUser interface{}) *MockRepoComponent_BatchGetRepoExtra_Call { + return &MockRepoComponent_BatchGetRepoExtra_Call{Call: _e.mock.On("BatchGetRepoExtra", ctx, repoIDs, currentUser)} +} + +func (_c *MockRepoComponent_BatchGetRepoExtra_Call) Run(run func(ctx context.Context, repoIDs []int64, currentUser string)) *MockRepoComponent_BatchGetRepoExtra_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].([]int64), args[2].(string)) + }) + return _c +} + +func (_c *MockRepoComponent_BatchGetRepoExtra_Call) Return(_a0 []types.RepoExtraItem, _a1 error) *MockRepoComponent_BatchGetRepoExtra_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockRepoComponent_BatchGetRepoExtra_Call) RunAndReturn(run func(context.Context, []int64, string) ([]types.RepoExtraItem, error)) *MockRepoComponent_BatchGetRepoExtra_Call { + _c.Call.Return(run) + return _c +} + // BatchMigrateRepoToHashedPath provides a mock function with given fields: ctx, auto, batchSize, lastID func (_m *MockRepoComponent) BatchMigrateRepoToHashedPath(ctx context.Context, auto bool, batchSize int, lastID int64) (int64, error) { ret := _m.Called(ctx, auto, batchSize, lastID) diff --git a/api/handler/repo.go b/api/handler/repo.go index d5a3be072..760daeb31 100644 --- a/api/handler/repo.go +++ b/api/handler/repo.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 diff --git a/api/handler/repo_test.go b/api/handler/repo_test.go index 8dc65b4f4..6f573c654 100644 --- a/api/handler/repo_test.go +++ b/api/handler/repo_test.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) + }) +} diff --git a/api/router/api.go b/api/router/api.go index 48af01952..14b2ec37d 100644 --- a/api/router/api.go +++ b/api/router/api.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 { diff --git a/builder/store/database/repository_statistics.go b/builder/store/database/repository_statistics.go index 369c721c5..a29858e1d 100644 --- a/builder/store/database/repository_statistics.go +++ b/builder/store/database/repository_statistics.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) diff --git a/builder/store/database/repository_statistics_test.go b/builder/store/database/repository_statistics_test.go index a75633b2b..2bfe65404 100644 --- a/builder/store/database/repository_statistics_test.go +++ b/builder/store/database/repository_statistics_test.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) + }) +} diff --git a/common/config/config.go b/common/config/config.go index 4b52c83ed..fcf096424 100644 --- a/common/config/config.go +++ b/common/config/config.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 { diff --git a/common/errorx/error_repo.go b/common/errorx/error_repo.go index 6d9392be0..c70a84ded 100644 --- a/common/errorx/error_repo.go +++ b/common/errorx/error_repo.go @@ -9,6 +9,7 @@ const ( codeRepoNotFoundErr codeRepoNoDefaultBranchErr codeCodeZipDownloadFailedErr + codeBatchGetRepoExtraFailedErr ) var ( @@ -77,6 +78,17 @@ var ( // // zh-HK: 下載代碼 zip 歸檔失敗 ErrCodeZipDownloadFailed error = CustomError{prefix: errRepoPrefix, code: codeCodeZipDownloadFailedErr} + + // Description: Failed to batch get repository extra information. + // + // Description_ZH: 批量获取仓库额外信息失败 + // + // en-US: Failed to batch get repository extra information + // + // zh-CN: 批量获取仓库额外信息失败 + // + // zh-HK: 批量獲取倉庫額外資訊失敗 + ErrBatchGetRepoExtraFailed error = CustomError{prefix: errRepoPrefix, code: codeBatchGetRepoExtraFailedErr} ) // RepoNotFound creates a REPO-ERR-3 error with context. @@ -93,3 +105,8 @@ func RepoNoDefaultBranch(ctx context) error { func CodeZipDownloadFailed(err error, ctx context) error { return CustomError{prefix: errRepoPrefix, code: codeCodeZipDownloadFailedErr, err: err, context: ctx} } + +// BatchGetRepoExtraFailed creates a REPO-ERR-6 error. +func BatchGetRepoExtraFailed(err error) error { + return CustomError{prefix: errRepoPrefix, code: codeBatchGetRepoExtraFailedErr, err: err} +} diff --git a/common/i18n/en-US/err_repo.json b/common/i18n/en-US/err_repo.json index db35365a7..66409f235 100644 --- a/common/i18n/en-US/err_repo.json +++ b/common/i18n/en-US/err_repo.json @@ -16,5 +16,8 @@ }, "error.REPO-ERR-5": { "other": "Failed to download code zip archive" + }, + "error.REPO-ERR-6": { + "other": "Failed to batch get repository extra information" } } diff --git a/common/i18n/zh-CN/err_repo.json b/common/i18n/zh-CN/err_repo.json index f4598986b..a6cd0a25f 100644 --- a/common/i18n/zh-CN/err_repo.json +++ b/common/i18n/zh-CN/err_repo.json @@ -16,5 +16,8 @@ }, "error.REPO-ERR-5": { "other": "下载代码 zip 归档失败" + }, + "error.REPO-ERR-6": { + "other": "批量获取仓库额外信息失败" } } diff --git a/common/i18n/zh-HK/err_repo.json b/common/i18n/zh-HK/err_repo.json index 4e31fd18f..9fcf2cfa7 100644 --- a/common/i18n/zh-HK/err_repo.json +++ b/common/i18n/zh-HK/err_repo.json @@ -16,5 +16,8 @@ }, "error.REPO-ERR-5": { "other": "下載代碼 zip 歸檔失敗" + }, + "error.REPO-ERR-6": { + "other": "批量獲取倉庫額外資訊失敗" } } diff --git a/common/types/repo.go b/common/types/repo.go index abe08a378..f0c1c0ee5 100644 --- a/common/types/repo.go +++ b/common/types/repo.go @@ -425,3 +425,12 @@ type DownloadCodeZipReq struct { Name string `json:"name"` Revision string `json:"revision"` } + +type BatchRepoExtraReq struct { + RepoIDs []int64 `json:"repo_ids" binding:"required"` +} + +type RepoExtraItem struct { + RepoID int64 `json:"repo_id"` + Size int64 `json:"size"` +} diff --git a/component/repo.go b/component/repo.go index 8e1eb4811..4ae0a35f2 100644 --- a/component/repo.go +++ b/component/repo.go @@ -3,6 +3,7 @@ package component import ( "bufio" "bytes" + "cmp" "context" "crypto/sha256" "database/sql" @@ -197,6 +198,8 @@ type RepoComponent interface { GetRepos(ctx context.Context, search, currentUser string, repoType types.RepositoryType) ([]string, error) // GetRepoSizeByBranch gets the repository size for a specific branch GetRepoSizeByBranch(ctx context.Context, repoType types.RepositoryType, namespace, name, branch, currentUser string) (int64, error) + // BatchGetRepoExtra gets the default branch size for multiple repositories + BatchGetRepoExtra(ctx context.Context, repoIDs []int64, currentUser string) ([]types.RepoExtraItem, error) // DownloadCodeZip downloads the code repository as a zip archive DownloadCodeZip(ctx context.Context, req types.DownloadCodeZipReq, currentUser string) ([]byte, error) advancedRepoInterface @@ -3305,6 +3308,82 @@ func (c *repoComponentImpl) GetRepoSizeByBranch(ctx context.Context, repoType ty return stats.TotalSize, nil } +// BatchGetRepoExtra gets the default branch size for multiple repositories by their IDs. +// Only repositories the current user has read permission for are included in the result. +func (c *repoComponentImpl) BatchGetRepoExtra(ctx context.Context, repoIDs []int64, currentUser string) ([]types.RepoExtraItem, error) { + repos, err := c.repoStore.FindByIds(ctx, repoIDs) + if err != nil { + return nil, errorx.BatchGetRepoExtraFailed(fmt.Errorf("failed to find repos by ids, error: %w", err)) + } + + // Determine which repos the user can read — same approach as PublicToUser: + // fetch user info once, then check ownership + privacy in-memory. + isAdmin := false + accessibleIDs := make(map[int64]bool) + if currentUser != "" { + user, err := c.userSvcClient.GetUserInfo(ctx, currentUser, currentUser) + if err != nil { + return nil, errorx.BatchGetRepoExtraFailed(fmt.Errorf("failed to get user info, error: %w", err)) + } + dbUser := &database.User{RoleMask: strings.Join(user.Roles, ",")} + if dbUser.CanAdmin() { + isAdmin = true + } else { + accessibleIDs[user.ID] = true + for _, org := range user.Orgs { + accessibleIDs[org.UserID] = true + } + } + } + + // Filter repos to those the user has read permission for + readableRepos := make([]*database.Repository, 0, len(repos)) + for _, repo := range repos { + if isAdmin || !repo.Private || accessibleIDs[repo.UserID] { + readableRepos = append(readableRepos, repo) + } + } + + // Collect repos that have a default branch for statistics lookup + repoDefaultBranches := make(map[int64]string, len(readableRepos)) + for _, repo := range readableRepos { + if repo.DefaultBranch != "" { + repoDefaultBranches[repo.ID] = repo.DefaultBranch + } + } + + sizeMap := make(map[int64]int64, len(readableRepos)) + + if len(repoDefaultBranches) > 0 { + repoIDsWithBranch := make([]int64, 0, len(repoDefaultBranches)) + for repoID := range repoDefaultBranches { + repoIDsWithBranch = append(repoIDsWithBranch, repoID) + } + slices.Sort(repoIDsWithBranch) + + stats, err := c.repoStatisticsStore.FindByRepositoryIDs(ctx, repoIDsWithBranch) + if err != nil { + return nil, errorx.BatchGetRepoExtraFailed(fmt.Errorf("failed to find repo statistics, error: %w", err)) + } + + for _, stat := range stats { + if defaultBranch, ok := repoDefaultBranches[stat.RepositoryID]; ok && stat.Branch == defaultBranch { + sizeMap[stat.RepositoryID] = stat.TotalSize + } + } + } + + result := make([]types.RepoExtraItem, 0, len(readableRepos)) + for _, repo := range readableRepos { + result = append(result, types.RepoExtraItem{RepoID: repo.ID, Size: sizeMap[repo.ID]}) + } + slices.SortFunc(result, func(a, b types.RepoExtraItem) int { + return cmp.Compare(a.RepoID, b.RepoID) + }) + + return result, nil +} + func (c *repoComponentImpl) DownloadCodeZip(ctx context.Context, req types.DownloadCodeZipReq, currentUser string) ([]byte, error) { repo, err := c.repoStore.FindByPath(ctx, types.CodeRepo, req.Namespace, req.Name) if err != nil { diff --git a/component/repo_test.go b/component/repo_test.go index 18d9cb178..a3b7bbc27 100644 --- a/component/repo_test.go +++ b/component/repo_test.go @@ -2667,6 +2667,198 @@ func TestRepoComponent_GetRepoSizeByBranch(t *testing.T) { }) } +func TestRepoComponent_BatchGetRepoExtra(t *testing.T) { + t.Run("admin user can read all repos", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "admin", "admin").Return(&rpc.User{ + ID: 1, + Roles: []string{"admin"}, + Orgs: []rpc.Organization{}, + }, nil) + + repos := []*database.Repository{ + {ID: 1, UserID: 10, Private: true, DefaultBranch: "main"}, + {ID: 2, UserID: 20, Private: true, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1, 2}).Return(repos, nil) + + stats := []*database.RepositoryStatistics{ + {RepositoryID: 1, Branch: "main", TotalSize: 1024}, + {RepositoryID: 2, Branch: "main", TotalSize: 2048}, + } + repoComp.mocks.stores.RepositoryStatisticsMock().EXPECT().FindByRepositoryIDs(ctx, []int64{1, 2}).Return(stats, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1, 2}, "admin") + require.Nil(t, err) + require.Equal(t, []types.RepoExtraItem{{RepoID: 1, Size: 1024}, {RepoID: 2, Size: 2048}}, result) + }) + + t.Run("regular user can read own repos and public repos", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(&rpc.User{ + ID: 1, + Roles: []string{}, + Orgs: []rpc.Organization{{UserID: 3}}, + }, nil) + + // Own repo (UserID=1), org repo (UserID=3), public repo (UserID=99), private unrelated repo (UserID=99) + repos := []*database.Repository{ + {ID: 1, UserID: 1, Private: true, DefaultBranch: "main"}, + {ID: 2, UserID: 3, Private: true, DefaultBranch: "main"}, + {ID: 3, UserID: 99, Private: false, DefaultBranch: "main"}, + {ID: 4, UserID: 99, Private: true, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1, 2, 3, 4}).Return(repos, nil) + + // Only repos 1, 2, 3 are readable + stats := []*database.RepositoryStatistics{ + {RepositoryID: 1, Branch: "main", TotalSize: 100}, + {RepositoryID: 2, Branch: "main", TotalSize: 200}, + {RepositoryID: 3, Branch: "main", TotalSize: 300}, + } + repoComp.mocks.stores.RepositoryStatisticsMock().EXPECT().FindByRepositoryIDs(ctx, []int64{1, 2, 3}).Return(stats, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1, 2, 3, 4}, "user1") + require.Nil(t, err) + require.Equal(t, []types.RepoExtraItem{{RepoID: 1, Size: 100}, {RepoID: 2, Size: 200}, {RepoID: 3, Size: 300}}, result) + }) + + t.Run("anonymous user can only read public repos", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repos := []*database.Repository{ + {ID: 1, UserID: 1, Private: true, DefaultBranch: "main"}, + {ID: 2, UserID: 2, Private: false, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1, 2}).Return(repos, nil) + + stats := []*database.RepositoryStatistics{ + {RepositoryID: 2, Branch: "main", TotalSize: 500}, + } + repoComp.mocks.stores.RepositoryStatisticsMock().EXPECT().FindByRepositoryIDs(ctx, []int64{2}).Return(stats, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1, 2}, "") + require.Nil(t, err) + require.Equal(t, []types.RepoExtraItem{{RepoID: 2, Size: 500}}, result) + }) + + t.Run("repo without default branch is included with size 0", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(&rpc.User{ + ID: 1, + Roles: []string{}, + Orgs: []rpc.Organization{}, + }, nil) + + repos := []*database.Repository{ + {ID: 1, UserID: 1, Private: true, DefaultBranch: ""}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return(repos, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Nil(t, err) + require.Equal(t, []types.RepoExtraItem{{RepoID: 1, Size: 0}}, result) + }) + + t.Run("repo with default branch but no statistics gets size 0", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(&rpc.User{ + ID: 1, + Roles: []string{}, + Orgs: []rpc.Organization{}, + }, nil) + + repos := []*database.Repository{ + {ID: 1, UserID: 1, Private: true, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return(repos, nil) + + stats := []*database.RepositoryStatistics{ + {RepositoryID: 1, Branch: "dev", TotalSize: 100}, + } + repoComp.mocks.stores.RepositoryStatisticsMock().EXPECT().FindByRepositoryIDs(ctx, []int64{1}).Return(stats, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Nil(t, err) + require.Equal(t, []types.RepoExtraItem{{RepoID: 1, Size: 0}}, result) + }) + + t.Run("all repos filtered out returns empty map", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(&rpc.User{ + ID: 1, + Roles: []string{}, + Orgs: []rpc.Organization{}, + }, nil) + + repos := []*database.Repository{ + {ID: 1, UserID: 99, Private: true, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return(repos, nil) + + result, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Nil(t, err) + require.Empty(t, result) + }) + + t.Run("FindByIds error", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return(nil, errors.New("db error")) + + _, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Error(t, err) + require.True(t, errors.Is(err, errorx.ErrBatchGetRepoExtraFailed)) + }) + + t.Run("GetUserInfo error", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return([]*database.Repository{ + {ID: 1, UserID: 1, Private: true}, + }, nil) + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(nil, errors.New("rpc error")) + + _, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Error(t, err) + require.True(t, errors.Is(err, errorx.ErrBatchGetRepoExtraFailed)) + }) + + t.Run("FindByRepositoryIDs error", func(t *testing.T) { + ctx := context.TODO() + repoComp := initializeTestRepoComponent(ctx, t) + + repoComp.mocks.userSvcClient.EXPECT().GetUserInfo(ctx, "user1", "user1").Return(&rpc.User{ + ID: 1, + Roles: []string{}, + Orgs: []rpc.Organization{}, + }, nil) + + repos := []*database.Repository{ + {ID: 1, UserID: 1, Private: true, DefaultBranch: "main"}, + } + repoComp.mocks.stores.RepoMock().EXPECT().FindByIds(ctx, []int64{1}).Return(repos, nil) + repoComp.mocks.stores.RepositoryStatisticsMock().EXPECT().FindByRepositoryIDs(ctx, []int64{1}).Return(nil, errors.New("db error")) + + _, err := repoComp.BatchGetRepoExtra(ctx, []int64{1}, "user1") + require.Error(t, err) + require.True(t, errors.Is(err, errorx.ErrBatchGetRepoExtraFailed)) + }) +} + func TestRepoComponent_DownloadCodeZip(t *testing.T) { t.Run("should return zip data when user has read permission and ref is provided", func(t *testing.T) { ctx := context.TODO() @@ -2682,10 +2874,10 @@ func TestRepoComponent_DownloadCodeZip(t *testing.T) { repoComp.mocks.stores.NamespaceMock().EXPECT().FindByPath(mock.Anything, ns.Path).Return(ns, nil) repo := &database.Repository{ - ID: 1, - Private: true, - User: user, - Path: fmt.Sprintf("%s/%s", ns.Path, "repo_name"), + ID: 1, + Private: true, + User: user, + Path: fmt.Sprintf("%s/%s", ns.Path, "repo_name"), DefaultBranch: "main", } repoComp.mocks.stores.RepoMock().EXPECT().FindByPath(mock.Anything, types.CodeRepo, ns.Path, repo.Name).Return(repo, nil)