aiforge/modules/cloudbrain/resty.go

package cloudbrain

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"

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

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/setting"
	"github.com/go-resty/resty/v2"
)

var (
	restyClient  *resty.Client
	HOST         string
	TOKEN        string
	ImagesUrlMap = map[string]string{Public: "/rest-server/api/v1/image/public/list/", Custom: "/rest-server/api/v1/image/list/"}
)

const (
	JobHasBeenStopped = "S410"
	Public            = "public"
	Custom            = "custom"
	LogPageSize		  = 500
	LogPageTokenExpired = "5m"
)

func getRestyClient() *resty.Client {
	if restyClient == nil {
		restyClient = resty.New()
	}
	return restyClient
}

func checkSetting() {
	if len(HOST) != 0 && len(TOKEN) != 0 && restyClient != nil {
		return
	}
	_ = loginCloudbrain()
}

func loginCloudbrain() error {
	conf := setting.GetCloudbrainConfig()

	username := conf.Username
	password := conf.Password
	HOST = conf.Host
	var loginResult models.CloudBrainLoginResult

	client := getRestyClient()

	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody(map[string]interface{}{"username": username, "password": password, "expiration": "604800"}).
		SetResult(&loginResult).
		Post(HOST + "/rest-server/api/v1/token")
	if err != nil {
		return fmt.Errorf("resty loginCloudbrain: %s", err)
	}

	if loginResult.Code != Success {
		return fmt.Errorf("%s: %s", loginResult.Msg, res.String())
	}

	TOKEN = loginResult.Payload["token"].(string)
	return nil
}

func CreateJob(jobName string, createJobParams models.CreateJobParams) (*models.CreateJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var jobResult models.CreateJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(createJobParams).
		SetResult(&jobResult).
		Post(HOST + "/rest-server/api/v1/jobs/")

	if err != nil {
		if res != nil {
			var response models.CloudBrainResult
			json.Unmarshal(res.Body(), &response)
			log.Error("code(%s), msg(%s)", response.Code, response.Msg)
			return nil, fmt.Errorf(response.Msg)
		}
		return nil, fmt.Errorf("resty create job: %s", err)
	}

	if jobResult.Code == "S401" && retry < 1 {
		retry++
		_ = loginCloudbrain()
		goto sendjob
	}

	if jobResult.Code != Success {
		return &jobResult, fmt.Errorf("jobResult err: %s", res.String())
	}

	return &jobResult, nil
}

func GetJob(jobID string) (*models.GetJobResult, error) {
	checkSetting()
	// http://192.168.204.24/rest-server/api/v1/jobs/90e26e500c4b3011ea0a251099a987938b96
	client := getRestyClient()
	var getJobResult models.GetJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&getJobResult).
		Get(HOST + "/rest-server/api/v1/jobs/" + jobID)

	if err != nil {
		return nil, fmt.Errorf("resty GetJob: %v", err)
	}

	if getJobResult.Code == "S401" && retry < 1 {
		retry++
		_ = loginCloudbrain()
		goto sendjob
	}

	if getJobResult.Code != Success {
		return &getJobResult, fmt.Errorf("jobResult GetJob err: %s", res.String())
	}

	return &getJobResult, nil
}

func GetImages() (*models.GetImagesResult, error) {

	return GetImagesPageable(1, 100, Custom, "")

}

func GetPublicImages() (*models.GetImagesResult, error) {
	return GetImagesPageable(1, 100, Public, "")
}

func GetImagesPageable(page int, size int, imageType string, name string) (*models.GetImagesResult, error) {
	checkSetting()
	client := getRestyClient()
	var getImagesResult models.GetImagesResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetQueryString(getQueryString(page, size, name)).
		SetResult(&getImagesResult).
		Get(HOST + ImagesUrlMap[imageType])

	if err != nil {
		return nil, fmt.Errorf("resty GetImages: %v", err)
	}

	var response models.CloudBrainResult
	err = json.Unmarshal(res.Body(), &response)
	if err != nil {
		log.Error("json.Unmarshal failed: %s", err.Error())
		return &getImagesResult, fmt.Errorf("json.Unmarshal failed: %s", err.Error())
	}

	if response.Code == "S401" && retry < 1 {
		retry++
		_ = loginCloudbrain()
		goto sendjob
	}

	if getImagesResult.Code != Success {
		return &getImagesResult, fmt.Errorf("getImagesResult err: %s", res.String())
	}

	getImagesResult.Payload.TotalPages = getTotalPages(getImagesResult, size)
	return &getImagesResult, nil
}

func getTotalPages(getImagesResult models.GetImagesResult, size int) int {
	totalCount := getImagesResult.Payload.Count
	var totalPages int
	if totalCount%size != 0 {
		totalPages = totalCount/size + 1
	} else {
		totalPages = totalCount / size
	}
	return totalPages
}

func getQueryString(page int, size int, name string) string {
	if strings.TrimSpace(name) == "" {
		return fmt.Sprintf("pageIndex=%d&pageSize=%d", page, size)
	}
	return fmt.Sprintf("pageIndex=%d&pageSize=%d&name=%s", page, size, name)
}

func CommitImage(jobID string, params models.CommitImageParams) error {
	checkSetting()
	client := getRestyClient()
	var result models.CommitImageResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(params).
		SetResult(&result).
		Post(HOST + "/rest-server/api/v1/jobs/" + jobID + "/commitImage")

	if err != nil {
		return fmt.Errorf("resty CommitImage: %v", err)
	}

	if result.Code == "S401" && retry < 1 {
		retry++
		_ = loginCloudbrain()
		goto sendjob
	}

	if result.Code != Success {
		return fmt.Errorf("CommitImage err: %s", res.String())
	}

	return nil
}

func StopJob(jobID string) error {
	checkSetting()
	client := getRestyClient()
	var result models.CloudBrainResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&result).
		Delete(HOST + "/rest-server/api/v1/jobs/" + jobID)

	if err != nil {
		return fmt.Errorf("resty StopJob: %v", err)
	}

	if result.Code == "S401" && retry < 1 {
		retry++
		_ = loginCloudbrain()
		goto sendjob
	}

	if result.Code != Success {
		if result.Code == JobHasBeenStopped {
			log.Info("StopJob(%s) failed:%s", jobID, result.Msg)
		} else {
			return fmt.Errorf("StopJob err: %s", res.String())
		}
	}

	return nil
}

func GetJobLog(jobID string) (*models.GetJobLogResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetJobLogResult
	req := models.GetJobLogParams{
		Size: strconv.Itoa(LogPageSize),
		Sort: "log.offset",
		QueryInfo: models.QueryInfo{
			MatchInfo: models.MatchInfo{
				PodName: jobID + "-task1-0",
			},
		},
	}

	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(req).
		SetResult(&result).
		Post(HOST + "es/_search?_source=message&scroll=" + LogPageTokenExpired)

	if err != nil {
		log.Error("GetJobLog failed: %v", err)
		return &result, fmt.Errorf("resty GetJobLog: %v, %s", err, res.String())
	}

	if !strings.Contains(res.Status(), strconv.Itoa(http.StatusOK)) {
		log.Error("res.Status(): %s, response: %s", res.Status(), res.String())
		return &result, errors.New(res.String())
	}

	return &result, nil
}

func GetJobAllLog(scrollID string) (*models.GetJobLogResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetJobLogResult
	req := models.GetAllJobLogParams{
		Scroll: LogPageTokenExpired,
		ScrollID: scrollID,
	}

	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(req).
		SetResult(&result).
		Post(HOST + "es/_search/scroll")

	if err != nil {
		log.Error("GetJobAllLog failed: %v", err)
		return &result, fmt.Errorf("resty GetJobAllLog: %v, %s", err, res.String())
	}

	if !strings.Contains(res.Status(), strconv.Itoa(http.StatusOK)) {
		log.Error("res.Status(): %s, response: %s", res.Status(), res.String())
		return &result, errors.New(res.String())
	}

	return &result, nil
}

func DeleteJobLogToken(scrollID string) (error) {
	checkSetting()
	client := getRestyClient()
	var result models.DeleteJobLogTokenResult
	req := models.DeleteJobLogTokenParams{
		ScrollID: scrollID,
	}

	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(req).
		SetResult(&result).
		Delete(HOST + "es/_search/scroll")

	if err != nil {
		log.Error("DeleteJobLogToken failed: %v", err)
		return fmt.Errorf("resty DeleteJobLogToken: %v, %s", err, res.String())
	}

	if !strings.Contains(res.Status(), strconv.Itoa(http.StatusOK)) {
		log.Error("res.Status(): %s, response: %s", res.Status(), res.String())
		return errors.New(res.String())
	}

	if !result.Succeeded {
		log.Error("DeleteJobLogToken failed")
		return errors.New("DeleteJobLogToken failed")
	}

	return nil
}