JCS-pub/common/pkgs/db2/storage.go

package db2

import (
	"fmt"

	"gitlink.org.cn/cloudream/common/pkgs/logger"
	cdssdk "gitlink.org.cn/cloudream/common/sdks/storage"
	stgmod "gitlink.org.cn/cloudream/storage/common/models"
	"gitlink.org.cn/cloudream/storage/common/pkgs/db2/model"
	"gorm.io/gorm"
)

type StorageDB struct {
	*DB
}

func (db *DB) Storage() *StorageDB {
	return &StorageDB{DB: db}
}

func (db *StorageDB) GetByID(ctx SQLContext, stgID cdssdk.StorageID) (model.Storage, error) {
	var stg model.Storage
	err := ctx.Table("Storage").First(&stg, stgID).Error
	return stg, err
}

func (StorageDB) GetAllIDs(ctx SQLContext) ([]cdssdk.StorageID, error) {
	var stgs []cdssdk.StorageID
	err := ctx.Table("Storage").Select("StorageID").Find(&stgs).Error
	return stgs, err
}

func (db *StorageDB) BatchGetByID(ctx SQLContext, stgIDs []cdssdk.StorageID) ([]model.Storage, error) {
	var stgs []model.Storage
	err := ctx.Table("Storage").Find(&stgs, "StorageID IN (?)", stgIDs).Error
	return stgs, err
}

func (db *StorageDB) GetUserStorages(ctx SQLContext, userID cdssdk.UserID) ([]model.Storage, error) {
	var stgs []model.Storage
	err := ctx.Table("Storage").Select("Storage.*").
		Joins("inner join UserStorage on Storage.StorageID = UserStorage.StorageID").
		Where("UserID = ?", userID).Find(&stgs).Error
	return stgs, err
}

func (db *StorageDB) BatchGetAllStorageIDs(ctx SQLContext, start int, count int) ([]cdssdk.StorageID, error) {
	var ret []cdssdk.StorageID
	err := ctx.Table("Storage").Select("StorageID").Find(&ret).Limit(count).Offset(start).Error
	return ret, err
}

func (db *StorageDB) IsAvailable(ctx SQLContext, userID cdssdk.UserID, storageID cdssdk.StorageID) (bool, error) {
	rows, err := ctx.Table("Storage").Select("Storage.StorageID").
		Joins("inner join UserStorage on Storage.StorageID = UserStorage.StorageID").
		Where("UserID = ? and StorageID = ?", userID, storageID).Rows()
	if err != nil {
		return false, fmt.Errorf("execute sql: %w", err)
	}
	defer rows.Close()

	return rows.Next(), nil
}

func (db *StorageDB) GetUserStorage(ctx SQLContext, userID cdssdk.UserID, storageID cdssdk.StorageID) (model.Storage, error) {
	var stg model.Storage
	err := ctx.Table("Storage").Select("Storage.*").
		Joins("inner join UserStorage on Storage.StorageID = UserStorage.StorageID").
		Where("UserID = ? and StorageID = ?", userID, storageID).First(&stg).Error

	return stg, err
}

func (db *StorageDB) GetUserStorageByName(ctx SQLContext, userID cdssdk.UserID, name string) (model.Storage, error) {
	var stg model.Storage
	err := ctx.Table("Storage").Select("Storage.*").
		Joins("inner join UserStorage on Storage.StorageID = UserStorage.StorageID").
		Where("UserID = ? and Name = ?", userID, name).First(&stg).Error

	return stg, err
}

func (db *StorageDB) GetHubStorages(ctx SQLContext, hubID cdssdk.HubID) ([]model.Storage, error) {
	var stgs []model.Storage
	err := ctx.Table("Storage").Select("Storage.*").Find(&stgs, "MasterHub = ?", hubID).Error
	return stgs, err
}

func (db *StorageDB) FillDetails(ctx SQLContext, details []stgmod.StorageDetail) error {
	stgsMp := make(map[cdssdk.StorageID]*stgmod.StorageDetail)
	var masterHubIDs []cdssdk.HubID
	for i := range details {
		stgsMp[details[i].Storage.StorageID] = &details[i]
		masterHubIDs = append(masterHubIDs, details[i].Storage.MasterHub)
	}

	// 获取监护Hub信息
	masterHubs, err := db.Hub().BatchGetByID(ctx, masterHubIDs)
	if err != nil && err != gorm.ErrRecordNotFound {
		return fmt.Errorf("getting master hub: %w", err)
	}
	masterHubMap := make(map[cdssdk.HubID]cdssdk.Hub)
	for _, hub := range masterHubs {
		masterHubMap[hub.HubID] = hub
	}
	for _, stg := range stgsMp {
		if stg.Storage.MasterHub != 0 {
			hub, ok := masterHubMap[stg.Storage.MasterHub]
			if !ok {
				logger.Warnf("master hub %v of storage %v not found, this storage will not be add to result", stg.Storage.MasterHub, stg.Storage)
				delete(stgsMp, stg.Storage.StorageID)
				continue
			}

			stg.MasterHub = &hub
		}
	}

	return nil
}