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

package db

import (
	"fmt"
	"strings"
	"time"

	"github.com/samber/lo"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"gitlink.org.cn/cloudream/storage2/client/types"
)

type ObjectDB struct {
	*DB
}

func (db *DB) Object() *ObjectDB {
	return &ObjectDB{DB: db}
}

func (db *ObjectDB) GetByID(ctx SQLContext, objectID types.ObjectID) (types.Object, error) {
	var ret types.Object
	err := ctx.Table("Object").Where("ObjectID = ?", objectID).First(&ret).Error
	return ret, err
}

func (db *ObjectDB) GetByPath(ctx SQLContext, packageID types.PackageID, path string) (types.Object, error) {
	var ret types.Object
	err := ctx.Table("Object").Where("PackageID = ? AND Path = ?", packageID, path).First(&ret).Error
	return ret, err
}

func (db *ObjectDB) GetByFullPath(ctx SQLContext, bktName string, pkgName string, path string) (types.Object, error) {
	var ret types.Object
	err := ctx.Table("Object").
		Joins("join Package on Package.PackageID = Object.PackageID and Package.Name = ?", pkgName).
		Joins("join Bucket on Bucket.BucketID = Package.BucketID and Bucket.Name = ?", bktName).
		Where("Object.Path = ?", path).First(&ret).Error
	return ret, err
}

func (db *ObjectDB) GetWithPathPrefix(ctx SQLContext, packageID types.PackageID, pathPrefix string) ([]types.Object, error) {
	var ret []types.Object
	err := ctx.Table("Object").Where("PackageID = ? AND Path LIKE ?", packageID, escapeLike("", "%", pathPrefix)).Order("ObjectID ASC").Find(&ret).Error
	return ret, err
}

// 查询结果将按照Path升序，而不是ObjectID升序
func (db *ObjectDB) GetWithPathPrefixPaged(ctx SQLContext, packageID types.PackageID, pathPrefix string, startPath string, limit int) ([]types.Object, error) {
	var ret []types.Object
	err := ctx.Table("Object").Where("PackageID = ? AND Path > ? AND Path LIKE ?", packageID, startPath, pathPrefix+"%").Order("Path ASC").Limit(limit).Find(&ret).Error
	return ret, err
}

func (db *ObjectDB) GetByPrefixGrouped(ctx SQLContext, packageID types.PackageID, pathPrefix string) (objs []types.Object, commonPrefixes []string, err error) {
	type ObjectOrDir struct {
		types.Object
		IsObject bool   `gorm:"IsObject"`
		Prefix   string `gorm:"Prefix"`
	}

	sepCnt := strings.Count(pathPrefix, types.ObjectPathSeparator) + 1

	prefixStatm := fmt.Sprintf("Substring_Index(Path, '%s', %d)", types.ObjectPathSeparator, sepCnt)
	grouping := ctx.Table("Object").
		Select(fmt.Sprintf("%s as Prefix, Max(ObjectID) as ObjectID, %s = Path as IsObject", prefixStatm, prefixStatm)).
		Where("PackageID = ?", packageID).
		Where("Path like ?", pathPrefix+"%").
		Group("Prefix, IsObject").
		Order("Prefix ASC")

	var ret []ObjectOrDir
	err = ctx.Table("Object").
		Select("Grouped.IsObject, Grouped.Prefix, Object.*").
		Joins("right join (?) as Grouped on Object.ObjectID = Grouped.ObjectID and Grouped.IsObject = 1", grouping).
		Find(&ret).Error
	if err != nil {
		return
	}

	for _, o := range ret {
		if o.IsObject {
			objs = append(objs, o.Object)
		} else {
			commonPrefixes = append(commonPrefixes, o.Prefix+types.ObjectPathSeparator)
		}
	}

	return
}

func (db *ObjectDB) GetByPrefixGroupedPaged(ctx SQLContext, packageID types.PackageID, pathPrefix string, startPath string, limit int) (objs []types.Object, commonPrefixes []string, nextStartPath string, err error) {
	type ObjectOrDir struct {
		types.Object
		IsObject bool   `gorm:"IsObject"`
		Prefix   string `gorm:"Prefix"`
	}

	sepCnt := strings.Count(pathPrefix, types.ObjectPathSeparator) + 1

	prefixStatm := fmt.Sprintf("Substring_Index(Path, '%s', %d)", types.ObjectPathSeparator, sepCnt)
	grouping := ctx.Table("Object").
		Select(fmt.Sprintf("%s as Prefix, Max(ObjectID) as ObjectID, %s = Path as IsObject", prefixStatm, prefixStatm)).
		Where("PackageID = ?", packageID).
		Where("Path like ?", pathPrefix+"%").
		Group("Prefix, IsObject").
		Having("Prefix > ?", startPath).
		Limit(limit).
		Order("Prefix ASC")

	var ret []ObjectOrDir
	err = ctx.Table("Object").
		Select("Grouped.IsObject, Grouped.Prefix, Object.*").
		Joins("right join (?) as Grouped on Object.ObjectID = Grouped.ObjectID and Grouped.IsObject = 1", grouping).
		Find(&ret).Error
	if err != nil {
		return
	}

	for _, o := range ret {
		if o.IsObject {
			objs = append(objs, o.Object)
		} else {
			commonPrefixes = append(commonPrefixes, o.Prefix+types.ObjectPathSeparator)
		}
		nextStartPath = o.Prefix
	}

	return
}

func (db *ObjectDB) HasObjectWithPrefix(ctx SQLContext, packageID types.PackageID, pathPrefix string) (bool, error) {
	var obj types.Object
	err := ctx.Table("Object").Where("PackageID = ? AND Path LIKE ?", packageID, escapeLike("", "%", pathPrefix)).First(&obj).Error
	if err == nil {
		return true, nil
	}

	if err == gorm.ErrRecordNotFound {
		return false, nil
	}

	return false, err
}

func (db *ObjectDB) BatchTestObjectID(ctx SQLContext, objectIDs []types.ObjectID) (map[types.ObjectID]bool, error) {
	if len(objectIDs) == 0 {
		return make(map[types.ObjectID]bool), nil
	}

	var avaiIDs []types.ObjectID
	err := ctx.Table("Object").Where("ObjectID IN ?", objectIDs).Pluck("ObjectID", &avaiIDs).Error
	if err != nil {
		return nil, err
	}

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

	return avaiIDMap, nil
}

func (db *ObjectDB) BatchGet(ctx SQLContext, objectIDs []types.ObjectID) ([]types.Object, error) {
	if len(objectIDs) == 0 {
		return nil, nil
	}

	var objs []types.Object
	err := ctx.Table("Object").Where("ObjectID IN ?", objectIDs).Order("ObjectID ASC").Find(&objs).Error
	if err != nil {
		return nil, err
	}

	return objs, nil
}

func (db *ObjectDB) BatchGetByPackagePath(ctx SQLContext, pkgID types.PackageID, pathes []string) ([]types.Object, error) {
	if len(pathes) == 0 {
		return nil, nil
	}

	var objs []types.Object
	err := ctx.Table("Object").Where("PackageID = ? AND Path IN ?", pkgID, pathes).Find(&objs).Error
	if err != nil {
		return nil, err
	}

	return objs, nil
}

func (db *ObjectDB) GetDetail(ctx SQLContext, objectID types.ObjectID) (types.ObjectDetail, error) {
	var obj types.Object
	err := ctx.Table("Object").Where("ObjectID = ?", objectID).First(&obj).Error
	if err != nil {
		return types.ObjectDetail{}, fmt.Errorf("getting object: %w", err)
	}

	// 获取所有的 ObjectBlock
	var allBlocks []types.ObjectBlock
	err = ctx.Table("ObjectBlock").Where("ObjectID = ?", objectID).Order("`Index` ASC").Find(&allBlocks).Error
	if err != nil {
		return types.ObjectDetail{}, fmt.Errorf("getting all object blocks: %w", err)
	}

	// 获取所有的 PinnedObject
	var allPinnedObjs []types.PinnedObject
	err = ctx.Table("PinnedObject").Where("ObjectID = ?", objectID).Order("ObjectID ASC").Find(&allPinnedObjs).Error
	if err != nil {
		return types.ObjectDetail{}, fmt.Errorf("getting all pinned objects: %w", err)
	}

	pinnedAt := make([]types.UserSpaceID, len(allPinnedObjs))
	for i, po := range allPinnedObjs {
		pinnedAt[i] = po.UserSpaceID
	}

	return types.ObjectDetail{
		Object:   obj,
		Blocks:   allBlocks,
		PinnedAt: pinnedAt,
	}, nil
}

// 仅返回查询到的对象
func (db *ObjectDB) BatchGetDetails(ctx SQLContext, objectIDs []types.ObjectID) ([]types.ObjectDetail, error) {
	var objs []types.Object

	err := ctx.Table("Object").Where("ObjectID IN ?", objectIDs).Order("ObjectID ASC").Find(&objs).Error
	if err != nil {
		return nil, err
	}

	// 获取所有的 ObjectBlock
	var allBlocks []types.ObjectBlock
	err = ctx.Table("ObjectBlock").Where("ObjectID IN ?", objectIDs).Order("ObjectID, `Index` ASC").Find(&allBlocks).Error
	if err != nil {
		return nil, err
	}

	// 获取所有的 PinnedObject
	var allPinnedObjs []types.PinnedObject
	err = ctx.Table("PinnedObject").Where("ObjectID IN ?", objectIDs).Order("ObjectID ASC").Find(&allPinnedObjs).Error
	if err != nil {
		return nil, err
	}

	details := make([]types.ObjectDetail, len(objs))
	for i, obj := range objs {
		details[i] = types.ObjectDetail{
			Object: obj,
		}
	}

	types.DetailsFillObjectBlocks(details, allBlocks)
	types.DetailsFillPinnedAt(details, allPinnedObjs)
	return details, nil
}

func (db *ObjectDB) Create(ctx SQLContext, obj types.Object) (types.ObjectID, error) {
	err := ctx.Table("Object").Create(&obj).Error
	if err != nil {
		return 0, fmt.Errorf("insert object failed, err: %w", err)
	}
	return obj.ObjectID, nil
}

// 批量创建对象，创建完成后会填充ObjectID。
func (db *ObjectDB) BatchCreate(ctx SQLContext, objs *[]types.Object) error {
	if len(*objs) == 0 {
		return nil
	}

	return ctx.Table("Object").Create(objs).Error
}

// 批量更新对象所有属性，objs中的对象必须包含ObjectID
func (db *ObjectDB) BatchUpdate(ctx SQLContext, objs []types.Object) error {
	if len(objs) == 0 {
		return nil
	}

	return ctx.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "ObjectID"}},
		UpdateAll: true,
	}).Create(objs).Error
}

// 批量更新对象指定属性，objs中的对象只需设置需要更新的属性即可，但：
//  1. 必须包含ObjectID
//  2. 日期类型属性不能设置为0值
func (db *ObjectDB) BatchUpdateColumns(ctx SQLContext, objs []types.Object, columns []string) error {
	if len(objs) == 0 {
		return nil
	}

	return ctx.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "ObjectID"}},
		DoUpdates: clause.AssignmentColumns(columns),
	}).Create(objs).Error
}

func (db *ObjectDB) GetPackageObjects(ctx SQLContext, packageID types.PackageID) ([]types.Object, error) {
	var ret []types.Object
	err := ctx.Table("Object").Where("PackageID = ?", packageID).Order("ObjectID ASC").Find(&ret).Error
	return ret, err
}

func (db *ObjectDB) GetPackageObjectDetails(ctx SQLContext, packageID types.PackageID) ([]types.ObjectDetail, error) {
	var objs []types.Object
	err := ctx.Table("Object").Where("PackageID = ?", packageID).Order("ObjectID ASC").Find(&objs).Error
	if err != nil {
		return nil, fmt.Errorf("getting objects: %w", err)
	}

	// 获取所有的 ObjectBlock
	var allBlocks []types.ObjectBlock
	err = ctx.Table("ObjectBlock").
		Select("ObjectBlock.*").
		Joins("JOIN Object ON ObjectBlock.ObjectID = Object.ObjectID").
		Where("Object.PackageID = ?", packageID).
		Order("ObjectBlock.ObjectID, `Index` ASC").
		Find(&allBlocks).Error
	if err != nil {
		return nil, fmt.Errorf("getting all object blocks: %w", err)
	}

	// 获取所有的 PinnedObject
	var allPinnedObjs []types.PinnedObject
	err = ctx.Table("PinnedObject").
		Select("PinnedObject.*").
		Joins("JOIN Object ON PinnedObject.ObjectID = Object.ObjectID").
		Where("Object.PackageID = ?", packageID).
		Order("PinnedObject.ObjectID").
		Find(&allPinnedObjs).Error
	if err != nil {
		return nil, fmt.Errorf("getting all pinned objects: %w", err)
	}

	details := make([]types.ObjectDetail, len(objs))
	for i, obj := range objs {
		details[i] = types.ObjectDetail{
			Object: obj,
		}
	}

	types.DetailsFillObjectBlocks(details, allBlocks)
	types.DetailsFillPinnedAt(details, allPinnedObjs)
	return details, nil
}

func (db *ObjectDB) GetObjectsIfAnyBlockOnStorage(ctx SQLContext, spaceID types.UserSpaceID) ([]types.Object, error) {
	var objs []types.Object
	err := ctx.Table("Object").Where("ObjectID IN (SELECT ObjectID FROM ObjectBlock WHERE UserSpaceID = ?)", spaceID).Order("ObjectID ASC").Find(&objs).Error
	if err != nil {
		return nil, fmt.Errorf("getting objects: %w", err)
	}

	return objs, nil
}

type AddObjectEntry struct {
	Path         string              `json:"path"`
	Size         int64               `json:"size,string"`
	FileHash     types.FileHash      `json:"fileHash"`
	UploadTime   time.Time           `json:"uploadTime"` // 开始上传文件的时间
	UserSpaceIDs []types.UserSpaceID `json:"userSpaceIDs"`
}

func (db *ObjectDB) BatchAdd(ctx SQLContext, packageID types.PackageID, adds []AddObjectEntry) ([]types.Object, error) {
	if len(adds) == 0 {
		return nil, nil
	}

	// 收集所有路径
	pathes := make([]string, 0, len(adds))
	for _, add := range adds {
		pathes = append(pathes, add.Path)
	}

	// 先查询要更新的对象，不存在也没关系
	existsObjs, err := db.BatchGetByPackagePath(ctx, packageID, pathes)
	if err != nil {
		return nil, fmt.Errorf("batch get object by path: %w", err)
	}

	existsObjsMap := make(map[string]types.Object)
	for _, obj := range existsObjs {
		existsObjsMap[obj.Path] = obj
	}

	var updatingObjs []types.Object
	var addingObjs []types.Object
	for i := range adds {
		o := types.Object{
			PackageID:  packageID,
			Path:       adds[i].Path,
			Size:       adds[i].Size,
			FileHash:   adds[i].FileHash,
			Redundancy: types.NewNoneRedundancy(), // 首次上传默认使用不分块的none模式
			CreateTime: adds[i].UploadTime,
			UpdateTime: adds[i].UploadTime,
		}

		e, ok := existsObjsMap[adds[i].Path]
		if ok {
			o.ObjectID = e.ObjectID
			o.CreateTime = e.CreateTime
			updatingObjs = append(updatingObjs, o)

		} else {
			addingObjs = append(addingObjs, o)
		}
	}

	// 先进行更新
	err = db.BatchUpdate(ctx, updatingObjs)
	if err != nil {
		return nil, fmt.Errorf("batch update objects: %w", err)
	}

	// 再执行插入，Create函数插入后会填充ObjectID
	err = db.BatchCreate(ctx, &addingObjs)
	if err != nil {
		return nil, fmt.Errorf("batch create objects: %w", err)
	}

	// 按照add参数的顺序返回结果
	affectedObjsMp := make(map[string]types.Object)
	for _, o := range updatingObjs {
		affectedObjsMp[o.Path] = o
	}
	for _, o := range addingObjs {
		affectedObjsMp[o.Path] = o
	}
	affectedObjs := make([]types.Object, 0, len(affectedObjsMp))
	affectedObjIDs := make([]types.ObjectID, 0, len(affectedObjsMp))
	for i := range adds {
		obj := affectedObjsMp[adds[i].Path]
		affectedObjs = append(affectedObjs, obj)
		affectedObjIDs = append(affectedObjIDs, obj.ObjectID)
	}

	if len(affectedObjIDs) > 0 {
		// 批量删除 ObjectBlock
		if err := db.ObjectBlock().BatchDeleteByObjectID(ctx, affectedObjIDs); err != nil {
			return nil, fmt.Errorf("batch delete object blocks: %w", err)
		}

		// 批量删除 PinnedObject
		if err := db.PinnedObject().BatchDeleteByObjectID(ctx, affectedObjIDs); err != nil {
			return nil, fmt.Errorf("batch delete pinned objects: %w", err)
		}
	}

	// 创建 ObjectBlock
	objBlocks := make([]types.ObjectBlock, 0, len(adds))
	for i, add := range adds {
		for _, spaceID := range add.UserSpaceIDs {
			objBlocks = append(objBlocks, types.ObjectBlock{
				ObjectID:    affectedObjIDs[i],
				Index:       0,
				UserSpaceID: spaceID,
				FileHash:    add.FileHash,
				Size:        add.Size,
			})
		}
	}
	if err := db.ObjectBlock().BatchCreate(ctx, objBlocks); err != nil {
		return nil, fmt.Errorf("batch create object blocks: %w", err)
	}

	return affectedObjs, nil
}

func (db *ObjectDB) BatchDelete(ctx SQLContext, ids []types.ObjectID) error {
	if len(ids) == 0 {
		return nil
	}

	return ctx.Table("Object").Where("ObjectID IN ?", ids).Delete(&types.Object{}).Error
}

func (db *ObjectDB) DeleteInPackage(ctx SQLContext, packageID types.PackageID) error {
	return ctx.Table("Object").Where("PackageID = ?", packageID).Delete(&types.Object{}).Error
}

type UpdatingObjectRedundancy struct {
	ObjectID   types.ObjectID      `json:"objectID"`
	FileHash   types.FileHash      `json:"fileHash"`
	Size       int64               `json:"size"`
	Redundancy types.Redundancy    `json:"redundancy"`
	PinnedAt   []types.UserSpaceID `json:"pinnedAt"`
	Blocks     []types.ObjectBlock `json:"blocks"`
}

func (db *ObjectDB) BatchUpdateRedundancy(ctx SQLContext, updates []UpdatingObjectRedundancy) error {
	objs := updates

	nowTime := time.Now()
	objIDs := make([]types.ObjectID, 0, len(objs))
	for _, obj := range objs {
		objIDs = append(objIDs, obj.ObjectID)
	}

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

	// 过滤掉已经不存在的对象。
	// 注意，objIDs没有被过滤，因为后续逻辑不过滤也不会出错
	objs = lo.Filter(objs, func(obj UpdatingObjectRedundancy, _ int) bool {
		return avaiIDs[obj.ObjectID]
	})

	dummyObjs := make([]types.Object, 0, len(objs))
	for _, obj := range objs {
		dummyObjs = append(dummyObjs, types.Object{
			ObjectID:   obj.ObjectID,
			FileHash:   obj.FileHash,
			Size:       obj.Size,
			Redundancy: obj.Redundancy,
			CreateTime: nowTime, // 实际不会更新，只因为不能是0值
			UpdateTime: nowTime,
		})
	}

	err = db.Object().BatchUpdateColumns(ctx, dummyObjs, []string{"FileHash", "Size", "Redundancy", "UpdateTime"})
	if err != nil {
		return fmt.Errorf("batch update object redundancy: %w", err)
	}

	// 删除原本所有的编码块记录，重新添加
	err = db.ObjectBlock().BatchDeleteByObjectID(ctx, objIDs)
	if err != nil {
		return fmt.Errorf("batch delete object blocks: %w", err)
	}

	// 删除原本Pin住的Object。暂不考虑FileHash没有变化的情况
	err = db.PinnedObject().BatchDeleteByObjectID(ctx, objIDs)
	if err != nil {
		return fmt.Errorf("batch delete pinned object: %w", err)
	}

	blocks := make([]types.ObjectBlock, 0, len(objs))
	for _, obj := range objs {
		blocks = append(blocks, obj.Blocks...)
	}
	err = db.ObjectBlock().BatchCreate(ctx, blocks)
	if err != nil {
		return fmt.Errorf("batch create object blocks: %w", err)
	}

	pinneds := make([]types.PinnedObject, 0, len(objs))
	for _, obj := range objs {
		for _, p := range obj.PinnedAt {
			pinneds = append(pinneds, types.PinnedObject{
				ObjectID:    obj.ObjectID,
				UserSpaceID: p,
				CreateTime:  nowTime,
			})
		}
	}
	err = db.PinnedObject().BatchTryCreate(ctx, pinneds)
	if err != nil {
		return fmt.Errorf("batch create pinned objects: %w", err)
	}

	return nil
}

func (db *ObjectDB) DeleteByPath(ctx SQLContext, packageID types.PackageID, path string) error {
	return ctx.Table("Object").Where("PackageID = ? AND Path = ?", packageID, path).Delete(&types.Object{}).Error
}

func (db *ObjectDB) MoveByPrefix(ctx SQLContext, oldPkgID types.PackageID, oldPrefix string, newPkgID types.PackageID, newPrefix string) error {
	return ctx.Table("Object").Where("PackageID = ? AND Path LIKE ?", oldPkgID, escapeLike("", "%", oldPrefix)).
		Updates(map[string]any{
			"PackageID": newPkgID,
			"Path":      gorm.Expr("concat(?, substring(Path, ?))", newPrefix, len(oldPrefix)+1),
		}).Error
}