pcm-coordinator/adaptor/pcm_slurm/service/slurmdb_user.go

package slurmer

/*
#cgo LDFLAGS: -lslurmdb

#include <stdio.h>
#include <slurm/slurm.h>
#include <slurm/slurmdb.h>
#include <memory.h>
#include <malloc.h>

typedef struct user_info_msg {
    uint32_t record_count;
    slurmdb_user_rec_t *user_array;
} user_info_msg_t;

typedef struct slurmdb_user_rec{
    uint16_t admin_level;
    List assoc_list;
    List coord_accts;
    char *default_acct;
    char *default_wckey;
    char *name;
    char *old_name;
    uint32_t uid;
    List wckey_list;
} slurmdb_user_rec_pcm;

struct user_info_msg get_user_info() {
    struct user_info_msg userinfo;
    List userList = NULL;
    slurmdb_user_cond_t *user_cond = NULL;
    void *db_conn;
    db_conn = slurmdb_connection_get();
    userList = slurmdb_users_get(db_conn, user_cond);
    slurmdb_connection_close(&db_conn);

    slurmdb_user_rec_t *rec = NULL;
    ListIterator itr = slurm_list_iterator_create(userList);
    int i = 0;
    uint32_t length;
    length = slurm_list_count(userList);
    userinfo.record_count = length;
    userinfo.user_array = malloc(length * sizeof(slurmdb_user_rec_t));
    while ((rec = slurm_list_next(itr))) {
        userinfo.user_array[i] = *rec;
        i++;
    }
    return userinfo;
}

struct slurmdb_user_rec *user_from_list(struct user_info_msg *list, int i) {
    return (struct slurmdb_user_rec *) &list->user_array[i];
}

*/
import "C"
import (
	pbslurm "code.gitlink.org.cn/JCCE/PCM.git/adaptor/pcm_slurm/gen/idl"
	ssh "code.gitlink.org.cn/JCCE/PCM.git/common/ssh"
	"context"
	"strings"
)

type UserInfoMsg struct {
	LastUpdate   int64
	RecordCount  uint32
	UserInfoList []pbslurm.UserInfo
}

func UserDescriptorConvertCToGo(cStruct *C.struct_slurmdb_user_rec) pbslurm.UserInfo {
	var goStruct pbslurm.UserInfo
	goStruct.Name = C.GoString(cStruct.name)
	return goStruct
}

func GetUserInfo() UserInfoMsg {
	var goUserBuffer UserInfoMsg
	cUserBuffer := C.get_user_info()
	goUserBuffer.RecordCount = uint32(cUserBuffer.record_count)
	goUserBuffer.UserInfoList = make([]pbslurm.UserInfo, cUserBuffer.record_count, cUserBuffer.record_count)

	for i := uint32(0); i < goUserBuffer.RecordCount; i++ {
		user := C.user_from_list(&cUserBuffer, C.int(i))
		goUser := UserDescriptorConvertCToGo(user)
		goUserBuffer.UserInfoList[i] = goUser
	}
	return goUserBuffer
}

func (slurmStruct SlurmStruct) ListUsers(ctx context.Context, req *pbslurm.ListUsersReq) (*pbslurm.ListUsersResp, error) {

	userList := GetUserInfo()

	resp := pbslurm.ListUsersResp{}
	for _, user := range userList.UserInfoList {
		userInfoResult := user
		//userInfoResult.Name = user.Name

		resp.UserInfos = append(resp.UserInfos, &userInfoResult)
	}

	return &resp, nil
}

func (slurmStruct SlurmStruct) GetUser(ctx context.Context, req *pbslurm.GetUserReq) (*pbslurm.GetUserResp, error) {

	userList := GetUserInfo()

	resp := pbslurm.GetUserResp{}
	for _, user := range userList.UserInfoList {
		if strings.Contains(user.Name, req.UserName) {
			userInfoResult := user
			resp.UserInfo = append(resp.UserInfo, &userInfoResult)
		}

	}

	return &resp, nil
}

func (slurmStruct SlurmStruct) AddUser(ctx context.Context, req *pbslurm.AddUserReq) (*pbslurm.AddUserResp, error) {

	cmd := "/usr/local/bin/sacctmgr add user "

	cmd = cmd + req.Names
	if len(req.Accounts) != 0 {
		cmd = cmd + " Accounts=" + req.Accounts
	}
	if len(req.AdminLevel) != 0 {
		cmd = cmd + " AdminLevel=" + req.AdminLevel
	}
	if len(req.Clusters) != 0 {
		cmd = cmd + " Clusters=" + req.Clusters
	}
	if len(req.DefaultAccount) != 0 {
		cmd = cmd + " DefaultAccount=" + req.DefaultAccount
	}
	if len(req.DefaultQos) != 0 {
		cmd = cmd + " DefaultQOS=" + req.DefaultQos
	}
	if len(req.DefaultWckey) != 0 {
		cmd = cmd + " DefaultWCKey=" + req.DefaultWckey
	}
	if len(req.FairShare) != 0 {
		cmd = cmd + " Fairshare=" + req.FairShare
	}
	if len(req.MaxCpuMins) != 0 {
		cmd = cmd + " MaxCPUMins=" + req.MaxCpuMins
	}
	if len(req.MaxCpus) != 0 {
		cmd = cmd + " MaxCPUs=" + req.MaxCpus
	}
	if len(req.MaxJobs) != 0 {
		cmd = cmd + " MaxJobs=" + req.MaxJobs
	}
	if len(req.MaxNodes) != 0 {
		cmd = cmd + " MaxNodes=" + req.MaxNodes
	}
	if len(req.MaxSubmitJobs) != 0 {
		cmd = cmd + " MaxSubmitJobs=" + req.MaxSubmitJobs
	}
	if len(req.MaxWall) != 0 {
		cmd = cmd + " MaxWall=" + req.MaxWall
	}
	if len(req.Partitions) != 0 {
		cmd = cmd + " Partitions=" + req.Partitions
	}
	if len(req.QosLevel) != 0 {
		cmd = cmd + " QosLevel=" + req.QosLevel
	}

	cmd = cmd + " -i"

	result := ssh.ExecCommand(cmd)

	resp := pbslurm.AddUserResp{}
	resp.Result = result

	return &resp, nil
}

func (slurmStruct SlurmStruct) DeleteUser(ctx context.Context, req *pbslurm.DeleteUserReq) (*pbslurm.DeleteUserResp, error) {

	cmd := "/usr/local/bin/sacctmgr delete user "
	cmd = cmd + req.Names
	cmd = cmd + " -i"

	result := ssh.ExecCommand(cmd)

	resp := pbslurm.DeleteUserResp{}
	resp.Result = result

	return &resp, nil
}