Repository: hyperledger/firefly
Pull request #1284 (2023-04-18 19:49:22)

v1.2: Replace "UpsertBatch" with "InsertOrGetBatch"
When persisting an off-chain batch initially received from data exchange or shared storage, it should always be an insert. If the batch exists, just move on. Otherwise we may accidentally revert a confirmed batch to its previous state if it happens to get redelivered.

Backport of #1208, #1209, and #1204


Patch:
diff --git a/docs/swagger/swagger.yaml b/docs/swagger/swagger.yaml
index 44fbd5f34..d4e8b6b7e 100644
--- a/docs/swagger/swagger.yaml
+++ b/docs/swagger/swagger.yaml
@@ -22993,7 +22993,8 @@ paths:
   /namespaces/{ns}/pins/rewind:
     post:
       description: Force a rewind of the event aggregator to a previous position,
-        to re-evaluate all unconfirmed pins since that point
+        to re-evaluate (and possibly dispatch) that pin and others after it. Only
+        accepts a sequence or batch ID for a currently undispatched pin
       operationId: postPinsRewindNamespace
       parameters:
       - description: The namespace which scopes this request
@@ -30539,7 +30540,8 @@ paths:
   /pins/rewind:
     post:
       description: Force a rewind of the event aggregator to a previous position,
-        to re-evaluate all unconfirmed pins since that point
+        to re-evaluate (and possibly dispatch) that pin and others after it. Only
+        accepts a sequence or batch ID for a currently undispatched pin
       operationId: postPinsRewind
       parameters:
       - description: Server-side request timeout (milliseconds, or set a custom suffix
diff --git a/internal/batch/batch_processor.go b/internal/batch/batch_processor.go
index 2887b8745..c5ca52f32 100644
--- a/internal/batch/batch_processor.go
+++ b/internal/batch/batch_processor.go
@@ -1,4 +1,4 @@
-// Copyright © 2022 Kaleido, Inc.
+// Copyright © 2023 Kaleido, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -577,7 +577,8 @@ func (bp *batchProcessor) sealBatch(state *DispatchState) (err error) {
 			log.L(ctx).Debugf("Batch %s sealed. Hash=%s", state.Persisted.ID, state.Persisted.Hash)
 
 			// At this point the manifest of the batch is finalized. We write it to the database
-			return bp.database.UpsertBatch(ctx, &state.Persisted)
+			_, err = bp.database.InsertOrGetBatch(ctx, &state.Persisted)
+			return err
 		})
 	})
 	if err != nil {
diff --git a/internal/coremsgs/en_api_translations.go b/internal/coremsgs/en_api_translations.go
index b62c69259..1bd482b42 100644
--- a/internal/coremsgs/en_api_translations.go
+++ b/internal/coremsgs/en_api_translations.go
@@ -169,7 +169,7 @@ var (
 	APIEndpointsPostNewOrganization             = ffm("api.endpoints.postNewOrganization", "Registers a new org in the network")
 	APIEndpointsPostNewSubscription             = ffm("api.endpoints.postNewSubscription", "Creates a new subscription for an application to receive events from FireFly")
 	APIEndpointsPostOpRetry                     = ffm("api.endpoints.postOpRetry", "Retries a failed operation")
-	APIEndpointsPostPinsRewind                  = ffm("api.endpoints.postPinsRewind", "Force a rewind of the event aggregator to a previous position, to re-evaluate all unconfirmed pins since that point")
+	APIEndpointsPostPinsRewind                  = ffm("api.endpoints.postPinsRewind", "Force a rewind of the event aggregator to a previous position, to re-evaluate (and possibly dispatch) that pin and others after it. Only accepts a sequence or batch ID for a currently undispatched pin")
 	APIEndpointsPostTokenApproval               = ffm("api.endpoints.postTokenApproval", "Creates a token approval")
 	APIEndpointsPostTokenBurn                   = ffm("api.endpoints.postTokenBurn", "Burns some tokens")
 	APIEndpointsPostTokenMint                   = ffm("api.endpoints.postTokenMint", "Mints some tokens")
diff --git a/internal/database/sqlcommon/batch_sql.go b/internal/database/sqlcommon/batch_sql.go
index 6e9684f4f..007209e40 100644
--- a/internal/database/sqlcommon/batch_sql.go
+++ b/internal/database/sqlcommon/batch_sql.go
@@ -1,4 +1,4 @@
-// Copyright © 2022 Kaleido, Inc.
+// Copyright © 2023 Kaleido, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -57,87 +57,49 @@ var (
 
 const batchesTable = "batches"
 
-func (s *SQLCommon) UpsertBatch(ctx context.Context, batch *core.BatchPersisted) (err error) {
+func (s *SQLCommon) InsertOrGetBatch(ctx context.Context, batch *core.BatchPersisted) (existing *core.BatchPersisted, err error) {
 	ctx, tx, autoCommit, err := s.BeginOrUseTx(ctx)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	defer s.RollbackTx(ctx, tx, autoCommit)
 
-	// Do a select within the transaction to detemine if the UUID already exists
-	batchRows, _, err := s.QueryTx(ctx, batchesTable, tx,
-		sq.Select("hash").
-			From(batchesTable).
-			Where(sq.Eq{"id": batch.ID, "namespace": batch.Namespace}),
+	// Try the insert first
+	_, insertErr := s.InsertTxExt(ctx, batchesTable, tx,
+		sq.Insert(batchesTable).
+			Columns(batchColumns...).
+			Values(
+				batch.ID,
+				string(batch.Type),
+				batch.Namespace,
+				batch.Author,
+				batch.Key,
+				batch.Group,
+				batch.Created,
+				batch.Hash,
+				batch.Manifest,
+				batch.Confirmed,
+				batch.TX.Type,
+				batch.TX.ID,
+				batch.Node,
+			),
+		func() {
+			s.callbacks.UUIDCollectionNSEvent(database.CollectionBatches, core.ChangeEventTypeCreated, batch.Namespace, batch.ID)
+		},
+		true, /* we want a failure here we can progress past */
 	)
-	if err != nil {
-		return err
+	if insertErr == nil {
+		return nil, s.CommitTx(ctx, tx, autoCommit)
 	}
 
-	existing := batchRows.Next()
-	if existing {
-		var hash *fftypes.Bytes32
-		_ = batchRows.Scan(&hash)
-		if !fftypes.SafeHashCompare(hash, batch.Hash) {
-			batchRows.Close()
-			log.L(ctx).Errorf("Existing=%s New=%s", hash, batch.Hash)
-			return database.HashMismatch
-		}
-	}
-	batchRows.Close()
-
-	if existing {
-
-		// Update the batch
-		if _, err = s.UpdateTx(ctx, batchesTable, tx,
-			sq.Update(batchesTable).
-				Set("btype", string(batch.Type)).
-				Set("author", batch.Author).
-				Set("key", batch.Key).
-				Set("group_hash", batch.Group).
-				Set("created", batch.Created).
-				Set("hash", batch.Hash).
-				Set("manifest", batch.Manifest).
-				Set("confirmed", batch.Confirmed).
-				Set("tx_type", batch.TX.Type).
-				Set("tx_id", batch.TX.ID).
-				Set("node_id", batch.Node).
-				Where(sq.Eq{"id": batch.ID, "namespace": batch.Namespace}),
-			func() {
-				s.callbacks.UUIDCollectionNSEvent(database.CollectionBatches, core.ChangeEventTypeUpdated, batch.Namespace, batch.ID)
-			},
-		); err != nil {
-			return err
-		}
-	} else {
-
-		if _, err = s.InsertTx(ctx, batchesTable, tx,
-			sq.Insert(batchesTable).
-				Columns(batchColumns...).
-				Values(
-					batch.ID,
-					string(batch.Type),
-					batch.Namespace,
-					batch.Author,
-					batch.Key,
-					batch.Group,
-					batch.Created,
-					batch.Hash,
-					batch.Manifest,
-					batch.Confirmed,
-					batch.TX.Type,
-					batch.TX.ID,
-					batch.Node,
-				),
-			func() {
-				s.callbacks.UUIDCollectionNSEvent(database.CollectionBatches, core.ChangeEventTypeCreated, batch.Namespace, batch.ID)
-			},
-		); err != nil {
-			return err
-		}
+	// Do a select within the transaction to determine if the batch already exists
+	existing, err = s.GetBatchByID(ctx, batch.Namespace, batch.ID)
+	if err != nil || existing != nil {
+		return existing, err
 	}
 
-	return s.CommitTx(ctx, tx, autoCommit)
+	// Error was apparently not an ID conflict - must have been something else
+	return nil, insertErr
 }
 
 func (s *SQLCommon) batchResult(ctx context.Context, row *sql.Rows) (*core.BatchPersisted, error) {
diff --git a/internal/events/event_poller.go b/internal/events/event_poller.go
index 6ffcc4a9b..9d876b6c8 100644
--- a/internal/events/event_poller.go
+++ b/internal/events/event_poller.go
@@ -1,4 +1,4 @@
-// Copyright © 2022 Kaleido, Inc.
+// Copyright © 2023 Kaleido, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -283,9 +283,15 @@ func (ep *eventPoller) shoulderTap() {
 func (ep *eventPoller) waitForShoulderTapOrPollTimeout(lastEventCount int) bool {
 	l := log.L(ep.ctx)
 	longTimeoutDuration := ep.conf.eventPollTimeout
+
+	if lastEventCount >= ep.conf.eventBatchSize {
+		l.Tracef("Polling immediately due to full previous event batch")
+		return true
+	}
+
 	// For throughput optimized environments, we can set an eventBatchingTimeout to allow messages to arrive
 	// between polling cycles (at the cost of some dispatch latency)
-	if ep.conf.eventBatchTimeout > 0 && lastEventCount > 0 && lastEventCount < ep.conf.eventBatchSize {
+	if ep.conf.eventBatchTimeout > 0 && lastEventCount > 0 {
 		shortTimeout := time.NewTimer(ep.conf.eventBatchTimeout)
 		select {
 		case <-shortTimeout.C:
diff --git a/internal/events/persist_batch.go b/internal/events/persist_batch.go
index 4d0cff6fc..99229c2c5 100644
--- a/internal/events/persist_batch.go
+++ b/internal/events/persist_batch.go
@@ -69,13 +69,9 @@ func (em *eventManager) persistBatch(ctx context.Context, batch *core.Batch) (pe
 		}
 	}
 
-	// Upsert the batch
-	err = em.database.UpsertBatch(ctx, persistedBatch)
+	// Insert the batch
+	existing, err := em.database.InsertOrGetBatch(ctx, persistedBatch)
 	if err != nil {
-		if err == database.HashMismatch {
-			l.Errorf("Invalid batch '%s'. Batch hash mismatch with existing record", batch.ID)
-			return nil, false, nil // This is not retryable. skip this batch
-		}
 		l.Errorf("Failed to insert batch '%s': %s", batch.ID, err)
 		return nil, false, err // a persistence failure here is considered retryable (so returned)
 	}
@@ -84,6 +80,11 @@ func (em *eventManager) persistBatch(ctx context.Context, batch *core.Batch) (pe
 	if err != nil || !valid {
 		return nil, valid, err
 	}
+
+	if existing != nil {
+		l.Infof("Skipped insert of batch '%s' (already exists)", batch.ID)
+		return existing, true, nil
+	}
 	em.aggregator.cacheBatch(em.aggregator.getBatchCacheKey(persistedBatch.ID, persistedBatch.Hash), persistedBatch, manifest)
 	return persistedBatch, true, err
 }
diff --git a/internal/orchestrator/orchestrator.go b/internal/orchestrator/orchestrator.go
index c744e5afc..60e3ab26d 100644
--- a/internal/orchestrator/orchestrator.go
+++ b/internal/orchestrator/orchestrator.go
@@ -552,7 +552,7 @@ func (or *orchestrator) Authorize(ctx context.Context, authReq *fftypes.AuthReq)
 func (or *orchestrator) RewindPins(ctx context.Context, rewind *core.PinRewind) (*core.PinRewind, error) {
 	if rewind.Sequence > 0 {
 		fb := database.PinQueryFactory.NewFilter(ctx)
-		if pins, _, err := or.GetPins(ctx, fb.And(fb.Eq("seq", rewind.Sequence))); err != nil {
+		if pins, _, err := or.GetPins(ctx, fb.And(fb.Eq("sequence", rewind.Sequence))); err != nil {
 			return nil, err
 		} else if len(pins) > 0 {
 			rewind.Batch = pins[0].Batch
diff --git a/mocks/databasemocks/plugin.go b/mocks/databasemocks/plugin.go
index 3b594ea19..94d52f8f3 100644
--- a/mocks/databasemocks/plugin.go
+++ b/mocks/databasemocks/plugin.go
@@ -2231,6 +2231,32 @@ func (_m *Plugin) InsertOperation(ctx context.Context, operation *core.Operation
 	return r0
 }
 
+// InsertOrGetBatch provides a mock function with given fields: ctx, data
+func (_m *Plugin) InsertOrGetBatch(ctx context.Context, data *core.BatchPersisted) (*core.BatchPersisted, error) {
+	ret := _m.Called(ctx, data)
+
+	var r0 *core.BatchPersisted
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *core.BatchPersisted) (*core.BatchPersisted, error)); ok {
+		return rf(ctx, data)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *core.BatchPersisted) *core.BatchPersisted); ok {
+		r0 = rf(ctx, data)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*core.BatchPersisted)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *core.BatchPersisted) error); ok {
+		r1 = rf(ctx, data)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // InsertOrGetBlockchainEvent provides a mock function with given fields: ctx, event
 func (_m *Plugin) InsertOrGetBlockchainEvent(ctx context.Context, event *core.BlockchainEvent) (*core.BlockchainEvent, error) {
 	ret := _m.Called(ctx, event)
@@ -2511,20 +2537,6 @@ func (_m *Plugin) UpdateTransaction(ctx context.Context, namespace string, id *f
 	return r0
 }
 
-// UpsertBatch provides a mock function with given fields: ctx, data
-func (_m *Plugin) UpsertBatch(ctx context.Context, data *core.BatchPersisted) error {
-	ret := _m.Called(ctx, data)
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, *core.BatchPersisted) error); ok {
-		r0 = rf(ctx, data)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
 // UpsertContractAPI provides a mock function with given fields: ctx, cd
 func (_m *Plugin) UpsertContractAPI(ctx context.Context, cd *core.ContractAPI) error {
 	ret := _m.Called(ctx, cd)
diff --git a/pkg/database/plugin.go b/pkg/database/plugin.go
index 856b243e0..110ccef62 100644
--- a/pkg/database/plugin.go
+++ b/pkg/database/plugin.go
@@ -139,10 +139,10 @@ type iDataCollection interface {
 }
 
 type iBatchCollection interface {
-	// UpsertBatch - Upsert a batch - the hash cannot change
-	UpsertBatch(ctx context.Context, data *core.BatchPersisted) (err error)
+	// InsertBatch - Insert a new batch, or retrieve the existing one if it has already been recorded
+	InsertOrGetBatch(ctx context.Context, data *core.BatchPersisted) (existing *core.BatchPersisted, err error)
 
-	// UpdateBatch - Update data
+	// UpdateBatch - Update a batch
 	UpdateBatch(ctx context.Context, namespace string, id *fftypes.UUID, update ffapi.Update) (err error)
 
 	// GetBatchByID - Get a batch by ID


Test patch:
diff --git a/internal/batch/batch_manager_test.go b/internal/batch/batch_manager_test.go
index f563cd76a..407c00466 100644
--- a/internal/batch/batch_manager_test.go
+++ b/internal/batch/batch_manager_test.go
@@ -6,7 +6,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
@@ -132,7 +132,7 @@ func TestE2EDispatchBroadcast(t *testing.T) {
 	mdm.On("UpdateMessageIfCached", mock.Anything, mock.Anything).Return()
 	mdi.On("GetMessageIDs", mock.Anything, "ns1", mock.Anything).Return([]*core.IDAndSequence{{ID: *msg.Header.ID}}, nil).Once()
 	mdi.On("GetMessageIDs", mock.Anything, "ns1", mock.Anything).Return([]*core.IDAndSequence{}, nil)
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 	mdi.On("UpdateBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	mdi.On("UpdateMessage", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil) // pins
 	rag := mdi.On("RunAsGroup", mock.Anything, mock.Anything, mock.Anything).Return(nil)
@@ -255,7 +255,7 @@ func TestE2EDispatchPrivateUnpinned(t *testing.T) {
 	mdi.On("GetMessageIDs", mock.Anything, "ns1", mock.Anything).Return([]*core.IDAndSequence{{ID: *msg.Header.ID}}, nil).Once()
 	mdi.On("GetMessageIDs", mock.Anything, "ns1", mock.Anything).Return([]*core.IDAndSequence{}, nil)
 	mdi.On("UpdateMessage", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil) // pins
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 	mdi.On("UpdateBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	rag := mdi.On("RunAsGroup", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	rag.RunFn = func(a mock.Arguments) {
@@ -441,7 +441,7 @@ func TestMessageSequencerUpdateMessagesFail(t *testing.T) {
 	mdm.On("UpdateMessageIfCached", mock.Anything, mock.Anything).Return()
 	mdi.On("InsertTransaction", mock.Anything, mock.Anything).Return(nil)
 	mdi.On("InsertEvent", mock.Anything, mock.Anything).Return(nil) // transaction submit
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 	mdi.On("UpdateMessages", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(fmt.Errorf("fizzle"))
 	rag := mdi.On("RunAsGroup", mock.Anything, mock.Anything, mock.Anything)
 	rag.RunFn = func(a mock.Arguments) {
@@ -538,7 +538,7 @@ func TestMessageSequencerUpdateBatchFail(t *testing.T) {
 	}
 	mdi.On("GetMessageIDs", mock.Anything, "ns1", mock.Anything).Return([]*core.IDAndSequence{{ID: *msg.Header.ID}}, nil)
 	mdm.On("GetMessageWithDataCached", mock.Anything, mock.Anything).Return(msg, core.DataArray{{ID: dataID}}, true, nil)
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(fmt.Errorf("fizzle"))
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, fmt.Errorf("fizzle"))
 	rag := mdi.On("RunAsGroup", mock.Anything, mock.Anything, mock.Anything)
 	rag.RunFn = func(a mock.Arguments) {
 		ctx := a.Get(0).(context.Context)
diff --git a/internal/batch/batch_processor_test.go b/internal/batch/batch_processor_test.go
index 6e9782325..8d692bb38 100644
--- a/internal/batch/batch_processor_test.go
+++ b/internal/batch/batch_processor_test.go
@@ -6,7 +6,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
@@ -85,7 +85,7 @@ func TestUnfilledBatch(t *testing.T) {
 
 	mockRunAsGroupPassthrough(mdi)
 	mdi.On("UpdateMessages", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil)
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 
 	mth := bp.txHelper.(*txcommonmocks.Helper)
 	mth.On("SubmitNewTransaction", mock.Anything, core.TransactionTypeBatchPin, core.IdempotencyKey("")).Return(fftypes.NewUUID(), nil)
@@ -134,7 +134,7 @@ func TestBatchSizeOverflow(t *testing.T) {
 	bp.conf.BatchMaxBytes = batchSizeEstimateBase + (&core.Message{}).EstimateSize(false) + 100
 	mockRunAsGroupPassthrough(mdi)
 	mdi.On("UpdateMessages", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil)
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 
 	mth := bp.txHelper.(*txcommonmocks.Helper)
 	mth.On("SubmitNewTransaction", mock.Anything, core.TransactionTypeBatchPin, core.IdempotencyKey("")).Return(fftypes.NewUUID(), nil)
@@ -419,7 +419,7 @@ func TestMarkMessageDispatchedUnpinnedOK(t *testing.T) {
 
 	mockRunAsGroupPassthrough(mdi)
 	mdi.On("UpdateMessages", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil)
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 	mdi.On("InsertEvent", mock.Anything, mock.Anything).Return(fmt.Errorf("pop")).Once()
 	mdi.On("InsertEvent", mock.Anything, mock.Anything).Return(nil)
 
@@ -476,7 +476,7 @@ func TestMaskContextsRetryAfterPinsAssigned(t *testing.T) {
 		return dbNonce.Nonce == 12347 // twice incremented
 	})).Return(nil).Once()
 	mdi.On("UpdateMessage", mock.Anything, "ns1", mock.Anything, mock.Anything).Return(nil).Twice()
-	mdi.On("UpsertBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
 
 	mdm := bp.data.(*datamocks.Manager)
 	mdm.On("UpdateMessageIfCached", mock.Anything, mock.Anything).Return()
diff --git a/internal/database/sqlcommon/batch_sql_test.go b/internal/database/sqlcommon/batch_sql_test.go
index 15f462d20..76d4a946d 100644
--- a/internal/database/sqlcommon/batch_sql_test.go
+++ b/internal/database/sqlcommon/batch_sql_test.go
@@ -63,10 +63,10 @@ func TestBatch2EWithDB(t *testing.T) {
 	}
 
 	s.callbacks.On("UUIDCollectionNSEvent", database.CollectionBatches, core.ChangeEventTypeCreated, "ns1", batchID, mock.Anything).Return()
-	s.callbacks.On("UUIDCollectionNSEvent", database.CollectionBatches, core.ChangeEventTypeUpdated, "ns1", batchID, mock.Anything).Return()
 
-	err := s.UpsertBatch(ctx, batch)
+	existing, err := s.InsertOrGetBatch(ctx, batch)
 	assert.NoError(t, err)
+	assert.Nil(t, existing)
 
 	// Check we get the exact same batch back
 	batchRead, err := s.GetBatchByID(ctx, "ns1", batchID)
@@ -76,58 +76,17 @@ func TestBatch2EWithDB(t *testing.T) {
 	batchReadJson, _ := json.Marshal(&batchRead)
 	assert.Equal(t, string(batchJson), string(batchReadJson))
 
-	// Update the batch (this is testing what's possible at the database layer,
-	// and does not account for the verification that happens at the higher level)
-	txid := fftypes.NewUUID()
-	msgID2 := fftypes.NewUUID()
-	batchUpdated := &core.BatchPersisted{
-		BatchHeader: core.BatchHeader{
-			ID:   batchID,
-			Type: core.BatchTypePrivate,
-			SignerRef: core.SignerRef{
-				Key:    "0x12345",
-				Author: "did:firefly:org/abcd",
-			},
-			Namespace: "ns1",
-			Node:      fftypes.NewUUID(),
-			Created:   fftypes.Now(),
-		},
-		Hash: fftypes.NewRandB32(),
-		TX: core.TransactionRef{
-			ID:   txid,
-			Type: core.TransactionTypeBatchPin,
-		},
-		Manifest: fftypes.JSONAnyPtr((&core.BatchManifest{
-			Messages: []*core.MessageManifestEntry{
-				{MessageRef: core.MessageRef{ID: msgID1}},
-				{MessageRef: core.MessageRef{ID: msgID2}},
-			},
-		}).String()),
-		Confirmed: fftypes.Now(),
-	}
-
-	// Rejects hash change
-	err = s.UpsertBatch(context.Background(), batchUpdated)
-	assert.Equal(t, database.HashMismatch, err)
-
-	batchUpdated.Hash = batch.Hash
-	err = s.UpsertBatch(context.Background(), batchUpdated)
+	// Try to insert again - should get back the existing row
+	existing, err = s.InsertOrGetBatch(ctx, batch)
 	assert.NoError(t, err)
-
-	// Check we get the exact same message back - note the removal of one of the batch elements
-	batchRead, err = s.GetBatchByID(ctx, "ns1", batchID)
-	assert.NoError(t, err)
-	batchJson, _ = json.Marshal(&batchUpdated)
-	batchReadJson, _ = json.Marshal(&batchRead)
-	assert.Equal(t, string(batchJson), string(batchReadJson))
+	assert.NotNil(t, existing)
 
 	// Query back the batch
 	fb := database.BatchQueryFactory.NewFilter(ctx)
 	filter := fb.And(
-		fb.Eq("id", batchUpdated.ID.String()),
-		fb.Eq("author", batchUpdated.Author),
+		fb.Eq("id", batch.ID.String()),
+		fb.Eq("author", batch.Author),
 		fb.Gt("created", "0"),
-		fb.Gt("confirmed", "0"),
 	)
 	batches, _, err := s.GetBatches(ctx, "ns1", filter)
 	assert.NoError(t, err)
@@ -137,7 +96,7 @@ func TestBatch2EWithDB(t *testing.T) {
 
 	// Negative test on filter
 	filter = fb.And(
-		fb.Eq("id", batchUpdated.ID.String()),
+		fb.Eq("id", batch.ID.String()),
 		fb.Eq("created", "0"),
 	)
 	batches, _, err = s.GetBatches(ctx, "ns1", filter)
@@ -152,7 +111,7 @@ func TestBatch2EWithDB(t *testing.T) {
 
 	// Test find updated value
 	filter = fb.And(
-		fb.Eq("id", batchUpdated.ID.String()),
+		fb.Eq("id", batch.ID.String()),
 		fb.Eq("author", author2),
 	)
 	batches, res, err := s.GetBatches(ctx, "ns1", filter.Count(true))
@@ -166,55 +125,30 @@ func TestBatch2EWithDB(t *testing.T) {
 func TestUpsertBatchFailBegin(t *testing.T) {
 	s, mock := newMockProvider().init()
 	mock.ExpectBegin().WillReturnError(fmt.Errorf("pop"))
-	err := s.UpsertBatch(context.Background(), &core.BatchPersisted{})
+	_, err := s.InsertOrGetBatch(context.Background(), &core.BatchPersisted{})
 	assert.Regexp(t, "FF00175", err)
 	assert.NoError(t, mock.ExpectationsWereMet())
 }
 
-func TestUpsertBatchFailSelect(t *testing.T) {
-	s, mock := newMockProvider().init()
-	mock.ExpectBegin()
-	mock.ExpectQuery("SELECT .*").WillReturnError(fmt.Errorf("pop"))
-	mock.ExpectRollback()
-	batchID := fftypes.NewUUID()
-	err := s.UpsertBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}})
-	assert.Regexp(t, "FF00176", err)
-	assert.NoError(t, mock.ExpectationsWereMet())
-}
-
 func TestUpsertBatchFailInsert(t *testing.T) {
 	s, mock := newMockProvider().init()
 	mock.ExpectBegin()
-	mock.ExpectQuery("SELECT .*").WillReturnRows(sqlmock.NewRows([]string{}))
 	mock.ExpectExec("INSERT .*").WillReturnError(fmt.Errorf("pop"))
+	mock.ExpectQuery("SELECT .*").WillReturnRows(sqlmock.NewRows([]string{}))
 	mock.ExpectRollback()
 	batchID := fftypes.NewUUID()
-	err := s.UpsertBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}})
+	_, err := s.InsertOrGetBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}})
 	assert.Regexp(t, "FF00177", err)
 	assert.NoError(t, mock.ExpectationsWereMet())
 }
 
-func TestUpsertBatchFailUpdate(t *testing.T) {
-	s, mock := newMockProvider().init()
-	batchID := fftypes.NewUUID()
-	hash := fftypes.NewRandB32()
-	mock.ExpectBegin()
-	mock.ExpectQuery("SELECT .*").WillReturnRows(sqlmock.NewRows([]string{"hash"}).AddRow(hash))
-	mock.ExpectExec("UPDATE .*").WillReturnError(fmt.Errorf("pop"))
-	mock.ExpectRollback()
-	err := s.UpsertBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}, Hash: hash})
-	assert.Regexp(t, "FF00178", err)
-	assert.NoError(t, mock.ExpectationsWereMet())
-}
-
 func TestUpsertBatchFailCommit(t *testing.T) {
 	s, mock := newMockProvider().init()
 	batchID := fftypes.NewUUID()
 	mock.ExpectBegin()
-	mock.ExpectQuery("SELECT .*").WillReturnRows(sqlmock.NewRows([]string{"id"}))
 	mock.ExpectExec("INSERT .*").WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectCommit().WillReturnError(fmt.Errorf("pop"))
-	err := s.UpsertBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}})
+	_, err := s.InsertOrGetBatch(context.Background(), &core.BatchPersisted{BatchHeader: core.BatchHeader{ID: batchID}})
 	assert.Regexp(t, "FF00180", err)
 	assert.NoError(t, mock.ExpectationsWereMet())
 }
diff --git a/internal/events/batch_pin_complete_test.go b/internal/events/batch_pin_complete_test.go
index 755792ab6..5fb4a6b22 100644
--- a/internal/events/batch_pin_complete_test.go
+++ b/internal/events/batch_pin_complete_test.go
@@ -469,20 +469,6 @@ func TestPersistBatchMismatchChainHash(t *testing.T) {
 	assert.False(t, valid)
 }
 
-func TestPersistBatchUpsertBatchMismatchHash(t *testing.T) {
-	em := newTestEventManager(t)
-	defer em.cleanup(t)
-	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
-	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
-
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(database.HashMismatch)
-
-	bp, valid, err := em.persistBatch(context.Background(), batch)
-	assert.False(t, valid)
-	assert.Nil(t, bp)
-	assert.NoError(t, err)
-}
-
 func TestPersistBatchBadHash(t *testing.T) {
 	em := newTestEventManager(t)
 	defer em.cleanup(t)
@@ -528,7 +514,7 @@ func TestPersistBatchUpsertBatchFail(t *testing.T) {
 	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
 	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
 
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(fmt.Errorf("pop"))
+	em.mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything).Return(nil, fmt.Errorf("pop"))
 
 	bp, valid, err := em.persistBatch(context.Background(), batch)
 	assert.Nil(t, bp)
@@ -559,7 +545,7 @@ func TestPersistBatchSwallowBadData(t *testing.T) {
 	}
 	batch.Hash = batch.Payload.Hash()
 
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything).Return(nil, nil)
 
 	bp, valid, err := em.persistBatch(context.Background(), batch)
 	assert.False(t, valid)
@@ -573,7 +559,7 @@ func TestPersistBatchGoodDataUpsertOptimizeFail(t *testing.T) {
 	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
 	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
 
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", mock.Anything, mock.Anything).Return(fmt.Errorf("optimzation miss"))
 	em.mdi.On("UpsertData", mock.Anything, mock.Anything, database.UpsertOptimizationExisting).Return(fmt.Errorf("pop"))
 
@@ -591,7 +577,7 @@ func TestPersistBatchGoodDataMessageFail(t *testing.T) {
 	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
 	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
 
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", mock.Anything, mock.Anything).Return(nil)
 	em.mdi.On("InsertMessages", mock.Anything, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(fmt.Errorf("optimzation miss"))
 	em.mdi.On("UpsertMessage", mock.Anything, mock.Anything, database.UpsertOptimizationExisting, mock.AnythingOfType("database.PostCompletionHook")).Return(fmt.Errorf("pop"))
@@ -614,7 +600,7 @@ func TestPersistBatchGoodMessageAuthorMismatch(t *testing.T) {
 	batch.Payload.Messages[0].Hash = batch.Payload.Messages[0].Header.Hash()
 	batch.Hash = batch.Payload.Hash()
 
-	em.mdi.On("UpsertBatch", mock.Anything, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", mock.Anything, mock.Anything).Return(nil, nil)
 
 	bp, valid, err := em.persistBatch(context.Background(), batch)
 	assert.Nil(t, bp)
diff --git a/internal/events/dx_callbacks_test.go b/internal/events/dx_callbacks_test.go
index c34529a2f..eacbb763f 100644
--- a/internal/events/dx_callbacks_test.go
+++ b/internal/events/dx_callbacks_test.go
@@ -153,7 +153,7 @@ func TestPinnedReceiveOK(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(nil, nil).Run(func(args mock.Arguments) {
 		args[2].(database.PostCompletionHook)()
@@ -224,7 +224,7 @@ func TestMessageReceivePersistBatchError(t *testing.T) {
 	}).Return(node1, nil)
 	em.mim.On("CachedIdentityLookupMustExist", em.ctx, "signingOrg").Return(org1, false, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(fmt.Errorf("pop"))
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, fmt.Errorf("pop"))
 
 	// no ack as we are simulating termination mid retry
 	mde := newMessageReceivedNoAck("peer1", b)
@@ -607,7 +607,7 @@ func TestMessageReceiveMessagePersistMessageFail(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(fmt.Errorf("optimization fail"))
 	em.mdi.On("UpsertMessage", em.ctx, mock.Anything, database.UpsertOptimizationExisting, mock.AnythingOfType("database.PostCompletionHook")).Return(fmt.Errorf("pop"))
@@ -648,7 +648,7 @@ func TestMessageReceiveMessagePersistDataFail(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(fmt.Errorf("optimization miss"))
 	em.mdi.On("UpsertData", em.ctx, mock.Anything, database.UpsertOptimizationExisting).Return(fmt.Errorf("pop"))
 
@@ -687,7 +687,7 @@ func TestMessageReceiveUnpinnedBatchOk(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(nil, nil).Run(func(args mock.Arguments) {
 		args[2].(database.PostCompletionHook)()
@@ -730,7 +730,7 @@ func TestMessageReceiveUnpinnedBatchConfirmMessagesFail(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(nil, nil).Run(func(args mock.Arguments) {
 		args[2].(database.PostCompletionHook)()
@@ -773,7 +773,7 @@ func TestMessageReceiveUnpinnedBatchPersistEventFail(t *testing.T) {
 	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
 	em.mim.On("ValidateNodeOwner", em.ctx, mock.Anything, mock.Anything).Return(true, nil)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(nil, nil).Run(func(args mock.Arguments) {
 		args[2].(database.PostCompletionHook)()
diff --git a/internal/events/event_poller_test.go b/internal/events/event_poller_test.go
index 68f39047f..9d9366525 100644
--- a/internal/events/event_poller_test.go
+++ b/internal/events/event_poller_test.go
@@ -280,12 +280,21 @@ func TestWaitForShoulderTapOrExitCloseBatch(t *testing.T) {
 	assert.False(t, ep.waitForShoulderTapOrPollTimeout(1))
 }
 
-func TestWaitForShoulderTapOrExitClosePoll(t *testing.T) {
+func TestWaitForShoulderTapImmediateRepollOnFullBatch(t *testing.T) {
 	mdi := &databasemocks.Plugin{}
 	ep, cancel := newTestEventPoller(t, mdi, nil, nil)
 	cancel()
 	ep.conf.eventBatchTimeout = 1 * time.Minute
 	ep.conf.eventBatchSize = 1
+	assert.True(t, ep.waitForShoulderTapOrPollTimeout(1))
+}
+
+func TestWaitForShoulderTapOrExitClosePoll(t *testing.T) {
+	mdi := &databasemocks.Plugin{}
+	ep, cancel := newTestEventPoller(t, mdi, nil, nil)
+	cancel()
+	ep.conf.eventBatchTimeout = 1 * time.Minute
+	ep.conf.eventBatchSize = 2
 	assert.False(t, ep.waitForShoulderTapOrPollTimeout(1))
 }
 
diff --git a/internal/events/persist_batch_test.go b/internal/events/persist_batch_test.go
index 49c142da6..8ce8d2c49 100644
--- a/internal/events/persist_batch_test.go
+++ b/internal/events/persist_batch_test.go
@@ -33,7 +33,7 @@ func TestPersistBatch(t *testing.T) {
 	em := newTestEventManager(t)
 	defer em.cleanup(t)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(fmt.Errorf(("pop")))
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, fmt.Errorf(("pop")))
 
 	org := newTestOrg("org1")
 	orgBytes, err := json.Marshal(&org)
@@ -88,12 +88,33 @@ func TestPersistBatch(t *testing.T) {
 
 }
 
+func TestPersistBatchAlreadyExisting(t *testing.T) {
+
+	em := newTestEventManager(t)
+	defer em.cleanup(t)
+
+	existing := &core.BatchPersisted{}
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(existing, nil)
+	em.mim.On("GetLocalNode", mock.Anything).Return(testNode, nil)
+	em.mdi.On("InsertDataArray", mock.Anything, mock.Anything).Return(nil)
+	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.Anything).Return(nil, nil)
+
+	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
+	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
+
+	result, valid, err := em.persistBatch(em.ctx, batch)
+	assert.True(t, valid)
+	assert.NoError(t, err)
+	assert.Equal(t, existing, result)
+
+}
+
 func TestPersistBatchNoCacheDataNotInBatch(t *testing.T) {
 
 	em := newTestEventManager(t)
 	defer em.cleanup(t)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 
 	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
 	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
@@ -113,7 +134,7 @@ func TestPersistBatchExtraDataInBatch(t *testing.T) {
 	em := newTestEventManager(t)
 	defer em.cleanup(t)
 
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 
 	data := &core.Data{ID: fftypes.NewUUID(), Value: fftypes.JSONAnyPtr(`"test"`)}
 	batch := sampleBatch(t, core.BatchTypeBroadcast, core.TransactionTypeBatchPin, core.DataArray{data})
diff --git a/internal/events/ss_callbacks_test.go b/internal/events/ss_callbacks_test.go
index 0a56dba9c..de5ab0186 100644
--- a/internal/events/ss_callbacks_test.go
+++ b/internal/events/ss_callbacks_test.go
@@ -39,7 +39,7 @@ func TestSharedStorageBatchDownloadedOk(t *testing.T) {
 	b, _ := json.Marshal(&batch)
 
 	mss := &sharedstoragemocks.Plugin{}
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(nil, nil)
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertDataArray", em.ctx, mock.Anything).Return(nil, nil)
 	em.mdi.On("InsertMessages", em.ctx, mock.Anything, mock.AnythingOfType("database.PostCompletionHook")).Return(nil, nil).Run(func(args mock.Arguments) {
 		args[2].(database.PostCompletionHook)()
@@ -71,7 +71,7 @@ func TestSharedStorageBatchDownloadedPersistFail(t *testing.T) {
 	b, _ := json.Marshal(&batch)
 
 	mss := &sharedstoragemocks.Plugin{}
-	em.mdi.On("UpsertBatch", em.ctx, mock.Anything).Return(fmt.Errorf("pop"))
+	em.mdi.On("InsertOrGetBatch", em.ctx, mock.Anything).Return(nil, fmt.Errorf("pop"))
 	mss.On("Name").Return("utdx").Maybe()
 
 	_, err := em.SharedStorageBatchDownloaded(mss, "payload1", b)