backend/implements/bunDatabase.go

package implements

import (
	"context"

	"gitea.konchin.com/go2025/backend/models"
	"gitea.konchin.com/go2025/backend/tracing"
	"gitea.konchin.com/go2025/backend/utils"
	"github.com/uptrace/bun"
	"go.uber.org/zap"
)

type BunDatabase struct {
	db *bun.DB
}

func NewBunDatabase(db *bun.DB) *BunDatabase {
	return &BunDatabase{db: db}
}

func (self *BunDatabase) GetSessionByLoginToken(
	ctx context.Context,
	loginToken string,
) (models.Session, error) {
	ret := models.Session{
		LoginToken: loginToken,
	}
	err := self.db.NewSelect().
		Model(&ret).
		Where("login_token = ?", loginToken).
		Scan(ctx)
	if err != nil {
		return models.Session{}, err
	}
	return ret, nil
}

func (self *BunDatabase) GetSessionByUserId(
	ctx context.Context,
	userId string,
) (models.Session, error) {
	ret := models.Session{
		UserId: userId,
	}
	err := self.db.NewSelect().
		Model(&ret).
		Where("user_id = ?", userId).
		Scan(ctx)
	if err != nil {
		return models.Session{}, err
	}
	return ret, nil
}

func (self *BunDatabase) UpdateRefreshToken(
	ctx context.Context,
	userId string,
) (models.Session, error) {
	ret := models.Session{
		UserId: userId,
	}

	err := self.db.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
		err := tx.NewSelect().
			Model(&ret).
			WherePK().
			Scan(ctx)
		if err != nil {
			return err
		}

		if err := ret.RotateRefreshToken(); err != nil {
			tracing.Logger.Ctx(ctx).
				Error("failed to rotate refresh token",
					zap.Error(err))
			return err
		}

		_, err = tx.NewUpdate().
			Model((*models.Session)(nil)).
			Set("refresh_token = ?", ret.RefreshToken).
			Where("user_id = ?", userId).
			Exec(ctx)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return models.Session{}, err
	}
	return ret, nil
}

func (self *BunDatabase) UpsertLoginToken(
	ctx context.Context,
	userId string,
) (string, error) {
	token, err := utils.RandomString(24)
	if err != nil {
		return "", err
	}

	session := models.Session{
		UserId:     userId,
		LoginToken: token,
		IsValid:    true,
	}
	_, err = self.db.NewInsert().
		Model(&session).
		On("CONFLICT (user_id) DO UPDATE").
		Set("login_token = EXCLUDED.login_token").
		Exec(ctx)
	if err != nil {
		return "", err
	}
	return token, nil
}

func (self *BunDatabase) GetAliases(
	ctx context.Context,
) ([]models.Alias, error) {
	var ret []models.Alias
	err := self.db.NewSelect().
		Model(&ret).
		Scan(ctx)
	if err != nil {
		return []models.Alias{}, err
	}
	return ret, nil
}

func (self *BunDatabase) GetImages(
	ctx context.Context,
	imageIds []int64,
	aliasIds []int64,
) ([]models.Image, error) {
	if len(aliasIds) > 0 {
		var rels []models.AliasImage
		err := self.db.NewSelect().
			Model(&rels).
			Where("alias_id IN (?)", bun.In(aliasIds)).
			Scan(ctx)
		if err != nil {
			return []models.Image{}, err
		}
		for _, rel := range rels {
			imageIds = append(imageIds, rel.ImageId)
		}
	}

	var ret []models.Image
	if len(imageIds) > 0 {
		err := self.db.NewSelect().
			Model(&ret).
			Where("id IN (?)", bun.In(imageIds)).
			Relation("Aliases").
			Scan(ctx)
		if err != nil {
			return []models.Image{}, err
		}
	}
	return ret, nil
}

func (self *BunDatabase) UpdateAliases(
	ctx context.Context,
	imageId int64,
	aliasNames []string,
) error {
	return self.db.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
		var aliases []models.Alias
		for _, ali := range aliasNames {
			aliases = append(aliases, models.Alias{
				Name: ali,
			})
		}
		_, err := tx.NewInsert().
			Model(&aliases).
			On("CONFLICT (name) DO NOTHING").
			Exec(ctx)
		if err != nil {
			return err
		}

		err = tx.NewSelect().
			Model(&aliases).
			Where("name IN (?)", bun.In(aliasNames)).
			Scan(ctx)
		if err != nil {
			return err
		}

		var rels []models.AliasImage
		for _, alias := range aliases {
			rels = append(rels, models.AliasImage{
				AliasId: alias.Id,
				ImageId: imageId,
			})
		}
		_, err = tx.NewInsert().
			Model(&rels).
			On(`CONFLICT ("alias_id", "image_id") DO NOTHING`).
			Exec(ctx)
		if err != nil {
			return err
		}

		return nil
	})
}

func (self *BunDatabase) InsertImage(
	ctx context.Context,
	image *models.Image,
) error {
	_, err := self.db.NewInsert().
		Model(image).
		Exec(ctx)
	return err
}

func (self *BunDatabase) DeleteImage(
	ctx context.Context,
	imageId int64,
) error {
	return self.db.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
		_, err := tx.NewDelete().
			Model((*models.Image)(nil)).
			Where("id = ?", imageId).
			Exec(ctx)
		if err != nil {
			return err
		}
		_, err = tx.NewDelete().
			Model((*models.AliasImage)(nil)).
			Where("image_id = ?", imageId).
			Exec(ctx)
		if err != nil {
			return err
		}
		return nil
	})
}

func (self *BunDatabase) DeleteAlias(
	ctx context.Context,
	aliasId int64,
) error {
	return self.db.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
		_, err := tx.NewDelete().
			Model((*models.Alias)(nil)).
			Where("id = ?", aliasId).
			Exec(ctx)
		if err != nil {
			return err
		}
		_, err = tx.NewDelete().
			Model((*models.AliasImage)(nil)).
			Where("alias_id = ?", aliasId).
			Exec(ctx)
		if err != nil {
			return err
		}
		return nil
	})
}