rust_rocksdb/
write_batch_with_index.rs

1use crate::db::DBInner;
2use crate::{
3    ffi, AsColumnFamilyRef, DBAccess, DBCommon, DBPinnableSlice, DBRawIteratorWithThreadMode,
4    Error, Options, ReadOptions, ThreadMode,
5};
6use libc::{c_char, c_uchar, size_t};
7
8pub struct WriteBatchWithIndex {
9    pub(crate) inner: *mut ffi::rocksdb_writebatch_wi_t,
10}
11
12impl WriteBatchWithIndex {
13    pub fn new(reserved_bytes: usize, overwrite_key: bool) -> Self {
14        Self {
15            inner: unsafe {
16                ffi::rocksdb_writebatch_wi_create(
17                    reserved_bytes as size_t,
18                    c_uchar::from(overwrite_key),
19                )
20            },
21        }
22    }
23
24    pub fn len(&self) -> usize {
25        unsafe { ffi::rocksdb_writebatch_wi_count(self.inner) as usize }
26    }
27
28    /// Return WriteBatch serialized size (in bytes).
29    pub fn size_in_bytes(&self) -> usize {
30        unsafe {
31            let mut batch_size: size_t = 0;
32            ffi::rocksdb_writebatch_wi_data(self.inner, &mut batch_size);
33            batch_size
34        }
35    }
36
37    /// Return a reference to a byte array which represents a serialized version of the batch.
38    pub fn data(&self) -> &[u8] {
39        unsafe {
40            let mut batch_size: size_t = 0;
41            let batch_data = ffi::rocksdb_writebatch_wi_data(self.inner, &mut batch_size);
42            std::slice::from_raw_parts(batch_data as _, batch_size)
43        }
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.len() == 0
48    }
49
50    pub fn get_from_batch<K>(&self, key: K, options: &Options) -> Result<Option<Vec<u8>>, Error>
51    where
52        K: AsRef<[u8]>,
53    {
54        let key = key.as_ref();
55        unsafe {
56            let mut value_size: size_t = 0;
57            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_from_batch(
58                self.inner,
59                options.inner,
60                key.as_ptr() as *const c_char,
61                key.len() as size_t,
62                &mut value_size
63            ));
64
65            if value_data.is_null() {
66                Ok(None)
67            } else {
68                Ok(Some(Vec::from_raw_parts(
69                    value_data as *mut u8,
70                    value_size,
71                    value_size,
72                )))
73            }
74        }
75    }
76
77    pub fn get_from_batch_cf<K>(
78        &self,
79        cf: &impl AsColumnFamilyRef,
80        key: K,
81        options: &Options,
82    ) -> Result<Option<Vec<u8>>, Error>
83    where
84        K: AsRef<[u8]>,
85    {
86        let key = key.as_ref();
87        unsafe {
88            let mut value_size: size_t = 0;
89            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_from_batch_cf(
90                self.inner,
91                options.inner,
92                cf.inner(),
93                key.as_ptr() as *const c_char,
94                key.len() as size_t,
95                &mut value_size
96            ));
97
98            if value_data.is_null() {
99                Ok(None)
100            } else {
101                Ok(Some(Vec::from_raw_parts(
102                    value_data as *mut u8,
103                    value_size,
104                    value_size,
105                )))
106            }
107        }
108    }
109
110    pub fn get_from_batch_and_db<T, I, K>(
111        &self,
112        db: &DBCommon<T, I>,
113        key: K,
114        readopts: &ReadOptions,
115    ) -> Result<Option<Vec<u8>>, Error>
116    where
117        T: ThreadMode,
118        I: DBInner,
119        K: AsRef<[u8]>,
120    {
121        if readopts.inner.is_null() {
122            return Err(Error::new(
123                "Unable to create RocksDB read options. This is a fairly trivial call, and its \
124                 failure may be indicative of a mis-compiled or mis-loaded RocksDB library."
125                    .to_owned(),
126            ));
127        }
128
129        let key = key.as_ref();
130        unsafe {
131            let mut value_size: size_t = 0;
132            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_from_batch_and_db(
133                self.inner,
134                db.inner.inner(),
135                readopts.inner,
136                key.as_ptr() as *const c_char,
137                key.len() as size_t,
138                &mut value_size
139            ));
140
141            if value_data.is_null() {
142                Ok(None)
143            } else {
144                Ok(Some(Vec::from_raw_parts(
145                    value_data as *mut u8,
146                    value_size,
147                    value_size,
148                )))
149            }
150        }
151    }
152
153    pub fn get_pinned_from_batch_and_db<T, I, K>(
154        &'_ self,
155        db: &DBCommon<T, I>,
156        key: K,
157        readopts: &ReadOptions,
158    ) -> Result<Option<DBPinnableSlice<'_>>, Error>
159    where
160        T: ThreadMode,
161        I: DBInner,
162        K: AsRef<[u8]>,
163    {
164        if readopts.inner.is_null() {
165            return Err(Error::new(
166                "Unable to create RocksDB read options. This is a fairly trivial call, and its \
167                 failure may be indicative of a mis-compiled or mis-loaded RocksDB library."
168                    .to_owned(),
169            ));
170        }
171
172        let key = key.as_ref();
173        unsafe {
174            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_pinned_from_batch_and_db(
175                self.inner,
176                db.inner.inner(),
177                readopts.inner,
178                key.as_ptr() as *const c_char,
179                key.len() as size_t,
180            ));
181
182            if value_data.is_null() {
183                Ok(None)
184            } else {
185                Ok(Some(DBPinnableSlice::from_c(value_data)))
186            }
187        }
188    }
189
190    pub fn get_from_batch_and_db_cf<T, I, K>(
191        &self,
192        db: &DBCommon<T, I>,
193        cf: &impl AsColumnFamilyRef,
194        key: K,
195        readopts: &ReadOptions,
196    ) -> Result<Option<Vec<u8>>, Error>
197    where
198        T: ThreadMode,
199        I: DBInner,
200        K: AsRef<[u8]>,
201    {
202        if readopts.inner.is_null() {
203            return Err(Error::new(
204                "Unable to create RocksDB read options. This is a fairly trivial call, and its \
205                 failure may be indicative of a mis-compiled or mis-loaded RocksDB library."
206                    .to_owned(),
207            ));
208        }
209
210        let key = key.as_ref();
211        unsafe {
212            let mut value_size: size_t = 0;
213            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_from_batch_and_db_cf(
214                self.inner,
215                db.inner.inner(),
216                readopts.inner,
217                cf.inner(),
218                key.as_ptr() as *const c_char,
219                key.len() as size_t,
220                &mut value_size
221            ));
222
223            if value_data.is_null() {
224                Ok(None)
225            } else {
226                Ok(Some(Vec::from_raw_parts(
227                    value_data as *mut u8,
228                    value_size,
229                    value_size,
230                )))
231            }
232        }
233    }
234
235    pub fn get_pinned_from_batch_and_db_cf<T, I, K>(
236        &'_ self,
237        db: &DBCommon<T, I>,
238        cf: &impl AsColumnFamilyRef,
239        key: K,
240        readopts: &ReadOptions,
241    ) -> Result<Option<DBPinnableSlice<'_>>, Error>
242    where
243        T: ThreadMode,
244        I: DBInner,
245        K: AsRef<[u8]>,
246    {
247        if readopts.inner.is_null() {
248            return Err(Error::new(
249                "Unable to create RocksDB read options. This is a fairly trivial call, and its \
250                 failure may be indicative of a mis-compiled or mis-loaded RocksDB library."
251                    .to_owned(),
252            ));
253        }
254
255        let key = key.as_ref();
256        unsafe {
257            let value_data = ffi_try!(ffi::rocksdb_writebatch_wi_get_pinned_from_batch_and_db_cf(
258                self.inner,
259                db.inner.inner(),
260                readopts.inner,
261                cf.inner(),
262                key.as_ptr() as *const c_char,
263                key.len() as size_t,
264            ));
265
266            if value_data.is_null() {
267                Ok(None)
268            } else {
269                Ok(Some(DBPinnableSlice::from_c(value_data)))
270            }
271        }
272    }
273
274    /// Insert a value into the database under the given key.
275    pub fn put<K, V>(&mut self, key: K, value: V)
276    where
277        K: AsRef<[u8]>,
278        V: AsRef<[u8]>,
279    {
280        let key = key.as_ref();
281        let value = value.as_ref();
282
283        unsafe {
284            ffi::rocksdb_writebatch_wi_put(
285                self.inner,
286                key.as_ptr() as *const c_char,
287                key.len() as size_t,
288                value.as_ptr() as *const c_char,
289                value.len() as size_t,
290            );
291        }
292    }
293
294    pub fn put_cf<K, V>(&mut self, cf: &impl AsColumnFamilyRef, key: K, value: V)
295    where
296        K: AsRef<[u8]>,
297        V: AsRef<[u8]>,
298    {
299        let key = key.as_ref();
300        let value = value.as_ref();
301
302        unsafe {
303            ffi::rocksdb_writebatch_wi_put_cf(
304                self.inner,
305                cf.inner(),
306                key.as_ptr() as *const c_char,
307                key.len() as size_t,
308                value.as_ptr() as *const c_char,
309                value.len() as size_t,
310            );
311        }
312    }
313
314    pub fn merge<K, V>(&mut self, key: K, value: V)
315    where
316        K: AsRef<[u8]>,
317        V: AsRef<[u8]>,
318    {
319        let key = key.as_ref();
320        let value = value.as_ref();
321
322        unsafe {
323            ffi::rocksdb_writebatch_wi_merge(
324                self.inner,
325                key.as_ptr() as *const c_char,
326                key.len() as size_t,
327                value.as_ptr() as *const c_char,
328                value.len() as size_t,
329            );
330        }
331    }
332
333    pub fn merge_cf<K, V>(&mut self, cf: &impl AsColumnFamilyRef, key: K, value: V)
334    where
335        K: AsRef<[u8]>,
336        V: AsRef<[u8]>,
337    {
338        let key = key.as_ref();
339        let value = value.as_ref();
340
341        unsafe {
342            ffi::rocksdb_writebatch_wi_merge_cf(
343                self.inner,
344                cf.inner(),
345                key.as_ptr() as *const c_char,
346                key.len() as size_t,
347                value.as_ptr() as *const c_char,
348                value.len() as size_t,
349            );
350        }
351    }
352
353    /// Removes the database entry for key. Does nothing if the key was not found.
354    pub fn delete<K: AsRef<[u8]>>(&mut self, key: K) {
355        let key = key.as_ref();
356
357        unsafe {
358            ffi::rocksdb_writebatch_wi_delete(
359                self.inner,
360                key.as_ptr() as *const c_char,
361                key.len() as size_t,
362            );
363        }
364    }
365
366    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf: &impl AsColumnFamilyRef, key: K) {
367        let key = key.as_ref();
368
369        unsafe {
370            ffi::rocksdb_writebatch_wi_delete_cf(
371                self.inner,
372                cf.inner(),
373                key.as_ptr() as *const c_char,
374                key.len() as size_t,
375            );
376        }
377    }
378
379    /// Clear all updates buffered in this batch.
380    pub fn clear(&mut self) {
381        unsafe {
382            ffi::rocksdb_writebatch_wi_clear(self.inner);
383        }
384    }
385
386    pub fn iterator_with_base<'a, D>(
387        &self,
388        base_iterator: DBRawIteratorWithThreadMode<'a, D>,
389    ) -> DBRawIteratorWithThreadMode<'a, D>
390    where
391        D: DBAccess,
392    {
393        let (base_iterator_inner, readopts) = base_iterator.into_inner();
394
395        let iterator = unsafe {
396            ffi::rocksdb_writebatch_wi_create_iterator_with_base_readopts(
397                self.inner,
398                base_iterator_inner.as_ptr(),
399                readopts.inner,
400            )
401        };
402
403        DBRawIteratorWithThreadMode::from_inner(iterator, readopts)
404    }
405
406    pub fn iterator_with_base_cf<'a, D>(
407        &self,
408        base_iterator: DBRawIteratorWithThreadMode<'a, D>,
409        cf: &impl AsColumnFamilyRef,
410    ) -> DBRawIteratorWithThreadMode<'a, D>
411    where
412        D: DBAccess,
413    {
414        let (base_iterator_inner, readopts) = base_iterator.into_inner();
415
416        let iterator = unsafe {
417            ffi::rocksdb_writebatch_wi_create_iterator_with_base_cf_readopts(
418                self.inner,
419                base_iterator_inner.as_ptr(),
420                cf.inner(),
421                readopts.inner,
422            )
423        };
424
425        DBRawIteratorWithThreadMode::from_inner(iterator, readopts)
426    }
427}
428
429impl Drop for WriteBatchWithIndex {
430    fn drop(&mut self) {
431        unsafe {
432            ffi::rocksdb_writebatch_wi_destroy(self.inner);
433        }
434    }
435}
436
437unsafe impl Send for WriteBatchWithIndex {}