myxce/src/frontend/xcecli/cmd/jobs/controller.go

package jobs

import (
	"encoding/json"
	"fmt"
	"strings"
	"xcecli/common/format"
	"xcecli/common/request"
	"xcecli/config"
	"xcecli/models"

	"github.com/spf13/cobra"
)

// JobController defines job control
type JobController struct {
	configManager *config.ConfigManager
	action        string
	filter        models.JobControlFilter
	jobIDs        []string
}

// NewJobController create job controller
func NewJobController(configManager *config.ConfigManager, action string) *JobController {
	return &JobController{
		configManager: configManager,
		action:        action,
	}
}

// Run execute job control
func (c *JobController) RunJobsControl() error {
	// Get server information
	serverInfo, _ := c.configManager.GetServerInfo()

	// Create API client
	apiClient := request.NewAPIClient(serverInfo.URL)

	if len(c.jobIDs) > 0 {
		// if job IDs are specified, use job ID interface
		return c.controlJobsByIDs(apiClient, serverInfo.Token)
	} else {
		// otherwise use filter conditions
		return c.controlJobsByFilter(apiClient, serverInfo.Token)
	}
}

// control jobs by job IDs
func (c *JobController) controlJobsByIDs(apiClient *request.APIClient, token string) error {
	apiIDs := make([]string, 0, len(c.jobIDs))
	for _, id := range c.jobIDs {
		apiID, err := format.ToAPIJobID(strings.TrimSpace(id), true)
		if err != nil {
			return fmt.Errorf("invalid job ID %q: %w", id, err)
		}
		apiIDs = append(apiIDs, apiID)
	}
	url := fmt.Sprintf("/xce/v1/jobs/%s/control", strings.Join(apiIDs, ","))

	req := models.JobControlJobIDRequest{
		Action: c.action,
	}
	// encode request data to JSON bytes
	return c.printJobControlResponse(req, apiClient, token, url)
}

// control jobs by filter conditions
func (c *JobController) controlJobsByFilter(apiClient *request.APIClient, token string) error {
	url := "/xce/v1/jobs/control"

	req := models.JobControlRequest{
		Action: c.action,
		Filter: c.filter,
	}
	// encode request data to JSON bytes
	return c.printJobControlResponse(req, apiClient, token, url)
}

func (c *JobController) printJobControlResponse(req interface{}, apiClient *request.APIClient, token string, url string) error {
	reqData, err := json.Marshal(req)
	if err != nil {
		errMsg := request.ExtractErrorMessage(err)
		return fmt.Errorf("failed to encode request data: [%s]", errMsg)
	}
	var resp models.JobControlResponse
	err = apiClient.PostRequestToAPIServer(
		token,
		url,
		reqData,
		&resp,
	)

	if err != nil {
		return err
	}
	for _, data := range resp.Data.Accepted {
		fmt.Printf("job <%s> %s\n", format.FormatArrayChildID(data.JobID, data.ArrayIdx), strings.TrimPrefix(data.Msg, "job "))
	}

	for _, data := range resp.Data.Failed {
		fmt.Printf("job <%s> %s\n", format.FormatArrayChildID(data.JobID, data.ArrayIdx), strings.TrimPrefix(data.Msg, "job "))
	}
	return nil
}

// Parse parse command line arguments
// Parse command line arguments
func (c *JobController) ParseCmdArgs(cmd *cobra.Command, args []string) error {
	if len(args) > 0 {
		if len(args) == 1 && args[0] == "0" {
			c.filter.AllJobs = true
		} else {
			// parse job ID parameters
			c.jobIDs = strings.Split(args[0], ",")
			return nil
		}
	}
	// check if any filter conditions are set
	if c.filter.JobName == "" && c.filter.QueueName == "" && c.filter.Users == "" && !c.filter.AllJobs {
		// if no job IDs or filter conditions are specified, default to processing all jobs of the current user
		return fmt.Errorf("no job IDs or filter conditions specified")
	}
	return nil
}

// NewKillJobsCmd create kill job command
func NewKillJobsCmd(configManager *config.ConfigManager) *cobra.Command {
	controller := NewJobController(configManager, "kill")

	cmd := &cobra.Command{
		Use:   "kill",
		Short: "Kill jobs",
		Long: `Kill jobs by job IDs or filters.
Examples:
  xce jobs kill jobid1,jobid2,...      # kill jobs by job IDs
  xce jobs kill -J jobname     	   # kill jobs by job name
  xce jobs kill -q queue_name      # kill jobs in a specific queue
  xce jobs kill -u username        # kill jobs by user`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := controller.ParseCmdArgs(cmd, args); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			if err := controller.RunJobsControl(); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			return nil
		},
	}

	// add control options
	flags := cmd.Flags()
	flags.StringVarP(&controller.filter.JobName, "job_name", "J", "", "filter by job name")
	flags.StringVarP(&controller.filter.QueueName, "queue", "q", "", "filter by queue name")
	flags.StringVarP(&controller.filter.Users, "user", "u", "", "filter by user")
	flags.BoolVarP(&controller.filter.AllJobs, "all", "a", false, "select all jobs")

	return cmd
}

// NewStopJobsCmd create stop job command
func NewStopJobsCmd(configManager *config.ConfigManager) *cobra.Command {
	controller := NewJobController(configManager, "stop")

	cmd := &cobra.Command{
		Use:   "stop",
		Short: "Stop jobs",
		Long: `Stop jobs by job IDs or filters.
Examples:
  xce jobs stop jobid1,jobid2,...      # stop jobs by job IDs
  xce jobs stop -J jobname     	   # stop jobs by job name
  xce jobs stop -q queue_name      # stop jobs in a specific queue
  xce jobs stop -u username        # stop jobs by user`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := controller.ParseCmdArgs(cmd, args); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			if err := controller.RunJobsControl(); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			return nil
		},
	}

	// add control options
	flags := cmd.Flags()
	flags.StringVarP(&controller.filter.JobName, "job_name", "J", "", "filter by job name")
	flags.StringVarP(&controller.filter.QueueName, "queue", "q", "", "filter by queue name")
	flags.StringVarP(&controller.filter.Users, "user", "u", "", "filter by user")
	flags.BoolVarP(&controller.filter.AllJobs, "all", "a", false, "select all jobs")

	return cmd
}

// NewResumeJobsCmd create resume job command
func NewResumeJobsCmd(configManager *config.ConfigManager) *cobra.Command {
	controller := NewJobController(configManager, "resume")

	cmd := &cobra.Command{
		Use:   "resume",
		Short: "Resume jobs",
		Long: `Resume jobs by job IDs or filters.
Examples:
  xce jobs resume jobid1,jobid2,...      # resume jobs by job IDs
  xce jobs resume -J jobname     	   # resume jobs by job name
  xce jobs resume -u username        # resume jobs by user`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := controller.ParseCmdArgs(cmd, args); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			if err := controller.RunJobsControl(); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			return nil
		},
	}

	// add control options
	flags := cmd.Flags()
	flags.StringVarP(&controller.filter.JobName, "job_name", "J", "", "filter by job name")
	flags.StringVarP(&controller.filter.Users, "user", "u", "", "filter by user")
	flags.BoolVarP(&controller.filter.AllJobs, "all", "a", false, "select all jobs")

	return cmd
}

func NewCheckPointJobsCmd(configManager *config.ConfigManager) *cobra.Command {
	controller := NewJobController(configManager, "checkpoint")

	cmd := &cobra.Command{
		Use:   "checkpoint",
		Short: "Checkpoint jobs",
		Long: `Checkpoint jobs by job IDs or filters.
Examples:
  xce jobs checkpoint jobid1,jobid2,...      # checkpoint jobs by job IDs`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := controller.ParseCmdArgs(cmd, args); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			if err := controller.RunJobsControl(); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			return nil
		},
	}

	// add control options
	return cmd
}

func NewRestartJobsCmd(configManager *config.ConfigManager) *cobra.Command {
	restarter := NewJobRestart(configManager)

	cmd := &cobra.Command{
		Use:   "restart",
		Short: "Restart jobs",
		Long: `Restart jobs by job ID.
Examples:
  xce jobs restart jobid    # restart job by jobid
  xce jobs restart -m hostname     	   # restart jobs by hostname`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := restarter.ParseCmdArgs(cmd, args); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			if err := restarter.RunJobsRestart(); err != nil {
				fmt.Println(strings.ToLower(err.Error()))
				return nil
			}
			return nil
		},
	}

	// add control options
	flags := cmd.Flags()
	flags.StringVarP(&restarter.askedHosts, "hostname", "m", "", "restart on the specific host name")
	return cmd
}