twitch-bot-streak/pkg/database/query.go

package database

import (
	"errors"
	"fmt"
	"time"

	"gorm.io/gorm"
)

var defaultMetaTime = time.Date(1970, 1, 1, 0, 0, 0, 0, time.UTC)

func (d DB) CountStreak(twitchID uint64, username string) (user StreakUser, err error) {
	if err = d.db.Transaction(func(tx *gorm.DB) (err error) {
		if err = tx.First(&user, "twitch_id = ?", twitchID).Error; err != nil {
			if !errors.Is(err, gorm.ErrRecordNotFound) {
				return fmt.Errorf("getting user: %w", err)
			}

			// User was not yet inserted
			user = StreakUser{
				TwitchID:      twitchID,
				Username:      username,
				StreamsCount:  0,
				CurrentStreak: 0,
				MaxStreak:     0,
				StreakStatus:  StatusBroken,
			}
		}

		switch user.StreakStatus {
		case StatusActive:
			// User has an active streak, do nothing
			return nil

		case StatusBroken:
			// User needs a new streak
			user.CurrentStreak = 1

		case StatusPending:
			// User can prolong their streak
			user.CurrentStreak += 1
		}

		// In any case set the streak active and count the current stream
		user.StreamsCount++
		user.StreakStatus = StatusActive
		if user.CurrentStreak > user.MaxStreak {
			user.MaxStreak = user.CurrentStreak
		}

		if err = tx.Save(&user).Error; err != nil {
			return fmt.Errorf("saving user: %w", err)
		}

		return nil
	}); err != nil {
		return user, fmt.Errorf("counting streak for user: %w", err)
	}

	return user, nil
}

func (d DB) SetStreamOffline() (err error) {
	return d.storeTimeToMeta(d.db, "stream_offline", time.Now())
}

func (d DB) StartStream(streamOfflineGrace time.Duration) (err error) {
	if err = d.db.Transaction(func(tx *gorm.DB) (err error) {
		lastOffline, err := d.getTimeFromMeta(tx, "stream_offline")
		if err != nil {
			return fmt.Errorf("getting offline time: %w", err)
		}

		lastOnline, err := d.getTimeFromMeta(tx, "stream_online")
		if err != nil {
			return fmt.Errorf("getting online time: %w", err)
		}

		if err = d.storeTimeToMeta(tx, "stream_online", time.Now()); err != nil {
			return fmt.Errorf("storing stream start: %w", err)
		}

		if time.Since(lastOffline) < streamOfflineGrace || lastOnline.After(lastOffline) {
			// We only had a short break or the stream was already started
			return nil
		}

		if err = tx.Model(&StreakUser{}).
			Where("streak_status = ?", StatusPending).
			Update("streak_status", StatusBroken).
			Error; err != nil {
			return fmt.Errorf("breaking streaks for pending users: %w", err)
		}

		if err = tx.Model(&StreakUser{}).
			Where("streak_status = ?", StatusActive).
			Update("streak_status", StatusPending).
			Error; err != nil {
			return fmt.Errorf("breaking streaks for pending users: %w", err)
		}

		return nil
	}); err != nil {
		return fmt.Errorf("starting stream: %w", err)
	}

	return nil
}

func (d DB) getTimeFromMeta(db *gorm.DB, key string) (t time.Time, err error) {
	var meta StreakMeta
	if err = db.First(&meta, "name = ?", key).Error; err != nil {
		if !errors.Is(err, gorm.ErrRecordNotFound) {
			return t, fmt.Errorf("getting last %s time: %w", key, err)
		}

		meta.Value = defaultMetaTime.Format(time.RFC3339Nano)
	}

	t, err = time.Parse(time.RFC3339Nano, meta.Value)
	if err != nil {
		return t, fmt.Errorf("parsing %s time: %w", key, err)
	}

	return t, nil
}

func (d DB) storeTimeToMeta(db *gorm.DB, key string, t time.Time) (err error) {
	if err = db.Save(&StreakMeta{
		Name:  key,
		Value: t.Format(time.RFC3339Nano),
	}).Error; err != nil {
		return fmt.Errorf("updating stream meta: %w", err)
	}

	return nil
}