aiforge/modules/cloudbrain/resty.go

package cloudbrain

import (
	"code.gitea.io/gitea/modules/log"
	"fmt"

	"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
)

const (
	JobHasBeenStopped = "S410"
)

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 {
		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) {
	checkSetting()
	client := getRestyClient()
	var getImagesResult models.GetImagesResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetQueryString("pageIndex=1&pageSize=100").
		SetResult(&getImagesResult).
		Get(HOST + "/rest-server/api/v1/image/list/")

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

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

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

	return &getImagesResult, nil
}

func GetPublicImages() (*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("pageIndex=1&pageSize=100").
		SetResult(&getImagesResult).
		Get(HOST + "/rest-server/api/v1/image/public/list/")

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

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

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

	return &getImagesResult, nil
}

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.StopJobResult

	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
}