Optimize GetTwoFaStatus

Author: sapk

models/userlist.go

@@ -4,6 +4,10 @@
 
 package models
 
+import (
+	"fmt"
+)
+
 //UserList is a list of user.
 // This type provide valuable methods to retrieve information for a group of users efficiently.
 type UserList []*User
@@ -29,9 +33,31 @@ func (users UserList) IsUserOrgOwner(orgID int64) map[int64]bool {
 // GetTwoFaStatus return state of 2FA enrollement
 func (users UserList) GetTwoFaStatus() map[int64]bool {
 	results := make(map[int64]bool, len(users))
-	//TODO use directly xorm
-	for _, u := range users {
-		results[u.ID] = u.IsTwoFaEnrolled()
+	for _, user := range users {
+		results[user.ID] = false //Set default to false
+	}
+	tokenMaps, err := users.loadTwoFactorStatus(x)
+	if err == nil {
+		for _, token := range tokenMaps {
+			results[token.UID] = true
+		}
 	}
+
 	return results
 }
+
+func (users UserList) loadTwoFactorStatus(e Engine) (map[int64]*TwoFactor, error) {
+	if len(users) == 0 {
+		return nil, nil
+	}
+
+	userIDs := users.getUserIDs()
+	tokenMaps := make(map[int64]*TwoFactor, len(userIDs))
+	err := e.
+		In("uid", userIDs).
+		Find(&tokenMaps)
+	if err != nil {
+		return nil, fmt.Errorf("find two factor: %v", err)
+	}
+	return tokenMaps, nil
+}

models/userlist_test.go

@@ -76,10 +76,10 @@ func TestUserListIsTwoFaEnrolled(t *testing.T) {
 		})
 	}
 }
+
 func testUserListIsTwoFaEnrolled(t *testing.T, orgID int64, expected map[int64]bool) {
 	org, err := GetUserByID(orgID)
 	assert.NoError(t, err)
 	assert.NoError(t, org.GetMembers())
 	assert.Equal(t, expected, org.Members.GetTwoFaStatus())
-
 }