diff --git a/cmd/armadactl/cmd/commands.go b/cmd/armadactl/cmd/commands.go index 35817ed6b35..1e4dad188a4 100644 --- a/cmd/armadactl/cmd/commands.go +++ b/cmd/armadactl/cmd/commands.go @@ -27,6 +27,7 @@ func createCmd(a *armadactl.App) *cobra.Command { } cmd.Flags().Bool("dry-run", false, "Validate the input file and exit without making any changes.") cmd.AddCommand(queueCreateCmd()) + cmd.AddCommand(retryPolicyCreateCmd()) return cmd } @@ -37,6 +38,7 @@ func deleteCmd() *cobra.Command { Long: "Delete Armada resource. Supported: queue", } cmd.AddCommand(queueDeleteCmd()) + cmd.AddCommand(retryPolicyDeleteCmd()) return cmd } @@ -47,6 +49,7 @@ func updateCmd() *cobra.Command { Long: "Update Armada resource. Supported: queue", } cmd.AddCommand(queueUpdateCmd()) + cmd.AddCommand(retryPolicyUpdateCmd()) return cmd } @@ -59,6 +62,8 @@ func getCmd() *cobra.Command { cmd.AddCommand( queueGetCmd(), queuesGetCmd(), + retryPolicyGetCmd(), + retryPolicyGetAllCmd(), getSchedulingReportCmd(armadactl.New()), getQueueSchedulingReportCmd(armadactl.New()), getJobSchedulingReportCmd(armadactl.New()), diff --git a/cmd/armadactl/cmd/params.go b/cmd/armadactl/cmd/params.go index 1bffacfcb84..135474d293d 100644 --- a/cmd/armadactl/cmd/params.go +++ b/cmd/armadactl/cmd/params.go @@ -8,6 +8,7 @@ import ( ce "github.com/armadaproject/armada/pkg/client/executor" cn "github.com/armadaproject/armada/pkg/client/node" cq "github.com/armadaproject/armada/pkg/client/queue" + crp "github.com/armadaproject/armada/pkg/client/retrypolicy" ) // initParams initialises the command parameters, flags, and a configuration file. @@ -33,6 +34,12 @@ func initParams(cmd *cobra.Command, params *armadactl.Params) error { params.QueueAPI.Preempt = cq.Preempt(client.ExtractCommandlineArmadaApiConnectionDetails) params.QueueAPI.Cancel = cq.Cancel(client.ExtractCommandlineArmadaApiConnectionDetails) + params.RetryPolicyAPI.Create = crp.Create(client.ExtractCommandlineArmadaApiConnectionDetails) + params.RetryPolicyAPI.Delete = crp.Delete(client.ExtractCommandlineArmadaApiConnectionDetails) + params.RetryPolicyAPI.Get = crp.Get(client.ExtractCommandlineArmadaApiConnectionDetails) + params.RetryPolicyAPI.GetAll = crp.GetAll(client.ExtractCommandlineArmadaApiConnectionDetails) + params.RetryPolicyAPI.Update = crp.Update(client.ExtractCommandlineArmadaApiConnectionDetails) + params.ExecutorAPI.Cordon = ce.CordonExecutor(client.ExtractCommandlineArmadaApiConnectionDetails) params.ExecutorAPI.Uncordon = ce.UncordonExecutor(client.ExtractCommandlineArmadaApiConnectionDetails) diff --git a/cmd/armadactl/cmd/queue.go b/cmd/armadactl/cmd/queue.go index 92a74995bd1..86ded8ae7c8 100644 --- a/cmd/armadactl/cmd/queue.go +++ b/cmd/armadactl/cmd/queue.go @@ -12,6 +12,8 @@ import ( "github.com/armadaproject/armada/pkg/client/queue" ) +const retryPoliciesFlag = "retry-policies" + func queueCreateCmd() *cobra.Command { return queueCreateCmdWithApp(armadactl.New()) } @@ -63,6 +65,11 @@ Job priority is evaluated inside queue, queue has its own priority. Any labels return fmt.Errorf("error converting queue labels to map: %s", err) } + retryPolicies, err := cmd.Flags().GetStringSlice(retryPoliciesFlag) + if err != nil { + return fmt.Errorf("error reading retry-policies: %s", err) + } + newQueue, err := queue.NewQueue(&api.Queue{ Name: name, PriorityFactor: priorityFactor, @@ -70,6 +77,7 @@ Job priority is evaluated inside queue, queue has its own priority. Any labels GroupOwners: groups, Cordoned: cordoned, Labels: labelsAsMap, + RetryPolicies: retryPolicies, }) if err != nil { return fmt.Errorf("invalid queue data: %s", err) @@ -83,6 +91,7 @@ Job priority is evaluated inside queue, queue has its own priority. Any labels cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.") cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.") cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.") + cmd.Flags().StringSlice(retryPoliciesFlag, []string{}, "Comma separated list of retry policy names to assign to this queue, in evaluation order. Defaults to empty list.") return cmd } @@ -192,8 +201,14 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "queue ", Short: "Update an existing queue", - Long: "Update settings of an existing queue", - Args: cobra.ExactArgs(1), + Long: `Update settings of an existing queue. + +This is a full replace, not a partial patch. Every queue attribute is set from +the flags on this command, and any flag you omit resets that attribute to its +default. If the queue has retry policies attached, pass --retry-policies on +every update, otherwise the attachment is cleared and the queue falls back to +the default retry behaviour.`, + Args: cobra.ExactArgs(1), PreRunE: func(cmd *cobra.Command, args []string) error { return initParams(cmd, a.Params) }, @@ -230,6 +245,11 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command { return fmt.Errorf("error converting queue labels to map: %s", err) } + retryPolicies, err := cmd.Flags().GetStringSlice(retryPoliciesFlag) + if err != nil { + return fmt.Errorf("error reading retry-policies: %s", err) + } + newQueue, err := queue.NewQueue(&api.Queue{ Name: name, PriorityFactor: priorityFactor, @@ -237,6 +257,7 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command { GroupOwners: groups, Cordoned: cordoned, Labels: labelsAsMap, + RetryPolicies: retryPolicies, }) if err != nil { return fmt.Errorf("invalid queue data: %s", err) @@ -251,5 +272,6 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command { cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.") cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.") cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.") + cmd.Flags().StringSlice(retryPoliciesFlag, []string{}, "Comma separated list of retry policy names to assign to this queue, in evaluation order. Defaults to empty list.") return cmd } diff --git a/cmd/armadactl/cmd/retrypolicy.go b/cmd/armadactl/cmd/retrypolicy.go new file mode 100644 index 00000000000..52337faf6de --- /dev/null +++ b/cmd/armadactl/cmd/retrypolicy.go @@ -0,0 +1,98 @@ +package cmd + +import ( + "github.com/spf13/cobra" + + "github.com/armadaproject/armada/internal/armadactl" +) + +func retryPolicyCreateCmd() *cobra.Command { + a := armadactl.New() + return retryPolicyFileCmd(a, + "Create a retry policy from a YAML/JSON file", + "Create a retry policy that defines rules for whether failed jobs should be retried.", + a.CreateRetryPolicyFromFile) +} + +func retryPolicyUpdateCmd() *cobra.Command { + a := armadactl.New() + return retryPolicyFileCmd(a, + "Update a retry policy from a YAML/JSON file", + "Update an existing retry policy with the definition from a YAML/JSON file.", + a.UpdateRetryPolicyFromFile) +} + +func retryPolicyGetCmd() *cobra.Command { + a := armadactl.New() + return retryPolicyNameCmd(a, + "Get a retry policy by name", + "Get the definition of a retry policy by its name.", + a.GetRetryPolicy) +} + +func retryPolicyDeleteCmd() *cobra.Command { + a := armadactl.New() + return retryPolicyNameCmd(a, + "Delete a retry policy by name", + "Delete an existing retry policy by its name.", + a.DeleteRetryPolicy) +} + +func retryPolicyGetAllCmd() *cobra.Command { + a := armadactl.New() + return &cobra.Command{ + Use: "retry-policies", + Short: "List all retry policies", + Long: "List all retry policies defined in the system.", + Args: cobra.NoArgs, + PreRunE: func(cmd *cobra.Command, args []string) error { + return initParams(cmd, a.Params) + }, + RunE: func(cmd *cobra.Command, args []string) error { + return a.GetAllRetryPolicies() + }, + } +} + +// retryPolicyFileCmd builds a command that reads a retry policy from a +// YAML/JSON file and applies it via run. +func retryPolicyFileCmd(a *armadactl.App, short, long string, run func(fileName string) error) *cobra.Command { + cmd := &cobra.Command{ + Use: "retry-policy", + Short: short, + Long: long, + Args: cobra.NoArgs, + PreRunE: func(cmd *cobra.Command, args []string) error { + return initParams(cmd, a.Params) + }, + RunE: func(cmd *cobra.Command, args []string) error { + filePath, err := cmd.Flags().GetString("file") + if err != nil { + return err + } + return run(filePath) + }, + } + cmd.Flags().StringP("file", "f", "", "Path to YAML/JSON file defining the retry policy.") + if err := cmd.MarkFlagRequired("file"); err != nil { + panic(err) + } + return cmd +} + +// retryPolicyNameCmd builds a command that takes a single retry policy name +// argument and applies it via run. +func retryPolicyNameCmd(a *armadactl.App, short, long string, run func(name string) error) *cobra.Command { + return &cobra.Command{ + Use: "retry-policy ", + Short: short, + Long: long, + Args: cobra.ExactArgs(1), + PreRunE: func(cmd *cobra.Command, args []string) error { + return initParams(cmd, a.Params) + }, + RunE: func(cmd *cobra.Command, args []string) error { + return run(args[0]) + }, + } +} diff --git a/internal/armadactl/app.go b/internal/armadactl/app.go index 24caa113a02..962f77f1c1d 100644 --- a/internal/armadactl/app.go +++ b/internal/armadactl/app.go @@ -21,6 +21,7 @@ import ( "github.com/armadaproject/armada/pkg/client/executor" "github.com/armadaproject/armada/pkg/client/node" "github.com/armadaproject/armada/pkg/client/queue" + "github.com/armadaproject/armada/pkg/client/retrypolicy" ) type App struct { @@ -41,6 +42,7 @@ type App struct { type Params struct { ApiConnectionDetails *client.ApiConnectionDetails QueueAPI *QueueAPI + RetryPolicyAPI *RetryPolicyAPI ExecutorAPI *ExecutorAPI NodeAPI *NodeAPI } @@ -69,6 +71,14 @@ type ExecutorAPI struct { PreemptOnExecutor executor.PreemptAPI } +type RetryPolicyAPI struct { + Create retrypolicy.CreateAPI + Delete retrypolicy.DeleteAPI + Get retrypolicy.GetAPI + GetAll retrypolicy.GetAllAPI + Update retrypolicy.UpdateAPI +} + type NodeAPI struct { PreemptOnNode node.PreemptAPI CancelOnNode node.CancelAPI @@ -79,9 +89,10 @@ type NodeAPI struct { func New() *App { return &App{ Params: &Params{ - QueueAPI: &QueueAPI{}, - ExecutorAPI: &ExecutorAPI{}, - NodeAPI: &NodeAPI{}, + QueueAPI: &QueueAPI{}, + RetryPolicyAPI: &RetryPolicyAPI{}, + ExecutorAPI: &ExecutorAPI{}, + NodeAPI: &NodeAPI{}, }, Out: os.Stdout, Random: rand.Reader, diff --git a/internal/armadactl/queue.go b/internal/armadactl/queue.go index b5b7cf75be1..aa6132a5289 100644 --- a/internal/armadactl/queue.go +++ b/internal/armadactl/queue.go @@ -55,6 +55,15 @@ func (a *App) CreateResource(fileName string, dryRun bool) error { if !dryRun { return a.Params.QueueAPI.Create(queue) } + case client.ResourceKindRetryPolicy: + // Parse the file even on a dry run so a malformed policy is caught. + policy, err := retryPolicyFromFile(fileName) + if err != nil { + return err + } + if !dryRun { + return a.CreateRetryPolicy(policy) + } default: return errors.Errorf("invalid resource kind: %s", resource.Kind) } diff --git a/internal/armadactl/retrypolicy.go b/internal/armadactl/retrypolicy.go new file mode 100644 index 00000000000..76789a40950 --- /dev/null +++ b/internal/armadactl/retrypolicy.go @@ -0,0 +1,131 @@ +package armadactl + +import ( + "fmt" + "os" + + "github.com/pkg/errors" + "sigs.k8s.io/yaml" + + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client" +) + +func (a *App) CreateRetryPolicy(policy *api.RetryPolicy) error { + if err := a.Params.RetryPolicyAPI.Create(policy); err != nil { + return errors.Errorf("error creating retry policy %s: %s", policy.Name, err) + } + fmt.Fprintf(a.Out, "Created retry policy %s\n", policy.Name) + return nil +} + +func (a *App) CreateRetryPolicyFromFile(fileName string) error { + policy, err := retryPolicyFromFile(fileName) + if err != nil { + return err + } + return a.CreateRetryPolicy(policy) +} + +func (a *App) UpdateRetryPolicy(policy *api.RetryPolicy) error { + if err := a.Params.RetryPolicyAPI.Update(policy); err != nil { + return errors.Errorf("error updating retry policy %s: %s", policy.Name, err) + } + fmt.Fprintf(a.Out, "Updated retry policy %s\n", policy.Name) + return nil +} + +func (a *App) UpdateRetryPolicyFromFile(fileName string) error { + policy, err := retryPolicyFromFile(fileName) + if err != nil { + return err + } + return a.UpdateRetryPolicy(policy) +} + +// retryPolicyDocument is the on-disk form of a policy: the resource envelope +// followed by the policy's own fields, written flat rather than nested. +type retryPolicyDocument struct { + client.Resource + api.RetryPolicy +} + +// retryPolicyListBody is the body of a policy list document. It exists because +// api.RetryPolicyList tags the slice omitempty, which drops the key entirely +// for an empty list. +type retryPolicyListBody struct { + RetryPolicies []*api.RetryPolicy `json:"retryPolicies"` +} + +func retryPolicyFromFile(fileName string) (*api.RetryPolicy, error) { + data, err := os.ReadFile(fileName) + if err != nil { + return nil, errors.Errorf("file %s error: %s", fileName, err) + } + + // Strict, because a silently dropped retryLimit leaves it at 0, which means + // never retry. + doc := &retryPolicyDocument{} + if err := yaml.UnmarshalStrict(data, doc); err != nil { + return nil, errors.Errorf("file %s error: %s", fileName, err) + } + + // Without this a queue definition would pass as a policy, since both carry a name. + if doc.Version != client.APIVersionV1 { + return nil, errors.Errorf("file %s error: apiVersion must be %q", fileName, client.APIVersionV1) + } + if doc.Kind != client.ResourceKindRetryPolicy { + return nil, errors.Errorf("file %s error: kind must be %q", fileName, client.ResourceKindRetryPolicy) + } + + return &doc.RetryPolicy, nil +} + +func (a *App) DeleteRetryPolicy(name string) error { + if err := a.Params.RetryPolicyAPI.Delete(name); err != nil { + return errors.Errorf("error deleting retry policy %s: %s", name, err) + } + fmt.Fprintf(a.Out, "Deleted retry policy %s (or it did not exist)\n", name) + return nil +} + +func (a *App) GetRetryPolicy(name string) error { + policy, err := a.Params.RetryPolicyAPI.Get(name) + if err != nil { + return errors.Errorf("error getting retry policy %s: %s", name, err) + } + b, err := yaml.Marshal(policy) + if err != nil { + return errors.Errorf("error marshalling retry policy %s: %s", name, err) + } + fmt.Fprint(a.Out, retryPolicyHeaderYaml()+string(b)) + return nil +} + +func (a *App) GetAllRetryPolicies() error { + policies, err := a.Params.RetryPolicyAPI.GetAll() + if err != nil { + return errors.Errorf("error getting retry policies: %s", err) + } + if policies == nil { + policies = []*api.RetryPolicy{} + } + // A mapping, so that it follows the mapping header the document opens with. + b, err := yaml.Marshal(retryPolicyListBody{RetryPolicies: policies}) + if err != nil { + return errors.Errorf("error marshalling retry policies: %s", err) + } + fmt.Fprint(a.Out, retryPolicyHeaderYaml()+string(b)) + return nil +} + +func retryPolicyHeaderYaml() string { + b, err := yaml.Marshal(client.Resource{ + Version: client.APIVersionV1, + Kind: client.ResourceKindRetryPolicy, + }) + if err != nil { + panic(err) + } + return string(b) +} diff --git a/internal/armadactl/retrypolicy_test.go b/internal/armadactl/retrypolicy_test.go new file mode 100644 index 00000000000..762f3792645 --- /dev/null +++ b/internal/armadactl/retrypolicy_test.go @@ -0,0 +1,253 @@ +package armadactl + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "sigs.k8s.io/yaml" + + "github.com/armadaproject/armada/pkg/api" +) + +func newTestApp() (*App, *bytes.Buffer) { + out := &bytes.Buffer{} + a := New() + a.Out = out + return a, out +} + +func writeRetryPolicyFile(t *testing.T, contents string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "policy.yaml") + require.NoError(t, os.WriteFile(path, []byte(contents), 0o600)) + return path +} + +func fileWith(contents string) func(t *testing.T) string { + return func(t *testing.T) string { + return writeRetryPolicyFile(t, contents) + } +} + +const validPolicyFile = `apiVersion: armadaproject.io/v1beta1 +kind: RetryPolicy +name: p1 +retryLimit: 3 +defaultAction: Retry +rules: + - action: Fail + onCategory: UserError +` + +func TestGetRetryPolicy_RendersFriendlyActionStrings(t *testing.T) { + a, out := newTestApp() + a.Params.RetryPolicyAPI.Get = func(name string) (*api.RetryPolicy, error) { + return &api.RetryPolicy{ + Name: name, + RetryLimit: 3, + DefaultAction: api.RetryAction_RETRY_ACTION_RETRY, + Rules: []*api.RetryRule{ + {Action: api.RetryAction_RETRY_ACTION_FAIL, OnCategory: "OutOfMemory"}, + }, + }, nil + } + + require.NoError(t, a.GetRetryPolicy("p1")) + + got := out.String() + assert.Contains(t, got, "kind: RetryPolicy") + // Actions must render as friendly aliases, not raw enum integers. + assert.Contains(t, got, "defaultAction: Retry") + assert.Contains(t, got, "action: Fail") + assert.NotContains(t, got, "action: 1") + assert.NotContains(t, got, "defaultAction: 2") +} + +func TestGetAllRetryPolicies_EmitsValidYaml(t *testing.T) { + tests := map[string]struct { + policies []*api.RetryPolicy + }{ + "a populated store": { + policies: []*api.RetryPolicy{ + {Name: "p1", DefaultAction: api.RetryAction_RETRY_ACTION_RETRY}, + {Name: "p2", DefaultAction: api.RetryAction_RETRY_ACTION_FAIL}, + }, + }, + // Its own case because "{}" after the header does not parse. + "an empty store": {policies: nil}, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + a, out := newTestApp() + a.Params.RetryPolicyAPI.GetAll = func() ([]*api.RetryPolicy, error) { + return tc.policies, nil + } + + require.NoError(t, a.GetAllRetryPolicies()) + + // Unmarshalling the whole output proves it is one mapping, not a + // header followed by a bare sequence. + var doc struct { + APIVersion string `json:"apiVersion"` + Kind string `json:"kind"` + RetryPolicies []*api.RetryPolicy `json:"retryPolicies"` + } + require.NoError(t, yaml.Unmarshal(out.Bytes(), &doc), "get-list output must be valid YAML: %s", out.String()) + + assert.Equal(t, "RetryPolicy", doc.Kind) + require.Len(t, doc.RetryPolicies, len(tc.policies)) + for i, want := range tc.policies { + assert.Equal(t, want.Name, doc.RetryPolicies[i].Name) + assert.Equal(t, want.DefaultAction, doc.RetryPolicies[i].DefaultAction) + } + }) + } +} + +// Each write command must hand the API what it was given and report what it did. +func TestRetryPolicyWriteCommands_CallAPIAndReport(t *testing.T) { + tests := map[string]struct { + stub func(a *App, gotName *string) + call func(a *App) error + wantOut string + }{ + "create": { + stub: func(a *App, gotName *string) { + a.Params.RetryPolicyAPI.Create = func(policy *api.RetryPolicy) error { + *gotName = policy.Name + return nil + } + }, + call: func(a *App) error { return a.CreateRetryPolicy(&api.RetryPolicy{Name: "p1"}) }, + wantOut: "Created retry policy p1", + }, + "update": { + stub: func(a *App, gotName *string) { + a.Params.RetryPolicyAPI.Update = func(policy *api.RetryPolicy) error { + *gotName = policy.Name + return nil + } + }, + call: func(a *App) error { return a.UpdateRetryPolicy(&api.RetryPolicy{Name: "p1"}) }, + wantOut: "Updated retry policy p1", + }, + "delete": { + stub: func(a *App, gotName *string) { + a.Params.RetryPolicyAPI.Delete = func(name string) error { + *gotName = name + return nil + } + }, + call: func(a *App) error { return a.DeleteRetryPolicy("p1") }, + wantOut: "Deleted retry policy p1", + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + a, out := newTestApp() + var gotName string + tc.stub(a, &gotName) + + require.NoError(t, tc.call(a)) + assert.Equal(t, "p1", gotName) + assert.Contains(t, out.String(), tc.wantOut) + }) + } +} + +func TestRetryPolicyFromFile_Parses(t *testing.T) { + tests := map[string]struct { + path func(t *testing.T) string + wantLimit uint32 + }{ + "a handwritten document": { + path: fileWith(validPolicyFile), + wantLimit: 3, + }, + // get output must be usable as create input. + "the output of get": { + path: func(t *testing.T) string { + a, out := newTestApp() + a.Params.RetryPolicyAPI.Get = func(name string) (*api.RetryPolicy, error) { + return &api.RetryPolicy{ + Name: name, + RetryLimit: 2, + DefaultAction: api.RetryAction_RETRY_ACTION_RETRY, + Rules: []*api.RetryRule{{Action: api.RetryAction_RETRY_ACTION_FAIL, OnCategory: "UserError"}}, + }, nil + } + require.NoError(t, a.GetRetryPolicy("p1")) + return writeRetryPolicyFile(t, out.String()) + }, + wantLimit: 2, + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + policy, err := retryPolicyFromFile(tc.path(t)) + require.NoError(t, err) + + assert.Equal(t, "p1", policy.Name) + assert.Equal(t, tc.wantLimit, policy.RetryLimit) + assert.Equal(t, api.RetryAction_RETRY_ACTION_RETRY, policy.DefaultAction) + require.Len(t, policy.Rules, 1) + assert.Equal(t, "UserError", policy.Rules[0].OnCategory) + assert.Equal(t, api.RetryAction_RETRY_ACTION_FAIL, policy.Rules[0].Action) + }) + } +} + +func TestRetryPolicyFromFile_Rejects(t *testing.T) { + tests := map[string]struct { + path func(t *testing.T) string + wantErr string + }{ + "a mistyped field": { + path: fileWith(`apiVersion: armadaproject.io/v1beta1 +kind: RetryPolicy +name: p1 +retry_limit: 5 +defaultAction: Retry +`), + wantErr: "unknown field", + }, + "a file describing another resource": { + path: fileWith(`apiVersion: armadaproject.io/v1beta1 +kind: Queue +name: q1 +`), + wantErr: "kind must be", + }, + "a bare document with no envelope": { + path: fileWith("name: p1\nretryLimit: 3\ndefaultAction: Retry\n"), + wantErr: "apiVersion must be", + }, + "an unknown kind": { + path: fileWith("apiVersion: armadaproject.io/v1beta1\nkind: Nonsense\nname: p1\n"), + wantErr: "invalid kind", + }, + "an unknown apiVersion": { + path: fileWith("apiVersion: armadaproject.io/v2\nkind: RetryPolicy\nname: p1\n"), + wantErr: "invalid version", + }, + "malformed yaml": { + path: fileWith("apiVersion: [unclosed\n"), + wantErr: "converting YAML to JSON", + }, + "a missing file": { + path: func(t *testing.T) string { return filepath.Join(t.TempDir(), "absent.yaml") }, + wantErr: "no such file", + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + _, err := retryPolicyFromFile(tc.path(t)) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} diff --git a/pkg/client/resource.go b/pkg/client/resource.go index 6d08f7a5900..5dab8d2de38 100644 --- a/pkg/client/resource.go +++ b/pkg/client/resource.go @@ -3,6 +3,7 @@ package client import ( "encoding/json" "fmt" + "slices" ) type Resource struct { @@ -13,16 +14,18 @@ type Resource struct { type ResourceKind string const ( - ResourceKindQueue ResourceKind = "Queue" + ResourceKindQueue ResourceKind = "Queue" + ResourceKindRetryPolicy ResourceKind = "RetryPolicy" ) func NewResourceKind(in string) (ResourceKind, error) { - validValues := []ResourceKind{ResourceKindQueue} - if in != string(ResourceKindQueue) { + validValues := []ResourceKind{ResourceKindQueue, ResourceKindRetryPolicy} + kind := ResourceKind(in) + if !slices.Contains(validValues, kind) { return "", fmt.Errorf("invalid kind: %s. Valid values: %v", in, validValues) } - return ResourceKind(in), nil + return kind, nil } func (kind *ResourceKind) UnmarshalJSON(data []byte) error { diff --git a/pkg/client/retrypolicy/create.go b/pkg/client/retrypolicy/create.go new file mode 100644 index 00000000000..a97fd8ff95e --- /dev/null +++ b/pkg/client/retrypolicy/create.go @@ -0,0 +1,35 @@ +package retrypolicy + +import ( + "fmt" + + "github.com/armadaproject/armada/internal/common" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client" +) + +type CreateAPI func(policy *api.RetryPolicy) error + +func Create(getConnectionDetails client.ConnectionDetails) CreateAPI { + return func(policy *api.RetryPolicy) error { + connectionDetails, err := getConnectionDetails() + if err != nil { + return fmt.Errorf("failed to obtain api connection details: %s", err) + } + conn, err := client.CreateApiConnection(connectionDetails) + if err != nil { + return fmt.Errorf("failed to connect to api because %s", err) + } + defer conn.Close() + + ctx, cancel := common.ContextWithDefaultTimeout() + defer cancel() + + c := api.NewRetryPolicyServiceClient(conn) + if _, err := c.CreateRetryPolicy(ctx, policy); err != nil { + return fmt.Errorf("create retry policy request failed: %s", err) + } + + return nil + } +} diff --git a/pkg/client/retrypolicy/delete.go b/pkg/client/retrypolicy/delete.go new file mode 100644 index 00000000000..9d1fbce7c33 --- /dev/null +++ b/pkg/client/retrypolicy/delete.go @@ -0,0 +1,35 @@ +package retrypolicy + +import ( + "fmt" + + "github.com/armadaproject/armada/internal/common" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client" +) + +type DeleteAPI func(name string) error + +func Delete(getConnectionDetails client.ConnectionDetails) DeleteAPI { + return func(name string) error { + connectionDetails, err := getConnectionDetails() + if err != nil { + return fmt.Errorf("failed to obtain api connection details: %s", err) + } + conn, err := client.CreateApiConnection(connectionDetails) + if err != nil { + return fmt.Errorf("failed to connect to api because %s", err) + } + defer conn.Close() + + ctx, cancel := common.ContextWithDefaultTimeout() + defer cancel() + + c := api.NewRetryPolicyServiceClient(conn) + if _, err = c.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: name}); err != nil { + return fmt.Errorf("delete retry policy request failed: %s", err) + } + + return nil + } +} diff --git a/pkg/client/retrypolicy/get.go b/pkg/client/retrypolicy/get.go new file mode 100644 index 00000000000..1e43b02d59e --- /dev/null +++ b/pkg/client/retrypolicy/get.go @@ -0,0 +1,63 @@ +package retrypolicy + +import ( + "fmt" + + "github.com/armadaproject/armada/internal/common" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client" +) + +type GetAPI func(name string) (*api.RetryPolicy, error) + +func Get(getConnectionDetails client.ConnectionDetails) GetAPI { + return func(name string) (*api.RetryPolicy, error) { + connectionDetails, err := getConnectionDetails() + if err != nil { + return nil, fmt.Errorf("failed to obtain api connection details: %s", err) + } + conn, err := client.CreateApiConnection(connectionDetails) + if err != nil { + return nil, fmt.Errorf("failed to connect to api because %s", err) + } + defer conn.Close() + + ctx, cancel := common.ContextWithDefaultTimeout() + defer cancel() + + c := api.NewRetryPolicyServiceClient(conn) + policy, err := c.GetRetryPolicy(ctx, &api.RetryPolicyGetRequest{Name: name}) + if err != nil { + return nil, fmt.Errorf("get retry policy request failed: %s", err) + } + + return policy, nil + } +} + +type GetAllAPI func() ([]*api.RetryPolicy, error) + +func GetAll(getConnectionDetails client.ConnectionDetails) GetAllAPI { + return func() ([]*api.RetryPolicy, error) { + connectionDetails, err := getConnectionDetails() + if err != nil { + return nil, fmt.Errorf("failed to obtain api connection details: %s", err) + } + conn, err := client.CreateApiConnection(connectionDetails) + if err != nil { + return nil, fmt.Errorf("failed to connect to api because %s", err) + } + defer conn.Close() + + ctx, cancel := common.ContextWithDefaultTimeout() + defer cancel() + + c := api.NewRetryPolicyServiceClient(conn) + list, err := c.GetRetryPolicies(ctx, &api.RetryPolicyListRequest{}) + if err != nil { + return nil, fmt.Errorf("get retry policies request failed: %s", err) + } + + return list.RetryPolicies, nil + } +} diff --git a/pkg/client/retrypolicy/update.go b/pkg/client/retrypolicy/update.go new file mode 100644 index 00000000000..967b0ad053f --- /dev/null +++ b/pkg/client/retrypolicy/update.go @@ -0,0 +1,34 @@ +package retrypolicy + +import ( + "fmt" + + "github.com/armadaproject/armada/internal/common" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client" +) + +type UpdateAPI func(policy *api.RetryPolicy) error + +func Update(getConnectionDetails client.ConnectionDetails) UpdateAPI { + return func(policy *api.RetryPolicy) error { + connectionDetails, err := getConnectionDetails() + if err != nil { + return fmt.Errorf("failed to obtain api connection details: %s", err) + } + conn, err := client.CreateApiConnection(connectionDetails) + if err != nil { + return fmt.Errorf("failed to connect to api because %s", err) + } + defer conn.Close() + + ctx, cancel := common.ContextWithDefaultTimeout() + defer cancel() + + c := api.NewRetryPolicyServiceClient(conn) + if _, err = c.UpdateRetryPolicy(ctx, policy); err != nil { + return fmt.Errorf("update retry policy request failed: %s", err) + } + return nil + } +}