From deb1f242b1feb755082172c728de4f539ae53779 Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Mon, 8 Nov 2021 17:38:11 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E6=B5=8B=E8=AF=95=E4=BB=A3?=
 =?UTF-8?q?=E7=A0=81=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zouap <zouap@pcl.ac.cn>
---
 modules/storage/obs.go          | 47 +++++++++++++++++++++++++++------
 routers/repo/ai_model_manage.go |  4 +--
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/modules/storage/obs.go b/modules/storage/obs.go
index b6083af7b..123ff4094 100755
--- a/modules/storage/obs.go
+++ b/modules/storage/obs.go
@@ -239,6 +239,39 @@ func ObsCopyFile(srcBucket string, srcKeyName string, destBucket string, destKey
 	return nil
 }
 
+func GetAllObsListObjectUnderDir(bucket string, prefix string) ([]FileInfo, error) {
+	input := &obs.ListObjectsInput{}
+	input.Bucket = bucket
+	input.Prefix = prefix
+	output, err := ObsCli.ListObjects(input)
+	fileInfos := make([]FileInfo, 0)
+	if err == nil {
+		for _, val := range output.Contents {
+			var isDir bool
+			if strings.HasSuffix(val.Key, "/") {
+				continue
+			} else {
+				isDir = false
+			}
+			fileInfo := FileInfo{
+				ModTime:  val.LastModified.Format("2006-01-02 15:04:05"),
+				FileName: val.Key[len(prefix):],
+				Size:     val.Size,
+				IsDir:    isDir,
+				ParenDir: "",
+			}
+			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 GetObsListObjectByBucketAndPrefix(bucket string, prefix string) ([]FileInfo, error) {
 	input := &obs.ListObjectsInput{}
 	input.Bucket = bucket
@@ -246,17 +279,18 @@ func GetObsListObjectByBucketAndPrefix(bucket string, prefix string) ([]FileInfo
 	output, err := ObsCli.ListObjects(input)
 	fileInfos := make([]FileInfo, 0)
 	if err == nil {
+		var nextParentDir string
 		for _, val := range output.Contents {
 			str1 := strings.Split(val.Key, "/")
 			var isDir bool
-			var fileName, nextParentDir string
+			var fileName string
 			if strings.HasSuffix(val.Key, "/") {
-				fileName = str1[len(str1)-2]
-				isDir = true
-				nextParentDir = fileName
-				if (fileName + "/") == setting.OutPutPath {
+				if prefix == val.Key {
 					continue
 				}
+				fileName = str1[len(str1)-2]
+				isDir = true
+				nextParentDir = nextParentDir + fileName + "/"
 			} else {
 				fileName = str1[len(str1)-1]
 				isDir = false
@@ -292,12 +326,10 @@ func GetObsListObject(jobName, parentDir string) ([]FileInfo, error) {
 			str1 := strings.Split(val.Key, "/")
 			var isDir bool
 			var fileName, nextParentDir string
-			log.Info("val.Key=" + val.Key)
 			if strings.HasSuffix(val.Key, "/") {
 				fileName = str1[len(str1)-2]
 				isDir = true
 				nextParentDir = fileName
-				log.Info("nextParentDir1=" + nextParentDir)
 				if fileName == parentDir || (fileName+"/") == setting.OutPutPath {
 					continue
 				}
@@ -305,7 +337,6 @@ func GetObsListObject(jobName, parentDir string) ([]FileInfo, error) {
 				fileName = str1[len(str1)-1]
 				isDir = false
 			}
-			log.Info("nextParentDir2=" + nextParentDir)
 			fileInfo := FileInfo{
 				ModTime:  val.LastModified.Format("2006-01-02 15:04:05"),
 				FileName: fileName,
diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go
index 692c34774..b63763fa3 100644
--- a/routers/repo/ai_model_manage.go
+++ b/routers/repo/ai_model_manage.go
@@ -99,8 +99,8 @@ func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir
 	dataActualPath := setting.Bucket + "/" + Model_prefix +
 		models.AttachmentRelativePath(modelUUID) +
 		"/"
-
-	modelDbResult, err := storage.GetObsListObject(jobName, parentDir)
+	objectkey := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir), "/")
+	modelDbResult, err := storage.GetAllObsListObjectUnderDir(setting.Bucket, objectkey)
 	if err != nil {
 		log.Info("get TrainJobListModel failed:", err)
 		return "", 0, err