zcash_client_sqlite: Add Orchard support to get_received_memo

zcash · Mar 10, 2024 · 4ef9452 · 4ef9452
1 parent 9ccedf2
commit 4ef9452
Showing 1 changed file with 19 additions and 19 deletions.
diff --git a/zcash_client_sqlite/src/wallet.rs b/zcash_client_sqlite/src/wallet.rs
@@ -1196,26 +1196,26 @@ pub(crate) fn get_received_memo(
     conn: &rusqlite::Connection,
     note_id: NoteId,
 ) -> Result<Option<Memo>, SqliteClientError> {
-    let memo_bytes: Option<Vec<_>> = match note_id.protocol() {
-        ShieldedProtocol::Sapling => conn
-            .query_row(
-                "SELECT memo FROM sapling_received_notes
-                JOIN transactions ON sapling_received_notes.tx = transactions.id_tx
+    let fetch_memo = |table_prefix: &'static str, output_col: &'static str| {
+        conn.query_row(
+            &format!(
+                "SELECT memo FROM {table_prefix}_received_notes
+                JOIN transactions ON {table_prefix}_received_notes.tx = transactions.id_tx
                 WHERE transactions.txid = :txid
-                AND sapling_received_notes.output_index = :output_index",
-                named_params![
-                    ":txid": note_id.txid().as_ref(),
-                    ":output_index": note_id.output_index()
-                ],
-                |row| row.get(0),
-            )
-            .optional()?
-            .flatten(),
-        _ => {
-            return Err(SqliteClientError::UnsupportedPoolType(PoolType::Shielded(
-                note_id.protocol(),
-            )))
-        }
+                AND {table_prefix}_received_notes.{output_col} = :output_index"
+            ),
+            named_params![
+                ":txid": note_id.txid().as_ref(),
+                ":output_index": note_id.output_index()
+            ],
+            |row| row.get(0),
+        )
+        .optional()
+    };
+
+    let memo_bytes: Option<Vec<_>> = match note_id.protocol() {
+        ShieldedProtocol::Sapling => fetch_memo(SAPLING_TABLES_PREFIX, "output_index")?.flatten(),
+        ShieldedProtocol::Orchard => fetch_memo(ORCHARD_TABLES_PREFIX, "action_index")?.flatten(),
     };
 
     memo_bytes