1 files changed, 75 insertions, 11 deletions

diff --git a/models/org_team.go b/models/org_team.go
index 5ea6e76cd9..3b37f936ff 100644
--- a/models/org_team.go
+++ b/models/org_team.go
@@ -1,3 +1,4 @@
+// Copyright 2018 The Gitea Authors. All rights reserved.
 // Copyright 2016 The Gogs Authors. All rights reserved.
 // Use of this source code is governed by a MIT-style
 // license that can be found in the LICENSE file.
@@ -10,7 +11,6 @@ import (
 	"strings"
 
 	"code.gitea.io/gitea/modules/log"
-
 	"github.com/go-xorm/xorm"
 )
 
@@ -28,15 +28,16 @@ type Team struct {
 	Members     []*User       `xorm:"-"`
 	NumRepos    int
 	NumMembers  int
-	UnitTypes   []UnitType `xorm:"json"`
+	Units       []*TeamUnit `xorm:"-"`
 }
 
-// GetUnitTypes returns unit types the team owned, empty means all the unit types
-func (t *Team) GetUnitTypes() []UnitType {
-	if len(t.UnitTypes) == 0 {
-		return allRepUnitTypes
+func (t *Team) getUnits(e Engine) (err error) {
+	if t.Units != nil {
+		return nil
 	}
-	return t.UnitTypes
+
+	t.Units, err = getUnitsByTeamID(e, t.ID)
+	return err
 }
 
 // HasWriteAccess returns true if team has at least write level access mode.
@@ -214,11 +215,12 @@ func (t *Team) RemoveRepository(repoID int64) error {
 
 // UnitEnabled returns if the team has the given unit type enabled
 func (t *Team) UnitEnabled(tp UnitType) bool {
-	if len(t.UnitTypes) == 0 {
-		return true
+	if err := t.getUnits(x); err != nil {
+		log.Warn("Error loading repository (ID: %d) units: %s", t.ID, err.Error())
 	}
-	for _, u := range t.UnitTypes {
-		if u == tp {
+
+	for _, unit := range t.Units {
+		if unit.Type == tp {
 			return true
 		}
 	}
@@ -275,6 +277,17 @@ func NewTeam(t *Team) (err error) {
 		return err
 	}
 
+	// insert units for team
+	if len(t.Units) > 0 {
+		for _, unit := range t.Units {
+			unit.TeamID = t.ID
+		}
+		if _, err = sess.Insert(&t.Units); err != nil {
+			sess.Rollback()
+			return err
+		}
+	}
+
 	// Update organization number of teams.
 	if _, err = sess.Exec("UPDATE `user` SET num_teams=num_teams+1 WHERE id = ?", t.OrgID); err != nil {
 		sess.Rollback()
@@ -424,6 +437,13 @@ func DeleteTeam(t *Team) error {
 		return err
 	}
 
+	// Delete team-unit.
+	if _, err := sess.
+		Where("team_id=?", t.ID).
+		Delete(new(TeamUnit)); err != nil {
+		return err
+	}
+
 	// Delete team.
 	if _, err := sess.ID(t.ID).Delete(new(Team)); err != nil {
 		return err
@@ -695,3 +715,47 @@ func GetTeamsWithAccessToRepo(orgID, repoID int64, mode AccessMode) ([]*Team, er
 		And("team_repo.repo_id = ?", repoID).
 		Find(&teams)
 }
+
+// ___________                    ____ ___      .__  __
+// \__    ___/___ _____    _____ |    |   \____ |__|/  |_
+//   |    |_/ __ \\__  \  /     \|    |   /    \|  \   __\
+//   |    |\  ___/ / __ \|  Y Y  \    |  /   |  \  ||  |
+//   |____| \___  >____  /__|_|  /______/|___|  /__||__|
+//              \/     \/      \/             \/
+
+// TeamUnit describes all units of a repository
+type TeamUnit struct {
+	ID     int64    `xorm:"pk autoincr"`
+	OrgID  int64    `xorm:"INDEX"`
+	TeamID int64    `xorm:"UNIQUE(s)"`
+	Type   UnitType `xorm:"UNIQUE(s)"`
+}
+
+// Unit returns Unit
+func (t *TeamUnit) Unit() Unit {
+	return Units[t.Type]
+}
+
+func getUnitsByTeamID(e Engine, teamID int64) (units []*TeamUnit, err error) {
+	return units, e.Where("team_id = ?", teamID).Find(&units)
+}
+
+// UpdateTeamUnits updates a teams's units
+func UpdateTeamUnits(team *Team, units []TeamUnit) (err error) {
+	sess := x.NewSession()
+	defer sess.Close()
+	if err = sess.Begin(); err != nil {
+		return err
+	}
+
+	if _, err = sess.Where("team_id = ?", team.ID).Delete(new(TeamUnit)); err != nil {
+		return err
+	}
+
+	if _, err = sess.Insert(units); err != nil {
+		sess.Rollback()
+		return err
+	}
+
+	return sess.Commit()
+}