From 950d6343dbfe6c36d2a7e4ef11625ce1617076ea Mon Sep 17 00:00:00 2001
From: Harry <109690906+harry-signal@users.noreply.github.com>
Date: Tue, 1 Oct 2024 12:56:58 -0700
Subject: [PATCH] Add method to fetch attachments by media name (has UNIQUE
 constraint and therefore has index)

---
 .../V2/AttachmentStore/AttachmentStore.swift  |  6 +++++
 .../AttachmentStore/AttachmentStoreImpl.swift | 22 +++++++++++++++++++
 .../AttachmentStore/AttachmentStoreMock.swift |  7 ++++++
 3 files changed, 35 insertions(+)

diff --git a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStore.swift b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStore.swift
index 4ef799a537..f43249fbe7 100644
--- a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStore.swift
+++ b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStore.swift
@@ -33,6 +33,12 @@ public protocol AttachmentStore {
         tx: DBReadTransaction
     ) -> Attachment?
 
+    /// Fetch attachment by media name. There can be only one match.
+    func fetchAttachment(
+        mediaName: String,
+        tx: DBReadTransaction
+    ) -> Attachment?
+
     /// Enumerate all references to a given attachment id, calling the block for each one.
     /// Blocks until all references have been enumerated.
     func enumerateAllReferences(
diff --git a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreImpl.swift b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreImpl.swift
index 7be9ce5043..806309e5e0 100644
--- a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreImpl.swift
+++ b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreImpl.swift
@@ -33,6 +33,17 @@ public class AttachmentStoreImpl: AttachmentStore {
         )
     }
 
+    public func fetchAttachment(
+        mediaName: String,
+        tx: DBReadTransaction
+    ) -> Attachment? {
+        try? fetchAttachment(
+            mediaName: mediaName,
+            db: SDSDB.shimOnlyBridge(tx).unwrapGrdbRead.database,
+            tx: tx
+        )
+    }
+
     public func enumerateAllReferences(
         toAttachmentId attachmentId: Attachment.IDType,
         tx: DBReadTransaction,
@@ -306,6 +317,17 @@ public class AttachmentStoreImpl: AttachmentStore {
             .map(Attachment.init(record:))
     }
 
+    func fetchAttachment(
+        mediaName: String,
+        db: GRDB.Database,
+        tx: DBReadTransaction
+    ) throws -> Attachment? {
+        return try Attachment.Record
+            .filter(Column(Attachment.Record.CodingKeys.mediaName) == mediaName)
+            .fetchOne(db)
+            .map(Attachment.init(record:))
+    }
+
     func allQuotedReplyAttachments(
         forOriginalAttachmentId originalAttachmentId: Attachment.IDType,
         db: GRDB.Database,
diff --git a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreMock.swift b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreMock.swift
index 414a0eed5b..2e2c3b6ffb 100644
--- a/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreMock.swift
+++ b/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreMock.swift
@@ -31,6 +31,13 @@ open class AttachmentStoreMock: AttachmentStore {
         return nil
     }
 
+    open func fetchAttachment(
+        mediaName: String,
+        tx: DBReadTransaction
+    ) -> Attachment? {
+        return nil
+    }
+
     open func enumerateAllReferences(
         toAttachmentId: Attachment.IDType,
         tx: DBReadTransaction,