Files
hpc/internal/slurm/options.go
2026-04-09 10:33:31 +08:00

101 lines
2.3 KiB
Go

package slurm
import (
"fmt"
"net/http"
"time"
)
// ClientOption configures a Client via functional options.
type ClientOption func(*clientConfig) error
type clientConfig struct {
jwtKeyPath string
username string
ttl time.Duration
leeway time.Duration
httpClient *http.Client
}
func defaultClientConfig() *clientConfig {
return &clientConfig{
ttl: 30 * time.Minute,
leeway: 30 * time.Second,
}
}
// WithJWTKey specifies the path to the JWT key file.
func WithJWTKey(path string) ClientOption {
return func(c *clientConfig) error {
c.jwtKeyPath = path
return nil
}
}
// WithUsername specifies the Slurm username for JWT authentication.
func WithUsername(username string) ClientOption {
return func(c *clientConfig) error {
c.username = username
return nil
}
}
// WithTokenTTL sets the JWT token time-to-live (default: 30 minutes).
func WithTokenTTL(ttl time.Duration) ClientOption {
return func(c *clientConfig) error {
c.ttl = ttl
return nil
}
}
// WithTokenLeeway sets the JWT token refresh leeway (default: 30 seconds).
func WithTokenLeeway(leeway time.Duration) ClientOption {
return func(c *clientConfig) error {
c.leeway = leeway
return nil
}
}
// WithHTTPClient specifies a custom HTTP client.
func WithHTTPClient(client *http.Client) ClientOption {
return func(c *clientConfig) error {
c.httpClient = client
return nil
}
}
// NewClientWithOpts creates a new Slurm API client using functional options.
// If WithJWTKey and WithUsername are provided, JWT authentication is configured
// automatically. If no JWT options are provided, http.DefaultClient is used.
func NewClientWithOpts(baseURL string, opts ...ClientOption) (*Client, error) {
cfg := defaultClientConfig()
for _, opt := range opts {
if err := opt(cfg); err != nil {
return nil, err
}
}
var httpClient *http.Client
if cfg.jwtKeyPath != "" && cfg.username != "" {
key, err := ReadJWTKey(cfg.jwtKeyPath)
if err != nil {
return nil, fmt.Errorf("read JWT key: %w", err)
}
transportOpts := []JWTTransportOption{
WithTTL(cfg.ttl),
WithLeeway(cfg.leeway),
}
tr := NewJWTAuthTransport(cfg.username, key, transportOpts...)
httpClient = tr.Client()
} else if cfg.httpClient != nil {
httpClient = cfg.httpClient
} else {
httpClient = http.DefaultClient
}
return NewClient(baseURL, httpClient)
}