JCS-pub/client/internal/db/package.go

package db

import (
	"fmt"
	"time"

	"github.com/samber/lo"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
	"gorm.io/gorm"
)

type PackageDB struct {
	*DB
}

func (db *DB) Package() *PackageDB {
	return &PackageDB{DB: db}
}

func (db *PackageDB) GetByID(ctx SQLContext, packageID jcstypes.PackageID) (jcstypes.Package, error) {
	var ret jcstypes.Package
	err := ctx.Table("Package").Where("PackageID = ?", packageID).First(&ret).Error
	return ret, err
}

func (db *PackageDB) GetByName(ctx SQLContext, bucketID jcstypes.BucketID, name string) (jcstypes.Package, error) {
	var ret jcstypes.Package
	err := ctx.Table("Package").Where("BucketID = ? AND Name = ?", bucketID, name).First(&ret).Error
	return ret, err
}

func (db *PackageDB) GetDetail(ctx SQLContext, packageID jcstypes.PackageID) (jcstypes.PackageDetail, error) {
	var pkg jcstypes.Package
	err := ctx.Table("Package").Where("PackageID = ?", packageID).First(&pkg).Error
	if err != nil {
		return jcstypes.PackageDetail{}, err
	}

	var ret struct {
		ObjectCount int64
		TotalSize   int64
	}

	err = ctx.Table("Object").
		Select("COUNT(*) as ObjectCount, SUM(Size) as TotalSize").
		Where("PackageID = ?", packageID).
		Scan(&ret).
		Error
	if err != nil {
		return jcstypes.PackageDetail{}, err
	}

	return jcstypes.PackageDetail{
		Package:     pkg,
		ObjectCount: ret.ObjectCount,
		TotalSize:   ret.TotalSize,
	}, nil
}

func (db *PackageDB) BatchGetIDPaged(ctx SQLContext, lastPkgID jcstypes.PackageID, count int) ([]jcstypes.PackageID, error) {
	var ret []jcstypes.PackageID
	err := ctx.Table("Package").
		Select("PackageID").
		Where("PackageID > ?", lastPkgID).
		Order("PackageID ASC").
		Limit(count).
		Find(&ret).Error
	return ret, err
}

func (db *PackageDB) BatchGetDetailPaged(ctx SQLContext, lastPkgID jcstypes.PackageID, count int) ([]jcstypes.PackageDetail, error) {
	var pkgs []jcstypes.Package
	err := ctx.Table("Package").
		Where("PackageID > ?", lastPkgID).
		Order("PackageID ASC").
		Limit(count).
		Find(&pkgs).Error
	if err != nil {
		return nil, err
	}

	if len(pkgs) == 0 {
		return nil, nil
	}

	minPkgID := pkgs[0].PackageID
	maxPkgID := pkgs[len(pkgs)-1].PackageID

	type detail struct {
		ObjectCount int64
		TotalSize   int64
	}

	var details []detail
	err = ctx.
		Table("Package").
		Select("COUNT(ObjectID) as ObjectCount, SUM(Size) as TotalSize").
		Joins("left join Object o on Package.PackageID = o.PackageID").
		Where("Package.PackageID >= ? AND Package.PackageID <= ?", minPkgID, maxPkgID).
		Group("Package.PackageID").
		Order("Package.PackageID ASC").
		Find(&details).Error
	if err != nil {
		return nil, err
	}

	ret := make([]jcstypes.PackageDetail, len(pkgs))
	for i := range pkgs {
		ret[i] = jcstypes.PackageDetail{
			Package:     pkgs[i],
			ObjectCount: details[i].ObjectCount,
			TotalSize:   details[i].TotalSize,
		}
	}

	return ret, nil
}

func (db *PackageDB) BatchTestPackageID(ctx SQLContext, pkgIDs []jcstypes.PackageID) (map[jcstypes.PackageID]bool, error) {
	if len(pkgIDs) == 0 {
		return make(map[jcstypes.PackageID]bool), nil
	}

	var avaiIDs []jcstypes.PackageID
	err := ctx.Table("Package").
		Select("PackageID").
		Where("PackageID IN ?", pkgIDs).
		Find(&avaiIDs).Error
	if err != nil {
		return nil, err
	}

	avaiIDMap := make(map[jcstypes.PackageID]bool)
	for _, pkgID := range avaiIDs {
		avaiIDMap[pkgID] = true
	}

	return avaiIDMap, nil
}

func (*PackageDB) BatchGetAllPackageIDs(ctx SQLContext, start int, count int) ([]jcstypes.PackageID, error) {
	var ret []jcstypes.PackageID
	err := ctx.Table("Package").Select("PackageID").Limit(count).Offset(start).Find(&ret).Error
	return ret, err
}

func (db *PackageDB) GetBucketPackages(ctx SQLContext, bucketID jcstypes.BucketID) ([]jcstypes.Package, error) {
	var ret []jcstypes.Package
	err := ctx.Table("Package").
		Select("Package.*").
		Where("BucketID = ?", bucketID).
		Find(&ret).Error
	return ret, err
}

func (db *PackageDB) GetBucketPackagesByName(ctx SQLContext, bucketName string) ([]jcstypes.Package, error) {
	var ret []jcstypes.Package
	err := ctx.Table("Package").
		Select("Package.*").
		Joins("JOIN Bucket ON Package.BucketID = Bucket.BucketID").
		Where("Bucket.Name = ?", bucketName).
		Find(&ret).Error
	return ret, err
}

// 在指定名称的Bucket中查找指定名称的Package
func (*PackageDB) GetByFullName(ctx SQLContext, bucketName string, packageName string) (jcstypes.Package, error) {
	var ret jcstypes.Package
	err := ctx.Table("Package").
		Select("Package.*").
		Joins("JOIN Bucket ON Package.BucketID = Bucket.BucketID").
		Where("Package.Name = ? AND Bucket.Name = ?", packageName, bucketName).
		First(&ret).Error
	return ret, err
}

func (db *PackageDB) Create(ctx SQLContext, bucketID jcstypes.BucketID, name string, createTime time.Time) (jcstypes.Package, error) {
	var packageID int64
	err := ctx.Table("Package").
		Select("PackageID").
		Where("Name = ? AND BucketID = ?", name, bucketID).
		Scan(&packageID).Error

	if err != nil {
		return jcstypes.Package{}, err
	}
	if packageID != 0 {
		return jcstypes.Package{}, gorm.ErrDuplicatedKey
	}

	newPackage := jcstypes.Package{Name: name, BucketID: bucketID, CreateTime: createTime}
	if err := ctx.Create(&newPackage).Error; err != nil {
		return jcstypes.Package{}, fmt.Errorf("insert package failed, err: %w", err)
	}

	return newPackage, nil
}

func (*PackageDB) Delete(ctx SQLContext, packageID jcstypes.PackageID) error {
	err := ctx.Delete(&jcstypes.Package{}, "PackageID = ?", packageID).Error
	return err
}

// 删除与Package相关的所有数据
func (db *PackageDB) DeleteComplete(ctx SQLContext, packageID jcstypes.PackageID) error {
	if err := db.Package().Delete(ctx, packageID); err != nil {
		return fmt.Errorf("delete package state: %w", err)
	}

	if err := db.ObjectAccessStat().DeleteInPackage(ctx, packageID); err != nil {
		return fmt.Errorf("delete from object access stat: %w", err)
	}

	if err := db.ObjectBlock().DeleteInPackage(ctx, packageID); err != nil {
		return fmt.Errorf("delete from object block failed, err: %w", err)
	}

	if err := db.PinnedObject().DeleteInPackage(ctx, packageID); err != nil {
		return fmt.Errorf("deleting pinned objects in package: %w", err)
	}

	if err := db.Object().DeleteInPackage(ctx, packageID); err != nil {
		return fmt.Errorf("deleting objects in package: %w", err)
	}

	if err := db.PackageAccessStat().DeleteByPackageID(ctx, packageID); err != nil {
		return fmt.Errorf("deleting package access stat: %w", err)
	}

	return nil
}

func (*PackageDB) ChangeState(ctx SQLContext, packageID jcstypes.PackageID, state string) error {
	err := ctx.Exec("UPDATE Package SET State = ? WHERE PackageID = ?", state, packageID).Error
	return err
}

// 返回ErrRecordNotFound表示没有找到指定名称的Bucket，nil表示找到了
func (*PackageDB) HasPackageIn(ctx SQLContext, bucketID jcstypes.BucketID) error {
	var pkg jcstypes.Package
	return ctx.Table("Package").Where("BucketID = ?", bucketID).First(&pkg).Error
}

func (*PackageDB) Move(ctx SQLContext, packageID jcstypes.PackageID, newBktID jcstypes.BucketID, newName string) error {
	err := ctx.Table("Package").Where("PackageID = ?", packageID).Update("BucketID", newBktID).Update("Name", newName).Error
	return err
}

type AddAccessStatEntry struct {
	ObjectID    jcstypes.ObjectID    `json:"objectID"`
	PackageID   jcstypes.PackageID   `json:"packageID"`
	UserSpaceID jcstypes.UserSpaceID `json:"userSpaceID"`
	Counter     float64              `json:"counter"`
}

func (db *PackageDB) BatchAddPackageAccessStat(ctx SQLContext, entries []AddAccessStatEntry) error {
	pkgIDs := make(map[jcstypes.PackageID]bool)
	objIDs := make(map[jcstypes.ObjectID]bool)
	for _, e := range entries {
		pkgIDs[e.PackageID] = true
		objIDs[e.ObjectID] = true
	}

	avaiPkgIDs, err := db.Package().BatchTestPackageID(ctx, lo.Keys(pkgIDs))
	if err != nil {
		return fmt.Errorf("batch test package id: %w", err)
	}

	avaiObjIDs, err := db.Object().BatchTestObjectID(ctx, lo.Keys(objIDs))
	if err != nil {
		return fmt.Errorf("batch test object id: %w", err)
	}

	var willAdds []AddAccessStatEntry
	for _, e := range entries {
		if avaiPkgIDs[e.PackageID] && avaiObjIDs[e.ObjectID] {
			willAdds = append(willAdds, e)
		}
	}

	if len(willAdds) > 0 {
		err := db.PackageAccessStat().BatchAddCounter(ctx, willAdds)
		if err != nil {
			return fmt.Errorf("batch add package access stat counter: %w", err)
		}

		err = db.ObjectAccessStat().BatchAddCounter(ctx, willAdds)
		if err != nil {
			return fmt.Errorf("batch add object access stat counter: %w", err)
		}
	}

	return nil
}

// 尝试创建指定名称的Bucket和Package，如果Bucket不存在，则创建Bucket，如果Package已存在，则直接返回已有的Package
func (db *PackageDB) TryCreateAll(ctx SQLContext, bktName string, pkgName string) (jcstypes.Package, error) {
	bkt, err := db.Bucket().GetByName(ctx, bktName)
	if err == gorm.ErrRecordNotFound {
		bkt, err = db.Bucket().Create(ctx, bktName, time.Now())
		if err != nil {
			return jcstypes.Package{}, fmt.Errorf("create bucket: %w", err)
		}
	} else if err != nil {
		return jcstypes.Package{}, fmt.Errorf("get bucket by name: %w", err)
	}

	pkg, err := db.GetByName(ctx, bkt.BucketID, pkgName)
	if err == nil {
		return pkg, nil
	}
	if err != gorm.ErrRecordNotFound {
		return jcstypes.Package{}, fmt.Errorf("get package by name: %w", err)
	}

	pkg, err = db.Create(ctx, bkt.BucketID, pkgName, time.Now())
	if err != nil {
		return jcstypes.Package{}, fmt.Errorf("create package: %w", err)
	}

	return pkg, nil
}

func (db *PackageDB) SetPinned(ctx SQLContext, packageID jcstypes.PackageID, pinned bool) error {
	err := ctx.Table("Package").Where("PackageID = ?", packageID).Update("Pinned", pinned).Error
	return err
}