website/pkg/models/friend.go

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
}