diff --git a/Session/Signal/MediaGalleryViewController.swift b/Session/Signal/MediaGalleryViewController.swift
index a8153b0c9..89925344d 100644
--- a/Session/Signal/MediaGalleryViewController.swift
+++ b/Session/Signal/MediaGalleryViewController.swift
@@ -801,14 +801,10 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
 
     func buildGalleryItem(attachment: TSAttachment, transaction: YapDatabaseReadTransaction) -> MediaGalleryItem? {
         guard let attachmentStream = attachment as? TSAttachmentStream else {
-            // Avoid crash in debug mode
-            // owsFailDebug("gallery doesn't yet support showing undownloaded attachments")
             return nil
         }
 
         guard let message = attachmentStream.fetchAlbumMessage(with: transaction) else {
-            // Avoid crash in debug mode
-            // owsFailDebug("message was unexpectedly nil")
             return nil
         }
 
diff --git a/SessionMessagingKit/Database/Storage+Messaging.swift b/SessionMessagingKit/Database/Storage+Messaging.swift
index e725a0c4a..dae3b1265 100644
--- a/SessionMessagingKit/Database/Storage+Messaging.swift
+++ b/SessionMessagingKit/Database/Storage+Messaging.swift
@@ -41,6 +41,10 @@ extension Storage {
             let thread = TSThread.fetch(uniqueId: threadID, transaction: transaction) else { return nil }
         let message = TSIncomingMessage.from(message, quotedMessage: quotedMessage, linkPreview: linkPreview, associatedWith: thread)
         message.save(with: transaction)
+        message.attachments(with: transaction).forEach { attachment in
+            attachment.albumMessageId = message.uniqueId!
+            attachment.save(with: transaction)
+        }
         DispatchQueue.main.async { message.touch() } // FIXME: Hack for a thread updating issue
         return message.uniqueId!
     }
diff --git a/SessionMessagingKit/Sending & Receiving/MessageSender.swift b/SessionMessagingKit/Sending & Receiving/MessageSender.swift
index f1952c505..6c028235b 100644
--- a/SessionMessagingKit/Sending & Receiving/MessageSender.swift	
+++ b/SessionMessagingKit/Sending & Receiving/MessageSender.swift	
@@ -43,7 +43,7 @@ public final class MessageSender : NSObject {
     public static let shared = MessageSender() // FIXME: Remove once requestSenderKey is static
 
     // MARK: Preparation
-    private static func prep(_ attachments: [SignalAttachment], for message: VisibleMessage, using transaction: YapDatabaseReadWriteTransaction) {
+    public static func prep(_ attachments: [SignalAttachment], for message: VisibleMessage, using transaction: YapDatabaseReadWriteTransaction) {
         guard let tsMessage = TSOutgoingMessage.find(withTimestamp: message.sentTimestamp!) else {
             #if DEBUG
             preconditionFailure()