diff --git a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
index e46a31528..8e5184917 100644
--- a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
+++ b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
@@ -82,10 +82,13 @@ NS_ASSUME_NONNULL_BEGIN
     return (preKeyRecord != nil);
 }
 
-- (void)removePreKey:(int)preKeyId
+- (void)removePreKey:(int)preKeyId protocolContext:(nullable id)protocolContext
 {
-    [self.dbReadWriteConnection removeObjectForKey:[self keyFromInt:preKeyId]
-                                      inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    if ([protocolContext isKindOfClass:YapDatabaseReadWriteTransaction.class]) {
+        [(YapDatabaseReadWriteTransaction *)protocolContext removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    } else {
+        [self.dbReadWriteConnection removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    }
 }
 
 - (int)nextPreKeyId:(int)batchSize
diff --git a/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m b/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
index d87c0f475..ad6888ee0 100644
--- a/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
+++ b/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
@@ -54,7 +54,7 @@
     PreKeyRecord *lastPreKeyRecord = [generatedKeys lastObject];
     PreKeyRecord *firstPreKeyRecord = [generatedKeys firstObject];
 
-    [[OWSPrimaryStorage sharedManager] removePreKey:lastPreKeyRecord.Id];
+    [[OWSPrimaryStorage sharedManager] removePreKey:lastPreKeyRecord.Id protocolContext:nil];
 
     XCTAssertThrows([[OWSPrimaryStorage sharedManager] throws_loadPreKey:lastPreKeyRecord.Id]);
     XCTAssertNoThrow([[OWSPrimaryStorage sharedManager] throws_loadPreKey:firstPreKeyRecord.Id]);