From a6c339e80862b8b72ef922bba0c96ade60672f78 Mon Sep 17 00:00:00 2001 From: Dejan Zele Pejchev Date: Fri, 24 Jul 2026 21:58:13 +0200 Subject: [PATCH 1/2] Move requireGrpcCode test helper into shared servertest package Signed-off-by: Dejan Zele Pejchev --- .../server/executor/executor_server_test.go | 33 ++++++++----------- internal/server/node/node_test.go | 21 ++++-------- internal/server/queue/queue_service_test.go | 25 +++++--------- internal/server/servertest/grpc.go | 17 ++++++++++ 4 files changed, 46 insertions(+), 50 deletions(-) create mode 100644 internal/server/servertest/grpc.go diff --git a/internal/server/executor/executor_server_test.go b/internal/server/executor/executor_server_test.go index 6d96b55dead..c50db5f5af6 100644 --- a/internal/server/executor/executor_server_test.go +++ b/internal/server/executor/executor_server_test.go @@ -9,7 +9,6 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "google.golang.org/grpc/codes" - grpcstatus "google.golang.org/grpc/status" clocktesting "k8s.io/utils/clock/testing" "github.com/armadaproject/armada/internal/common/armadacontext" @@ -20,6 +19,7 @@ import ( protoutil "github.com/armadaproject/armada/internal/common/proto" servermocks "github.com/armadaproject/armada/internal/server/mocks" "github.com/armadaproject/armada/internal/server/permissions" + "github.com/armadaproject/armada/internal/server/servertest" "github.com/armadaproject/armada/pkg/api" "github.com/armadaproject/armada/pkg/controlplaneevents" ) @@ -40,13 +40,6 @@ func newExecutorTestServer(t *testing.T) (*Server, *executorTestMocks) { return s, m } -func requireGrpcCode(t *testing.T, err error, code codes.Code) { - t.Helper() - st, ok := grpcstatus.FromError(err) - require.True(t, ok, "expected gRPC status error") - assert.Equal(t, code, st.Code()) -} - func TestUpsertExecutorSettings_PermissionDenied(t *testing.T) { s, m := newExecutorTestServer(t) grpcCtx := armadacontext.Background() @@ -59,7 +52,7 @@ func TestUpsertExecutorSettings_PermissionDenied(t *testing.T) { _, err := s.UpsertExecutorSettings(grpcCtx, &api.ExecutorSettingsUpsertRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestUpsertExecutorSettings_AuthorizeErrorUnavailable(t *testing.T) { @@ -74,7 +67,7 @@ func TestUpsertExecutorSettings_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.UpsertExecutorSettings(grpcCtx, &api.ExecutorSettingsUpsertRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Unavailable) + servertest.RequireGrpcCode(t, err, codes.Unavailable) } func TestUpsertExecutorSettings_ValidationName(t *testing.T) { @@ -125,7 +118,7 @@ func TestUpsertExecutorSettings_PublishErrorInternal(t *testing.T) { _, err := s.UpsertExecutorSettings(grpcCtx, &api.ExecutorSettingsUpsertRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestUpsertExecutorSettings_SuccessPublishesExpectedEvent(t *testing.T) { @@ -186,7 +179,7 @@ func TestDeleteExecutorSettings_PermissionDenied(t *testing.T) { _, err := s.DeleteExecutorSettings(grpcCtx, &api.ExecutorSettingsDeleteRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestDeleteExecutorSettings_AuthorizeErrorUnavailable(t *testing.T) { @@ -201,7 +194,7 @@ func TestDeleteExecutorSettings_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.DeleteExecutorSettings(grpcCtx, &api.ExecutorSettingsDeleteRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Unavailable) + servertest.RequireGrpcCode(t, err, codes.Unavailable) } func TestDeleteExecutorSettings_ValidationName(t *testing.T) { @@ -237,7 +230,7 @@ func TestDeleteExecutorSettings_PublishErrorInternal(t *testing.T) { _, err := s.DeleteExecutorSettings(grpcCtx, &api.ExecutorSettingsDeleteRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestDeleteExecutorSettings_SuccessPublishesExpectedEvent(t *testing.T) { @@ -288,7 +281,7 @@ func TestPreemptOnExecutor_PermissionDenied(t *testing.T) { _, err := s.PreemptOnExecutor(grpcCtx, &api.ExecutorPreemptRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestPreemptOnExecutor_AuthorizeErrorUnavailable(t *testing.T) { @@ -303,7 +296,7 @@ func TestPreemptOnExecutor_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.PreemptOnExecutor(grpcCtx, &api.ExecutorPreemptRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Unavailable) + servertest.RequireGrpcCode(t, err, codes.Unavailable) } func TestPreemptOnExecutor_ValidationName(t *testing.T) { @@ -339,7 +332,7 @@ func TestPreemptOnExecutor_PublishErrorInternal(t *testing.T) { _, err := s.PreemptOnExecutor(grpcCtx, &api.ExecutorPreemptRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestPreemptOnExecutor_SuccessPublishesExpectedEvent(t *testing.T) { @@ -392,7 +385,7 @@ func TestCancelOnExecutor_PermissionDenied(t *testing.T) { _, err := s.CancelOnExecutor(grpcCtx, &api.ExecutorCancelRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestCancelOnExecutor_AuthorizeErrorUnavailable(t *testing.T) { @@ -407,7 +400,7 @@ func TestCancelOnExecutor_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.CancelOnExecutor(grpcCtx, &api.ExecutorCancelRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Unavailable) + servertest.RequireGrpcCode(t, err, codes.Unavailable) } func TestCancelOnExecutor_ValidationName(t *testing.T) { @@ -443,7 +436,7 @@ func TestCancelOnExecutor_PublishErrorInternal(t *testing.T) { _, err := s.CancelOnExecutor(grpcCtx, &api.ExecutorCancelRequest{Name: "executor-1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestCancelOnExecutor_SuccessPublishesExpectedEvent(t *testing.T) { diff --git a/internal/server/node/node_test.go b/internal/server/node/node_test.go index 4929053e873..e4c66ab4fce 100644 --- a/internal/server/node/node_test.go +++ b/internal/server/node/node_test.go @@ -9,7 +9,6 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "google.golang.org/grpc/codes" - grpcstatus "google.golang.org/grpc/status" clocktesting "k8s.io/utils/clock/testing" "github.com/armadaproject/armada/internal/common/armadacontext" @@ -19,6 +18,7 @@ import ( protoutil "github.com/armadaproject/armada/internal/common/proto" servermocks "github.com/armadaproject/armada/internal/server/mocks" "github.com/armadaproject/armada/internal/server/permissions" + "github.com/armadaproject/armada/internal/server/servertest" "github.com/armadaproject/armada/pkg/api" "github.com/armadaproject/armada/pkg/controlplaneevents" ) @@ -39,13 +39,6 @@ func newTestServer(t *testing.T) (*Server, *testMocks) { return s, m } -func requireGrpcCode(t *testing.T, err error, code codes.Code) { - t.Helper() - st, ok := grpcstatus.FromError(err) - require.True(t, ok, "expected gRPC status error") - assert.Equal(t, code, st.Code()) -} - func TestPreemptOnNode_PermissionDenied(t *testing.T) { s, m := newTestServer(t) ctx := armadacontext.Background() @@ -58,7 +51,7 @@ func TestPreemptOnNode_PermissionDenied(t *testing.T) { _, err := s.PreemptOnNode(ctx, &api.NodePreemptRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestPreemptOnNode_AuthorizeErrorUnavailable(t *testing.T) { @@ -73,7 +66,7 @@ func TestPreemptOnNode_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.PreemptOnNode(ctx, &api.NodePreemptRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestPreemptOnNode_Validation(t *testing.T) { @@ -109,7 +102,7 @@ func TestPreemptOnNode_PublishErrorInternal(t *testing.T) { _, err := s.PreemptOnNode(ctx, &api.NodePreemptRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestPreemptOnNode_SuccessPublishesExpectedEvent(t *testing.T) { @@ -168,7 +161,7 @@ func TestCancelOnNode_PermissionDenied(t *testing.T) { _, err := s.CancelOnNode(ctx, &api.NodeCancelRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestCancelOnNode_AuthorizeErrorUnavailable(t *testing.T) { @@ -183,7 +176,7 @@ func TestCancelOnNode_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.CancelOnNode(ctx, &api.NodeCancelRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestCancelOnNode_Validation(t *testing.T) { @@ -219,7 +212,7 @@ func TestCancelOnNode_PublishErrorInternal(t *testing.T) { _, err := s.CancelOnNode(ctx, &api.NodeCancelRequest{Name: "executor-1", Executor: "executor-id"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestCancelOnNode_SuccessPublishesExpectedEvent(t *testing.T) { diff --git a/internal/server/queue/queue_service_test.go b/internal/server/queue/queue_service_test.go index f0dc7ff2885..71497e279f0 100644 --- a/internal/server/queue/queue_service_test.go +++ b/internal/server/queue/queue_service_test.go @@ -12,7 +12,6 @@ import ( "go.uber.org/mock/gomock" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" - grpcstatus "google.golang.org/grpc/status" clocktesting "k8s.io/utils/clock/testing" "github.com/armadaproject/armada/internal/common/armadacontext" @@ -23,6 +22,7 @@ import ( protoutil "github.com/armadaproject/armada/internal/common/proto" servermocks "github.com/armadaproject/armada/internal/server/mocks" "github.com/armadaproject/armada/internal/server/permissions" + "github.com/armadaproject/armada/internal/server/servertest" "github.com/armadaproject/armada/pkg/api" "github.com/armadaproject/armada/pkg/client/queue" "github.com/armadaproject/armada/pkg/controlplaneevents" @@ -46,13 +46,6 @@ func newTestQueueServer(t *testing.T) (*Server, *queueServiceTestMocks) { return s, m } -func requireGrpcCode(t *testing.T, err error, code codes.Code) { - t.Helper() - st, ok := grpcstatus.FromError(err) - require.True(t, ok, "expected gRPC status error") - assert.Equal(t, code, st.Code()) -} - func TestCreateQueue_PermissionDenied(t *testing.T) { s, m := newTestQueueServer(t) ctx := armadacontext.Background() @@ -65,7 +58,7 @@ func TestCreateQueue_PermissionDenied(t *testing.T) { _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.PermissionDenied) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } func TestCreateQueue_AuthorizeErrorUnavailable(t *testing.T) { @@ -80,7 +73,7 @@ func TestCreateQueue_AuthorizeErrorUnavailable(t *testing.T) { _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.Unavailable) + servertest.RequireGrpcCode(t, err, codes.Unavailable) } func TestCreateQueue_DefaultsUserOwnerFromPrincipal(t *testing.T) { @@ -132,7 +125,7 @@ func TestCreateQueue_ValidationInvalidArgument(t *testing.T) { _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1", PriorityFactor: 1, Labels: map[string]string{"k": ""}}) require.Error(t, err) - requireGrpcCode(t, err, codes.InvalidArgument) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) } func TestCreateQueue_AlreadyExists(t *testing.T) { @@ -153,7 +146,7 @@ func TestCreateQueue_AlreadyExists(t *testing.T) { _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1", PriorityFactor: 1}) require.Error(t, err) - requireGrpcCode(t, err, codes.AlreadyExists) + servertest.RequireGrpcCode(t, err, codes.AlreadyExists) } func TestUpdateQueue_NotFound(t *testing.T) { @@ -174,7 +167,7 @@ func TestUpdateQueue_NotFound(t *testing.T) { _, err := s.UpdateQueue(ctx, &api.Queue{Name: "q1", PriorityFactor: 1}) require.Error(t, err) - requireGrpcCode(t, err, codes.NotFound) + servertest.RequireGrpcCode(t, err, codes.NotFound) } func TestDeleteQueue_RepoErrorInvalidArgument(t *testing.T) { @@ -195,7 +188,7 @@ func TestDeleteQueue_RepoErrorInvalidArgument(t *testing.T) { _, err := s.DeleteQueue(ctx, &api.QueueDeleteRequest{Name: "q1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.InvalidArgument) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) } func TestGetQueue_NotFound(t *testing.T) { @@ -210,7 +203,7 @@ func TestGetQueue_NotFound(t *testing.T) { _, err := s.GetQueue(ctx, &api.QueueGetRequest{Name: "q1"}) require.Error(t, err) - requireGrpcCode(t, err, codes.NotFound) + servertest.RequireGrpcCode(t, err, codes.NotFound) } type fakeQueueStream struct { @@ -289,7 +282,7 @@ func TestCancelOnQueue_PublishErrorInternal(t *testing.T) { _, err := s.CancelOnQueue(ctx, &api.QueueCancelRequest{Name: "q1", JobStates: []api.JobState{api.JobState_RUNNING}}) require.Error(t, err) - requireGrpcCode(t, err, codes.Internal) + servertest.RequireGrpcCode(t, err, codes.Internal) } func TestCancelOnQueue_SuccessPublishesExpectedEvent(t *testing.T) { diff --git a/internal/server/servertest/grpc.go b/internal/server/servertest/grpc.go new file mode 100644 index 00000000000..0c19679847a --- /dev/null +++ b/internal/server/servertest/grpc.go @@ -0,0 +1,17 @@ +package servertest + +import ( + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// RequireGrpcCode asserts that err is a gRPC status error with the given code. +func RequireGrpcCode(t *testing.T, err error, code codes.Code) { + t.Helper() + st, ok := status.FromError(err) + require.True(t, ok, "expected gRPC status error") + require.Equal(t, code, st.Code()) +} From e493f4d7515e3be05f4b44de102175b9bdee75a3 Mon Sep 17 00:00:00 2001 From: Dejan Zele Pejchev Date: Fri, 24 Jul 2026 21:58:39 +0200 Subject: [PATCH 2/2] Add RetryPolicy CRUD service and queue attachment validation Signed-off-by: Dejan Zele Pejchev --- cmd/server/main.go | 1 + .../migrations/035_create_retry_policy.sql | 5 + internal/server/mocks/generate.go | 1 + .../mocks/mock_retry_policy_repository.go | 108 +++++ internal/server/permissions/permissions.go | 3 + internal/server/queryapi/database/models.go | 5 + internal/server/queue/queue_service.go | 31 ++ internal/server/queue/queue_service_test.go | 63 ++- internal/server/retrypolicy/repository.go | 153 +++++++ .../server/retrypolicy/repository_test.go | 119 ++++++ internal/server/retrypolicy/service.go | 180 ++++++++ internal/server/retrypolicy/service_test.go | 384 ++++++++++++++++++ internal/server/retrypolicy/validation.go | 58 +++ .../server/retrypolicy/validation_test.go | 121 ++++++ internal/server/server.go | 7 +- pkg/client/queue/queue.go | 7 +- 16 files changed, 1236 insertions(+), 10 deletions(-) create mode 100644 internal/lookout/schema/migrations/035_create_retry_policy.sql create mode 100644 internal/server/mocks/mock_retry_policy_repository.go create mode 100644 internal/server/retrypolicy/repository.go create mode 100644 internal/server/retrypolicy/repository_test.go create mode 100644 internal/server/retrypolicy/service.go create mode 100644 internal/server/retrypolicy/service_test.go create mode 100644 internal/server/retrypolicy/validation.go create mode 100644 internal/server/retrypolicy/validation_test.go diff --git a/cmd/server/main.go b/cmd/server/main.go index cab9c7bc0b7..c0da47334ec 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -102,6 +102,7 @@ func main() { config.CorsAllowedOrigins, api.SwaggerJsonTemplate(), api.RegisterSubmitHandler, + api.RegisterRetryPolicyServiceHandler, api.RegisterEventHandler, api.RegisterJobsHandler, schedulerobjects.RegisterSchedulerReportingHandler, diff --git a/internal/lookout/schema/migrations/035_create_retry_policy.sql b/internal/lookout/schema/migrations/035_create_retry_policy.sql new file mode 100644 index 00000000000..3455aceadf0 --- /dev/null +++ b/internal/lookout/schema/migrations/035_create_retry_policy.sql @@ -0,0 +1,5 @@ +CREATE TABLE IF NOT EXISTS retry_policy +( + name text NOT NULL PRIMARY KEY, + definition bytea NOT NULL +) diff --git a/internal/server/mocks/generate.go b/internal/server/mocks/generate.go index 2efeadc2e72..de0fff32137 100644 --- a/internal/server/mocks/generate.go +++ b/internal/server/mocks/generate.go @@ -4,3 +4,4 @@ package mocks //go:generate mockgen -destination=./mock_deduplicator.go -package=mocks "github.com/armadaproject/armada/internal/server/submit" Deduplicator //go:generate mockgen -destination=./mock_authorizer.go -package=mocks "github.com/armadaproject/armada/internal/common/auth" ActionAuthorizer //go:generate mockgen -destination=./mock_repository.go -package=mocks "github.com/armadaproject/armada/internal/server/queue" QueueRepository +//go:generate mockgen -destination=./mock_retry_policy_repository.go -package=mocks "github.com/armadaproject/armada/internal/server/retrypolicy" RetryPolicyRepository diff --git a/internal/server/mocks/mock_retry_policy_repository.go b/internal/server/mocks/mock_retry_policy_repository.go new file mode 100644 index 00000000000..9bab4bac379 --- /dev/null +++ b/internal/server/mocks/mock_retry_policy_repository.go @@ -0,0 +1,108 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/armadaproject/armada/internal/server/retrypolicy (interfaces: RetryPolicyRepository) + +package mocks + +import ( + reflect "reflect" + + armadacontext "github.com/armadaproject/armada/internal/common/armadacontext" + api "github.com/armadaproject/armada/pkg/api" + gomock "go.uber.org/mock/gomock" +) + +// MockRetryPolicyRepository is a mock of RetryPolicyRepository interface. +type MockRetryPolicyRepository struct { + ctrl *gomock.Controller + recorder *MockRetryPolicyRepositoryMockRecorder + isgomock struct{} +} + +// MockRetryPolicyRepositoryMockRecorder is the mock recorder for MockRetryPolicyRepository. +type MockRetryPolicyRepositoryMockRecorder struct { + mock *MockRetryPolicyRepository +} + +// NewMockRetryPolicyRepository creates a new mock instance. +func NewMockRetryPolicyRepository(ctrl *gomock.Controller) *MockRetryPolicyRepository { + mock := &MockRetryPolicyRepository{ctrl: ctrl} + mock.recorder = &MockRetryPolicyRepositoryMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRetryPolicyRepository) EXPECT() *MockRetryPolicyRepositoryMockRecorder { + return m.recorder +} + +// CreateRetryPolicy mocks base method. +func (m *MockRetryPolicyRepository) CreateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CreateRetryPolicy", ctx, policy) + ret0, _ := ret[0].(error) + return ret0 +} + +// CreateRetryPolicy indicates an expected call of CreateRetryPolicy. +func (mr *MockRetryPolicyRepositoryMockRecorder) CreateRetryPolicy(ctx, policy any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateRetryPolicy", reflect.TypeOf((*MockRetryPolicyRepository)(nil).CreateRetryPolicy), ctx, policy) +} + +// DeleteRetryPolicy mocks base method. +func (m *MockRetryPolicyRepository) DeleteRetryPolicy(ctx *armadacontext.Context, name string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteRetryPolicy", ctx, name) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteRetryPolicy indicates an expected call of DeleteRetryPolicy. +func (mr *MockRetryPolicyRepositoryMockRecorder) DeleteRetryPolicy(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteRetryPolicy", reflect.TypeOf((*MockRetryPolicyRepository)(nil).DeleteRetryPolicy), ctx, name) +} + +// GetAllRetryPolicies mocks base method. +func (m *MockRetryPolicyRepository) GetAllRetryPolicies(ctx *armadacontext.Context) ([]*api.RetryPolicy, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAllRetryPolicies", ctx) + ret0, _ := ret[0].([]*api.RetryPolicy) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAllRetryPolicies indicates an expected call of GetAllRetryPolicies. +func (mr *MockRetryPolicyRepositoryMockRecorder) GetAllRetryPolicies(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllRetryPolicies", reflect.TypeOf((*MockRetryPolicyRepository)(nil).GetAllRetryPolicies), ctx) +} + +// GetRetryPolicy mocks base method. +func (m *MockRetryPolicyRepository) GetRetryPolicy(ctx *armadacontext.Context, name string) (*api.RetryPolicy, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetRetryPolicy", ctx, name) + ret0, _ := ret[0].(*api.RetryPolicy) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetRetryPolicy indicates an expected call of GetRetryPolicy. +func (mr *MockRetryPolicyRepositoryMockRecorder) GetRetryPolicy(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRetryPolicy", reflect.TypeOf((*MockRetryPolicyRepository)(nil).GetRetryPolicy), ctx, name) +} + +// UpdateRetryPolicy mocks base method. +func (m *MockRetryPolicyRepository) UpdateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateRetryPolicy", ctx, policy) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpdateRetryPolicy indicates an expected call of UpdateRetryPolicy. +func (mr *MockRetryPolicyRepositoryMockRecorder) UpdateRetryPolicy(ctx, policy any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateRetryPolicy", reflect.TypeOf((*MockRetryPolicyRepository)(nil).UpdateRetryPolicy), ctx, policy) +} diff --git a/internal/server/permissions/permissions.go b/internal/server/permissions/permissions.go index 21b06426b26..3a11183c4ad 100644 --- a/internal/server/permissions/permissions.go +++ b/internal/server/permissions/permissions.go @@ -17,4 +17,7 @@ const ( CordonNodes = "cordon_nodes" ExecuteJobs = "execute_jobs" UpdateExecutorSettings = "update_executor_settings" + CreateRetryPolicy = "create_retry_policy" + UpdateRetryPolicy = "update_retry_policy" + DeleteRetryPolicy = "delete_retry_policy" ) diff --git a/internal/server/queryapi/database/models.go b/internal/server/queryapi/database/models.go index 69ea05982e3..bca1d8217f5 100644 --- a/internal/server/queryapi/database/models.go +++ b/internal/server/queryapi/database/models.go @@ -74,3 +74,8 @@ type Queue struct { Name string `db:"name"` Definition []byte `db:"definition"` } + +type RetryPolicy struct { + Name string `db:"name"` + Definition []byte `db:"definition"` +} diff --git a/internal/server/queue/queue_service.go b/internal/server/queue/queue_service.go index 70ded541b1b..f22f427c914 100644 --- a/internal/server/queue/queue_service.go +++ b/internal/server/queue/queue_service.go @@ -23,9 +23,15 @@ import ( "github.com/armadaproject/armada/pkg/controlplaneevents" ) +// RetryPolicyExistenceChecker reports whether a named retry policy exists. +type RetryPolicyExistenceChecker interface { + RetryPolicyExists(ctx *armadacontext.Context, name string) (bool, error) +} + type Server struct { publisher pulsarutils.Publisher[*controlplaneevents.Event] queueRepository QueueRepository + retryPolicies RetryPolicyExistenceChecker authorizer auth.ActionAuthorizer clock clock.Clock } @@ -33,16 +39,33 @@ type Server struct { func NewServer( publisher pulsarutils.Publisher[*controlplaneevents.Event], queueRepository QueueRepository, + retryPolicies RetryPolicyExistenceChecker, authorizer auth.ActionAuthorizer, ) *Server { return &Server{ publisher: publisher, queueRepository: queueRepository, + retryPolicies: retryPolicies, authorizer: authorizer, clock: clock.RealClock{}, } } +// validateRetryPolicy checks that every retry policy referenced by q exists, +// so a typo cannot silently disable retries for the queue. +func (s *Server) validateRetryPolicy(ctx *armadacontext.Context, q queue.Queue) error { + for _, name := range q.RetryPolicies { + exists, err := s.retryPolicies.RetryPolicyExists(ctx, name) + if err != nil { + return status.Errorf(codes.Unavailable, "error validating retry policy %q: %s", name, err) + } + if !exists { + return status.Errorf(codes.InvalidArgument, "retry policy %q does not exist", name) + } + } + return nil +} + func (s *Server) CreateQueue(grpcCtx context.Context, req *api.Queue) (*types.Empty, error) { ctx := armadacontext.FromGrpcCtx(grpcCtx) err := s.authorizer.AuthorizeAction(ctx, permissions.CreateQueue) @@ -63,6 +86,10 @@ func (s *Server) CreateQueue(grpcCtx context.Context, req *api.Queue) (*types.Em return nil, status.Errorf(codes.InvalidArgument, "error validating queue: %s", err) } + if err := s.validateRetryPolicy(ctx, queue); err != nil { + return nil, err + } + err = s.queueRepository.CreateQueue(ctx, queue) var eq *ErrQueueAlreadyExists if errors.As(err, &eq) { @@ -108,6 +135,10 @@ func (s *Server) UpdateQueue(grpcCtx context.Context, req *api.Queue) (*types.Em return nil, status.Errorf(codes.InvalidArgument, "error: %s", err) } + if err := s.validateRetryPolicy(ctx, queue); err != nil { + return nil, err + } + err = s.queueRepository.UpdateQueue(ctx, queue) var e *ErrQueueNotFound if errors.As(err, &e) { diff --git a/internal/server/queue/queue_service_test.go b/internal/server/queue/queue_service_test.go index 71497e279f0..e3598d0275a 100644 --- a/internal/server/queue/queue_service_test.go +++ b/internal/server/queue/queue_service_test.go @@ -28,21 +28,31 @@ import ( "github.com/armadaproject/armada/pkg/controlplaneevents" ) +type fakeRetryPolicyChecker struct { + exists bool +} + +func (f *fakeRetryPolicyChecker) RetryPolicyExists(_ *armadacontext.Context, _ string) (bool, error) { + return f.exists, nil +} + type queueServiceTestMocks struct { - publisher *commonMocks.MockPublisher[*controlplaneevents.Event] - authorizer *servermocks.MockActionAuthorizer - repo *servermocks.MockQueueRepository + publisher *commonMocks.MockPublisher[*controlplaneevents.Event] + authorizer *servermocks.MockActionAuthorizer + repo *servermocks.MockQueueRepository + retryPolicies *fakeRetryPolicyChecker } func newTestQueueServer(t *testing.T) (*Server, *queueServiceTestMocks) { t.Helper() ctrl := gomock.NewController(t) m := &queueServiceTestMocks{ - publisher: commonMocks.NewMockPublisher[*controlplaneevents.Event](ctrl), - authorizer: servermocks.NewMockActionAuthorizer(ctrl), - repo: servermocks.NewMockQueueRepository(ctrl), + publisher: commonMocks.NewMockPublisher[*controlplaneevents.Event](ctrl), + authorizer: servermocks.NewMockActionAuthorizer(ctrl), + repo: servermocks.NewMockQueueRepository(ctrl), + retryPolicies: &fakeRetryPolicyChecker{exists: true}, } - s := NewServer(m.publisher, m.repo, m.authorizer) + s := NewServer(m.publisher, m.repo, m.retryPolicies, m.authorizer) return s, m } @@ -61,6 +71,45 @@ func TestCreateQueue_PermissionDenied(t *testing.T) { servertest.RequireGrpcCode(t, err, codes.PermissionDenied) } +func TestCreateQueue_UnknownRetryPolicyRejected(t *testing.T) { + s, m := newTestQueueServer(t) + ctx := armadacontext.Background() + + m.authorizer. + EXPECT(). + AuthorizeAction(ctx, permission.Permission(permissions.CreateQueue)). + Return(nil). + Times(1) + // No repo expectation is set: gomock fails if CreateQueue reaches the repository. + m.retryPolicies.exists = false + + _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1", PriorityFactor: 1, RetryPolicies: []string{"ghost"}}) + require.Error(t, err) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) +} + +func TestCreateQueue_MultipleRetryPoliciesAccepted(t *testing.T) { + s, m := newTestQueueServer(t) + ctx := armadacontext.Background() + + m.authorizer. + EXPECT(). + AuthorizeAction(ctx, permission.Permission(permissions.CreateQueue)). + Return(nil). + Times(1) + var created queue.Queue + m.repo. + EXPECT(). + CreateQueue(ctx, gomock.Any()). + Do(func(_ *armadacontext.Context, q queue.Queue) { created = q }). + Return(nil). + Times(1) + + _, err := s.CreateQueue(ctx, &api.Queue{Name: "q1", PriorityFactor: 1, RetryPolicies: []string{"a", "b"}}) + require.NoError(t, err) + assert.Equal(t, []string{"a", "b"}, created.RetryPolicies) +} + func TestCreateQueue_AuthorizeErrorUnavailable(t *testing.T) { s, m := newTestQueueServer(t) ctx := armadacontext.Background() diff --git a/internal/server/retrypolicy/repository.go b/internal/server/retrypolicy/repository.go new file mode 100644 index 00000000000..dd41a86d09a --- /dev/null +++ b/internal/server/retrypolicy/repository.go @@ -0,0 +1,153 @@ +package retrypolicy + +import ( + "fmt" + + "github.com/gogo/protobuf/proto" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/pkg/errors" + + "github.com/armadaproject/armada/internal/common/armadacontext" + "github.com/armadaproject/armada/pkg/api" +) + +type ErrRetryPolicyNotFound struct { + Name string +} + +func (err *ErrRetryPolicyNotFound) Error() string { + return fmt.Sprintf("could not find retry policy %q", err.Name) +} + +type ErrRetryPolicyAlreadyExists struct { + Name string +} + +func (err *ErrRetryPolicyAlreadyExists) Error() string { + return fmt.Sprintf("retry policy %s already exists", err.Name) +} + +type RetryPolicyRepository interface { + GetAllRetryPolicies(ctx *armadacontext.Context) ([]*api.RetryPolicy, error) + GetRetryPolicy(ctx *armadacontext.Context, name string) (*api.RetryPolicy, error) + CreateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error + UpdateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error + DeleteRetryPolicy(ctx *armadacontext.Context, name string) error +} + +type PostgresRetryPolicyRepository struct { + db *pgxpool.Pool +} + +func NewPostgresRetryPolicyRepository(db *pgxpool.Pool) *PostgresRetryPolicyRepository { + return &PostgresRetryPolicyRepository{db: db} +} + +func (r *PostgresRetryPolicyRepository) GetAllRetryPolicies(ctx *armadacontext.Context) ([]*api.RetryPolicy, error) { + rows, err := r.db.Query(ctx, "SELECT definition FROM retry_policy ORDER BY name") + if err != nil { + return nil, errors.WithStack(err) + } + defer rows.Close() + + var policies []*api.RetryPolicy + for rows.Next() { + var definitionBytes []byte + err := rows.Scan(&definitionBytes) + if err != nil { + return nil, errors.WithStack(err) + } + policy, err := unmarshalRetryPolicy(definitionBytes) + if err != nil { + return nil, err + } + policies = append(policies, policy) + } + + if err := rows.Err(); err != nil { + return nil, errors.WithStack(err) + } + return policies, nil +} + +func (r *PostgresRetryPolicyRepository) GetRetryPolicy(ctx *armadacontext.Context, name string) (*api.RetryPolicy, error) { + var definitionBytes []byte + query := "SELECT definition FROM retry_policy WHERE name = $1" + + err := r.db.QueryRow(ctx, query, name).Scan(&definitionBytes) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, &ErrRetryPolicyNotFound{Name: name} + } + return nil, errors.WithStack(err) + } + + return unmarshalRetryPolicy(definitionBytes) +} + +// RetryPolicyExists reports whether a retry policy with the given name exists. +func (r *PostgresRetryPolicyRepository) RetryPolicyExists(ctx *armadacontext.Context, name string) (bool, error) { + _, err := r.GetRetryPolicy(ctx, name) + if err != nil { + var notFound *ErrRetryPolicyNotFound + if errors.As(err, ¬Found) { + return false, nil + } + return false, err + } + return true, nil +} + +func (r *PostgresRetryPolicyRepository) CreateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error { + data, err := proto.Marshal(policy) + if err != nil { + return errors.WithStack(err) + } + + query := "INSERT INTO retry_policy (name, definition) VALUES ($1, $2) ON CONFLICT (name) DO NOTHING" + result, err := r.db.Exec(ctx, query, policy.Name, data) + if err != nil { + return errors.WithStack(err) + } + if result.RowsAffected() == 0 { + return &ErrRetryPolicyAlreadyExists{Name: policy.Name} + } + return nil +} + +func (r *PostgresRetryPolicyRepository) UpdateRetryPolicy(ctx *armadacontext.Context, policy *api.RetryPolicy) error { + data, err := proto.Marshal(policy) + if err != nil { + return errors.WithStack(err) + } + + query := "UPDATE retry_policy SET definition = $2 WHERE name = $1" + result, err := r.db.Exec(ctx, query, policy.Name, data) + if err != nil { + return errors.WithStack(err) + } + if result.RowsAffected() == 0 { + return &ErrRetryPolicyNotFound{Name: policy.Name} + } + return nil +} + +// DeleteRetryPolicy removes a retry policy by name. +// Deleting a non-existent policy is a no-op, consistent with DeleteQueue. +func (r *PostgresRetryPolicyRepository) DeleteRetryPolicy(ctx *armadacontext.Context, name string) error { + query := "DELETE FROM retry_policy WHERE name = $1" + _, err := r.db.Exec(ctx, query, name) + if err != nil { + return errors.WithStack(err) + } + return nil +} + +func unmarshalRetryPolicy(definitionBytes []byte) (*api.RetryPolicy, error) { + policy := &api.RetryPolicy{} + if err := proto.Unmarshal(definitionBytes, policy); err != nil { + return nil, errors.WithStack(err) + } + return policy, nil +} diff --git a/internal/server/retrypolicy/repository_test.go b/internal/server/retrypolicy/repository_test.go new file mode 100644 index 00000000000..6bef8ed53de --- /dev/null +++ b/internal/server/retrypolicy/repository_test.go @@ -0,0 +1,119 @@ +package retrypolicy + +import ( + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/armadaproject/armada/internal/common/armadacontext" + "github.com/armadaproject/armada/internal/common/database/lookout" + "github.com/armadaproject/armada/pkg/api" +) + +func withRetryPolicyRepo(t *testing.T, action func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository)) { + t.Helper() + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 5*time.Second) + defer cancel() + err := lookout.WithLookoutDb(func(db *pgxpool.Pool) error { + action(ctx, NewPostgresRetryPolicyRepository(db)) + return nil + }) + require.NoError(t, err) +} + +func policyFixture(name string) *api.RetryPolicy { + return &api.RetryPolicy{ + Name: name, + RetryLimit: 3, + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + Rules: []*api.RetryRule{ + {Action: api.RetryAction_RETRY_ACTION_RETRY, OnCategory: "gpu", OnSubcategory: "transient"}, + }, + } +} + +func TestPostgresRetryPolicyRepository_CreateThenGetRoundTrips(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + policy := policyFixture("p1") + require.NoError(t, repo.CreateRetryPolicy(ctx, policy)) + + got, err := repo.GetRetryPolicy(ctx, "p1") + require.NoError(t, err) + assert.Equal(t, "p1", got.Name) + assert.Equal(t, uint32(3), got.RetryLimit) + assert.Equal(t, api.RetryAction_RETRY_ACTION_FAIL, got.DefaultAction) + require.Len(t, got.Rules, 1) + assert.Equal(t, "gpu", got.Rules[0].OnCategory) + assert.Equal(t, "transient", got.Rules[0].OnSubcategory) + }) +} + +func TestPostgresRetryPolicyRepository_CreateDuplicateReturnsAlreadyExists(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + require.NoError(t, repo.CreateRetryPolicy(ctx, policyFixture("dup"))) + err := repo.CreateRetryPolicy(ctx, policyFixture("dup")) + require.Error(t, err) + assert.IsType(t, &ErrRetryPolicyAlreadyExists{}, err) + }) +} + +func TestPostgresRetryPolicyRepository_GetMissingReturnsNotFound(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + _, err := repo.GetRetryPolicy(ctx, "absent") + require.Error(t, err) + assert.IsType(t, &ErrRetryPolicyNotFound{}, err) + }) +} + +func TestPostgresRetryPolicyRepository_UpdateReplacesDefinition(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + require.NoError(t, repo.CreateRetryPolicy(ctx, policyFixture("p1"))) + updated := policyFixture("p1") + updated.RetryLimit = 9 + require.NoError(t, repo.UpdateRetryPolicy(ctx, updated)) + + got, err := repo.GetRetryPolicy(ctx, "p1") + require.NoError(t, err) + assert.Equal(t, uint32(9), got.RetryLimit) + }) +} + +func TestPostgresRetryPolicyRepository_UpdateMissingReturnsNotFound(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + err := repo.UpdateRetryPolicy(ctx, policyFixture("absent")) + require.Error(t, err) + assert.IsType(t, &ErrRetryPolicyNotFound{}, err) + }) +} + +func TestPostgresRetryPolicyRepository_GetAllIsOrderedByNameAndEmptyWhenNone(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + all, err := repo.GetAllRetryPolicies(ctx) + require.NoError(t, err) + assert.Empty(t, all) + + require.NoError(t, repo.CreateRetryPolicy(ctx, policyFixture("b"))) + require.NoError(t, repo.CreateRetryPolicy(ctx, policyFixture("a"))) + + all, err = repo.GetAllRetryPolicies(ctx) + require.NoError(t, err) + require.Len(t, all, 2) + assert.Equal(t, "a", all[0].Name) + assert.Equal(t, "b", all[1].Name) + }) +} + +func TestPostgresRetryPolicyRepository_DeleteIsIdempotent(t *testing.T) { + withRetryPolicyRepo(t, func(ctx *armadacontext.Context, repo *PostgresRetryPolicyRepository) { + require.NoError(t, repo.CreateRetryPolicy(ctx, policyFixture("p1"))) + require.NoError(t, repo.DeleteRetryPolicy(ctx, "p1")) + + _, err := repo.GetRetryPolicy(ctx, "p1") + assert.IsType(t, &ErrRetryPolicyNotFound{}, err) + + require.NoError(t, repo.DeleteRetryPolicy(ctx, "p1")) + }) +} diff --git a/internal/server/retrypolicy/service.go b/internal/server/retrypolicy/service.go new file mode 100644 index 00000000000..3f3d6c16798 --- /dev/null +++ b/internal/server/retrypolicy/service.go @@ -0,0 +1,180 @@ +package retrypolicy + +import ( + "context" + "fmt" + "slices" + "strings" + + "github.com/gogo/protobuf/types" + "github.com/pkg/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/armadaproject/armada/internal/common/armadacontext" + "github.com/armadaproject/armada/internal/common/armadaerrors" + "github.com/armadaproject/armada/internal/common/auth" + "github.com/armadaproject/armada/internal/common/auth/permission" + "github.com/armadaproject/armada/internal/server/permissions" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client/queue" +) + +// QueueLister supplies the queues checked before a retry policy is deleted. +type QueueLister interface { + GetAllQueues(ctx *armadacontext.Context) ([]queue.Queue, error) +} + +type Server struct { + repository RetryPolicyRepository + queueLister QueueLister + authorizer auth.ActionAuthorizer +} + +func NewServer(repository RetryPolicyRepository, queueLister QueueLister, authorizer auth.ActionAuthorizer) *Server { + return &Server{ + repository: repository, + queueLister: queueLister, + authorizer: authorizer, + } +} + +func (s *Server) authorize(ctx *armadacontext.Context, perm permission.Permission, verb, name string) error { + err := s.authorizer.AuthorizeAction(ctx, perm) + var ep *armadaerrors.ErrUnauthorized + if errors.As(err, &ep) { + return status.Errorf(codes.PermissionDenied, "error %s retry policy %s: %s", verb, name, ep) + } + if err != nil { + return status.Errorf(codes.Unavailable, "error checking permissions: %s", err) + } + return nil +} + +func (s *Server) CreateRetryPolicy(grpcCtx context.Context, req *api.RetryPolicy) (*types.Empty, error) { + ctx := armadacontext.FromGrpcCtx(grpcCtx) + if err := s.authorize(ctx, permissions.CreateRetryPolicy, "creating", req.Name); err != nil { + return nil, err + } + + if err := ValidatePolicy(req); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "invalid retry policy: %s", err) + } + + err := s.repository.CreateRetryPolicy(ctx, req) + var ea *ErrRetryPolicyAlreadyExists + if errors.As(err, &ea) { + return nil, status.Errorf(codes.AlreadyExists, "error creating retry policy: %s", err) + } + if err != nil { + return nil, status.Errorf(codes.Unavailable, "error creating retry policy: %s", err) + } + + return &types.Empty{}, nil +} + +func (s *Server) UpdateRetryPolicy(grpcCtx context.Context, req *api.RetryPolicy) (*types.Empty, error) { + ctx := armadacontext.FromGrpcCtx(grpcCtx) + if err := s.authorize(ctx, permissions.UpdateRetryPolicy, "updating", req.Name); err != nil { + return nil, err + } + + if err := ValidatePolicy(req); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "invalid retry policy: %s", err) + } + + err := s.repository.UpdateRetryPolicy(ctx, req) + var enf *ErrRetryPolicyNotFound + if errors.As(err, &enf) { + return nil, status.Errorf(codes.NotFound, "error: %s", err) + } + if err != nil { + return nil, status.Errorf(codes.Unavailable, "error updating retry policy %q: %s", req.Name, err) + } + + return &types.Empty{}, nil +} + +func (s *Server) DeleteRetryPolicy(grpcCtx context.Context, req *api.RetryPolicyDeleteRequest) (*types.Empty, error) { + ctx := armadacontext.FromGrpcCtx(grpcCtx) + if err := s.authorize(ctx, permissions.DeleteRetryPolicy, "deleting", req.Name); err != nil { + return nil, err + } + + if req.Name == "" { + return nil, status.Errorf(codes.InvalidArgument, "retry policy name must not be empty") + } + + // Reject the delete while queues still reference the policy, so a queue is + // never left pointing at a policy that no longer exists. + // + // This check and the delete below are not atomic: the queue and policy + // stores are separate, so a queue could attach the policy in between. The + // race is accepted because closing it would require a transaction spanning + // both stores. + referencing, err := s.queuesReferencingPolicy(ctx, req.Name) + if err != nil { + return nil, status.Errorf(codes.Unavailable, "error checking queues referencing retry policy %s: %s", req.Name, err) + } + if len(referencing) > 0 { + shown := referencing[:min(len(referencing), maxReportedReferencingQueues)] + return nil, status.Errorf( + codes.FailedPrecondition, + "retry policy %s is still referenced by %d queue(s), including: %s", + req.Name, len(referencing), strings.Join(shown, ", "), + ) + } + + if err := s.repository.DeleteRetryPolicy(ctx, req.Name); err != nil { + return nil, status.Errorf(codes.Unavailable, "error deleting retry policy %s: %s", req.Name, err) + } + return &types.Empty{}, nil +} + +// maxReportedReferencingQueues caps the queue names listed in the delete error message. +const maxReportedReferencingQueues = 5 + +func (s *Server) queuesReferencingPolicy(ctx *armadacontext.Context, policyName string) ([]string, error) { + queues, err := s.queueLister.GetAllQueues(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list queues: %w", err) + } + var names []string + for _, q := range queues { + if slices.Contains(q.RetryPolicies, policyName) { + names = append(names, q.Name) + } + } + return names, nil +} + +// GetRetryPolicy returns a single retry policy by name. +// Reads require no permission, consistent with GetQueue. +func (s *Server) GetRetryPolicy(grpcCtx context.Context, req *api.RetryPolicyGetRequest) (*api.RetryPolicy, error) { + ctx := armadacontext.FromGrpcCtx(grpcCtx) + + if req.Name == "" { + return nil, status.Errorf(codes.InvalidArgument, "retry policy name must not be empty") + } + + policy, err := s.repository.GetRetryPolicy(ctx, req.Name) + var enf *ErrRetryPolicyNotFound + if errors.As(err, &enf) { + return nil, status.Errorf(codes.NotFound, "error: %s", err) + } + if err != nil { + return nil, status.Errorf(codes.Unavailable, "error getting retry policy %q: %s", req.Name, err) + } + return policy, nil +} + +// GetRetryPolicies returns all retry policies. +// Reads require no permission, consistent with GetQueue. +func (s *Server) GetRetryPolicies(grpcCtx context.Context, _ *api.RetryPolicyListRequest) (*api.RetryPolicyList, error) { + ctx := armadacontext.FromGrpcCtx(grpcCtx) + policies, err := s.repository.GetAllRetryPolicies(ctx) + if err != nil { + return nil, status.Errorf(codes.Unavailable, "error getting retry policies: %s", err) + } + return &api.RetryPolicyList{RetryPolicies: policies}, nil +} diff --git a/internal/server/retrypolicy/service_test.go b/internal/server/retrypolicy/service_test.go new file mode 100644 index 00000000000..aebfe513059 --- /dev/null +++ b/internal/server/retrypolicy/service_test.go @@ -0,0 +1,384 @@ +package retrypolicy + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "google.golang.org/grpc/codes" + + "github.com/armadaproject/armada/internal/common/armadacontext" + "github.com/armadaproject/armada/internal/common/armadaerrors" + "github.com/armadaproject/armada/internal/common/auth/permission" + servermocks "github.com/armadaproject/armada/internal/server/mocks" + "github.com/armadaproject/armada/internal/server/permissions" + "github.com/armadaproject/armada/internal/server/servertest" + "github.com/armadaproject/armada/pkg/api" + "github.com/armadaproject/armada/pkg/client/queue" +) + +type fakeQueueLister struct { + queues []queue.Queue + err error +} + +func (f *fakeQueueLister) GetAllQueues(_ *armadacontext.Context) ([]queue.Queue, error) { + return f.queues, f.err +} + +type testMocks struct { + authorizer *servermocks.MockActionAuthorizer + repo *servermocks.MockRetryPolicyRepository + queueLister *fakeQueueLister +} + +func (m *testMocks) expectAuthorizeAction(ctx *armadacontext.Context, perm string, authErr error) { + m.authorizer. + EXPECT(). + AuthorizeAction(ctx, permission.Permission(perm)). + Return(authErr). + Times(1) +} + +func newTestServer(t *testing.T) (*Server, *testMocks) { + t.Helper() + ctrl := gomock.NewController(t) + m := &testMocks{ + authorizer: servermocks.NewMockActionAuthorizer(ctrl), + repo: servermocks.NewMockRetryPolicyRepository(ctrl), + queueLister: &fakeQueueLister{}, + } + s := NewServer(m.repo, m.queueLister, m.authorizer) + return s, m +} + +// validPolicy returns a minimal policy that passes ValidatePolicy. +func validPolicy(name string) *api.RetryPolicy { + return &api.RetryPolicy{ + Name: name, + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + } +} + +func TestCreateRetryPolicy_PermissionDenied(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, &armadaerrors.ErrUnauthorized{Principal: "alice", Permission: "create_retry_policy"}) + + _, err := s.CreateRetryPolicy(ctx, &api.RetryPolicy{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) +} + +func TestCreateRetryPolicy_AuthorizeErrorUnavailable(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, errors.New("authorizer down")) + + _, err := s.CreateRetryPolicy(ctx, &api.RetryPolicy{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.Unavailable) +} + +func TestCreateRetryPolicy_AlreadyExists(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, nil) + + m.repo. + EXPECT(). + CreateRetryPolicy(gomock.Any(), gomock.Any()). + Return(&ErrRetryPolicyAlreadyExists{Name: "p1"}). + Times(1) + + _, err := s.CreateRetryPolicy(ctx, validPolicy("p1")) + servertest.RequireGrpcCode(t, err, codes.AlreadyExists) +} + +func TestCreateRetryPolicy_InvalidPolicyRejected(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, nil) + + // The policy lacks a default action, so ValidatePolicy rejects it. No repo + // expectation is set: gomock fails the test if the repository is called. + _, err := s.CreateRetryPolicy(ctx, &api.RetryPolicy{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) +} + +func TestCreateRetryPolicy_Success(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, nil) + + var persisted *api.RetryPolicy + m.repo. + EXPECT(). + CreateRetryPolicy(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ *armadacontext.Context, p *api.RetryPolicy) error { + persisted = p + return nil + }). + Times(1) + + policy := validPolicy("p1") + policy.RetryLimit = 3 + _, err := s.CreateRetryPolicy(ctx, policy) + require.NoError(t, err) + require.NotNil(t, persisted) + assert.Equal(t, "p1", persisted.Name) + assert.Equal(t, uint32(3), persisted.RetryLimit) +} + +func TestUpdateRetryPolicy_PermissionDenied(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.UpdateRetryPolicy, &armadaerrors.ErrUnauthorized{Principal: "alice", Permission: "update_retry_policy"}) + + _, err := s.UpdateRetryPolicy(ctx, &api.RetryPolicy{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) +} + +func TestUpdateRetryPolicy_NotFound(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.UpdateRetryPolicy, nil) + + m.repo. + EXPECT(). + UpdateRetryPolicy(gomock.Any(), gomock.Any()). + Return(&ErrRetryPolicyNotFound{Name: "p1"}). + Times(1) + + _, err := s.UpdateRetryPolicy(ctx, validPolicy("p1")) + servertest.RequireGrpcCode(t, err, codes.NotFound) +} + +func TestUpdateRetryPolicy_InvalidPolicyRejected(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.UpdateRetryPolicy, nil) + + _, err := s.UpdateRetryPolicy(ctx, &api.RetryPolicy{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) +} + +func TestUpdateRetryPolicy_Success(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.UpdateRetryPolicy, nil) + + var persisted *api.RetryPolicy + m.repo. + EXPECT(). + UpdateRetryPolicy(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ *armadacontext.Context, p *api.RetryPolicy) error { + persisted = p + return nil + }). + Times(1) + + policy := validPolicy("p1") + policy.RetryLimit = 5 + _, err := s.UpdateRetryPolicy(ctx, policy) + require.NoError(t, err) + require.NotNil(t, persisted) + assert.Equal(t, "p1", persisted.Name) + assert.Equal(t, uint32(5), persisted.RetryLimit) +} + +func TestDeleteRetryPolicy_EmptyName(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.DeleteRetryPolicy, nil) + + _, err := s.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: ""}) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) +} + +func TestDeleteRetryPolicy_PermissionDenied(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.DeleteRetryPolicy, &armadaerrors.ErrUnauthorized{Principal: "alice", Permission: "delete_retry_policy"}) + + _, err := s.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.PermissionDenied) +} + +func TestDeleteRetryPolicy_Success(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.DeleteRetryPolicy, nil) + + m.repo. + EXPECT(). + DeleteRetryPolicy(gomock.Any(), "p1"). + Return(nil). + Times(1) + + _, err := s.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: "p1"}) + require.NoError(t, err) +} + +func TestDeleteRetryPolicy_ReferencedByQueues(t *testing.T) { + tests := map[string]struct { + queues []queue.Queue + wantInMessage []string + notInMessage []string + }{ + "single referencing queue": { + queues: []queue.Queue{ + {Name: "queue-a", RetryPolicies: []string{"p1"}}, + {Name: "queue-b", RetryPolicies: []string{"other"}}, + {Name: "queue-c"}, + }, + wantInMessage: []string{"queue-a"}, + notInMessage: []string{"queue-b", "queue-c"}, + }, + "more referencing queues than the reporting cap": { + queues: []queue.Queue{ + {Name: "q1", RetryPolicies: []string{"p1"}}, + {Name: "q2", RetryPolicies: []string{"p1"}}, + {Name: "q3", RetryPolicies: []string{"p1"}}, + {Name: "q4", RetryPolicies: []string{"p1"}}, + {Name: "q5", RetryPolicies: []string{"p1"}}, + {Name: "q6", RetryPolicies: []string{"p1"}}, + }, + wantInMessage: []string{"6 queue(s)", "q1", "q5"}, + notInMessage: []string{"q6"}, + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.DeleteRetryPolicy, nil) + + // No repo delete expectation: gomock fails the test if the + // repository is called despite live references. + m.queueLister.queues = tc.queues + + _, err := s.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: "p1"}) + require.Error(t, err) + servertest.RequireGrpcCode(t, err, codes.FailedPrecondition) + for _, want := range tc.wantInMessage { + assert.Contains(t, err.Error(), want) + } + for _, notWant := range tc.notInMessage { + assert.NotContains(t, err.Error(), notWant) + } + }) + } +} + +func TestDeleteRetryPolicy_QueueListError(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.DeleteRetryPolicy, nil) + + m.queueLister.err = errors.New("postgres down") + + _, err := s.DeleteRetryPolicy(ctx, &api.RetryPolicyDeleteRequest{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.Unavailable) +} + +func TestGetRetryPolicy_EmptyName(t *testing.T) { + s, _ := newTestServer(t) + ctx := armadacontext.Background() + + _, err := s.GetRetryPolicy(ctx, &api.RetryPolicyGetRequest{Name: ""}) + servertest.RequireGrpcCode(t, err, codes.InvalidArgument) +} + +func TestGetRetryPolicy_NotFound(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.repo. + EXPECT(). + GetRetryPolicy(gomock.Any(), "p1"). + Return(nil, &ErrRetryPolicyNotFound{Name: "p1"}). + Times(1) + + _, err := s.GetRetryPolicy(ctx, &api.RetryPolicyGetRequest{Name: "p1"}) + servertest.RequireGrpcCode(t, err, codes.NotFound) +} + +func TestGetRetryPolicy_Success(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + expected := &api.RetryPolicy{Name: "p1", RetryLimit: 5} + m.repo. + EXPECT(). + GetRetryPolicy(gomock.Any(), "p1"). + Return(expected, nil). + Times(1) + + result, err := s.GetRetryPolicy(ctx, &api.RetryPolicyGetRequest{Name: "p1"}) + require.NoError(t, err) + assert.Equal(t, expected, result) +} + +func TestGetRetryPolicies_Success(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + expected := []*api.RetryPolicy{ + {Name: "p1", RetryLimit: 3}, + {Name: "p2", RetryLimit: 5}, + } + m.repo. + EXPECT(). + GetAllRetryPolicies(gomock.Any()). + Return(expected, nil). + Times(1) + + result, err := s.GetRetryPolicies(ctx, &api.RetryPolicyListRequest{}) + require.NoError(t, err) + assert.Equal(t, expected, result.RetryPolicies) +} + +func TestCreateRetryPolicy_RepoErrorUnavailable(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.expectAuthorizeAction(ctx, permissions.CreateRetryPolicy, nil) + + m.repo. + EXPECT(). + CreateRetryPolicy(gomock.Any(), gomock.Any()). + Return(errors.New("db down")). + Times(1) + + _, err := s.CreateRetryPolicy(ctx, validPolicy("p1")) + servertest.RequireGrpcCode(t, err, codes.Unavailable) +} + +func TestGetRetryPolicies_RepoErrorUnavailable(t *testing.T) { + s, m := newTestServer(t) + ctx := armadacontext.Background() + + m.repo. + EXPECT(). + GetAllRetryPolicies(gomock.Any()). + Return(nil, errors.New("db down")). + Times(1) + + _, err := s.GetRetryPolicies(ctx, &api.RetryPolicyListRequest{}) + servertest.RequireGrpcCode(t, err, codes.Unavailable) +} diff --git a/internal/server/retrypolicy/validation.go b/internal/server/retrypolicy/validation.go new file mode 100644 index 00000000000..24945a5fd1d --- /dev/null +++ b/internal/server/retrypolicy/validation.go @@ -0,0 +1,58 @@ +package retrypolicy + +import ( + "fmt" + "regexp" + + "github.com/armadaproject/armada/pkg/api" +) + +// Policy names follow RFC 1123 label rules, like queue names, because they +// may end up in Kubernetes labels. +var policyNamePattern = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]*[a-z0-9])?$`) + +const maxPolicyNameLength = 63 + +// ValidatePolicy checks that a retry policy is structurally valid, so that +// malformed policies are rejected at write time. +func ValidatePolicy(p *api.RetryPolicy) error { + if p == nil { + return fmt.Errorf("retry policy must not be nil") + } + if p.Name == "" { + return fmt.Errorf("retry policy name must not be empty") + } + if len(p.Name) > maxPolicyNameLength { + return fmt.Errorf("retry policy name %q must be at most %d characters", p.Name, maxPolicyNameLength) + } + if !policyNamePattern.MatchString(p.Name) { + return fmt.Errorf( + "retry policy name %q is invalid: must consist of lowercase alphanumeric characters or '-', and must start and end with an alphanumeric character", + p.Name, + ) + } + for i, rule := range p.Rules { + if err := validateRule(rule); err != nil { + return fmt.Errorf("retry policy %q rule %d: %w", p.Name, i, err) + } + } + if p.DefaultAction != api.RetryAction_RETRY_ACTION_FAIL && p.DefaultAction != api.RetryAction_RETRY_ACTION_RETRY { + return fmt.Errorf("retry policy %q must set a default action (Fail or Retry)", p.Name) + } + return nil +} + +func validateRule(r *api.RetryRule) error { + if r == nil { + return fmt.Errorf("rule must not be nil") + } + if r.Action != api.RetryAction_RETRY_ACTION_FAIL && r.Action != api.RetryAction_RETRY_ACTION_RETRY { + return fmt.Errorf("action must be Fail or Retry") + } + // on_subcategory only narrows an on_category match, so it does not count + // as a matcher on its own. + if r.OnCategory == "" { + return fmt.Errorf("on_category must be set") + } + return nil +} diff --git a/internal/server/retrypolicy/validation_test.go b/internal/server/retrypolicy/validation_test.go new file mode 100644 index 00000000000..279427b4fbb --- /dev/null +++ b/internal/server/retrypolicy/validation_test.go @@ -0,0 +1,121 @@ +package retrypolicy + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/armadaproject/armada/pkg/api" +) + +func TestValidatePolicy(t *testing.T) { + tests := map[string]struct { + policy *api.RetryPolicy + wantErr string // empty means the policy is expected to validate + }{ + "nil policy": { + policy: nil, + wantErr: "must not be nil", + }, + "empty name": { + policy: &api.RetryPolicy{DefaultAction: api.RetryAction_RETRY_ACTION_FAIL}, + wantErr: "name must not be empty", + }, + "name with uppercase characters": { + policy: &api.RetryPolicy{ + Name: "MyPolicy", + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + }, + wantErr: "is invalid", + }, + "name with leading dash": { + policy: &api.RetryPolicy{ + Name: "-policy", + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + }, + wantErr: "is invalid", + }, + "name with trailing dash": { + policy: &api.RetryPolicy{ + Name: "policy-", + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + }, + wantErr: "is invalid", + }, + "name too long": { + policy: &api.RetryPolicy{ + Name: strings.Repeat("a", maxPolicyNameLength+1), + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + }, + wantErr: "at most 63 characters", + }, + "nil rule": { + policy: &api.RetryPolicy{ + Name: "p1", + Rules: []*api.RetryRule{nil}, + }, + wantErr: "rule must not be nil", + }, + "rule with unspecified action": { + policy: &api.RetryPolicy{ + Name: "p1", + Rules: []*api.RetryRule{ + {}, + }, + }, + wantErr: "action must be Fail or Retry", + }, + "rule with no matchers": { + policy: &api.RetryPolicy{ + Name: "p1", + Rules: []*api.RetryRule{ + {Action: api.RetryAction_RETRY_ACTION_RETRY}, + }, + }, + wantErr: "on_category must be set", + }, + "rule with only subcategory": { + policy: &api.RetryPolicy{ + Name: "p1", + Rules: []*api.RetryRule{ + { + Action: api.RetryAction_RETRY_ACTION_RETRY, + OnSubcategory: "oom", + }, + }, + }, + wantErr: "on_category must be set", + }, + "valid policy with default action only": { + policy: &api.RetryPolicy{ + Name: "p1", + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + }, + }, + "valid policy with on_category rule": { + policy: &api.RetryPolicy{ + Name: "p1", + DefaultAction: api.RetryAction_RETRY_ACTION_FAIL, + Rules: []*api.RetryRule{ + { + Action: api.RetryAction_RETRY_ACTION_RETRY, + OnCategory: "transient", + }, + }, + }, + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + err := ValidatePolicy(tc.policy) + if tc.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index b0932998b4b..14906e8682b 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -31,6 +31,7 @@ import ( "github.com/armadaproject/armada/internal/server/node" "github.com/armadaproject/armada/internal/server/queryapi" "github.com/armadaproject/armada/internal/server/queue" + "github.com/armadaproject/armada/internal/server/retrypolicy" "github.com/armadaproject/armada/internal/server/submit" "github.com/armadaproject/armada/pkg/api" "github.com/armadaproject/armada/pkg/api/schedulerobjects" @@ -179,7 +180,10 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt } defer controlPlaneEventsPublisher.Close() - queueServer := queue.NewServer(controlPlaneEventsPublisher, queueRepository, authorizer) + retryPolicyRepo := retrypolicy.NewPostgresRetryPolicyRepository(dbPool) + + queueServer := queue.NewServer(controlPlaneEventsPublisher, queueRepository, retryPolicyRepo, authorizer) + retryPolicyServer := retrypolicy.NewServer(retryPolicyRepo, queueRepository, authorizer) submitServer := submit.NewServer( queueServer, @@ -209,6 +213,7 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt api.RegisterSubmitServer(grpcServer, submitServer) api.RegisterEventServer(grpcServer, eventServer) api.RegisterQueueServiceServer(grpcServer, queueServer) + api.RegisterRetryPolicyServiceServer(grpcServer, retryPolicyServer) api.RegisterExecutorServer(grpcServer, executorServer) api.RegisterNodeServer(grpcServer, nodeServer) diff --git a/pkg/client/queue/queue.go b/pkg/client/queue/queue.go index effb16b6422..09f739d0cb7 100644 --- a/pkg/client/queue/queue.go +++ b/pkg/client/queue/queue.go @@ -14,6 +14,7 @@ type Queue struct { ResourceLimitsByPriorityClassName map[string]api.PriorityClassResourceLimits Cordoned bool `json:"cordoned"` Labels map[string]string `json:"labels"` + RetryPolicies []string `json:"retryPolicies"` } // NewQueue returns new Queue using the in parameter. Error is returned if @@ -62,6 +63,7 @@ func NewQueue(in *api.Queue) (Queue, error) { ResourceLimitsByPriorityClassName: resourceLimitsByPriorityClassName, Cordoned: in.Cordoned, Labels: in.Labels, + RetryPolicies: in.RetryPolicies, }, nil } @@ -75,8 +77,9 @@ func (q Queue) ToAPI() *api.Queue { func(p api.PriorityClassResourceLimits) *api.PriorityClassResourceLimits { return &p }), - Cordoned: q.Cordoned, - Labels: q.Labels, + Cordoned: q.Cordoned, + Labels: q.Labels, + RetryPolicies: q.RetryPolicies, } for _, permission := range q.Permissions { rv.Permissions = append(rv.Permissions, permission.ToAPI())