Add team retrieval and counting methods for user membership

- Introduced GetTeamsByMember and CountTeamsByMember methods in the DefaultUser struct to retrieve teams associated with a specific member and count the total number of teams, respectively.
- Updated the UserProvider interface to include these new methods.
- Enhanced the login response structure to include a status indicating whether team selection is required based on the user's team membership.
- Added error handling for team retrieval in the authentication process.

openapi/oauth/providers/user/team.go

@@ -251,6 +251,80 @@ func (u *DefaultUser) GetTeamsByOwner(ctx context.Context, ownerID string) ([]ma
 	return u.GetTeams(ctx, param)
 }
 
+// GetTeamsByMember retrieves teams by member_id
+func (u *DefaultUser) GetTeamsByMember(ctx context.Context, memberID string) ([]maps.MapStr, error) {
+
+	// Set default select fields if not provided
+	param := model.QueryParam{
+		Select: []interface{}{"team_id", "user_id", "member_type"},
+		Wheres: []model.QueryWhere{
+			{Column: "user_id", Value: memberID},
+			{Column: "member_type", Value: "user"},
+			{Column: "status", Value: "active"},
+		},
+	}
+
+	if param.Select == nil {
+		param.Select = u.memberFields
+	}
+
+	m := model.Select(u.memberModel)
+	members, err := m.Get(param)
+	if err != nil {
+		return nil, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	if len(members) == 0 {
+		return []maps.MapStr{}, nil
+	}
+
+	// Get team ids
+	teamIDs := []string{}
+	for _, member := range members {
+		teamIDs = append(teamIDs, member["team_id"].(string))
+	}
+
+	// Get teams
+	teams, err := u.GetTeams(ctx, model.QueryParam{
+		Select: u.teamFields,
+		Wheres: []model.QueryWhere{
+			{Column: "team_id", Value: teamIDs, Method: "in"},
+		},
+	})
+	if err != nil {
+		return nil, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	return teams, nil
+}
+
+// CountTeamsByMember returns total count of teams by member_id
+func (u *DefaultUser) CountTeamsByMember(ctx context.Context, memberID string) (int64, error) {
+
+	param := model.QueryParam{
+		Select: []interface{}{"team_id", "user_id", "member_type"},
+		Wheres: []model.QueryWhere{
+			{Column: "user_id", Value: memberID},
+			{Column: "member_type", Value: "user"},
+			{Column: "status", Value: "active"},
+		},
+	}
+	// Use Paginate with a small page size to get the total count
+	// This is more reliable than manual COUNT(*) queries
+	m := model.Select(u.memberModel)
+	result, err := m.Paginate(param, 1, 1) // Get first page with 1 item to get total
+	if err != nil {
+		return 0, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	// Extract total from pagination result using utility function
+	if totalInterface, ok := result["total"]; ok {
+		return parseIntFromDB(totalInterface)
+	}
+
+	return 0, fmt.Errorf("total not found in pagination result")
+}
+
 // GetTeamsByStatus retrieves teams by status
 func (u *DefaultUser) GetTeamsByStatus(ctx context.Context, status string) ([]maps.MapStr, error) {
 	param := model.QueryParam{

openapi/oauth/types/interfaces.go

@@ -266,7 +266,9 @@ type UserProvider interface {
 
 	// Team Query Methods
 	GetTeamsByOwner(ctx context.Context, ownerID string) ([]maps.MapStr, error)
+	GetTeamsByMember(ctx context.Context, memberID string) ([]maps.MapStr, error)
 	GetTeamsByStatus(ctx context.Context, status string) ([]maps.MapStr, error)
+	CountTeamsByMember(ctx context.Context, memberID string) (int64, error)
 
 	// Team Management
 	UpdateTeamStatus(ctx context.Context, teamID string, status string) error

openapi/response/response.go

@@ -113,6 +113,7 @@ var (
 	ErrInvalidSoftwareStatement = &ErrorResponse{Code: "invalid_software_statement", ErrorDescription: "The software statement is invalid or cannot be verified."}
 	ErrUnapprovedSoftware       = &ErrorResponse{Code: "unapproved_software", ErrorDescription: "The software statement represents software that has been replaced or is otherwise invalid."}
 	ErrMFARequired              = &ErrorResponse{Code: "mfa_required", ErrorDescription: "Multi-factor authentication is required to access this resource."}
+	ErrTeamSelectionRequired    = &ErrorResponse{Code: "team_selection_required", ErrorDescription: "Team selection is required to access this resource."}
 
 	// Configuration and service errors
 	ErrInvalidConfiguration     = types.ErrInvalidConfiguration

openapi/user/oauth.go

@@ -1,6 +1,7 @@
 package user
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/http"
@@ -13,6 +14,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/yaoapp/gou/session"
 	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/kun/maps"
 	"github.com/yaoapp/yao/openapi/oauth"
 	"github.com/yaoapp/yao/openapi/response"
 	"github.com/yaoapp/yao/openapi/utils"
@@ -195,6 +197,22 @@ func authback(c *gin.Context) {
 	// Send all login cookies (access token, refresh token, and session ID)
 	SendLoginCookies(c, loginResponse, sid)
 
+	// Get Teams
+	numTeams, err := countUserTeams(c.Request.Context(), loginResponse.UserID)
+	if err != nil {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrInvalidRequest.Code,
+			ErrorDescription: "Failed to count teams: " + err.Error(),
+		}
+		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
+		return
+	}
+
+	status := LoginStatusSuccess
+	if numTeams > 0 {
+		status = LoginStatusTeamSelection
+	}
+
 	// Send IDToken to the client
 	response.RespondWithSuccess(c, response.StatusOK, LoginSuccessResponse{
 		SessionID:             sid,
@@ -204,6 +222,7 @@ func authback(c *gin.Context) {
 		ExpiresIn:             loginResponse.ExpiresIn,
 		RefreshTokenExpiresIn: loginResponse.RefreshTokenExpiresIn,
 		MFAEnabled:            loginResponse.MFAEnabled,
+		Status:                status,
 	})
 }
 
@@ -453,6 +472,24 @@ func getUserInfo(providerID, state string) (string, error) {
 	return value.(string), nil
 }
 
+// getUserTeams gets the user teams
+func getUserTeams(ctx context.Context, userID string) ([]maps.MapStr, error) {
+	userProvider, err := oauth.OAuth.GetUserProvider()
+	if err != nil {
+		return nil, err
+	}
+	return userProvider.GetTeamsByMember(ctx, userID)
+}
+
+// countUserTeams counts the number of teams a user is a member of
+func countUserTeams(ctx context.Context, userID string) (int64, error) {
+	userProvider, err := oauth.OAuth.GetUserProvider()
+	if err != nil {
+		return 0, err
+	}
+	return userProvider.CountTeamsByMember(ctx, userID)
+}
+
 // removeUserInfo removes the user info from cache
 func removeUserInfo(providerID, state string) error {
 	key := userInfoKey(providerID, state)

openapi/user/types.go

@@ -2,6 +2,16 @@ package user
 
 import (
 	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
+	"github.com/yaoapp/yao/openapi/response"
+)
+
+const (
+	// LoginStatusSuccess is the success status
+	LoginStatusSuccess = "ok"
+	// LoginStatusMFA is the MFA status
+	LoginStatusMFA = "mfa_required"
+	// LoginStatusTeamSelection is the team selection status
+	LoginStatusTeamSelection = "team_selection_required"
 )
 
 // Config represents the signin page configuration
@@ -177,13 +187,15 @@ type LoginResponse struct {
 
 // LoginSuccessResponse represents the response for login success
 type LoginSuccessResponse struct {
-	IDToken               string `json:"id_token,omitempty"`
-	AccessToken           string `json:"access_token,omitempty"`
-	SessionID             string `json:"session_id,omitempty"`
-	RefreshToken          string `json:"refresh_token,omitempty"`
-	ExpiresIn             int    `json:"expires_in,omitempty"`
-	MFAEnabled            bool   `json:"mfa_enabled"`
-	RefreshTokenExpiresIn int    `json:"refresh_token_expires_in,omitempty"`
+	IDToken               string                  `json:"id_token,omitempty"`
+	AccessToken           string                  `json:"access_token,omitempty"`
+	SessionID             string                  `json:"session_id,omitempty"`
+	RefreshToken          string                  `json:"refresh_token,omitempty"`
+	ExpiresIn             int                     `json:"expires_in,omitempty"`
+	MFAEnabled            bool                    `json:"mfa_enabled"`
+	RefreshTokenExpiresIn int                     `json:"refresh_token_expires_in,omitempty"`
+	Status                string                  `json:"status,omitempty"`
+	Error                 *response.ErrorResponse `json:"error,omitempty"`
 }
 
 // Built-in preset mapping types