aiforge/modules/storage/obs.go

// Copyright 2020 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package storage

import (
	"errors"
	"io"
	"net/url"
	"path"
	"sort"
	"strconv"
	"strings"

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

	"github.com/unknwon/com"
)

type FileInfo struct {
	FileName string `json:"FileName"`
	ModTime  string `json:"ModTime"`
	IsDir    bool   `json:"IsDir"`
	Size     int64  `json:"Size"`
	ParenDir string `json:"ParenDir"`
	UUID     string `json:"UUID"`
}
type FileInfoList []FileInfo

const MAX_LIST_PARTS = 1000

func (ulist FileInfoList) Swap(i, j int) { ulist[i], ulist[j] = ulist[j], ulist[i] }
func (ulist FileInfoList) Len() int      { return len(ulist) }
func (ulist FileInfoList) Less(i, j int) bool {
	return strings.Compare(ulist[i].FileName, ulist[j].FileName) > 0
}

//check if has the object
func ObsHasObject(path string) (bool, error) {
	hasObject := false

	input := &obs.GetObjectMetadataInput{}
	input.Bucket = setting.Bucket
	input.Key = path
	_, err := ObsCli.GetObjectMetadata(input)
	if err == nil {
		hasObject = true
	} else {
		if obsError, ok := err.(obs.ObsError); ok {
			log.Error("GetObjectMetadata failed(%d): %s", obsError.StatusCode, obsError.Message)
		} else {
			log.Error("%v", err.Error())
		}
	}

	return hasObject, nil
}

func listAllParts(uuid, uploadID, key string) (output *obs.ListPartsOutput, err error) {
	output = &obs.ListPartsOutput{}
	partNumberMarker := 0
	for {
		temp, err := ObsCli.ListParts(&obs.ListPartsInput{
			Bucket:           setting.Bucket,
			Key:              key,
			UploadId:         uploadID,
			MaxParts:         MAX_LIST_PARTS,
			PartNumberMarker: partNumberMarker,
		})
		if err != nil {
			log.Error("ListParts failed:", err.Error())
			return output, err
		}

		partNumberMarker = temp.NextPartNumberMarker
		log.Info("uuid:%s, MaxParts:%d, PartNumberMarker:%d, NextPartNumberMarker:%d, len:%d", uuid, temp.MaxParts, temp.PartNumberMarker, temp.NextPartNumberMarker, len(temp.Parts))

		for _, partInfo := range temp.Parts {
			output.Parts = append(output.Parts, obs.Part{
				PartNumber: partInfo.PartNumber,
				ETag:       partInfo.ETag,
			})
		}

		if !temp.IsTruncated {
			break
		} else {
			continue
		}

		break
	}

	return output, nil
}

func GetObsPartInfos(uuid, uploadID, fileName string) (string, error) {
	key := strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")

	allParts, err := listAllParts(uuid, uploadID, key)
	if err != nil {
		log.Error("listAllParts failed: %v", err)
		return "", err
	}

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

	return chunks, nil
}

func NewObsMultiPartUpload(uuid, fileName string) (string, error) {
	input := &obs.InitiateMultipartUploadInput{}
	input.Bucket = setting.Bucket
	input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")

	output, err := ObsCli.InitiateMultipartUpload(input)
	if err != nil {
		log.Error("InitiateMultipartUpload failed:", err.Error())
		return "", err
	}

	return output.UploadId, nil
}

func CompleteObsMultiPartUpload(uuid, uploadID, fileName string, totalChunks int) error {
	input := &obs.CompleteMultipartUploadInput{}
	input.Bucket = setting.Bucket
	input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")
	input.UploadId = uploadID

	allParts, err := listAllParts(uuid, uploadID, input.Key)
	if err != nil {
		log.Error("listAllParts failed: %v", err)
		return err
	}

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

	input.Parts = allParts.Parts

	output, err := ObsCli.CompleteMultipartUpload(input)
	if err != nil {
		log.Error("CompleteMultipartUpload failed:", err.Error())
		return err
	}

	log.Info("uuid:%s, RequestId:%s", uuid, output.RequestId)

	return nil
}

func ObsMultiPartUpload(uuid string, uploadId string, partNumber int, fileName string, putBody io.ReadCloser) error {
	input := &obs.UploadPartInput{}
	input.Bucket = setting.Bucket
	input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")
	input.UploadId = uploadId
	input.PartNumber = partNumber
	input.Body = putBody
	output, err := ObsCli.UploadPart(input)
	if err == nil {
		log.Info("RequestId:%s\n", output.RequestId)
		log.Info("ETag:%s\n", output.ETag)
		return nil
	} else {
		if obsError, ok := err.(obs.ObsError); ok {
			log.Info(obsError.Code)
			log.Info(obsError.Message)
			return obsError
		} else {
			log.Error("error:", err.Error())
			return err
		}
	}

}

//delete all file under the dir path
func ObsRemoveObject(bucket string, path string) error {
	log.Info("Bucket=" + bucket + "  path=" + path)
	if len(path) == 0 {
		return errors.New("path canot be null.")
	}
	input := &obs.ListObjectsInput{}
	input.Bucket = bucket
	// 设置每页100个对象
	input.MaxKeys = 100
	input.Prefix = path
	index := 1
	log.Info("prefix=" + input.Prefix)
	for {
		output, err := ObsCli.ListObjects(input)
		if err == nil {
			log.Info("Page:%d\n", index)
			index++
			for _, val := range output.Contents {
				log.Info("delete obs file:" + val.Key)
				delObj := &obs.DeleteObjectInput{}
				delObj.Bucket = setting.Bucket
				delObj.Key = val.Key
				ObsCli.DeleteObject(delObj)
			}
			if output.IsTruncated {
				input.Marker = output.NextMarker
			} else {
				break
			}
		} else {
			if obsError, ok := err.(obs.ObsError); ok {
				log.Info("Code:%s\n", obsError.Code)
				log.Info("Message:%s\n", obsError.Message)
			}
			return err
		}
	}
	return nil
}

func ObsDownloadAFile(bucket string, key string) (io.ReadCloser, error) {
	input := &obs.GetObjectInput{}
	input.Bucket = bucket
	input.Key = key
	output, err := ObsCli.GetObject(input)
	if err == nil {
		log.Info("StorageClass:%s, ETag:%s, ContentType:%s, ContentLength:%d, LastModified:%s\n",
			output.StorageClass, output.ETag, output.ContentType, output.ContentLength, output.LastModified)
		return output.Body, nil
	} else if obsError, ok := err.(obs.ObsError); ok {
		log.Error("Code:%s, Message:%s", obsError.Code, obsError.Message)
		return nil, obsError
	} else {
		return nil, err
	}
}

func ObsDownload(uuid string, fileName string) (io.ReadCloser, error) {

	return ObsDownloadAFile(setting.Bucket, strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/"))
}

func ObsModelDownload(JobName string, fileName string) (io.ReadCloser, error) {
	input := &obs.GetObjectInput{}
	input.Bucket = setting.Bucket
	input.Key = strings.TrimPrefix(path.Join(setting.TrainJobModelPath, JobName, setting.OutPutPath, fileName), "/")
	// input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid)), "/")
	output, err := ObsCli.GetObject(input)
	if err == nil {
		log.Info("StorageClass:%s, ETag:%s, ContentType:%s, ContentLength:%d, LastModified:%s\n",
			output.StorageClass, output.ETag, output.ContentType, output.ContentLength, output.LastModified)
		return output.Body, nil
	} else if obsError, ok := err.(obs.ObsError); ok {
		log.Error("Code:%s, Message:%s", obsError.Code, obsError.Message)
		return nil, obsError
	} else {
		return nil, err
	}
}

func ObsCopyManyFile(srcBucket string, srcPath string, destBucket string, destPath string, Files []string) (int64, error) {

	var fileTotalSize int64

	for _, file := range Files {
		srcKey := srcPath + file
		destKey := destPath + file
		log.Info("srcKey=" + srcKey + " destKey=" + destKey)
		out, err := ObsCli.GetObjectMetadata(&obs.GetObjectMetadataInput{
			Bucket: srcBucket,
			Key:    srcKey,
		})
		if err != nil {
			log.Info("Get File error, error=" + err.Error())
			continue
		}
		obsCopyFile(srcBucket, srcKey, destBucket, destKey)
		fileTotalSize += out.ContentLength
	}

	return fileTotalSize, nil
}

func ObsCopyAllFile(srcBucket string, srcPath string, destBucket string, destPath string) (int64, error) {
	input := &obs.ListObjectsInput{}
	input.Bucket = srcBucket
	// 设置每页100个对象
	input.MaxKeys = 100
	input.Prefix = srcPath
	index := 1
	length := len(srcPath)
	var fileTotalSize int64
	log.Info("prefix=" + input.Prefix)
	for {
		output, err := ObsCli.ListObjects(input)
		if err == nil {
			log.Info("Page:%d\n", index)
			index++
			for _, val := range output.Contents {
				destKey := destPath + val.Key[length:]
				obsCopyFile(srcBucket, val.Key, destBucket, destKey)
				fileTotalSize += val.Size
			}
			if output.IsTruncated {
				input.Marker = output.NextMarker
			} else {
				break
			}
		} else {
			if obsError, ok := err.(obs.ObsError); ok {
				log.Info("Code:%s\n", obsError.Code)
				log.Info("Message:%s\n", obsError.Message)
			}
			return 0, err
		}
	}
	return fileTotalSize, nil
}

func obsCopyFile(srcBucket string, srcKeyName string, destBucket string, destKeyName string) error {
	input := &obs.CopyObjectInput{}
	input.Bucket = destBucket
	input.Key = destKeyName
	input.CopySourceBucket = srcBucket
	input.CopySourceKey = srcKeyName
	_, err := ObsCli.CopyObject(input)
	if err == nil {
		log.Info("copy success,destBuckName:%s, destkeyname:%s", destBucket, destKeyName)
	} else {
		log.Info("copy failed,,destBuckName:%s, destkeyname:%s", destBucket, destKeyName)
		if obsError, ok := err.(obs.ObsError); ok {
			log.Info(obsError.Code)
			log.Info(obsError.Message)
		}
		return err
	}
	return nil
}

func GetOneLevelAllObjectUnderDir(bucket string, prefixRootPath string, relativePath string) ([]FileInfo, error) {
	input := &obs.ListObjectsInput{}
	input.Bucket = bucket
	input.Prefix = prefixRootPath + relativePath
	if !strings.HasSuffix(input.Prefix, "/") {
		input.Prefix += "/"
	}
	output, err := ObsCli.ListObjects(input)
	fileInfos := make([]FileInfo, 0)
	prefixLen := len(input.Prefix)
	fileMap := make(map[string]bool, 0)
	if err == nil {
		for _, val := range output.Contents {
			log.Info("val key=" + val.Key)
			var isDir bool
			var fileName string
			if val.Key == input.Prefix {
				continue
			}
			if strings.Contains(val.Key[prefixLen:len(val.Key)-1], "/") {
				fileName = val.Key[prefixLen : len(val.Key)-1]
				files := strings.Split(fileName, "/")
				fileName = files[0]
				isDir = true
				if fileMap[files[0]] {
					continue
				} else {
					fileMap[files[0]] = true
				}
			} else {
				if strings.HasSuffix(val.Key, "/") {
					isDir = true
					fileName = val.Key[prefixLen : len(val.Key)-1]
					relativePath += val.Key[prefixLen:]
				} else {
					isDir = false
					fileName = val.Key[prefixLen:]
				}
				fileMap[fileName] = true
			}

			fileInfo := FileInfo{
				ModTime:  val.LastModified.Local().Format("2006-01-02 15:04:05"),
				FileName: fileName,
				Size:     val.Size,
				IsDir:    isDir,
				ParenDir: relativePath,
			}
			fileInfos = append(fileInfos, fileInfo)
		}
		return fileInfos, err
	} else {
		if obsError, ok := err.(obs.ObsError); ok {
			log.Error("Code:%s, Message:%s", obsError.Code, obsError.Message)
		}
		return nil, err
	}

}

func GetAllObjectByBucketAndPrefix(bucket string, prefix string) ([]FileInfo, error) {
	input := &obs.ListObjectsInput{}
	input.Bucket = bucket
	// 设置每页100个对象
	input.MaxKeys = 100
	input.Prefix = prefix
	index := 1
	fileInfoList := FileInfoList{}

	prefixLen := len(prefix)
	log.Info("prefix=" + input.Prefix)
	for {
		output, err := ObsCli.ListObjects(input)
		if err == nil {
			log.Info("Page:%d\n", index)
			index++
			for _, val := range output.Contents {
				var isDir bool
				if prefixLen == len(val.Key) {
					continue
				}
				if strings.HasSuffix(val.Key, "/") {
					isDir = true
				} else {
					isDir = false
				}
				fileInfo := FileInfo{
					ModTime:  val.LastModified.Format("2006-01-02 15:04:05"),
					FileName: val.Key[prefixLen:],
					Size:     val.Size,
					IsDir:    isDir,
					ParenDir: "",
				}
				fileInfoList = append(fileInfoList, fileInfo)
			}
			if output.IsTruncated {
				input.Marker = output.NextMarker
			} else {
				break
			}
		} else {
			if obsError, ok := err.(obs.ObsError); ok {
				log.Info("Code:%s\n", obsError.Code)
				log.Info("Message:%s\n", obsError.Message)
			}
			return nil, err
		}
	}
	sort.Sort(fileInfoList)
	return fileInfoList, nil
}

func GetObsListObject(jobName, outPutPath, parentDir, versionName string) ([]FileInfo, error) {
	input := &obs.ListObjectsInput{}
	input.Bucket = setting.Bucket
	input.Prefix = strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, outPutPath, versionName, parentDir), "/")
	log.Info("bucket=" + input.Bucket + "  Prefix=" + input.Prefix)
	strPrefix := strings.Split(input.Prefix, "/")
	output, err := ObsCli.ListObjects(input)
	fileInfos := make([]FileInfo, 0)
	if err == nil {
		for _, val := range output.Contents {
			str1 := strings.Split(val.Key, "/")
			var isDir bool
			var fileName, nextParentDir string
			if strings.HasSuffix(val.Key, "/") {
				//dirs in next level dir
				if len(str1)-len(strPrefix) > 2 {
					continue
				}
				fileName = str1[len(str1)-2]
				isDir = true
				if parentDir == "" {
					nextParentDir = fileName
				} else {
					nextParentDir = parentDir + "/" + fileName
				}

				if fileName == strPrefix[len(strPrefix)-1] || (fileName+"/") == outPutPath {
					continue
				}
			} else {
				//files in next level dir
				if len(str1)-len(strPrefix) > 1 {
					continue
				}
				fileName = str1[len(str1)-1]
				isDir = false
				nextParentDir = parentDir
			}

			fileInfo := FileInfo{
				ModTime:  val.LastModified.Local().Format("2006-01-02 15:04:05"),
				FileName: fileName,
				Size:     val.Size,
				IsDir:    isDir,
				ParenDir: nextParentDir,
			}
			fileInfos = append(fileInfos, fileInfo)
		}
		sort.Slice(fileInfos, func(i, j int) bool {
			return fileInfos[i].ModTime > fileInfos[j].ModTime
		})
		return fileInfos, err
	} else {
		if obsError, ok := err.(obs.ObsError); ok {
			log.Error("Code:%s, Message:%s", obsError.Code, obsError.Message)
		}
		return nil, err
	}
}

func ObsGenMultiPartSignedUrl(uuid string, uploadId string, partNumber int, fileName string) (string, error) {

	input := &obs.CreateSignedUrlInput{}
	input.Bucket = setting.Bucket
	input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")
	input.Expires = 60 * 60
	input.Method = obs.HttpMethodPut

	input.QueryParams = map[string]string{
		"partNumber": com.ToStr(partNumber, 10),
		"uploadId":   uploadId,
		//"partSize":	com.ToStr(partSize,10),
	}

	output, err := ObsCli.CreateSignedUrl(input)
	if err != nil {
		log.Error("CreateSignedUrl failed:", err.Error())
		return "", err
	}

	return output.SignedUrl, nil
}

func GetObsCreateSignedUrlByBucketAndKey(bucket, key string) (string, error) {
	input := &obs.CreateSignedUrlInput{}
	input.Bucket = bucket
	input.Key = key

	input.Expires = 60 * 60
	input.Method = obs.HttpMethodGet
	comma := strings.LastIndex(key, "/")
	filename := key
	if comma != -1 {
		filename = key[comma+1:]
	}
	reqParams := make(map[string]string)
	filename = url.PathEscape(filename)
	reqParams["response-content-disposition"] = "attachment; filename=\"" + filename + "\""
	input.QueryParams = reqParams
	output, err := ObsCli.CreateSignedUrl(input)
	if err != nil {
		log.Error("CreateSignedUrl failed:", err.Error())
		return "", err
	}

	return output.SignedUrl, nil
}

func GetObsCreateSignedUrl(jobName, parentDir, fileName string) (string, error) {
	return GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir, fileName), "/"))
}

func ObsGetPreSignedUrl(uuid, fileName string) (string, error) {
	input := &obs.CreateSignedUrlInput{}
	input.Method = obs.HttpMethodGet
	input.Key = strings.TrimPrefix(path.Join(setting.BasePath, path.Join(uuid[0:1], uuid[1:2], uuid, fileName)), "/")
	input.Bucket = setting.Bucket
	input.Expires = 60 * 60

	fileName = url.PathEscape(fileName)
	reqParams := make(map[string]string)
	reqParams["response-content-disposition"] = "attachment; filename=\"" + fileName + "\""
	input.QueryParams = reqParams
	output, err := ObsCli.CreateSignedUrl(input)
	if err != nil {
		log.Error("CreateSignedUrl failed:", err.Error())
		return "", err
	}

	return output.SignedUrl, nil
}

func ObsCreateObject(path string) error {
	input := &obs.PutObjectInput{}
	input.Bucket = setting.Bucket
	input.Key = path

	_, err := ObsCli.PutObject(input)
	if err != nil {
		log.Error("PutObject failed:", err.Error())
		return err
	}

	return nil
}

func GetObsLogFileName(prefix string) (string, error) {
	input := &obs.ListObjectsInput{}
	input.Bucket = setting.Bucket
	input.Prefix = prefix

	output, err := ObsCli.ListObjects(input)
	if err != nil {
		log.Error("PutObject failed:", err.Error())
		return "", err
	}

	return output.Contents[0].Key, nil
}