new-api/model/channel_asset_credential.go

package model

import (
	"github.com/QuantumNous/new-api/common"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// ChannelAssetCredential stores the China Mobile asset credentials separately
// from the channel video-generation key.
type ChannelAssetCredential struct {
	Id        int    `json:"id"`
	ChannelId int    `json:"channel_id" gorm:"uniqueIndex;not null"`
	AccessKey string `json:"-" gorm:"not null;size:255"`
	SecretKey string `json:"-" gorm:"not null;size:255"`
	PoolID    string `json:"pool_id" gorm:"size:255"`
	CreatedAt int64  `json:"created_at" gorm:"bigint;not null"`
	UpdatedAt int64  `json:"updated_at" gorm:"bigint;not null"`
}

// ChannelAssetCredentialSummary is safe to expose in channel management APIs.
type ChannelAssetCredentialSummary struct {
	ChannelId int
	PoolID    string
}

func GetChannelAssetCredential(channelID int) (*ChannelAssetCredential, error) {
	var credential ChannelAssetCredential
	err := DB.Where("channel_id = ?", channelID).First(&credential).Error
	if err == gorm.ErrRecordNotFound {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return &credential, nil
}

func UpsertChannelAssetCredential(credential *ChannelAssetCredential) error {
	return DB.Transaction(func(tx *gorm.DB) error {
		return UpsertChannelAssetCredentialWithTx(tx, credential)
	})
}

func UpsertChannelAssetCredentialWithTx(tx *gorm.DB, credential *ChannelAssetCredential) error {
	now := common.GetTimestamp()
	credential.CreatedAt = now
	credential.UpdatedAt = now
	return tx.Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "channel_id"}},
		DoUpdates: clause.Assignments(map[string]any{
			"access_key": credential.AccessKey,
			"secret_key": credential.SecretKey,
			"pool_id":    credential.PoolID,
			"updated_at": credential.UpdatedAt,
		}),
	}).Create(credential).Error
}

func DeleteChannelAssetCredentialWithTx(tx *gorm.DB, channelID int) error {
	return tx.Where("channel_id = ?", channelID).Delete(&ChannelAssetCredential{}).Error
}

func DeleteChannelAssetCredential(channelID int) error {
	return DB.Transaction(func(tx *gorm.DB) error {
		return DeleteChannelAssetCredentialWithTx(tx, channelID)
	})
}

func DeleteChannelAssetCredentialsWithTx(tx *gorm.DB, channelIDs []int) error {
	if len(channelIDs) == 0 {
		return nil
	}
	return tx.Where("channel_id IN ?", channelIDs).Delete(&ChannelAssetCredential{}).Error
}

func DeleteChannelAssetCredentials(channelIDs []int) error {
	return DB.Transaction(func(tx *gorm.DB) error {
		return DeleteChannelAssetCredentialsWithTx(tx, channelIDs)
	})
}

func GetChannelAssetCredentialSummaries(channelIDs []int) (map[int]ChannelAssetCredentialSummary, error) {
	summaries := make(map[int]ChannelAssetCredentialSummary)
	if len(channelIDs) == 0 {
		return summaries, nil
	}
	var rows []ChannelAssetCredentialSummary
	if err := DB.Model(&ChannelAssetCredential{}).
		Select("channel_id", "pool_id").
		Where("channel_id IN ?", channelIDs).
		Find(&rows).Error; err != nil {
		return nil, err
	}
	for _, row := range rows {
		summaries[row.ChannelId] = row
	}
	return summaries, nil
}