hpc/internal/slurm/client.go

package slurm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const (
	// DefaultBaseURL is the default Slurm REST API endpoint.
	DefaultBaseURL = "http://localhost:6820/"
	// DefaultUserAgent is the default User-Agent header value.
	DefaultUserAgent = "slurm-go-sdk"
	// DefaultTimeout is the default HTTP request timeout.
	DefaultTimeout = 30 * time.Second
)

// Client manages communication with the Slurm REST API.
type Client struct {
	client    *http.Client
	baseURL   *url.URL
	UserAgent string

	common service

	Jobs         *JobsService
	Nodes        *NodesService
	Partitions   *PartitionsService
	Reservations *ReservationsService
	Diag         *DiagService
	Ping         *PingService
	Licenses     *LicensesService
	Reconfigure  *ReconfigureService
	Shares       *SharesService

	SlurmdbDiag      *SlurmdbDiagService
	SlurmdbConfig    *SlurmdbConfigService
	SlurmdbTres      *SlurmdbTresService
	SlurmdbQos       *SlurmdbQosService
	SlurmdbAssocs    *SlurmdbAssocsService
	SlurmdbInstances *SlurmdbInstancesService
	SlurmdbUsers     *SlurmdbUsersService
	SlurmdbClusters  *SlurmdbClustersService
	SlurmdbWckeys    *SlurmdbWckeysService
	SlurmdbAccounts  *SlurmdbAccountsService
	SlurmdbJobs      *SlurmdbJobsService
}

// service is the base struct for all API services.
type service struct {
	client *Client
}

type JobsService service
type NodesService service
type PartitionsService service
type ReservationsService service
type DiagService service
type PingService service
type LicensesService service
type ReconfigureService service
type SharesService service
type SlurmdbDiagService service
type SlurmdbConfigService service
type SlurmdbTresService service
type SlurmdbQosService service
type SlurmdbAssocsService service
type SlurmdbInstancesService service
type SlurmdbUsersService service
type SlurmdbClustersService service
type SlurmdbWckeysService service
type SlurmdbAccountsService service
type SlurmdbJobsService service

// Response wraps an http.Response and will later hold Meta/Errors/Warnings
// fields parsed from Slurm API responses.
type Response struct {
	*http.Response
}

// NewClient returns a new Slurm API client. If httpClient is nil,
// http.DefaultClient is used.
func NewClient(baseURL string, httpClient *http.Client) (*Client, error) {
	if httpClient == nil {
		httpClient = &http.Client{Timeout: DefaultTimeout}
	}

	parsedURL, err := url.Parse(baseURL)
	if err != nil {
		return nil, fmt.Errorf("invalid base URL %q: %w", baseURL, err)
	}
	if !strings.HasSuffix(parsedURL.Path, "/") {
		parsedURL.Path += "/"
	}

	c := &Client{
		client:    httpClient,
		baseURL:   parsedURL,
		UserAgent: DefaultUserAgent,
	}
	c.common.client = c

	c.Jobs = (*JobsService)(&c.common)
	c.Nodes = (*NodesService)(&c.common)
	c.Partitions = (*PartitionsService)(&c.common)
	c.Reservations = (*ReservationsService)(&c.common)
	c.Diag = (*DiagService)(&c.common)
	c.Ping = (*PingService)(&c.common)
	c.Licenses = (*LicensesService)(&c.common)
	c.Reconfigure = (*ReconfigureService)(&c.common)
	c.Shares = (*SharesService)(&c.common)

	c.SlurmdbDiag = (*SlurmdbDiagService)(&c.common)
	c.SlurmdbConfig = (*SlurmdbConfigService)(&c.common)
	c.SlurmdbTres = (*SlurmdbTresService)(&c.common)
	c.SlurmdbQos = (*SlurmdbQosService)(&c.common)
	c.SlurmdbAssocs = (*SlurmdbAssocsService)(&c.common)
	c.SlurmdbInstances = (*SlurmdbInstancesService)(&c.common)
	c.SlurmdbUsers = (*SlurmdbUsersService)(&c.common)
	c.SlurmdbClusters = (*SlurmdbClustersService)(&c.common)
	c.SlurmdbWckeys = (*SlurmdbWckeysService)(&c.common)
	c.SlurmdbAccounts = (*SlurmdbAccountsService)(&c.common)
	c.SlurmdbJobs = (*SlurmdbJobsService)(&c.common)

	return c, nil
}

// NewRequest creates an API request. A relative URL path can be provided in
// urlStr, which will be resolved against the base URL of the Client.
// If body is not nil, it will be JSON-encoded and set as the request body.
func (c *Client) NewRequest(method, urlStr string, body interface{}) (*http.Request, error) {
	u, err := c.baseURL.Parse(urlStr)
	if err != nil {
		return nil, fmt.Errorf("invalid relative URL %q: %w", urlStr, err)
	}

	var buf io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal request body: %w", err)
		}
		buf = bytes.NewReader(b)
	}

	req, err := http.NewRequest(method, u.String(), buf)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.UserAgent != "" {
		req.Header.Set("User-Agent", c.UserAgent)
	}

	return req, nil
}

// Do sends an API request and returns the API response. The API response is
// JSON-decoded into the value pointed to by v (if v is not nil).
func (c *Client) Do(ctx context.Context, req *http.Request, v interface{}) (*Response, error) {
	req = req.WithContext(ctx)

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if err := CheckResponse(resp); err != nil {
		return nil, err
	}

	response := &Response{Response: resp}

	if v != nil {
		if err := json.NewDecoder(resp.Body).Decode(v); err != nil {
			return response, fmt.Errorf("failed to decode response: %w", err)
		}
	}

	return response, nil
}