package models import ( "errors" "time" "gorm.io/gorm" ) // Friend table. type Friend struct { ID uint64 `gorm:"primaryKey"` SourceUserID uint64 `gorm:"index"` TargetUserID uint64 `gorm:"index"` Approved bool `gorm:"index"` CreatedAt time.Time UpdatedAt time.Time } // AddFriend sends a friend request or accepts one if there was already a pending one. func AddFriend(sourceUserID, targetUserID uint64) error { // Did we already send a friend request? f := &Friend{} forward := DB.Where( "source_user_id = ? AND target_user_id = ?", sourceUserID, targetUserID, ).First(&f).Error // Is there a reverse friend request pending? rev := &Friend{} reverse := DB.Where( "source_user_id = ? AND target_user_id = ?", targetUserID, sourceUserID, ).First(&rev).Error // If the reverse exists (requested us) but not the forward, this completes the friendship. if reverse == nil && forward != nil { // Approve the reverse. rev.Approved = true rev.Save() // Add the matching forward. f = &Friend{ SourceUserID: sourceUserID, TargetUserID: targetUserID, Approved: true, } return DB.Create(f).Error } // If the forward already existed, error. if forward == nil { if f.Approved { return errors.New("you are already friends") } return errors.New("a friend request had already been sent") } // Create the pending forward request. f = &Friend{ SourceUserID: sourceUserID, TargetUserID: targetUserID, Approved: false, } return DB.Create(f).Error } // AreFriends quickly checks if two user IDs are friends. func AreFriends(sourceUserID, targetUserID uint64) bool { f := &Friend{} DB.Where( "source_user_id = ? AND target_user_id = ?", sourceUserID, targetUserID, ).First(&f) return f.Approved } // FriendStatus returns an indicator of friendship status: "none", "pending", "approved" func FriendStatus(sourceUserID, targetUserID uint64) string { f := &Friend{} result := DB.Where( "source_user_id = ? AND target_user_id = ?", sourceUserID, targetUserID, ).First(&f) if result.Error == nil { if f.Approved { return "approved" } return "pending" } return "none" } // CountFriendRequests gets a count of pending requests for the user. func CountFriendRequests(userID uint64) (int64, error) { var count int64 result := DB.Where( "target_user_id = ? AND approved = ?", userID, false, ).Model(&Friend{}).Count(&count) return count, result.Error } // PaginateFriends gets a page of friends (or pending friend requests) as User objects ordered // by friendship date. func PaginateFriends(userID uint64, requests bool, pager *Pagination) ([]*User, error) { // We paginate over the Friend table. var ( fs = []*Friend{} userIDs = []uint64{} query *gorm.DB ) if requests { query = DB.Where( "target_user_id = ? AND approved = ?", userID, false, ) } else { query = DB.Where( "source_user_id = ? AND approved = ?", userID, true, ) } query = query.Order(pager.Sort) query.Model(&Friend{}).Count(&pager.Total) result := query.Offset(pager.GetOffset()).Limit(pager.PerPage).Find(&fs) if result.Error != nil { return nil, result.Error } // Now of these friends get their User objects. for _, friend := range fs { if requests { userIDs = append(userIDs, friend.SourceUserID) } else { userIDs = append(userIDs, friend.TargetUserID) } } return GetUsers(userIDs) } // GetFriendRequests returns all pending friend requests for a user. func GetFriendRequests(userID uint64) ([]*Friend, error) { var fs = []*Friend{} result := DB.Where( "target_user_id = ? AND approved = ?", userID, false, ).Find(&fs) return fs, result.Error } // RemoveFriend severs a friend connection both directions, used when // rejecting a request or removing a friend. func RemoveFriend(sourceUserID, targetUserID uint64) error { result := DB.Where( "(source_user_id = ? AND target_user_id = ?) OR "+ "(target_user_id = ? AND source_user_id = ?)", sourceUserID, targetUserID, sourceUserID, targetUserID, ).Delete(&Friend{}) return result.Error } // Save photo. func (f *Friend) Save() error { result := DB.Save(f) return result.Error }