aiforge/modules/storage/minio_ext.go

package storage

import (
	"encoding/xml"
	"errors"
	"path"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/minio_ext"
	"code.gitea.io/gitea/modules/setting"

	miniov6 "github.com/minio/minio-go/v6"
)

const (
	PresignedUploadPartUrlExpireTime = time.Hour * 24 * 7
)

type ComplPart struct {
	PartNumber int    `json:"partNumber"`
	ETag       string `json:"eTag"`
}

type CompleteParts struct {
	Data []ComplPart `json:"completedParts"`
}

// completedParts is a collection of parts sortable by their part numbers.
// used for sorting the uploaded parts before completing the multipart request.
type completedParts []miniov6.CompletePart

func (a completedParts) Len() int           { return len(a) }
func (a completedParts) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a completedParts) Less(i, j int) bool { return a[i].PartNumber < a[j].PartNumber }

// completeMultipartUpload container for completing multipart upload.
type completeMultipartUpload struct {
	XMLName xml.Name               `xml:"http://s3.amazonaws.com/doc/2006-03-01/ CompleteMultipartUpload" json:"-"`
	Parts   []miniov6.CompletePart `xml:"Part"`
}

var (
	adminClient *minio_ext.Client = nil
	coreClient  *miniov6.Core     = nil
)

var mutex *sync.Mutex

func init() {
	mutex = new(sync.Mutex)
}

func getClients() (*minio_ext.Client, *miniov6.Core, error) {
	var client *minio_ext.Client
	var core *miniov6.Core
	mutex.Lock()

	defer mutex.Unlock()

	if nil != adminClient && nil != coreClient {
		client = adminClient
		core = coreClient
		return client, core, nil
	}

	var err error
	minio := setting.Attachment.Minio
	if nil == adminClient {
		adminClient, err = minio_ext.New(
			minio.Endpoint,
			minio.AccessKeyID,
			minio.SecretAccessKey,
			minio.UseSSL,
		)
		if nil != err {
			return nil, nil, err
		}
	}

	client = adminClient

	if nil == coreClient {
		coreClient, err = miniov6.NewCore(
			minio.Endpoint,
			minio.AccessKeyID,
			minio.SecretAccessKey,
			minio.UseSSL,
		)
		if nil != err {
			return nil, nil, err
		}
	}

	core = coreClient

	return client, core, nil
}

func GenMultiPartSignedUrl(uuid string, uploadId string, partNumber int, partSize int64) (string, error) {
	minioClient, _, err := getClients()
	if err != nil {
		log.Error("getClients failed:", err.Error())
		return "", err
	}

	minio := setting.Attachment.Minio
	bucketName := minio.Bucket
	objectName := strings.TrimPrefix(path.Join(minio.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid)), "/")

	return minioClient.GenUploadPartSignedUrl(uploadId, bucketName, objectName, partNumber, partSize, PresignedUploadPartUrlExpireTime, setting.Attachment.Minio.Location)

}

func NewMultiPartUpload(uuid string) (string, error) {
	_, core, err := getClients()
	if err != nil {
		log.Error("getClients failed:", err.Error())
		return "", err
	}

	minio := setting.Attachment.Minio
	bucketName := minio.Bucket
	objectName := strings.TrimPrefix(path.Join(minio.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid)), "/")

	return core.NewMultipartUpload(bucketName, objectName, miniov6.PutObjectOptions{})
}

func CompleteMultiPartUpload(uuid string, uploadID string, totalChunks int) (string, error) {
	client, core, err := getClients()
	if err != nil {
		log.Error("getClients failed:", err.Error())
		return "", err
	}

	minio := setting.Attachment.Minio
	bucketName := minio.Bucket
	objectName := strings.TrimPrefix(path.Join(minio.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid)), "/")

	partInfos, err := client.ListObjectParts(bucketName, objectName, uploadID)
	if err != nil {
		log.Error("ListObjectParts failed:", err.Error())
		return "", err
	}

	if len(partInfos) != totalChunks {
		log.Error("ListObjectParts number(%d) is not equal the set total chunk number(%d)", len(partInfos), totalChunks)
		return "", errors.New("the parts is not complete")
	}

	var complMultipartUpload completeMultipartUpload
	for _, partInfo := range partInfos {
		complMultipartUpload.Parts = append(complMultipartUpload.Parts, miniov6.CompletePart{
			PartNumber: partInfo.PartNumber,
			ETag:       partInfo.ETag,
		})
	}

	// Sort all completed parts.
	sort.Sort(completedParts(complMultipartUpload.Parts))

	return core.CompleteMultipartUpload(bucketName, objectName, uploadID, complMultipartUpload.Parts)
}

func GetPartInfos(uuid string, uploadID string) (string, error) {
	minioClient, _, err := getClients()
	if err != nil {
		log.Error("getClients failed:", err.Error())
		return "", err
	}

	minio := setting.Attachment.Minio
	bucketName := minio.Bucket
	objectName := strings.TrimPrefix(path.Join(minio.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid)), "/")

	partInfos, err := minioClient.ListObjectParts(bucketName, objectName, uploadID)
	if err != nil {
		log.Error("ListObjectParts failed:", err.Error())
		return "", err
	}

	var chunks string
	for _, partInfo := range partInfos {
		chunks += strconv.Itoa(partInfo.PartNumber) + "-" + partInfo.ETag + ","
	}

	return chunks, nil
}