tamabox/internal/db/upload_images.go

// Copyright 2023 E99p1ant. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package db

import (
	"context"

	"github.com/pkg/errors"
	"github.com/samber/lo"
	"gorm.io/gorm"
)

var UploadImages UploadImagesStore

var _ UploadImagesStore = (*uploadImages)(nil)

type UploadImagesStore interface {
	Create(ctx context.Context, opts CreateUploadImageOptions) (*UploadImage, error)
	GetByQuestionID(ctx context.Context, questionID uint) ([]*UploadImage, error)
	// GetByQuestionIDs 批量获取多个提问的图片，按 question_id 分组返回。
	// 用于账号导出等场景，将 N 次单题查询合并为少量 IN 查询，避免 N+1。
	GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error)
	GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error)
	// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录（软删除，与提问保持一致）。
	DeleteByQuestionID(ctx context.Context, questionID uint) error
}

func NewUploadImagesStore(db *gorm.DB) UploadImagesStore {
	return &uploadImages{db}
}

type uploadImages struct {
	*gorm.DB
}

type UploadImage struct {
	gorm.Model
	UploaderUserID uint
	Name           string
	FileSize       int64
	Md5            string
	Key            string
}

type UploadImageQuestionType string

const (
	UploadImageQuestionTypeAsk    UploadImageQuestionType = "ask"
	UploadImageQuestionTypeAnswer UploadImageQuestionType = "answer"
)

type UploadImageQuestion struct {
	Type          UploadImageQuestionType
	UploadImageID uint `gorm:"primaryKey"`
	QuestionID    uint `gorm:"primaryKey"`
}

type CreateUploadImageOptions struct {
	Type               UploadImageQuestionType
	QuestionID         uint
	UploaderUserID     uint
	Name               string
	FileSize           int64
	Md5                string
	Key                string
	IsDeletingPrevious bool
}

func (db *uploadImages) Create(ctx context.Context, opts CreateUploadImageOptions) (*UploadImage, error) {
	image := &UploadImage{
		UploaderUserID: opts.UploaderUserID,
		Name:           opts.Name,
		FileSize:       opts.FileSize,
		Md5:            opts.Md5,
		Key:            opts.Key,
	}
	if err := db.Transaction(func(tx *gorm.DB) error {
		if err := tx.WithContext(ctx).Create(image).Error; err != nil {
			return errors.Wrap(err, "create image")
		}

		if opts.IsDeletingPrevious {
			if err := tx.WithContext(ctx).Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID).Delete(&UploadImageQuestion{}).Error; err != nil {
				return errors.Wrap(err, "delete previous image question link")
			}
		}
		if err := tx.WithContext(ctx).Create(&UploadImageQuestion{
			Type:          opts.Type,
			UploadImageID: image.ID,
			QuestionID:    opts.QuestionID,
		}).Error; err != nil {
			return errors.Wrap(err, "create image question link")
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return image, nil
}

func (db *uploadImages) getBy(ctx context.Context, where string, args ...interface{}) ([]*UploadImage, error) {
	var uploadImageQuestions []*UploadImageQuestion
	if err := db.WithContext(ctx).Model(&UploadImageQuestion{}).Where(where, args...).Find(&uploadImageQuestions).Error; err != nil {
		return nil, errors.Wrap(err, "get image questions")
	}

	uploadImageIDs := lo.Map(uploadImageQuestions, func(item *UploadImageQuestion, _ int) uint {
		return item.UploadImageID
	})

	var uploadImages []*UploadImage
	if err := db.WithContext(ctx).Model(&UploadImage{}).Where("id IN ?", uploadImageIDs).Order("id ASC").Find(&uploadImages).Error; err != nil {
		return nil, errors.Wrap(err, "get images")
	}
	return uploadImages, nil
}

func (db *uploadImages) GetByQuestionID(ctx context.Context, questionID uint) ([]*UploadImage, error) {
	return db.getBy(ctx, "question_id = ?", questionID)
}

// GetByQuestionIDs 批量获取多个提问的图片，按 question_id 分组返回（解决导出 N+1）。
// 实现：先按 IN 批量取关联行，再按 IN 批量取图片主体，最后在内存中按 question_id 归组，
// 将每个提问原本的 2 次查询降为整体 2 次查询。
func (db *uploadImages) GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error) {
	result := make(map[uint][]*UploadImage)
	if len(questionIDs) == 0 {
		return result, nil
	}

	var links []*UploadImageQuestion
	if err := db.WithContext(ctx).Model(&UploadImageQuestion{}).Where("question_id IN ?", questionIDs).Find(&links).Error; err != nil {
		return nil, errors.Wrap(err, "get image questions batch")
	}
	if len(links) == 0 {
		return result, nil
	}

	imageIDs := lo.Map(links, func(item *UploadImageQuestion, _ int) uint {
		return item.UploadImageID
	})
	var images []*UploadImage
	if err := db.WithContext(ctx).Model(&UploadImage{}).Where("id IN ?", imageIDs).Order("id ASC").Find(&images).Error; err != nil {
		return nil, errors.Wrap(err, "get images batch")
	}
	imageByID := lo.SliceToMap(images, func(img *UploadImage) (uint, *UploadImage) {
		return img.ID, img
	})

	for _, link := range links {
		if img, ok := imageByID[link.UploadImageID]; ok {
			result[link.QuestionID] = append(result[link.QuestionID], img)
		}
	}
	return result, nil
}

func (db *uploadImages) GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error) {
	return db.getBy(ctx, "type = ? AND question_id = ?", typ, questionID)
}

// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录（软删除，与提问保持一致）。
// 注意：question_id 列在关联表 upload_image_questions 上，不在 upload_images 表上，
// 需先查关联行拿到 imageID 列表，再软删图片主体并删除关联行（同一事务内完成）。
// 注：仅清理数据库记录，对象存储中的实际文件不在此处删除（避免误删复用资源）。
func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint) error {
	return db.Transaction(func(tx *gorm.DB) error {
		var links []*UploadImageQuestion
		if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Find(&links).Error; err != nil {
			return errors.Wrap(err, "get image questions for delete")
		}
		if len(links) == 0 {
			return nil
		}

		imageIDs := lo.Map(links, func(item *UploadImageQuestion, _ int) uint {
			return item.UploadImageID
		})
		if err := tx.WithContext(ctx).Where("id IN ?", imageIDs).Delete(&UploadImage{}).Error; err != nil {
			return errors.Wrap(err, "delete upload images by question id")
		}
		if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Delete(&UploadImageQuestion{}).Error; err != nil {
			return errors.Wrap(err, "delete image question links")
		}
		return nil
	})
}