1use async_trait::async_trait;
2use sea_orm::{
3 ActiveModelTrait, ActiveValue, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter,
4 TransactionError, TransactionTrait, prelude::Uuid,
5};
6use shield::{CreateEmailAddress, CreateUser, Storage, StorageError, UpdateUser};
7
8#[cfg(feature = "entity")]
9use crate::entities::entity;
10use crate::{
11 entities::{email_address, user},
12 user::User,
13};
14
15pub const SEA_ORM_STORAGE_ID: &str = "sea-orm";
16
17#[derive(Clone, Debug)]
18pub struct SeaOrmStorage {
19 pub(crate) database: DatabaseConnection,
20}
21
22impl SeaOrmStorage {
23 pub fn new(database: DatabaseConnection) -> Self {
24 Self { database }
25 }
26
27 pub(crate) fn parse_uuid(uuid: &str) -> Result<Uuid, StorageError> {
28 Uuid::parse_str(uuid).map_err(|err| StorageError::Validation(err.to_string()))
29 }
30}
31
32#[async_trait]
33impl Storage<User> for SeaOrmStorage {
34 fn id(&self) -> String {
35 SEA_ORM_STORAGE_ID.to_owned()
36 }
37
38 async fn user_by_id(&self, user_id: &str) -> Result<Option<User>, StorageError> {
39 #[cfg(feature = "entity")]
40 {
41 let user_and_entity = user::Entity::find_by_id(Self::parse_uuid(user_id)?)
42 .find_also_related(entity::Entity)
43 .one(&self.database)
44 .await
45 .map_err(|err| StorageError::Engine(err.to_string()))?;
46
47 match user_and_entity {
48 Some((user, Some(entity))) => {
49 Ok(Some(User::new(self.database.clone(), user, entity)))
50 }
51 Some((user, None)) => Err(StorageError::NotFound(
52 "Entity".to_owned(),
53 user.entity_id.to_string(),
54 )),
55 None => Ok(None),
56 }
57 }
58
59 #[cfg(not(feature = "entity"))]
60 {
61 user::Entity::find_by_id(Self::parse_uuid(user_id)?)
62 .one(&self.database)
63 .await
64 .map_err(|err| StorageError::Engine(err.to_string()))
65 .map(|user| user.map(|user| User::new(self.database.clone(), user)))
66 }
67 }
68
69 async fn user_by_email(&self, email: &str) -> Result<Option<User>, StorageError> {
70 #[cfg(feature = "entity")]
71 {
72 use sea_orm::{JoinType, QuerySelect, RelationTrait};
73
74 let user_and_entity = user::Entity::find()
75 .find_also_related(entity::Entity)
76 .join(JoinType::LeftJoin, entity::Relation::EmailAddress.def())
77 .filter(email_address::Column::Email.eq(email))
78 .one(&self.database)
79 .await
80 .map_err(|err| StorageError::Engine(err.to_string()))?;
81
82 match user_and_entity {
83 Some((user, Some(entity))) => {
84 Ok(Some(User::new(self.database.clone(), user, entity)))
85 }
86 Some((user, None)) => Err(StorageError::NotFound(
87 "Entity".to_owned(),
88 user.entity_id.to_string(),
89 )),
90 None => Ok(None),
91 }
92 }
93
94 #[cfg(not(feature = "entity"))]
95 {
96 user::Entity::find()
97 .left_join(email_address::Entity)
98 .filter(email_address::Column::Email.eq(email))
99 .one(&self.database)
100 .await
101 .map_err(|err| StorageError::Engine(err.to_string()))
102 .map(|user| user.map(|user| User::new(self.database.clone(), user)))
103 }
104 }
105
106 async fn create_user(
107 &self,
108 user: CreateUser,
109 email_address: CreateEmailAddress,
110 ) -> Result<User, StorageError> {
111 #[cfg(feature = "entity")]
112 type UserAndEntity = (user::Model, entity::Model);
113
114 #[cfg(not(feature = "entity"))]
115 type UserAndEntity = user::Model;
116
117 let user_and_entity = self
118 .database
119 .transaction::<_, UserAndEntity, StorageError>(|database_transaction| {
120 Box::pin(async move {
121 #[cfg(feature = "entity")]
122 {
123 let active_model = entity::ActiveModel {
124 name: ActiveValue::Set(user.name.unwrap_or_default()),
125 ..Default::default()
126 };
127
128 let entity = active_model
129 .insert(database_transaction)
130 .await
131 .map_err(|err| StorageError::Engine(err.to_string()))?;
132
133 let active_model = user::ActiveModel {
134 entity_id: ActiveValue::Set(entity.id),
135 ..Default::default()
136 };
137
138 let user = active_model
139 .insert(database_transaction)
140 .await
141 .map_err(|err| StorageError::Engine(err.to_string()))?;
142
143 let active_model = email_address::ActiveModel {
144 email: ActiveValue::Set(email_address.email),
145 is_primary: ActiveValue::Set(email_address.is_primary),
146 is_verified: ActiveValue::Set(email_address.is_verified),
147 verification_token: ActiveValue::Set(email_address.verification_token),
148 verification_token_expired_at: ActiveValue::Set(
149 email_address.verification_token_expired_at,
150 ),
151 verified_at: ActiveValue::Set(email_address.verified_at),
152 entity_id: ActiveValue::Set(entity.id),
153 ..Default::default()
154 };
155
156 active_model
157 .insert(database_transaction)
158 .await
159 .map_err(|err| StorageError::Engine(err.to_string()))?;
160
161 Ok((user, entity))
162 }
163
164 #[cfg(not(feature = "entity"))]
165 {
166 let active_model = user::ActiveModel {
167 name: ActiveValue::Set(user.name.unwrap_or_default()),
168 ..Default::default()
169 };
170
171 let user = active_model
172 .insert(database_transaction)
173 .await
174 .map_err(|err| StorageError::Engine(err.to_string()))?;
175
176 let active_model = email_address::ActiveModel {
177 email: ActiveValue::Set(email_address.email),
178 is_primary: ActiveValue::Set(email_address.is_primary),
179 is_verified: ActiveValue::Set(email_address.is_verified),
180 verification_token: ActiveValue::Set(email_address.verification_token),
181 verification_token_expired_at: ActiveValue::Set(
182 email_address.verification_token_expired_at,
183 ),
184 verified_at: ActiveValue::Set(email_address.verified_at),
185 user_id: ActiveValue::Set(user.id),
186 ..Default::default()
187 };
188
189 active_model
190 .insert(database_transaction)
191 .await
192 .map_err(|err| StorageError::Engine(err.to_string()))?;
193
194 Ok(user)
195 }
196 })
197 })
198 .await
199 .map_err(|err| match err {
200 TransactionError::Connection(err) => StorageError::Engine(err.to_string()),
201 TransactionError::Transaction(err) => err,
202 })?;
203
204 #[cfg(feature = "entity")]
205 {
206 let (user, entity) = user_and_entity;
207 Ok(User::new(self.database.clone(), user, entity))
208 }
209
210 #[cfg(not(feature = "entity"))]
211 {
212 let user = user_and_entity;
213 Ok(User::new(self.database.clone(), user))
214 }
215 }
216
217 async fn update_user(&self, user: UpdateUser) -> Result<User, StorageError> {
218 #[cfg(feature = "entity")]
219 type UserAndEntity = (user::Model, entity::Model);
220
221 #[cfg(not(feature = "entity"))]
222 type UserAndEntity = user::Model;
223
224 let user_and_entity = self
225 .database
226 .transaction::<_, UserAndEntity, StorageError>(|database_transaction| {
227 Box::pin(async move {
228 #[cfg(feature = "entity")]
229 {
230 use sea_orm::ModelTrait;
231
232 let user_entity = user::Entity::find()
233 .filter(user::Column::Id.eq(Self::parse_uuid(&user.id)?))
234 .one(database_transaction)
235 .await
236 .map_err(|err| StorageError::Engine(err.to_string()))?
237 .ok_or_else(|| {
238 StorageError::NotFound("User".to_owned(), user.id.clone())
239 })?;
240
241 let mut entity_active_model: entity::ActiveModel = user_entity
242 .find_related(entity::Entity)
243 .one(database_transaction)
244 .await
245 .map_err(|err| StorageError::Engine(err.to_string()))?
246 .ok_or_else(|| {
247 StorageError::NotFound(
248 "Entity".to_owned(),
249 user_entity.entity_id.to_string(),
250 )
251 })?
252 .into();
253
254 if let Some(Some(name)) = user.name {
255 entity_active_model.name = ActiveValue::Set(name);
256 }
257
258 let entity = entity_active_model
259 .update(database_transaction)
260 .await
261 .map_err(|err| StorageError::Engine(err.to_string()))?;
262
263 Ok((user_entity, entity))
264 }
265
266 #[cfg(not(feature = "entity"))]
267 {
268 let user_entity = user::Entity::find()
269 .filter(user::Column::Id.eq(Self::parse_uuid(&user.id)?))
270 .one(database_transaction)
271 .await
272 .map_err(|err| StorageError::Engine(err.to_string()))?
273 .ok_or_else(|| {
274 StorageError::NotFound("User".to_owned(), user.id.clone())
275 })?;
276
277 let mut user_active_model: user::ActiveModel = user_entity.into();
278
279 #[cfg(not(feature = "entity"))]
280 if let Some(Some(name)) = user.name {
281 user_active_model.name = ActiveValue::Set(name);
282 }
283
284 let user = user_active_model
285 .update(database_transaction)
286 .await
287 .map_err(|err| StorageError::Engine(err.to_string()))?;
288
289 Ok(user)
290 }
291 })
292 })
293 .await
294 .map_err(|err| match err {
295 TransactionError::Connection(err) => StorageError::Engine(err.to_string()),
296 TransactionError::Transaction(err) => err,
297 })?;
298
299 #[cfg(feature = "entity")]
300 {
301 let (user, entity) = user_and_entity;
302 Ok(User::new(self.database.clone(), user, entity))
303 }
304
305 #[cfg(not(feature = "entity"))]
306 {
307 let user = user_and_entity;
308 Ok(User::new(self.database.clone(), user))
309 }
310 }
311
312 async fn delete_user(&self, user_id: &str) -> Result<(), StorageError> {
313 user::Entity::delete_by_id(Self::parse_uuid(user_id)?)
314 .exec(&self.database)
315 .await
316 .map_err(|err| StorageError::Engine(err.to_string()))
317 .map(|_| ())
318 }
319}