daddy/internal/model/workgroup_repository.go

package model

import (
	"context"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/pkg/errors"
	errs "github.com/pkg/errors"
)

type WorkgroupRepository struct {
	db *gorm.DB
}

func (r *WorkgroupRepository) FindWorkgroups(ctx context.Context, criteria ...interface{}) ([]*Workgroup, error) {
	workgroups := make([]*Workgroup, 0)
	if err := r.db.Model(&Workgroup{}).Preload("Members").Find(&workgroups, criteria...).Error; err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroups, nil
}

func (r *WorkgroupRepository) UpdateWorkgroup(ctx context.Context, workgroupID uint, changes WorkgroupChanges) (*Workgroup, error) {
	workgroup := &Workgroup{
		Name: changes.Name,
	}
	workgroup.ID = workgroupID

	err := r.db.Model(workgroup).
		Update(workgroup).
		Error
	if err != nil {
		return nil, errors.WithStack(err)
	}

	err = r.db.Model(workgroup).Preload("Members").First(workgroup, "id = ?", workgroupID).Error
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroup, nil
}

func (r *WorkgroupRepository) CreateWorkgroup(ctx context.Context, changes WorkgroupChanges) (*Workgroup, error) {
	workgroup := &Workgroup{
		Name: changes.Name,
	}

	if err := r.db.Model(&Workgroup{}).Preload("Members").Create(workgroup).Error; err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroup, nil
}

func (r *WorkgroupRepository) CloseWorkgroup(ctx context.Context, workgroupID uint) (*Workgroup, error) {
	workgroup := &Workgroup{}

	err := r.db.Model(workgroup).
		Where("id = ?", workgroupID).
		UpdateColumn("closedAt", time.Now()).
		Error
	if err != nil {
		return nil, errors.WithStack(err)
	}

	err = r.db.Model(workgroup).Preload("Members").First(workgroup, "id = ?", workgroupID).Error
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroup, nil
}

func (r *WorkgroupRepository) AddUserToWorkgroup(ctx context.Context, userID, workgroupID uint) (*Workgroup, error) {
	user := &User{}

	err := r.db.Model(user).Preload("Workgroups").First(user, "id = ?", userID).Error
	if err != nil {
		return nil, errors.Wrap(err, "could not find user")
	}

	workgroup := &Workgroup{}
	workgroup.ID = workgroupID

	err = r.db.Model(user).
		Association("Workgroups").
		Append(workgroup).
		Error

	if err != nil {
		return nil, errors.Wrap(err, "could not add user to workgroup")
	}

	err = r.db.Model(workgroup).
		Preload("Members").
		First(workgroup, "id = ?", workgroupID).
		Error
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroup, nil
}

func (r *WorkgroupRepository) RemoveUserFromWorkgroup(ctx context.Context, userID, workgroupID uint) (*Workgroup, error) {
	workgroup := &Workgroup{}
	workgroup.ID = workgroupID

	err := r.db.Transaction(func(tx *gorm.DB) error {
		user := &User{}
		err := tx.First(user, "id = ?", userID).Error
		if err != nil {
			return errors.Wrap(err, "could not find user")
		}

		err = tx.Model(user).
			Association("Workgroups").
			Delete(workgroup).
			Error

		if err != nil {
			return errors.Wrap(err, "could not add user to workgroup")
		}

		err = tx.Model(workgroup).
			Preload("Members").
			First(workgroup, "id = ?", workgroupID).
			Error
		if err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return workgroup, nil
}

func (r *WorkgroupRepository) Find(ctx context.Context, id string) (*Workgroup, error) {
	wg := &Workgroup{}
	query := r.db.Model(wg).Preload("Members").Where("id = ?", id)

	if err := query.First(&wg).Error; err != nil {
		return nil, errs.WithStack(err)
	}

	return wg, nil
}

func NewWorkgroupRepository(db *gorm.DB) *WorkgroupRepository {
	return &WorkgroupRepository{db}
}