diff --git a/go/tasks/pluginmachinery/k8s/mocks/plugin_abort_override.go b/go/tasks/pluginmachinery/k8s/mocks/plugin_abort_override.go new file mode 100644 index 000000000..41463d103 --- /dev/null +++ b/go/tasks/pluginmachinery/k8s/mocks/plugin_abort_override.go @@ -0,0 +1,59 @@ +// Code generated by mockery v1.0.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + client "sigs.k8s.io/controller-runtime/pkg/client" + + core "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core" + + k8s "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/k8s" + + mock "github.com/stretchr/testify/mock" +) + +// PluginAbortOverride is an autogenerated mock type for the PluginAbortOverride type +type PluginAbortOverride struct { + mock.Mock +} + +type PluginAbortOverride_OnAbort struct { + *mock.Call +} + +func (_m PluginAbortOverride_OnAbort) Return(behavior k8s.AbortBehavior, err error) *PluginAbortOverride_OnAbort { + return &PluginAbortOverride_OnAbort{Call: _m.Call.Return(behavior, err)} +} + +func (_m *PluginAbortOverride) OnOnAbort(ctx context.Context, tCtx core.TaskExecutionContext, resource client.Object) *PluginAbortOverride_OnAbort { + c := _m.On("OnAbort", ctx, tCtx, resource) + return &PluginAbortOverride_OnAbort{Call: c} +} + +func (_m *PluginAbortOverride) OnOnAbortMatch(matchers ...interface{}) *PluginAbortOverride_OnAbort { + c := _m.On("OnAbort", matchers...) + return &PluginAbortOverride_OnAbort{Call: c} +} + +// OnAbort provides a mock function with given fields: ctx, tCtx, resource +func (_m *PluginAbortOverride) OnAbort(ctx context.Context, tCtx core.TaskExecutionContext, resource client.Object) (k8s.AbortBehavior, error) { + ret := _m.Called(ctx, tCtx, resource) + + var r0 k8s.AbortBehavior + if rf, ok := ret.Get(0).(func(context.Context, core.TaskExecutionContext, client.Object) k8s.AbortBehavior); ok { + r0 = rf(ctx, tCtx, resource) + } else { + r0 = ret.Get(0).(k8s.AbortBehavior) + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, core.TaskExecutionContext, client.Object) error); ok { + r1 = rf(ctx, tCtx, resource) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} diff --git a/go/tasks/pluginmachinery/k8s/plugin.go b/go/tasks/pluginmachinery/k8s/plugin.go index 0a746c991..8f83d4ee0 100644 --- a/go/tasks/pluginmachinery/k8s/plugin.go +++ b/go/tasks/pluginmachinery/k8s/plugin.go @@ -91,3 +91,72 @@ type Plugin interface { // Properties desired by the plugin GetProperties() PluginProperties } + +// An optional interface a Plugin can implement to override its default OnAbort finalizer (deletion of the underlying resource). +type PluginAbortOverride interface { + OnAbort(ctx context.Context, tCtx pluginsCore.TaskExecutionContext, resource client.Object) (behavior AbortBehavior, err error) +} + +// Defines the overridden OnAbort behavior. The resource (by default, the underlying resource, although this +// can be overridden) can be either patched, updated, or deleted. +type AbortBehavior struct { + // Optional override to the default k8s Resource being acted on. + Resource client.Object + DeleteResource bool + Update *UpdateResourceOperation + Patch *PatchResourceOperation + // Determines whether to delete the Resource if the specified operations return an error + DeleteOnErr bool +} + +// Defines a Patch operation on a Resource +type PatchResourceOperation struct { + Patch client.Patch + Options []client.PatchOption +} + +// Defines an Update operation on a Resource +type UpdateResourceOperation struct { + Options []client.UpdateOption +} + +// AbortBehavior that patches the default resource +func AbortBehaviorPatchDefaultResource(patchOperation PatchResourceOperation, deleteOnErr bool) AbortBehavior { + return AbortBehaviorPatch(patchOperation, deleteOnErr, nil) +} + +// AbortBehavior that patches the specified resource +func AbortBehaviorPatch(patchOperation PatchResourceOperation, deleteOnErr bool, resource client.Object) AbortBehavior { + return AbortBehavior{ + Resource: resource, + Patch: &patchOperation, + DeleteOnErr: deleteOnErr, + } +} + +// AbortBehavior that updates the default resource +func AbortBehaviorUpdateDefaultResource(updateOperation UpdateResourceOperation, deleteOnErr bool) AbortBehavior { + return AbortBehaviorUpdate(updateOperation, deleteOnErr, nil) +} + +// AbortBehavior that updates the specified resource +func AbortBehaviorUpdate(updateOperation UpdateResourceOperation, deleteOnErr bool, resource client.Object) AbortBehavior { + return AbortBehavior{ + Resource: resource, + Update: &updateOperation, + DeleteOnErr: deleteOnErr, + } +} + +// AbortBehavior that deletes the default resource +func AbortBehaviorDeleteDefaultResource() AbortBehavior { + return AbortBehaviorDelete(nil) +} + +// AbortBehavior that deletes the specified resource +func AbortBehaviorDelete(resource client.Object) AbortBehavior { + return AbortBehavior{ + Resource: resource, + DeleteResource: true, + } +}