diff --git a/go/tasks/pluginmachinery/catalog/client.go b/go/tasks/pluginmachinery/catalog/client.go index 7531a59d9..30b99bd3a 100644 --- a/go/tasks/pluginmachinery/catalog/client.go +++ b/go/tasks/pluginmachinery/catalog/client.go @@ -131,6 +131,9 @@ type Client interface { // GetOrExtendReservation tries to retrieve a (valid) reservation for the given key, creating a new one using the // specified owner ID if none was found or updating an existing one if it has expired. GetOrExtendReservation(ctx context.Context, key Key, ownerID string, heartbeatInterval time.Duration) (*datacatalog.Reservation, error) + // GetOrExtendReservationByArtifactTag tries to retrieve a (valid) reservation for the given dataset ID and artifact + // tag, creating a new one using the specified owner ID if none was found or updating an existing one if it has expired. + GetOrExtendReservationByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string, ownerID string, heartbeatInterval time.Duration) (*datacatalog.Reservation, error) // Put stores the given data using the specified key, creating artifact entries as required. // To update an existing artifact, use Update instead. Put(ctx context.Context, key Key, reader io.OutputReader, metadata Metadata) (Status, error) @@ -139,6 +142,11 @@ type Client interface { Update(ctx context.Context, key Key, reader io.OutputReader, metadata Metadata) (Status, error) // ReleaseReservation releases an acquired reservation for the given key and owner ID. ReleaseReservation(ctx context.Context, key Key, ownerID string) error + // Delete removes the artifact associated with the given key and deletes its underlying data from blob storage. + Delete(ctx context.Context, key Key) error + // DeleteByArtifactTag removes the artifact associated with the given dataset ID and artifact tag and deletes its + // underlying data from blob storage. + DeleteByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string) error } func IsNotFound(err error) bool { diff --git a/go/tasks/pluginmachinery/catalog/mocks/client.go b/go/tasks/pluginmachinery/catalog/mocks/client.go index fe9d75f1c..172b72177 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/client.go +++ b/go/tasks/pluginmachinery/catalog/mocks/client.go @@ -21,6 +21,70 @@ type Client struct { mock.Mock } +type Client_Delete struct { + *mock.Call +} + +func (_m Client_Delete) Return(_a0 error) *Client_Delete { + return &Client_Delete{Call: _m.Call.Return(_a0)} +} + +func (_m *Client) OnDelete(ctx context.Context, key catalog.Key) *Client_Delete { + c_call := _m.On("Delete", ctx, key) + return &Client_Delete{Call: c_call} +} + +func (_m *Client) OnDeleteMatch(matchers ...interface{}) *Client_Delete { + c_call := _m.On("Delete", matchers...) + return &Client_Delete{Call: c_call} +} + +// Delete provides a mock function with given fields: ctx, key +func (_m *Client) Delete(ctx context.Context, key catalog.Key) error { + ret := _m.Called(ctx, key) + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, catalog.Key) error); ok { + r0 = rf(ctx, key) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +type Client_DeleteByArtifactTag struct { + *mock.Call +} + +func (_m Client_DeleteByArtifactTag) Return(_a0 error) *Client_DeleteByArtifactTag { + return &Client_DeleteByArtifactTag{Call: _m.Call.Return(_a0)} +} + +func (_m *Client) OnDeleteByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string) *Client_DeleteByArtifactTag { + c_call := _m.On("DeleteByArtifactTag", ctx, datasetID, artifactTag) + return &Client_DeleteByArtifactTag{Call: c_call} +} + +func (_m *Client) OnDeleteByArtifactTagMatch(matchers ...interface{}) *Client_DeleteByArtifactTag { + c_call := _m.On("DeleteByArtifactTag", matchers...) + return &Client_DeleteByArtifactTag{Call: c_call} +} + +// DeleteByArtifactTag provides a mock function with given fields: ctx, datasetID, artifactTag +func (_m *Client) DeleteByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string) error { + ret := _m.Called(ctx, datasetID, artifactTag) + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.DatasetID, string) error); ok { + r0 = rf(ctx, datasetID, artifactTag) + } else { + r0 = ret.Error(0) + } + + return r0 +} + type Client_Get struct { *mock.Call } @@ -101,6 +165,47 @@ func (_m *Client) GetOrExtendReservation(ctx context.Context, key catalog.Key, o return r0, r1 } +type Client_GetOrExtendReservationByArtifactTag struct { + *mock.Call +} + +func (_m Client_GetOrExtendReservationByArtifactTag) Return(_a0 *datacatalog.Reservation, _a1 error) *Client_GetOrExtendReservationByArtifactTag { + return &Client_GetOrExtendReservationByArtifactTag{Call: _m.Call.Return(_a0, _a1)} +} + +func (_m *Client) OnGetOrExtendReservationByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string, ownerID string, heartbeatInterval time.Duration) *Client_GetOrExtendReservationByArtifactTag { + c_call := _m.On("GetOrExtendReservationByArtifactTag", ctx, datasetID, artifactTag, ownerID, heartbeatInterval) + return &Client_GetOrExtendReservationByArtifactTag{Call: c_call} +} + +func (_m *Client) OnGetOrExtendReservationByArtifactTagMatch(matchers ...interface{}) *Client_GetOrExtendReservationByArtifactTag { + c_call := _m.On("GetOrExtendReservationByArtifactTag", matchers...) + return &Client_GetOrExtendReservationByArtifactTag{Call: c_call} +} + +// GetOrExtendReservationByArtifactTag provides a mock function with given fields: ctx, datasetID, artifactTag, ownerID, heartbeatInterval +func (_m *Client) GetOrExtendReservationByArtifactTag(ctx context.Context, datasetID *datacatalog.DatasetID, artifactTag string, ownerID string, heartbeatInterval time.Duration) (*datacatalog.Reservation, error) { + ret := _m.Called(ctx, datasetID, artifactTag, ownerID, heartbeatInterval) + + var r0 *datacatalog.Reservation + if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.DatasetID, string, string, time.Duration) *datacatalog.Reservation); ok { + r0 = rf(ctx, datasetID, artifactTag, ownerID, heartbeatInterval) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*datacatalog.Reservation) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.DatasetID, string, string, time.Duration) error); ok { + r1 = rf(ctx, datasetID, artifactTag, ownerID, heartbeatInterval) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + type Client_Put struct { *mock.Call }