reth_storage_api/
database_provider.rs1use alloc::vec::Vec;
2use core::ops::{Bound, RangeBounds};
3use reth_db_api::{
4 common::KeyValue,
5 cursor::DbCursorRO,
6 database::Database,
7 table::Table,
8 transaction::{DbTx, DbTxMut},
9 DatabaseError,
10};
11use reth_prune_types::PruneModes;
12use reth_storage_errors::provider::ProviderResult;
13
14#[auto_impl::auto_impl(&)]
16pub trait DbTxProvider {
17 type Tx: DbTx;
19
20 fn tx(&self) -> &Self::Tx;
22}
23
24pub trait DBProvider: DbTxProvider + Sized {
26 fn tx_ref(&self) -> &Self::Tx {
28 self.tx()
29 }
30
31 fn tx_mut(&mut self) -> &mut Self::Tx;
33
34 fn into_tx(self) -> Self::Tx;
36
37 fn disable_long_read_transaction_safety(mut self) -> Self {
44 self.tx_mut().disable_long_read_transaction_safety();
45 self
46 }
47
48 fn commit(self) -> ProviderResult<()>;
50
51 fn prune_modes_ref(&self) -> &PruneModes;
53
54 fn table<T: Table>(&self) -> Result<Vec<KeyValue<T>>, DatabaseError>
56 where
57 T::Key: Default + Ord,
58 {
59 self.tx_ref()
60 .cursor_read::<T>()?
61 .walk(Some(T::Key::default()))?
62 .collect::<Result<Vec<_>, DatabaseError>>()
63 }
64
65 #[inline]
67 fn get<T: Table>(
68 &self,
69 range: impl RangeBounds<T::Key>,
70 ) -> Result<Vec<KeyValue<T>>, DatabaseError> {
71 self.tx_ref().cursor_read::<T>()?.walk_range(range)?.collect::<Result<Vec<_>, _>>()
72 }
73
74 fn cursor_read_collect<T: Table<Key = u64>>(
78 &self,
79 range: impl RangeBounds<T::Key>,
80 ) -> ProviderResult<Vec<T::Value>> {
81 let capacity = match range_size_hint(&range) {
82 Some(0) | None => return Ok(Vec::new()),
83 Some(capacity) => capacity,
84 };
85 let mut cursor = self.tx_ref().cursor_read::<T>()?;
86 self.cursor_collect_with_capacity(&mut cursor, range, capacity)
87 }
88
89 fn cursor_collect<T: Table<Key = u64>>(
91 &self,
92 cursor: &mut impl DbCursorRO<T>,
93 range: impl RangeBounds<T::Key>,
94 ) -> ProviderResult<Vec<T::Value>> {
95 let capacity = range_size_hint(&range).unwrap_or(0);
96 self.cursor_collect_with_capacity(cursor, range, capacity)
97 }
98
99 fn cursor_collect_with_capacity<T: Table<Key = u64>>(
102 &self,
103 cursor: &mut impl DbCursorRO<T>,
104 range: impl RangeBounds<T::Key>,
105 capacity: usize,
106 ) -> ProviderResult<Vec<T::Value>> {
107 let mut items = Vec::with_capacity(capacity);
108 for entry in cursor.walk_range(range)? {
109 items.push(entry?.1);
110 }
111 Ok(items)
112 }
113
114 #[inline]
116 fn remove<T: Table>(&self, range: impl RangeBounds<T::Key>) -> Result<usize, DatabaseError>
117 where
118 Self::Tx: DbTxMut,
119 {
120 let mut entries = 0;
121 let mut cursor_write = self.tx_ref().cursor_write::<T>()?;
122 let mut walker = cursor_write.walk_range(range)?;
123 while walker.next().transpose()?.is_some() {
124 walker.delete_current()?;
125 entries += 1;
126 }
127 Ok(entries)
128 }
129
130 #[inline]
132 fn take<T: Table>(
133 &self,
134 range: impl RangeBounds<T::Key>,
135 ) -> Result<Vec<KeyValue<T>>, DatabaseError>
136 where
137 Self::Tx: DbTxMut,
138 {
139 let mut cursor_write = self.tx_ref().cursor_write::<T>()?;
140 let mut walker = cursor_write.walk_range(range)?;
141 let mut items = Vec::new();
142 while let Some(i) = walker.next().transpose()? {
143 walker.delete_current()?;
144 items.push(i)
145 }
146 Ok(items)
147 }
148}
149
150#[auto_impl::auto_impl(&, Arc)]
152pub trait DatabaseProviderFactory: Send + Sync {
153 type DB: Database;
155
156 type Provider: DBProvider<Tx = <Self::DB as Database>::TX>;
158
159 type ProviderRW: DBProvider<Tx = <Self::DB as Database>::TXMut>;
161
162 fn database_provider_ro(&self) -> ProviderResult<Self::Provider>;
164
165 fn database_provider_rw(&self) -> ProviderResult<Self::ProviderRW>;
167}
168
169pub type FactoryTx<F> = <<F as DatabaseProviderFactory>::DB as Database>::TX;
171
172pub trait DatabaseProviderROFactory {
174 type Provider;
179
180 fn database_provider_ro(&self) -> ProviderResult<Self::Provider>;
182}
183
184impl<T> DatabaseProviderROFactory for T
185where
186 T: DatabaseProviderFactory,
187{
188 type Provider = T::Provider;
189
190 fn database_provider_ro(&self) -> ProviderResult<Self::Provider> {
191 <T as DatabaseProviderFactory>::database_provider_ro(self)
192 }
193}
194
195pub fn range_size_hint(range: &impl RangeBounds<u64>) -> Option<usize> {
197 let start = match range.start_bound().cloned() {
198 Bound::Included(start) => start,
199 Bound::Excluded(start) => start.checked_add(1)?,
200 Bound::Unbounded => 0,
201 };
202 let end = match range.end_bound().cloned() {
203 Bound::Included(end) => end.saturating_add(1),
204 Bound::Excluded(end) => end,
205 Bound::Unbounded => return None,
206 };
207 end.checked_sub(start).map(|x| x as _)
208}