From 61d4924e644b182247b672d5075ffd1bf127960d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 30 Nov 2023 13:45:53 +1000
Subject: [PATCH] Expose raw database methods with shielded-scan

---
 zebra-scan/src/storage/db/sapling.rs                        | 4 +++-
 zebra-state/src/service/finalized_state/zebra_db.rs         | 3 ++-
 .../src/service/finalized_state/zebra_db/arbitrary.rs       | 6 +++++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/zebra-scan/src/storage/db/sapling.rs b/zebra-scan/src/storage/db/sapling.rs
index 823d143a757..d67437a023c 100644
--- a/zebra-scan/src/storage/db/sapling.rs
+++ b/zebra-scan/src/storage/db/sapling.rs
@@ -113,7 +113,9 @@ impl Storage {
 
     /// Returns a handle to the `sapling_tx_ids` column family.
     pub(crate) fn sapling_tx_ids_cf(&self) -> impl AsColumnFamilyRef + '_ {
-        self.db.cf_handle(SAPLING_TX_IDS).unwrap()
+        self.db
+            .cf_handle(SAPLING_TX_IDS)
+            .expect("column family was created when database was created")
     }
 
     // Writing batches
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index f02de7d678d..6b7c6823dcd 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -35,7 +35,8 @@ pub mod metrics;
 pub mod shielded;
 pub mod transparent;
 
-#[cfg(any(test, feature = "proptest-impl"))]
+#[cfg(any(test, feature = "proptest-impl", feature = "shielded-scan"))]
+// TODO: when the database is split out of zebra-state, always expose these methods.
 pub mod arbitrary;
 
 /// Wrapper struct to ensure high-level `zebra-state` database access goes through the correct API.
diff --git a/zebra-state/src/service/finalized_state/zebra_db/arbitrary.rs b/zebra-state/src/service/finalized_state/zebra_db/arbitrary.rs
index 2b9f03526a8..8a83a2894c5 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/arbitrary.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/arbitrary.rs
@@ -1,5 +1,7 @@
 //! Arbitrary value generation and test harnesses for high-level typed database access.
 
+#![allow(unused_imports)]
+
 use std::ops::Deref;
 
 use zebra_chain::{amount::NonNegative, block::Block, sprout, value_balance::ValueBalance};
@@ -21,13 +23,14 @@ impl Deref for ZebraDb {
 impl ZebraDb {
     /// Returns the inner database.
     ///
-    /// This is a test-only method, because it allows write access
+    /// This is a test-only and shielded-scan-only method, because it allows write access
     /// and raw read access to the RocksDB instance.
     pub fn db(&self) -> &DiskDb {
         &self.db
     }
 
     /// Allow to set up a fake value pool in the database for testing purposes.
+    #[cfg(any(test, feature = "proptest-impl"))]
     pub fn set_finalized_value_pool(&self, fake_value_pool: ValueBalance<NonNegative>) {
         let mut batch = DiskWriteBatch::new();
         let value_pool_cf = self.db().cf_handle("tip_chain_value_pool").unwrap();
@@ -38,6 +41,7 @@ impl ZebraDb {
 
     /// Artificially prime the note commitment tree anchor sets with anchors
     /// referenced in a block, for testing purposes _only_.
+    #[cfg(any(test, feature = "proptest-impl"))]
     pub fn populate_with_anchors(&self, block: &Block) {
         let mut batch = DiskWriteBatch::new();