mongo_orm_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, Data, DeriveInput, Fields};
4
5#[proc_macro_derive(mongo_doc, attributes(object_id, primary_key))]
6pub fn mongo_entity_derive(input: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(input as DeriveInput);
8    let name = input.ident;
9
10    let data = match input.data {
11        Data::Struct(data) => data,
12        _ => panic!("MongoDoc can only be derived for structs"),
13    };
14
15    let mut object_id_conversions = Vec::new();
16    let mut object_id_reverse_conversions = Vec::new();
17    let mut primary_key_field = None;
18
19    if let Fields::Named(fields) = data.fields {
20        for field in fields.named {
21            let field_name = field.ident.clone().unwrap();
22            
23            let is_object_id = field.attrs.iter().any(|attr| attr.path().is_ident("object_id"));
24            let is_primary_key = field.attrs.iter().any(|attr| attr.path().is_ident("primary_key"));
25
26            if is_primary_key {
27                if primary_key_field.is_some() {
28                    panic!("Only one field can be marked as primary_key");
29                }
30                primary_key_field = Some(field_name.clone());
31            }
32
33            if is_object_id {
34                object_id_conversions.push(quote! {
35                    if let Some(value) = doc.get_mut(stringify!(#field_name)) {
36                        if let mongo_orm::bson::Bson::String(s) = value {
37                            if let Ok(oid) = mongo_orm::bson::oid::ObjectId::parse_str(s) {
38                                *value = mongo_orm::bson::Bson::ObjectId(oid);
39                            }
40                        }
41                    }
42                });
43
44                object_id_reverse_conversions.push(quote! {
45                    if let Some(value) = doc.get_mut(stringify!(#field_name)) {
46                        if let mongo_orm::bson::Bson::ObjectId(oid) = value {
47                            *value = mongo_orm::bson::Bson::String(oid.to_string());
48                        }
49                    }
50                });
51            }
52        }
53    }
54
55    let primary_key_handling = if let Some(pk_field) = primary_key_field.clone() {
56        quote! {
57            // When converting to document
58            if let Some(id_value) = doc.remove(stringify!(#pk_field)) {
59                // Try to convert string to ObjectId if possible
60                let converted_value = if let mongo_orm::bson::Bson::String(s) = &id_value {
61                    if let Ok(oid) = mongo_orm::bson::oid::ObjectId::parse_str(s) {
62                        mongo_orm::bson::Bson::ObjectId(oid)
63                    } else {
64                        id_value
65                    }
66                } else {
67                    id_value
68                };
69                doc.insert("_id", converted_value);
70            }
71        }
72    } else {
73        quote! {}
74    };
75
76    let primary_key_reverse_handling = if let Some(pk_field) = primary_key_field {
77        quote! {
78            // Handle primary key if it exists
79            if let Some(id_value) = doc.remove("_id") {
80                // Convert ObjectId back to String if it is an ObjectId
81                let converted_value = if let mongo_orm::bson::Bson::ObjectId(oid) = id_value {
82                    mongo_orm::bson::Bson::String(oid.to_string())
83                } else {
84                    id_value
85                };
86                doc.insert(stringify!(#pk_field), converted_value);
87            }
88        }
89    } else {
90        quote! {
91            // Remove _id field if no primary key is defined
92            doc.remove("_id");
93        }
94    };
95
96    let expanded = quote! {
97        impl mongo_orm::entity::MongoEntity for #name {
98            fn to_document(&self) -> mongo_orm::bson::Document {
99                let value = mongo_orm::serde_json::to_value(self).unwrap();
100                let mut doc = mongo_orm::bson::Document::from_iter(
101                    value.as_object().unwrap().iter().map(|(k, v)| {
102                        (k.clone(), mongo_orm::bson::Bson::try_from(v.clone()).unwrap())
103                    })
104                );
105
106                // Handle ObjectId conversions
107                #(#object_id_conversions)*
108
109                // Handle primary key
110                #primary_key_handling
111
112                doc
113            }
114
115            fn from_document(mut doc: mongo_orm::bson::Document) -> Self {
116                #primary_key_reverse_handling
117
118                // Convert ObjectId back to String before deserializing
119                #(#object_id_reverse_conversions)*
120
121                let bson_value = mongo_orm::bson::Bson::Document(doc);
122                mongo_orm::serde_json::from_value(mongo_orm::serde_json::Value::from(bson_value)).unwrap()
123            }
124        }
125    };
126
127    TokenStream::from(expanded)
128}
mongo_orm_macro/lib.rs

mongo_orm_macro/
lib.rs
