pbbson 0.1.8

Utilities for pbjson to BSON conversion
use bson::document::{Iter, IterMut, Keys, ValueAccessError, ValueAccessResult, Values};
use bson::spec::BinarySubtype;
use bson::{Bson, DateTime, Decimal128, Document, Timestamp};
use serde::{Deserialize, Serialize};
use std::borrow::Borrow;
use std::str::FromStr;

#[derive(Clone, Debug, Default, Deserialize)]
pub struct Model(Document);

impl Model {
    pub fn clear(&mut self) {
        self.0.clear()
    }

    pub fn contains_field(&self, field: impl AsRef<str>) -> bool {
        self.0.contains_key(field.as_ref())
    }

    pub fn get(&self, field: impl AsRef<str>) -> Option<&Bson> {
        self.0.get(field)
    }

    pub fn get_mut(&mut self, field: impl AsRef<str>) -> Option<&mut Bson> {
        self.0.get_mut(field.as_ref())
    }

    pub fn get_bool(&self, field: impl AsRef<str>) -> ValueAccessResult<bool> {
        self.0.get_bool(field)
    }

    pub fn get_datetime(&self, field: impl AsRef<str>) -> ValueAccessResult<&DateTime> {
        self.0.get_datetime(field)
    }

    pub fn get_decimal128(&self, field: impl AsRef<str>) -> ValueAccessResult<&Decimal128> {
        self.0.get_decimal128(field)
    }

    pub fn get_f64(&self, field: impl AsRef<str>) -> ValueAccessResult<f64> {
        self.0.get_f64(field)
    }

    pub fn get_i32(&self, field: impl AsRef<str>) -> ValueAccessResult<i32> {
        self.0.get_i32(field)
    }

    pub fn get_i64(&self, field: impl AsRef<str>) -> ValueAccessResult<i64> {
        self.0.get_i64(field)
    }

    pub fn get_str(&self, field: impl AsRef<str>) -> ValueAccessResult<&str> {
        self.0.get_str(field)
    }

    pub fn get_timestamp(&self, field: impl AsRef<str>) -> ValueAccessResult<Timestamp> {
        self.0.get_timestamp(field)
    }

    pub fn id(&self) -> ValueAccessResult<String> {
        match self.0.get("id") {
            None => Err(ValueAccessError::NotPresent),
            Some(Bson::Binary(binary)) => match binary.subtype {
                BinarySubtype::Md5 => {
                    use base64::{engine::general_purpose, Engine as _};
                    Ok(general_purpose::STANDARD.encode(&binary.bytes))
                }
                BinarySubtype::Uuid => {
                    let uuid = binary.to_uuid().map_err(|_e| ValueAccessError::UnexpectedType)?;
                    Ok(uuid.to_string())
                }
                _ => Err(ValueAccessError::UnexpectedType),
            },
            Some(Bson::ObjectId(object_id)) => try_from_object_id(object_id),
            Some(Bson::String(str)) => Ok(str.clone()),
            _ => Err(ValueAccessError::UnexpectedType),
        }
    }

    pub fn insert<KT: Into<String>, BT: Into<Bson>>(&mut self, field: KT, value: BT) -> Option<Bson> {
        self.0.insert(field, value)
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn is_null(&self, field: impl AsRef<str>) -> bool {
        self.0.is_null(field)
    }

    pub fn iter(&self) -> Iter {
        self.0.iter()
    }

    pub fn iter_mut(&mut self) -> IterMut {
        self.0.iter_mut()
    }

    pub fn keys(&self) -> Keys {
        self.0.keys()
    }

    pub fn len(&self) -> usize {
        self.0.len()
    }

    pub fn remove(&mut self, field: impl AsRef<str>) -> Option<Bson> {
        self.0.remove(field.as_ref())
    }

    pub fn set_datetime(&mut self, field: &str, value: DateTime) {
        self.0.insert(field, Some(value));
    }

    pub fn try_from<T: prost::Message + Serialize>(other: &T) -> ValueAccessResult<Self> {
        let buf = serde_json::to_vec(&other).map_err(|_e| ValueAccessError::UnexpectedType)?;
        let mut model: Model = serde_json::from_slice(&buf).map_err(|_e| ValueAccessError::UnexpectedType)?;
        for key in model.clone().keys() {
            if key == "id" {
                if let Ok(value) = model.get_str(key) {
                    if value.is_empty() {
                        model.remove(key);
                    } else if let Ok(object_id) = try_to_object_id(value) {
                        model.insert(key, object_id);
                    } else if let Ok(uuid) = try_to_uuid(value) {
                        model.insert(key, uuid);
                    } else {
                        model.insert(key, value.to_string());
                    }
                }
            }
            if key.ends_with("At") {
                if let Ok(value) = model.0.get_str(key) {
                    if let Ok(value) = bson::DateTime::parse_rfc3339_str(value) {
                        model.0.insert(key, value);
                    }
                }
            }
        }
        Ok(model)
    }

    pub fn try_into<T: prost::Message + Clone + Default + for<'a> Deserialize<'a>>(self) -> ValueAccessResult<T> {
        let mut this = self.0.clone();

        // Perform some translations
        for key in self.keys() {
            let maybe_value = this.get(key);
            if let Some(value) = maybe_value {
                let new_value = to_pbjson(value.clone())?;
                this.insert(key, new_value);
            }
        }

        let json = serde_json::to_vec(&this).map_err(|_e| ValueAccessError::UnexpectedType)?;
        let message: T = serde_json::from_slice::<T>(&json).map_err(|_e| ValueAccessError::UnexpectedType)?;
        Ok(message)
    }

    pub fn values(&self) -> Values {
        self.0.values()
    }
}

pub fn to_pbjson(value: Bson) -> ValueAccessResult<Bson> {
    match value {
        Bson::ObjectId(object_id) => Ok(Bson::String(try_from_object_id(&object_id)?)),
        Bson::DateTime(date_time) => Ok(Bson::String(
            date_time
                .try_to_rfc3339_string()
                .map_err(|_e| ValueAccessError::UnexpectedType)?
                .to_string(),
        )),
        Bson::Array(array) => {
            let mut new_array = Vec::with_capacity(array.len());
            for value in array {
                new_array.push(to_pbjson(value)?);
            }
            Ok(Bson::Array(new_array))
        }
        Bson::Binary(ref binary) => match binary.subtype {
            BinarySubtype::Md5 => {
                use base64::{engine::general_purpose, Engine as _};
                Ok(Bson::String(general_purpose::STANDARD.encode(&binary.bytes)))
            }
            BinarySubtype::Uuid => {
                let uuid = binary.to_uuid().map_err(|_e| ValueAccessError::UnexpectedType)?;
                Ok(Bson::String(uuid.to_string()))
            }
            _ => Ok(value),
        },
        _ => Ok(value),
    }
}

impl From<Document> for Model {
    fn from(other: Document) -> Self {
        Model(other)
    }
}

fn try_from_object_id(id: &bson::oid::ObjectId) -> ValueAccessResult<String> {
    let id = xid::Id::from_bytes(&id.bytes()).map_err(|_e| ValueAccessError::UnexpectedType)?;
    Ok(id.to_string())
}

fn try_to_object_id(id: &str) -> ValueAccessResult<bson::oid::ObjectId> {
    const WHAT: &str = "IDs";
    require_xid_base32hex(WHAT, id)?;
    let as_xid = match xid::Id::from_str(id) {
        Ok(xid) => xid,
        Err(_e) => {
            // let msg = format!("{WHAT} must be 12-byte XID's in 20-byte base32hex encoded form ({e:?})");
            return Err(ValueAccessError::UnexpectedType);
        }
    };
    Ok(bson::oid::ObjectId::from_bytes(*as_xid.as_bytes()))
}

fn try_to_uuid(id: &str) -> ValueAccessResult<bson::binary::Binary> {
    let uuid = bson::uuid::Uuid::parse_str(id).map_err(|_e| ValueAccessError::UnexpectedType)?;
    Ok(bson::binary::Binary::from_uuid(uuid))
}

fn require_xid_base32hex(_what: &str, id: &str) -> ValueAccessResult<()> {
    if id.len() != 20 {
        // return Err(Status::invalid_argument(format!(
        //     "{what} must be 12-byte XID's in 20-byte base32hex encoded form"
        // )));
        return Err(ValueAccessError::UnexpectedType);
    }
    Ok(())
}

impl Borrow<Document> for Model {
    fn borrow(&self) -> &Document {
        self.0.borrow()
    }
}

impl From<Model> for Document {
    fn from(other: Model) -> Self {
        other.0
    }
}

#[cfg(test)]
mod tests {
    use super::Model;
    use bson::spec::BinarySubtype;
    use bson::{Bson, Document};

    #[test]
    fn can_access_id_when_none() {
        let model = Model::default();
        assert!(model.id().is_err());
    }

    #[test]
    fn can_deser() {
        const JSON: &str = r#"{"id":"d05vu4r71n0pfhlalahg","firstName":"Joe","age":23}"#;
        let model: Model = serde_json::from_str(JSON).unwrap();
        assert_eq!(model.id(), Ok("d05vu4r71n0pfhlalahg".to_string()));
        assert_eq!(model.get_str("firstName").unwrap(), "Joe");
        assert_eq!(model.get_i32("age").unwrap(), 23);
    }

    #[test]
    fn can_convert_into_document() {
        let _doc: Document = Model::default().into();
    }

    #[test]
    fn can_get_id_for_string() {
        let mut model = Model::default();
        model.insert("id", Bson::String("123".to_string()));
        assert_eq!(model.id(), Ok("123".to_string()));
    }

    #[test]
    fn can_get_id_for_xid() {
        const JSON: &str = r#"{"id":"d05vu4r71n0pfhlalahg"}"#;
        let model: Model = serde_json::from_str(JSON).unwrap();
        assert_eq!(model.id(), Ok("d05vu4r71n0pfhlalahg".to_string()));
    }

    #[test]
    fn can_get_id_for_uuid() {
        const UUID: &str = "67e55044-10b1-426f-9247-bb680e5fe0c8";
        let uuid = bson::uuid::Uuid::parse_str(UUID).unwrap();
        let mut model = Model::default();
        model.insert("id", uuid);
        assert_eq!(model.id(), Ok(UUID.to_string()));
    }

    #[test]
    fn can_get_id_for_md5() {
        use base64::{engine::general_purpose, Engine as _};
        const MD5_AS_BASE64: &str = "uEzK4pS84wA+OqxuXYwZfQ==";
        let bytes = general_purpose::STANDARD.decode(MD5_AS_BASE64).unwrap();
        let mut model = Model::default();
        model.insert(
            "id",
            Bson::Binary(bson::Binary {
                subtype: BinarySubtype::Md5,
                bytes,
            }),
        );
        assert_eq!(model.id(), Ok(MD5_AS_BASE64.to_string()));
    }
}