Adding BucketExist check during Upload-Model

Issue-Id: AIMLFW-208
Change-Id: I8813f790a7d20e8b7d8cb3db160b560cd05c7a21
Signed-off-by: ashishj1729 <jain.ashish@samsung.com>

diff --git a/apis/mmes_apis.go b/apis/mmes_apis.go
index 0ec3ea8..e0aeda6 100644 (file)
--- a/apis/mmes_apis.go
+++ b/apis/mmes_apis.go
@@ -18,11 +18,11 @@ limitations under the License.
 package apis
 
 import (
+       "errors"
        "fmt"
        "io"
        "net/http"
        "os"
-       "errors"
 
        "gerrit.o-ran-sc.org/r/aiml-fw/awmf/modelmgmtservice/core"
        "gerrit.o-ran-sc.org/r/aiml-fw/awmf/modelmgmtservice/db"
@@ -110,7 +110,6 @@ func (m *MmeApiHandler) RegisterModel(cont *gin.Context) {
                "modelInfo": modelInfo,
        })
 
-
 }
 
 /*
@@ -131,7 +130,7 @@ func (m *MmeApiHandler) GetModelInfo(cont *gin.Context) {
                        logging.ERROR("error:", "Only allowed params are modelname and modelversion")
                        cont.JSON(http.StatusBadRequest, models.ProblemDetail{
                                Status: http.StatusBadRequest,
-                               Title: "Bad Request",
+                               Title:  "Bad Request",
                                Detail: fmt.Sprintf("Only allowed params are modelname and modelversion"),
                        })
                        return
@@ -164,7 +163,7 @@ func (m *MmeApiHandler) GetModelInfo(cont *gin.Context) {
                                logging.ERROR("Error occurred, send status code: ", statusCode)
                                cont.JSON(statusCode, models.ProblemDetail{
                                        Status: http.StatusInternalServerError,
-                                       Title: "Internal Server Error",
+                                       Title:  "Internal Server Error",
                                        Detail: fmt.Sprintf("Can't fetch the models due to , %s", err.Error()),
                                })
                                return
@@ -259,12 +258,18 @@ func (m *MmeApiHandler) UploadModel(cont *gin.Context) {
        fileHeader, _ := cont.FormFile("file")
        //TODO : Accept only .zip file for trained model
        file, _ := fileHeader.Open()
-       //TODO: Handle error response
        defer file.Close()
        byteFile, _ := io.ReadAll((file))
 
-       logging.INFO("Uploading model : ", modelName)
-       m.dbmgr.UploadFile(byteFile, modelName+os.Getenv("MODEL_FILE_POSTFIX"), modelName)
+       logging.INFO("Uploading model : " + modelName)
+       if err := m.dbmgr.UploadFile(byteFile, modelName+os.Getenv("MODEL_FILE_POSTFIX"), modelName); err != nil {
+               logging.ERROR("Failed to Upload Model : ", err)
+               cont.JSON(http.StatusInternalServerError, gin.H{
+                       "code":    http.StatusInternalServerError,
+                       "message": err.Error(),
+               })
+               return
+       }
        cont.JSON(http.StatusOK, gin.H{
                "code":    http.StatusOK,
                "message": string("Model uploaded successfully.."),
@@ -305,7 +310,7 @@ func (m *MmeApiHandler) UpdateModel(c *gin.Context) {
        }
        existingModelInfo, err := m.iDB.GetModelInfoById(id)
 
-       if err != nil || existingModelInfo.Id == ""{
+       if err != nil || existingModelInfo.Id == "" {
                statusCode := http.StatusNotFound
                logging.ERROR("Error occurred, send status code: ", statusCode)
                c.JSON(statusCode, gin.H{
@@ -314,7 +319,7 @@ func (m *MmeApiHandler) UpdateModel(c *gin.Context) {
                })
                return
        }
-       if existingModelInfo.ModelId.ModelName != modelInfo.ModelId.ModelName || existingModelInfo.ModelId.ModelVersion != modelInfo.ModelId.ModelVersion{
+       if existingModelInfo.ModelId.ModelName != modelInfo.ModelId.ModelName || existingModelInfo.ModelId.ModelVersion != modelInfo.ModelId.ModelVersion {
                statusCode := http.StatusBadRequest
                logging.ERROR("Error occurred, send status code: ", statusCode)
                c.JSON(statusCode, gin.H{

diff --git a/core/s3_manager.go b/core/s3_manager.go
index 867de59..558b56d 100644 (file)
--- a/core/s3_manager.go
+++ b/core/s3_manager.go
@@ -20,6 +20,7 @@ package core
 import (
        "bytes"
        "errors"
+       "fmt"
        "io"
        "os"
 
@@ -48,7 +49,7 @@ type DBMgr interface {
        GetBucketObject(objectName string, bucketName string) BucketObject
        DeleteBucket(client *s3.S3, objectName string, bucketName string)
        DeleteBucketObject(client *s3.S3, objectName string, bucketName string) bool
-       UploadFile(dataBytes []byte, file_name string, bucketName string)
+       UploadFile(dataBytes []byte, file_name string, bucketName string) error
        ListBucket(bucketObjPostfix string) ([]Bucket, error)
        GetBucketItems(bucketName string)
 }
@@ -107,7 +108,7 @@ func (s3manager *S3Manager) CreateBucket(bucketName string) (err error) {
                        return
                }
        }
-       println("Bucket created : ", bucketName)
+       logging.INFO("Bucket created : " + bucketName)
        return nil
 }
 
@@ -163,19 +164,54 @@ func (s3manager *S3Manager) DeleteBucketObject(client *s3.S3, objectName string,
        return true
 }
 
-func (s3manager *S3Manager) UploadFile(dataBytes []byte, file_name string, bucketName string) {
+func (s3manager *S3Manager) checkIfBucketExists(bucketName string) (bool, error) {
+       _, err := s3manager.S3Client.HeadBucket(&s3.HeadBucketInput{
+               Bucket: aws.String(bucketName),
+       })
+       if err != nil {
+               // Check if the error means the bucket doesn't exist
+               if aerr, ok := err.(awserr.Error); ok {
+                       switch aerr.Code() {
+                       case "NotFound":
+                               return false, nil
+                       default:
+                               return false, fmt.Errorf("error checking bucket: %v", aerr)
+                       }
+               }
+               return false, fmt.Errorf("error checking bucket: %v", err)
+       }
+       return true, nil
+}
+
+func (s3manager *S3Manager) UploadFile(dataBytes []byte, file_name string, bucketName string) error {
+
+       doesBucketExist, err := s3manager.checkIfBucketExists(bucketName)
+       if err != nil {
+               logging.DEBUG(fmt.Sprintf("unable to check bucket %s existence, Error : %v", bucketName, err))
+               return fmt.Errorf("unable to check bucket %s existence, Error : %v", bucketName, err)
+       }
+       if !doesBucketExist {
+               logging.INFO("Bucket " + bucketName + " Doesn't Exist, Creating One")
+               if err := s3manager.CreateBucket(bucketName); err != nil {
+                       logging.DEBUG(fmt.Sprintf("unable to create bucket for uploading-model, Error : %v", err))
+                       return fmt.Errorf("unable to create bucket for uploading-model, Error : %v", err)
+               }
+       }
 
        dataReader := bytes.NewReader(dataBytes) //bytes.Reader is type of io.ReadSeeker
+
        params := &s3.PutObjectInput{
                Bucket: aws.String(bucketName),
                Key:    aws.String(file_name),
                Body:   dataReader,
        }
-       _, err := s3manager.S3Client.PutObject(params)
+       _, err = s3manager.S3Client.PutObject(params)
        if err != nil {
                logging.ERROR("Error in uploading file to bucket ", err)
+               return err
        }
-       logging.INFO("File uploaded to bucket ", bucketName)
+       logging.INFO("File uploaded to bucket " + bucketName)
+       return nil
 }
 
 func (s3manager *S3Manager) ListBucket(bucketObjPostfix string) ([]Bucket, error) {