docxTemplater/worker/tariff.go

package worker

import (
	"context"
	"time"

	"go.uber.org/zap"
	"penahub.gitlab.yandexcloud.net/backend/templategen/broker/tariff"
	"penahub.gitlab.yandexcloud.net/backend/templategen/broker/tariff/models"
	"penahub.gitlab.yandexcloud.net/backend/templategen/dal"
	"penahub.gitlab.yandexcloud.net/backend/templategen/penadisk"
	"penahub.gitlab.yandexcloud.net/backend/templategen/privileges"
)

/* TODO: Сделать рефакторинг кода!
Какой то колхоз с привилегиями получился.
Сейчас у меня 3 разных модели привилегий из разных пакетов: boker\tariff\models, proto, privileges.
*/

type TariffWorkerDeps struct {
	Logger           *zap.Logger
	Dal              *dal.MongoDAL
	TariffConsumer   *tariff.Consumer
	PrivilegesDomain string
}

type TariffWorker struct {
	logger           *zap.Logger
	dal              *dal.MongoDAL
	tariffConsumer   *tariff.Consumer
	privilegesDomain string
}

func NewTariffWorker(deps TariffWorkerDeps) *TariffWorker {
	return &TariffWorker{
		logger:           deps.Logger,
		dal:              deps.Dal,
		tariffConsumer:   deps.TariffConsumer,
		privilegesDomain: deps.PrivilegesDomain,
	}
}

func (w *TariffWorker) Run(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Minute)

	for {
		select {
		case <-ctx.Done():
			w.logger.Info("tariff worker done")
			return
		case <-ticker.C:
			w.Do(ctx)
		}
	}
}

func (w *TariffWorker) Do(ctx context.Context) {
	// получаем актуальные привилегии
	getActualPrivileges, err := privileges.GetActualPrivileges(ctx, w.privilegesDomain)

	if err != nil {
		w.logger.Error("cannot get actual privileges")
		return
	}

	actualPrivileges := make(map[string]models.Privilege)

	for _, privilege := range getActualPrivileges {
		actualPrivileges[privilege.ID] = models.Privilege{
			ID:          privilege.ID,
			Amount:      0,
			PrivilegeID: privilege.PrivilegeID,
			Name:        privilege.Name,
			ServiceKey:  privilege.ServiceKey,
			Description: privilege.Description,
			Type:        privilege.Type,
			Value:       privilege.Value,
			Price:       privilege.Price,
			IsDeleted:   privilege.IsDeleted,
			CreatedAt:   privilege.CreatedAt,
			UpdatedAt:   privilege.UpdatedAt,
			DeletedAt:   privilege.DeletedAt,
		}
	}

	// проверяем новые тарифы
	tariffs := w.tariffConsumer.FetchTariffs(ctx)

	for _, tariff := range tariffs {
		penadisk := penadisk.NewClient(tariff.UserID)

		for _, newPrivilege := range tariff.Privileges {
			if newPrivilege.ServiceKey != models.ServiceKey {
				continue
			}

			privilege := actualPrivileges[newPrivilege.ID]
			privilege.Amount = newPrivilege.Amount

			if err = w.dal.Amo.AddPrivilege(ctx,
				tariff.UserID,
				privilege.PrivilegeID,
				privilege.Amount); err != nil {
				w.logger.Error("cannot add privilege",
					zap.Error(err),
					zap.String("user_id", tariff.UserID),
					zap.String("privilege_id", privilege.PrivilegeID))
			}

			// передаем в penahub disk запрос на изменение размера диска пользователя
			if privilege.PrivilegeID == models.PrivilegeTemplateStorage {
				if err = penadisk.SetBucketQuota(ctx, uint64(privilege.Amount)<<20); err != nil {
					w.logger.Error("cannot set bucket quota",
						zap.Error(err),
						zap.String("user_id", tariff.UserID),
						zap.String("privilege_id", privilege.PrivilegeID))
				}
			}
		}
	}

	// получаем пользователей
	amos, err := w.dal.Amo.GetAll(ctx)

	if err != nil {
		w.logger.Error("cannot get all amos", zap.Error(err))
		return
	}

	for _, amo := range amos {
		penadisk := penadisk.NewClient(amo.PenaID)

		for privilegeID, privilege := range amo.Privileges {
			// проверяем истекшие привилегии
			if privilegeID == models.PrivilegeTemplateUnlimTime &&
				privilege.CreatedAt.AddDate(0, 0, int(privilege.Amount)).After(time.Now()) {
				if err = w.dal.Amo.DeletePrivilege(ctx, amo.ID, privilegeID); err != nil {
					w.logger.Error("cannot delete privilege",
						zap.Error(err),
						zap.String("amo_id", amo.ID),
						zap.String("privilege_id", privilegeID))
				}
			}

			// обновляем базовую привилегию на количество, если она закончилась и прошел месяц
			if privilegeID == models.PrivilegeTemplateCount &&
				privilege.Amount == 0 &&
				privilege.CreatedAt.AddDate(0, 1, 0).Before(time.Now()) {
				if err = w.dal.Amo.AddPrivilege(ctx,
					amo.ID,
					privilegeID,
					models.BasicAmountPrivilegeTemplateCount); err != nil {
					w.logger.Error("cannot add privilege",
						zap.Error(err),
						zap.String("amo_ID", amo.ID),
						zap.String("privilege_id", privilegeID))
				}
			}
		}

		// добавляем базовые привилегии если у пользователя нет никаких
		if len(amo.Privileges) == 0 {
			// привилегия на количество генераций
			if err = w.dal.Amo.AddPrivilege(ctx,
				amo.ID,
				models.PrivilegeTemplateCount,
				models.BasicAmountPrivilegeTemplateCount); err != nil {
				w.logger.Error("cannot add privilege",
					zap.Error(err),
					zap.String("amo_ID", amo.ID),
					zap.String("privilege_id", models.PrivilegeTemplateCount))
			}

			// привилегия на объем диска
			if err = w.dal.Amo.AddPrivilege(ctx,
				amo.ID,
				models.PrivilegeTemplateStorage,
				models.BasicAmountPrivilegeTemplateStorage); err != nil {
				w.logger.Error("cannot add privilege",
					zap.Error(err),
					zap.String("amo_ID", amo.ID),
					zap.String("privilege_id", models.PrivilegeTemplateStorage))
			}

			if err := penadisk.SetBucketQuota(ctx, 100<<20); err != nil {
				w.logger.Error("cannot set bucket quota",
					zap.Error(err),
					zap.String("user_id", amo.PenaID),
					zap.String("privilege_id", models.PrivilegeTemplateStorage))
			}
		}
	}
}