Refactor service user handling, user cache lookup, and cache loading

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>

management/server/account.go

@@ -1441,20 +1441,8 @@ func isNil(i idp.Manager) bool {
 // addAccountIDToIDPAppMeta update user's  app metadata in idp manager
 func (am *DefaultAccountManager) addAccountIDToIDPAppMeta(ctx context.Context, userID string, accountID string) error {
 	if !isNil(am.idpManager) {
-		accountUsers, err := am.Store.GetAccountUsers(ctx, LockingStrengthShare, accountID)
-		if err != nil {
-			return err
-		}
-		cachedAccount := &Account{
-			Id:    accountID,
-			Users: make(map[string]*User),
-		}
-		for _, user := range accountUsers {
-			cachedAccount.Users[user.Id] = user
-		}
-
 		// user can be nil if it wasn't found (e.g., just created)
-		user, err := am.lookupUserInCache(ctx, userID, cachedAccount)
+		user, err := am.lookupUserInCache(ctx, userID, accountID)
 		if err != nil {
 			return err
 		}
@@ -1483,11 +1471,10 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID inte
 	log.WithContext(ctx).Debugf("account %s not found in cache, reloading", accountID)
 	accountIDString := fmt.Sprintf("%v", accountID)
 
-	account, err := am.Store.GetAccount(ctx, accountIDString)
+	accountUsers, err := am.Store.GetAccountUsers(ctx, LockingStrengthShare, accountIDString)
 	if err != nil {
 		return nil, err
 	}
-
 	userData, err := am.idpManager.GetAccount(ctx, accountIDString)
 	if err != nil {
 		return nil, err
@@ -1500,7 +1487,7 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID inte
 	}
 
 	matchedUserData := make([]*idp.UserData, 0)
-	for _, user := range account.Users {
+	for _, user := range accountUsers {
 		if user.IsServiceUser {
 			continue
 		}
@@ -1530,10 +1517,15 @@ func (am *DefaultAccountManager) lookupUserInCacheByEmail(ctx context.Context, e
 }
 
 // lookupUserInCache looks up user in the IdP cache and returns it. If the user wasn't found, the function returns nil
-func (am *DefaultAccountManager) lookupUserInCache(ctx context.Context, userID string, account *Account) (*idp.UserData, error) {
+func (am *DefaultAccountManager) lookupUserInCache(ctx context.Context, userID string, accountID string) (*idp.UserData, error) {
-	users := make(map[string]userLoggedInOnce, len(account.Users))
+	accountUsers, err := am.Store.GetAccountUsers(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	users := make(map[string]userLoggedInOnce, len(accountUsers))
 	// ignore service users and users provisioned by integrations than are never logged in
-	for _, user := range account.Users {
+	for _, user := range accountUsers {
 		if user.IsServiceUser {
 			continue
 		}
@@ -1542,8 +1534,8 @@ func (am *DefaultAccountManager) lookupUserInCache(ctx context.Context, userID s
 		}
 		users[user.Id] = userLoggedInOnce(!user.LastLogin.IsZero())
 	}
-	log.WithContext(ctx).Debugf("looking up user %s of account %s in cache", userID, account.Id)
+	log.WithContext(ctx).Debugf("looking up user %s of account %s in cache", userID, accountID)
-	userData, err := am.lookupCache(ctx, users, account.Id)
+	userData, err := am.lookupCache(ctx, users, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -1556,13 +1548,13 @@ func (am *DefaultAccountManager) lookupUserInCache(ctx context.Context, userID s
 
 	// add extra check on external cache manager. We may get to this point when the user is not yet findable in IDP,
 	// or it didn't have its metadata updated with am.addAccountIDToIDPAppMeta
-	user, err := account.FindUser(userID)
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
-		log.WithContext(ctx).Errorf("failed finding user %s in account %s", userID, account.Id)
+		log.WithContext(ctx).Errorf("failed finding user %s in account %s", userID, accountID)
 		return nil, err
 	}
 
-	key := user.IntegrationReference.CacheKey(account.Id, userID)
+	key := user.IntegrationReference.CacheKey(accountID, userID)
 	ud, err := am.externalCacheManager.Get(am.ctx, key)
 	if err != nil {
 		log.WithContext(ctx).Debugf("failed to get externalCache for key: %s, error: %s", key, err)
@@ -1827,12 +1819,7 @@ func (am *DefaultAccountManager) redeemInvite(ctx context.Context, accountID str
 		return nil
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	user, err := am.lookupUserInCache(ctx, userID, accountID)
-	if err != nil {
-		return err
-	}
-
-	user, err := am.lookupUserInCache(ctx, userID, account)
 	if err != nil {
 		return err
 	}
@@ -1842,17 +1829,17 @@ func (am *DefaultAccountManager) redeemInvite(ctx context.Context, accountID str
 	}
 
 	if user.AppMetadata.WTPendingInvite != nil && *user.AppMetadata.WTPendingInvite {
-		log.WithContext(ctx).Infof("redeeming invite for user %s account %s", userID, account.Id)
+		log.WithContext(ctx).Infof("redeeming invite for user %s account %s", userID, accountID)
 		// User has already logged in, meaning that IdP should have set wt_pending_invite to false.
 		// Our job is to just reload cache.
 		go func() {
-			_, err = am.refreshCache(ctx, account.Id)
+			_, err = am.refreshCache(ctx, accountID)
 			if err != nil {
-				log.WithContext(ctx).Warnf("failed reloading cache when redeeming user %s under account %s", userID, account.Id)
+				log.WithContext(ctx).Warnf("failed reloading cache when redeeming user %s under account %s", userID, accountID)
 				return
 			}
-			log.WithContext(ctx).Debugf("user %s of account %s redeemed invite", user.ID, account.Id)
+			log.WithContext(ctx).Debugf("user %s of account %s redeemed invite", user.ID, accountID)
-			am.StoreEvent(ctx, userID, userID, account.Id, activity.UserJoined, nil)
+			am.StoreEvent(ctx, userID, userID, accountID, activity.UserJoined, nil)
 		}()
 	}
 

management/server/sql_store.go

@@ -413,7 +413,8 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 func (s *SqlStore) SaveUser(ctx context.Context, lockStrength LockingStrength, user *User) error {
 	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
 	if result.Error != nil {
-		return status.Errorf(status.Internal, "failed to save user to store: %v", result.Error)
+		log.WithContext(ctx).Errorf("failed to save user to store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to save user to store")
 	}
 	return nil
 }
@@ -822,6 +823,21 @@ func (s *SqlStore) GetAccountSettings(ctx context.Context, lockStrength LockingS
 	return accountSettings.Settings, nil
 }
 
+func (s *SqlStore) GetAccountCreatedBy(ctx context.Context, lockStrength LockingStrength, accountID string) (string, error) {
+	var createdBy string
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+		Select("created_by").First(&createdBy, idQueryCondition, accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.NewAccountNotFoundError()
+		}
+		log.WithContext(ctx).Errorf("error when getting account creator from the store: %s", result.Error)
+		return "", status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return createdBy, nil
+}
+
 // SaveAccountSettings stores the account settings in DB.
 func (s *SqlStore) SaveAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *Settings) error {
 	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).

management/server/status/error.go

@@ -119,6 +119,19 @@ func NewGetUserFromStoreError() error {
 	return Errorf(Internal, "issue getting user from store")
 }
 
+func NewUnauthorizedToViewUsersError() error {
+	return Errorf(PermissionDenied, "only users with admin power can view users")
+}
+
+func NewUnauthorizedToViewServiceUsersError() error {
+	return Errorf(PermissionDenied, "only users with admin power can view service users")
+}
+
+// NewServiceUserRoleInvalidError creates a new Error with InvalidArgument type for creating a service user with owner role
+func NewServiceUserRoleInvalidError() error {
+	return Errorf(InvalidArgument, "can't create a service user with owner role")
+}
+
 // NewInvalidKeyIDError creates a new Error with InvalidArgument type for an issue getting a setup key
 func NewInvalidKeyIDError() error {
 	return Errorf(InvalidArgument, "invalid key ID")

management/server/store.go

@@ -57,6 +57,7 @@ type Store interface {
 	GetAccountIDByPrivateDomain(ctx context.Context, lockStrength LockingStrength, domain string) (string, error)
 	GetAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*Settings, error)
 	GetAccountDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*DNSSettings, error)
+	GetAccountCreatedBy(ctx context.Context, lockStrength LockingStrength, accountID string) (string, error)
 	SaveDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *DNSSettings) error
 	SaveAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *Settings) error
 	SaveAccount(ctx context.Context, account *Account) error

management/server/user.go

@@ -222,33 +222,25 @@ func NewOwnerUser(id string) *User {
 
 // createServiceUser creates a new service user under the given account.
 func (am *DefaultAccountManager) createServiceUser(ctx context.Context, accountID string, initiatorUserID string, role UserRole, serviceUserName string, nonDeletable bool, autoGroups []string) (*UserInfo, error) {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	initiatorUser, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, initiatorUserID)
-	defer unlock()
-
-	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return nil, status.Errorf(status.NotFound, "account %s doesn't exist", accountID)
+		return nil, err
 	}
 
-	executingUser := account.Users[initiatorUserID]
+	if !initiatorUser.HasAdminPower() {
-	if executingUser == nil {
+		return nil, status.NewUnauthorizedToViewServiceUsersError()
-		return nil, status.Errorf(status.NotFound, "user not found")
-	}
-	if !executingUser.HasAdminPower() {
-		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can create service users")
 	}
 
 	if role == UserRoleOwner {
-		return nil, status.Errorf(status.InvalidArgument, "can't create a service user with owner role")
+		return nil, status.NewServiceUserRoleInvalidError()
 	}
 
 	newUserID := uuid.New().String()
 	newUser := NewUser(newUserID, role, true, nonDeletable, serviceUserName, autoGroups, UserIssuedAPI)
+	newUser.AccountID = accountID
 	log.WithContext(ctx).Debugf("New User: %v", newUser)
-	account.Users[newUserID] = newUser
 
-	err = am.Store.SaveAccount(ctx, account)
+	if err = am.Store.SaveUser(ctx, LockingStrengthUpdate, newUser); err != nil {
-	if err != nil {
 		return nil, err
 	}
 
@@ -301,23 +293,22 @@ func (am *DefaultAccountManager) inviteNewUser(ctx context.Context, accountID, u
 	default:
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	initiatorUser, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.Errorf(status.NotFound, "account %s doesn't exist", accountID)
+		return nil, err
-	}
-
-	initiatorUser, err := account.FindUser(userID)
-	if err != nil {
-		return nil, status.Errorf(status.NotFound, "initiator user with ID %s doesn't exist", userID)
 	}
 
 	inviterID := userID
 	if initiatorUser.IsServiceUser {
-		inviterID = account.CreatedBy
+		createdBy, err := am.Store.GetAccountCreatedBy(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return nil, err
+		}
+		inviterID = createdBy
 	}
 
 	// inviterUser is the one who is inviting the new user
-	inviterUser, err := am.lookupUserInCache(ctx, inviterID, account)
+	inviterUser, err := am.lookupUserInCache(ctx, inviterID, accountID)
 	if err != nil || inviterUser == nil {
 		return nil, status.Errorf(status.NotFound, "inviter user with ID %s doesn't exist in IdP", inviterID)
 	}
@@ -348,27 +339,31 @@ func (am *DefaultAccountManager) inviteNewUser(ctx context.Context, accountID, u
 
 	newUser := &User{
 		Id:                   idpUser.ID,
+		AccountID:            accountID,
 		Role:                 invitedRole,
 		AutoGroups:           invite.AutoGroups,
 		Issued:               invite.Issued,
 		IntegrationReference: invite.IntegrationReference,
 		CreatedAt:            time.Now().UTC(),
 	}
-	account.Users[idpUser.ID] = newUser
 
-	err = am.Store.SaveAccount(ctx, account)
+	settings, err := am.Store.GetAccountSettings(ctx, LockingStrengthShare, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	_, err = am.refreshCache(ctx, account.Id)
+	if err = am.Store.SaveUser(ctx, LockingStrengthUpdate, newUser); err != nil {
+		return nil, err
+	}
+
+	_, err = am.refreshCache(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	am.StoreEvent(ctx, userID, newUser.Id, accountID, activity.UserInvited, nil)
 
-	return newUser.ToUserInfo(idpUser, account.Settings)
+	return newUser.ToUserInfo(idpUser, settings)
 }
 
 func (am *DefaultAccountManager) GetUserByID(ctx context.Context, id string) (*User, error) {
@@ -408,20 +403,7 @@ func (am *DefaultAccountManager) GetUser(ctx context.Context, claims jwtclaims.A
 // ListUsers returns lists of all users under the account.
 // It doesn't populate user information such as email or name.
 func (am *DefaultAccountManager) ListUsers(ctx context.Context, accountID string) ([]*User, error) {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	return am.Store.GetAccountUsers(ctx, LockingStrengthShare, accountID)
-	defer unlock()
-
-	account, err := am.Store.GetAccount(ctx, accountID)
-	if err != nil {
-		return nil, err
-	}
-
-	users := make([]*User, 0, len(account.Users))
-	for _, item := range account.Users {
-		users = append(users, item)
-	}
-
-	return users, nil
 }
 
 func (am *DefaultAccountManager) deleteServiceUser(ctx context.Context, account *Account, initiatorUserID string, targetUser *User) {
@@ -519,20 +501,12 @@ func (am *DefaultAccountManager) deleteUserPeers(ctx context.Context, initiatorU
 
 // InviteUser resend invitations to users who haven't activated their accounts prior to the expiration period.
 func (am *DefaultAccountManager) InviteUser(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	if am.idpManager == nil {
 		return status.Errorf(status.PreconditionFailed, "IdP manager must be enabled to send user invites")
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
-	if err != nil {
-		return status.Errorf(status.NotFound, "account %s doesn't exist", accountID)
-	}
-
 	// check if the user is already registered with this ID
-	user, err := am.lookupUserInCache(ctx, targetUserID, account)
+	user, err := am.lookupUserInCache(ctx, targetUserID, accountID)
 	if err != nil {
 		return err
 	}
@@ -559,9 +533,6 @@ func (am *DefaultAccountManager) InviteUser(ctx context.Context, accountID strin
 
 // CreatePAT creates a new PAT for the given user
 func (am *DefaultAccountManager) CreatePAT(ctx context.Context, accountID string, initiatorUserID string, targetUserID string, tokenName string, expiresIn int) (*PersonalAccessTokenGenerated, error) {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	if tokenName == "" {
 		return nil, status.Errorf(status.InvalidArgument, "token name can't be empty")
 	}
@@ -570,35 +541,31 @@ func (am *DefaultAccountManager) CreatePAT(ctx context.Context, accountID string
 		return nil, status.Errorf(status.InvalidArgument, "expiration has to be between 1 and 365")
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	initiatorUser, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, initiatorUserID)
 	if err != nil {
 		return nil, err
 	}
 
-	targetUser, ok := account.Users[targetUserID]
+	if initiatorUser.AccountID != accountID {
-	if !ok {
+		return nil, status.NewUserNotPartOfAccountError()
-		return nil, status.Errorf(status.NotFound, "user not found")
 	}
 
-	executingUser, ok := account.Users[initiatorUserID]
+	if initiatorUserID != targetUserID && initiatorUser.IsRegularUser() {
-	if !ok {
+		return nil, status.NewUnauthorizedToViewPATsError()
-		return nil, status.Errorf(status.NotFound, "user not found")
 	}
 
-	if !(initiatorUserID == targetUserID || (executingUser.HasAdminPower() && targetUser.IsServiceUser)) {
+	targetUser, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, initiatorUserID)
-		return nil, status.Errorf(status.PermissionDenied, "no permission to create PAT for this user")
+	if err != nil {
+		return nil, err
 	}
 
-	pat, err := CreateNewPAT(tokenName, expiresIn, targetUserID, executingUser.Id)
+	pat, err := CreateNewPAT(tokenName, expiresIn, targetUserID, initiatorUser.Id)
 	if err != nil {
 		return nil, status.Errorf(status.Internal, "failed to create PAT: %v", err)
 	}
 
-	targetUser.PATs[pat.ID] = &pat.PersonalAccessToken
+	if err = am.Store.SavePAT(ctx, LockingStrengthUpdate, &pat.PersonalAccessToken); err != nil {
-
+		return nil, err
-	err = am.Store.SaveAccount(ctx, account)
-	if err != nil {
-		return nil, status.Errorf(status.Internal, "failed to save account: %v", err)
 	}
 
 	meta := map[string]any{"name": pat.Name, "is_service_user": targetUser.IsServiceUser, "user_name": targetUser.ServiceUserName}
@@ -777,7 +744,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 		events := am.prepareUserUpdateEvents(ctx, initiatorUser.Id, oldUser, newUser, account, transferredOwnerRole)
 		eventsToStore = append(eventsToStore, events...)
 
-		updatedUserInfo, err := getUserInfo(ctx, am, newUser, account)
+		updatedUserInfo, err := getUserInfo(ctx, am, newUser, accountID)
 		if err != nil {
 			return nil, err
 		}
@@ -876,15 +843,20 @@ func handleOwnerRoleTransfer(account *Account, initiatorUser, update *User) bool
 // getUserInfo retrieves the UserInfo for a given User and Account.
 // If the AccountManager has a non-nil idpManager and the User is not a service user,
 // it will attempt to look up the UserData from the cache.
-func getUserInfo(ctx context.Context, am *DefaultAccountManager, user *User, account *Account) (*UserInfo, error) {
+func getUserInfo(ctx context.Context, am *DefaultAccountManager, user *User, accountID string) (*UserInfo, error) {
+	settings, err := am.Store.GetAccountSettings(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
+	}
+
 	if !isNil(am.idpManager) && !user.IsServiceUser {
-		userData, err := am.lookupUserInCache(ctx, user.Id, account)
+		userData, err := am.lookupUserInCache(ctx, user.Id, accountID)
 		if err != nil {
 			return nil, err
 		}
-		return user.ToUserInfo(userData, account.Settings)
+		return user.ToUserInfo(userData, settings)
 	}
-	return user.ToUserInfo(nil, account.Settings)
+	return user.ToUserInfo(nil, settings)
 }
 
 // validateUserUpdate validates the update operation for a user.