diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index 3f0500166a..095109c0f1 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -84,13 +84,7 @@ where chain_spec: Arc, snapshot_provider: &Option>, ) -> ProviderResult> { - let mut provider = DatabaseProvider::new_rw(tx, chain_spec.clone()); - - if let Some(snapshot_provider) = snapshot_provider { - provider = provider.with_snapshot_provider(snapshot_provider.clone()); - } - - Ok(DatabaseProviderRW(provider)) + DatabaseProvider::with_tx_mut(tx, chain_spec, snapshot_provider).map(DatabaseProviderRW) } /// Takes a function and passes a read-write transaction into it, making sure it's committed at @@ -145,11 +139,36 @@ pub struct DatabaseProvider { snapshot_provider: Option>, } +impl DatabaseProvider { + /// Creates a new [`Self`] with access to a [`SnapshotProvider`]. + pub fn with_snapshot_provider(mut self, snapshot_provider: Arc) -> Self { + self.snapshot_provider = Some(snapshot_provider); + self + } +} + impl DatabaseProvider { /// Creates a provider with an inner read-write transaction. pub fn new_rw(tx: TX, chain_spec: Arc) -> Self { Self { tx, chain_spec, snapshot_provider: None } } + + /// Returns a provider with the given `DbTxMut` inside, which allows fetching and updating + /// data from the database using different types of providers. + #[track_caller] + pub fn with_tx_mut( + tx: TX, + chain_spec: Arc, + snapshot_provider: &Option>, + ) -> ProviderResult> { + let mut provider = DatabaseProvider::new_rw(tx, chain_spec.clone()); + + if let Some(snapshot_provider) = snapshot_provider { + provider = provider.with_snapshot_provider(snapshot_provider.clone()); + } + + Ok(provider) + } } impl DatabaseProvider { @@ -254,12 +273,6 @@ impl DatabaseProvider { Self { tx, chain_spec, snapshot_provider: None } } - /// Creates a new [`Self`] with access to a [`SnapshotProvider`]. - pub fn with_snapshot_provider(mut self, snapshot_provider: Arc) -> Self { - self.snapshot_provider = Some(snapshot_provider); - self - } - /// Consume `DbTx` or `DbTxMut`. pub fn into_tx(self) -> TX { self.tx