Fix db upgrade from sapling

binding.h

@@ -50,9 +50,11 @@ void dart_post_cobject(DartPostCObjectFnType ptr);
 
 void deallocate_str(char *s);
 
-void init_wallet(char *db_path);
+CResult init_wallet(uint8_t coin, char *db_path);
 
-void migrate_db(char *db_path);
+CResult migrate_db(uint8_t coin, char *db_path);
+
+CResult migrate_data_db(uint8_t coin);
 
 void set_active(uint8_t active);
 

src/api/dart_ffi.rs

@@ -1,4 +1,4 @@
-use crate::coinconfig::{init_coin, migrate_coin, CoinConfig, MEMPOOL, MEMPOOL_RUNNER};
+use crate::coinconfig::{init_coin, CoinConfig, MEMPOOL, MEMPOOL_RUNNER};
 use crate::note_selection::TransactionReport;
 use crate::{ChainError, TransactionPlan, Tx};
 use allo_isolate::{ffi, IntoDart};
@@ -90,24 +90,25 @@ pub struct CResult<T> {
     error: *mut c_char,
 }
 
-const COIN_DBNAMES: &[&str] = &["zec.db", "yec.db"];
-
 #[no_mangle]
-pub unsafe extern "C" fn init_wallet(db_path: *mut c_char) {
+pub unsafe extern "C" fn init_wallet(coin: u8, db_path: *mut c_char) -> CResult<()> {
     try_init_logger();
     from_c_str!(db_path);
-    for (coin, filename) in COIN_DBNAMES.iter().enumerate() {
-        let _ = init_coin(coin as u8, &format!("{}/{}", &db_path, filename));
-    }
+    to_cresult(init_coin(coin, &db_path))
 }
 
 #[no_mangle]
-pub unsafe extern "C" fn migrate_db(db_path: *mut c_char) {
+pub unsafe extern "C" fn migrate_db(coin: u8, db_path: *mut c_char) -> CResult<()> {
     try_init_logger();
     from_c_str!(db_path);
-    for (coin, filename) in COIN_DBNAMES.iter().enumerate() {
-        let _ = migrate_coin(coin as u8, &format!("{}/{}", &db_path, filename));
-    }
+    to_cresult(crate::coinconfig::migrate_db(coin, &db_path))
+}
+
+#[no_mangle]
+#[tokio::main]
+pub async unsafe extern "C" fn migrate_data_db(coin: u8) -> CResult<()> {
+    try_init_logger();
+    to_cresult(crate::coinconfig::migrate_data(coin).await)
 }
 
 #[no_mangle]

src/coinconfig.rs

@@ -62,12 +62,20 @@ pub fn init_coin(coin: u8, db_path: &str) -> anyhow::Result<()> {
 
 /// Upgrade database schema for given coin and db path
 /// Used from ywallet
-pub fn migrate_coin(coin: u8, db_path: &str) -> anyhow::Result<()> {
+pub fn migrate_db(coin: u8, db_path: &str) -> anyhow::Result<()> {
     let chain = get_coin_chain(get_coin_type(coin));
     DbAdapter::migrate_db(chain.network(), db_path, chain.has_unified())?;
     Ok(())
 }
 
+pub async fn migrate_data(coin: u8) -> anyhow::Result<()> {
+    let c = CoinConfig::get(coin);
+    let db = c.db()?;
+    let mut client = c.connect_lwd().await?;
+    db.migrate_data(&mut client).await?;
+    Ok(())
+}
+
 #[derive(Clone)]
 pub struct CoinConfig {
     pub coin: u8,

src/db.rs

@@ -8,7 +8,7 @@ use crate::sync::tree::{CTree, TreeCheckpoint};
 use crate::taddr::{derive_tkeys, TBalance};
 use crate::transaction::{GetTransactionDetailRequest, TransactionDetails};
 use crate::unified::UnifiedAddressType;
-use crate::{sync, Hash};
+use crate::{sync, BlockId, CompactTxStreamerClient, Hash};
 use orchard::keys::FullViewingKey;
 use rusqlite::Error::QueryReturnedNoRows;
 use rusqlite::{params, Connection, OptionalExtension, Transaction};
@@ -16,6 +16,8 @@ use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use std::collections::HashMap;
 use std::convert::TryInto;
+use tonic::transport::Channel;
+use tonic::Request;
 use zcash_client_backend::encoding::decode_extended_full_viewing_key;
 use zcash_params::coin::{get_coin_chain, get_coin_id, CoinType};
 use zcash_primitives::consensus::{Network, NetworkUpgrade, Parameters};
@@ -119,6 +121,39 @@ impl DbAdapter {
         Ok(())
     }
 
+    pub async fn migrate_data(
+        &self,
+        client: &mut CompactTxStreamerClient<Channel>,
+    ) -> anyhow::Result<()> {
+        let mut stmt = self.connection.prepare("select s.height from sapling_tree s LEFT JOIN orchard_tree o ON s.height = o.height WHERE o.height IS NULL")?;
+        let rows = stmt.query_map([], |row| {
+            let height: u32 = row.get(0)?;
+            Ok(height)
+        })?;
+        let mut trees = HashMap::new();
+        for r in rows {
+            trees.insert(r?, vec![]);
+        }
+        for (height, tree) in trees.iter_mut() {
+            let tree_state = client
+                .get_tree_state(Request::new(BlockId {
+                    height: *height as u64,
+                    hash: vec![],
+                }))
+                .await?
+                .into_inner();
+            let orchard_tree = hex::decode(&tree_state.orchard_tree).unwrap();
+            tree.extend(orchard_tree);
+        }
+        for (height, tree) in trees.iter() {
+            self.connection.execute(
+                "INSERT INTO orchard_tree(height, tree) VALUES (?1, ?2) ON CONFLICT DO NOTHING",
+                params![height, tree],
+            )?;
+        }
+        Ok(())
+    }
+
     pub fn disable_wal(db_path: &str) -> anyhow::Result<()> {
         let connection = Connection::open(db_path)?;
         connection.query_row("PRAGMA journal_mode = OFF", [], |_| Ok(()))?;

src/db/migration.rs

@@ -34,6 +34,8 @@ pub fn reset_db(connection: &Connection) -> anyhow::Result<()> {
     Ok(())
 }
 
+const LATEST_VERSION: u32 = 5;
+
 pub fn init_db(connection: &Connection, network: &Network, has_ua: bool) -> anyhow::Result<()> {
     connection.execute(
         "CREATE TABLE IF NOT EXISTS schema_version (
@@ -265,8 +267,8 @@ pub fn init_db(connection: &Connection, network: &Network, has_ua: bool) -> anyh
         )?;
     }
 
-    if version != 5 {
-        update_schema_version(connection, 5)?;
+    if version != LATEST_VERSION {
+        update_schema_version(connection, LATEST_VERSION)?;
         connection.cache_flush()?;
         log::info!("Database migrated");
     }