add Database::get_or_put

heed/src/database.rs

@@ -1993,6 +1993,116 @@ impl<KC, DC, C> Database<KC, DC, C> {
  Ok(())
  }
 
+ /// Attempt to insert a key-value pair in this database, or if a value already exists for the
+ /// key, returns the previous value. The entry is written with no specific flag.
+ ///
+ /// ```
+ /// # use heed::EnvOpenOptions;
+ /// use heed::Database;
+ /// use heed::types::*;
+ /// use heed::byteorder::BigEndian;
+ ///
+ /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+ /// # let dir = tempfile::tempdir()?;
+ /// # let env = unsafe { EnvOpenOptions::new()
+ /// # .map_size(10 * 1024 * 1024) // 10MB
+ /// # .max_dbs(3000)
+ /// # .open(dir.path())?
+ /// # };
+ /// type BEI32 = I32<BigEndian>;
+ ///
+ /// let mut wtxn = env.write_txn()?;
+ /// let db: Database<BEI32, Str> = env.create_database(&mut wtxn, Some("iter-i32"))?;
+ ///
+ /// # db.clear(&mut wtxn)?;
+ /// assert_eq!(db.get_or_put(&mut wtxn, &42, "i-am-forty-two")?, None);
+ /// assert_eq!(db.get_or_put(&mut wtxn, &42, "the meaning of life")?, Some("i-am-forty-two"));
+ ///
+ /// let ret = db.get(&mut wtxn, &42)?;
+ /// assert_eq!(ret, Some("i-am-forty-two"));
+ ///
+ /// wtxn.commit()?;
+ /// # Ok(()) }
+ /// ```
+ pub fn get_or_put<'a, 'txn>(
+ &'txn self,
+ txn: &mut RwTxn,
+ key: &'a KC::EItem,
+ data: &'a DC::EItem,
+ ) -> Result<Option<DC::DItem>>
+ where
+ KC: BytesEncode<'a>,
+ DC: BytesEncode<'a> + BytesDecode<'a>,
+ {
+ self.get_or_put_with_flags(txn, PutFlags::empty(), key, data)
+ }
+
+ /// Attempt to insert a key-value pair in this database, or if a value already exists for the
+ /// key, returns the previous value. The entry is written with the specified flags.
+ ///
+ /// ```
+ /// # use heed::EnvOpenOptions;
+ /// use heed::{Database, PutFlags};
+ /// use heed::types::*;
+ /// use heed::byteorder::BigEndian;
+ ///
+ /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+ /// # let dir = tempfile::tempdir()?;
+ /// # let env = unsafe { EnvOpenOptions::new()
+ /// # .map_size(10 * 1024 * 1024) // 10MB
+ /// # .max_dbs(3000)
+ /// # .open(dir.path())?
+ /// # };
+ /// type BEI32 = I32<BigEndian>;
+ ///
+ /// let mut wtxn = env.write_txn()?;
+ /// let db: Database<BEI32, Str> = env.create_database(&mut wtxn, Some("iter-i32"))?;
+ ///
+ /// # db.clear(&mut wtxn)?;
+ /// assert_eq!(db.get_or_put_with_flags(&mut wtxn, PutFlags::empty(), &42, "i-am-forty-two")?, None);
+ /// assert_eq!(db.get_or_put_with_flags(&mut wtxn, PutFlags::empty(), &42, "the meaning of life")?, Some("i-am-forty-two"));
+ ///
+ /// let ret = db.get(&mut wtxn, &42)?;
+ /// assert_eq!(ret, Some("i-am-forty-two"));
+ ///
+ /// wtxn.commit()?;
+ /// # Ok(()) }
+ /// ```
+ pub fn get_or_put_with_flags<'a, 'txn>(
+ &'txn self,
+ txn: &mut RwTxn,
+ flags: PutFlags,
+ key: &'a KC::EItem,
+ data: &'a DC::EItem,
+ ) -> Result<Option<DC::DItem>>
+ where
+ KC: BytesEncode<'a>,
+ DC: BytesEncode<'a> + BytesDecode<'a>,
+ {
+ assert_eq_env_db_txn!(self, txn);
+
+ let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?;
+ let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Encoding)?;
+
+ let mut key_val = unsafe { crate::into_val(&key_bytes) };
+ let mut data_val = unsafe { crate::into_val(&data_bytes) };
+ let flags = (flags | PutFlags::NO_OVERWRITE).bits();
+
+ unsafe {
+ match ffi::mdb_put(txn.txn.txn, self.dbi, &mut key_val, &mut data_val, flags) {
+ lmdb_master_sys::MDB_KEYEXIST => {
+ let bytes = crate::from_val(data_val);
+ let data = DC::bytes_decode(bytes).map_err(Error::Decoding)?;
+ Ok(Some(data))
+ }
+ err_code => {
+ mdb_result(err_code)?;
+ Ok(None)
+ }
+ }
+ }
+ }
+
  /// Deletes an entry or every duplicate data items of a key
  /// if the database supports duplicate data items.
  ///
@@ -2355,3 +2465,27 @@ pub struct DatabaseStat {
  /// Number of data items.
  pub entries: usize,
 }
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use heed_types::*;
+
+ #[test]
+ fn put_overwrite() -> Result<()> {
+ let dir = tempfile::tempdir()?;
+ let env = unsafe { EnvOpenOptions::new().open(dir.path())? };
+ let mut txn = env.write_txn()?;
+ let db = env.create_database::<Bytes, Bytes>(&mut txn, None)?;
+
+ assert_eq!(db.get(&txn, b"hello").unwrap(), None);
+
+ db.put(&mut txn, b"hello", b"hi").unwrap();
+ assert_eq!(db.get(&txn, b"hello").unwrap(), Some(&b"hi"[..]));
+
+ db.put(&mut txn, b"hello", b"bye").unwrap();
+ assert_eq!(db.get(&txn, b"hello").unwrap(), Some(&b"bye"[..]));
+
+ Ok(())
+ }
+}