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 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 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 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 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 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 {}