aihub-dag-helmbe/dag-jupyter/internal/api/helm-wrapper.go

package api

import (
	"dag-jupyter/internal/types"
	"fmt"

	"github.com/parnurzeal/gorequest"
	"github.com/zeromicro/go-zero/core/logx"
)

type Resp struct {
	Code  int         `json:"code"`
	Data  interface{} `json:"data,omitempty"`
	Error string      `json:"error,omitempty"`
}

type StatusResp struct {
	Code  int                    `json:"code"`
	Data  []*types.ReleaseStatus `json:"data,omitempty"`
	Error string                 `json:"error,omitempty"`
}

const url = "http://localhost:8081"

// const url = "http://192.168.199.109:18080"

func genParam(req *types.CreateJupyterInfo) map[string]interface{} {
	settings := make([]string, 0, 16)
	settings = append(settings, fmt.Sprintf("ingress.jupyterlab.hosts[0].host=%s", req.Host))
	settings = append(settings, fmt.Sprintf("ingress.jupyterlab.hosts[0].paths[0].path=%s", req.Path))
	settings = append(settings, fmt.Sprintf("ingress.jupyterlab.hosts[0].paths[0].pathType=%s", req.PathType))
	logx.Info(req.IngressClass, "-------  len: ", len(req.IngressClass))
	if len(req.IngressClass) > 0 {

		settings = append(settings, fmt.Sprintf("ingress.jupyterlab.annotations.kubernetes\\.io/ingress\\.class=%s", req.IngressClass))
		settings = append(settings, fmt.Sprintf("ingress.jupyterlab.annotations.nginx\\.ingress\\.kubernetes\\.io/proxy-body-size=%s", "8192m"))

	}
	if req.NodeSelector != "" {
		settings = append(settings, fmt.Sprintf("jupyterlab.nodeSelector.kubernetes\\.io/hostname=%s", req.NodeSelector))
	}

	if req.ImagePullSecret != "" {
		settings = append(settings, fmt.Sprintf("jupyterlab.imagePullSecrets[0].name=%s", req.ImagePullSecret))
	}
	settings = append(settings, fmt.Sprintf("jupyterlab.image.repository=%s", req.Image))
	settings = append(settings, fmt.Sprintf("jupyterlab.image.tag=%s", req.Tag))
	settings = append(settings, fmt.Sprintf("jupyterlab.config.baseUrl=%s", req.BaseUrl))
	settings = append(settings, fmt.Sprintf("jupyterlab.config.workspace=%s", req.Workspace))
	settings = append(settings, fmt.Sprintf("jupyterlab.config.password=%s", req.Password))

	m := map[string]interface{}{
		"dry_run":          false,
		"create_namespace": false,
		"set":              settings,
	}

	return m
}

func Install(req *types.CreateJupyterInfo) error {
	m := genParam(req)

	request := gorequest.New()
	var resp Resp
	_, _, errs := request.Post(fmt.Sprintf("%s/api/namespaces/%s/releases/%s?chart=%s", url, req.Namespace, req.ReleaseName, req.Chart)).SendMap(m).EndStruct(&resp)

	if len(errs) != 0 {
		return fmt.Errorf("helm install error: %s", errs[0])
	}

	if resp.Code != 0 {
		return fmt.Errorf("helm install error: %s", resp.Error)
	}

	logx.Info("helm install success~~")

	return nil
}

func Uninstall(req *types.DeleteJupyterInfo) error {
	request := gorequest.New()
	var resp Resp
	_, _, errs := request.Delete(fmt.Sprintf("%s/api/namespaces/%s/releases/%s", url, req.Namespace, req.ReleaseName)).EndStruct(&resp)

	if len(errs) != 0 {
		return fmt.Errorf("helm uninstall error: %s", errs[0])
	}

	if resp.Code != 0 {
		return fmt.Errorf("helm uninstall error: %s", resp.Error)
	}

	logx.Info("helm uninstall success~~")

	return nil
}

func Upgrade(req *types.CreateJupyterInfo) error {
	m := genParam(req)

	request := gorequest.New()
	var resp Resp
	_, _, errs := request.Put(fmt.Sprintf("%s/api/namespaces/%s/releases/%s?chart=%s", url, req.Namespace, req.ReleaseName, req.Chart)).SendMap(m).EndStruct(&resp)

	if len(errs) != 0 {
		return fmt.Errorf("helm upgrade error: %s", errs[0])
	}

	if resp.Code != 0 {
		return fmt.Errorf("helm upgrade error: %s", resp.Error)
	}

	logx.Info("helm upgrade success~~")

	return nil
}

func Status(req *types.ReleaseInfo) ([]*types.ReleaseStatus, error) {
	var resp StatusResp
	request := gorequest.New()
	m := map[string]string{
		"filter": req.Filter,
	}
	_, _, errs := request.Get(fmt.Sprintf("%s/api/namespaces/%s/releases", url, req.Namespace)).SendMap(m).EndStruct(&resp)

	if len(errs) != 0 {
		return nil, fmt.Errorf("helm status error: %s", errs[0])
	}

	if resp.Code != 0 {
		return nil, fmt.Errorf("helm status error: %s", resp.Error)
	}

	logx.Info("helm status success~~")

	return resp.Data, nil

}