1use async_trait::async_trait;
2use sea_orm::{
3 prelude::Uuid, ActiveModelTrait, ActiveValue, ColumnTrait, DatabaseConnection, EntityTrait,
4 QueryFilter, TransactionError, TransactionTrait,
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, user::Relation::Entity.def())
77 .join(JoinType::LeftJoin, entity::Relation::EmailAddress.def())
78 .filter(email_address::Column::Email.eq(email))
79 .one(&self.database)
80 .await
81 .map_err(|err| StorageError::Engine(err.to_string()))?;
82
83 match user_and_entity {
84 Some((user, Some(entity))) => {
85 Ok(Some(User::new(self.database.clone(), user, entity)))
86 }
87 Some((user, None)) => Err(StorageError::NotFound(
88 "Entity".to_owned(),
89 user.entity_id.to_string(),
90 )),
91 None => Ok(None),
92 }
93 }
94
95 #[cfg(not(feature = "entity"))]
96 {
97 user::Entity::find()
98 .left_join(email_address::Entity)
99 .filter(email_address::Column::Email.eq(email))
100 .one(&self.database)
101 .await
102 .map_err(|err| StorageError::Engine(err.to_string()))
103 .map(|user| user.map(|user| User::new(self.database.clone(), user)))
104 }
105 }
106
107 async fn create_user(
108 &self,
109 user: CreateUser,
110 email_address: CreateEmailAddress,
111 ) -> Result<User, StorageError> {
112 #[cfg(feature = "entity")]
113 type UserAndEntity = (user::Model, entity::Model);
114
115 #[cfg(not(feature = "entity"))]
116 type UserAndEntity = user::Model;
117
118 let user_and_entity = self
119 .database
120 .transaction::<_, UserAndEntity, StorageError>(|database_transaction| {
121 Box::pin(async move {
122 #[cfg(feature = "entity")]
123 {
124 let active_model = entity::ActiveModel {
125 name: ActiveValue::Set(user.name.unwrap_or_default()),
126 ..Default::default()
127 };
128
129 let entity = active_model
130 .insert(database_transaction)
131 .await
132 .map_err(|err| StorageError::Engine(err.to_string()))?;
133
134 let active_model = user::ActiveModel {
135 entity_id: ActiveValue::Set(entity.id),
136 ..Default::default()
137 };
138
139 let user = active_model
140 .insert(database_transaction)
141 .await
142 .map_err(|err| StorageError::Engine(err.to_string()))?;
143
144 let active_model = email_address::ActiveModel {
145 email: ActiveValue::Set(email_address.email),
146 is_primary: ActiveValue::Set(email_address.is_primary),
147 is_verified: ActiveValue::Set(email_address.is_verified),
148 verification_token: ActiveValue::Set(email_address.verification_token),
149 verification_token_expired_at: ActiveValue::Set(
150 email_address.verification_token_expired_at,
151 ),
152 verified_at: ActiveValue::Set(email_address.verified_at),
153 entity_id: ActiveValue::Set(entity.id),
154 ..Default::default()
155 };
156
157 active_model
158 .insert(database_transaction)
159 .await
160 .map_err(|err| StorageError::Engine(err.to_string()))?;
161
162 Ok((user, entity))
163 }
164
165 #[cfg(not(feature = "entity"))]
166 {
167 let active_model = user::ActiveModel {
168 name: ActiveValue::Set(user.name.unwrap_or_default()),
169 ..Default::default()
170 };
171
172 let user = active_model
173 .insert(database_transaction)
174 .await
175 .map_err(|err| StorageError::Engine(err.to_string()))?;
176
177 let active_model = email_address::ActiveModel {
178 email: ActiveValue::Set(email_address.email),
179 is_primary: ActiveValue::Set(email_address.is_primary),
180 is_verified: ActiveValue::Set(email_address.is_verified),
181 verification_token: ActiveValue::Set(email_address.verification_token),
182 verification_token_expired_at: ActiveValue::Set(
183 email_address.verification_token_expired_at,
184 ),
185 verified_at: ActiveValue::Set(email_address.verified_at),
186 user_id: ActiveValue::Set(user.id),
187 ..Default::default()
188 };
189
190 active_model
191 .insert(database_transaction)
192 .await
193 .map_err(|err| StorageError::Engine(err.to_string()))?;
194
195 Ok(user)
196 }
197 })
198 })
199 .await
200 .map_err(|err| match err {
201 TransactionError::Connection(err) => StorageError::Engine(err.to_string()),
202 TransactionError::Transaction(err) => err,
203 })?;
204
205 #[cfg(feature = "entity")]
206 {
207 let (user, entity) = user_and_entity;
208 Ok(User::new(self.database.clone(), user, entity))
209 }
210
211 #[cfg(not(feature = "entity"))]
212 {
213 let user = user_and_entity;
214 Ok(User::new(self.database.clone(), user))
215 }
216 }
217
218 async fn update_user(&self, user: UpdateUser) -> Result<User, StorageError> {
219 #[cfg(feature = "entity")]
220 type UserAndEntity = (user::Model, entity::Model);
221
222 #[cfg(not(feature = "entity"))]
223 type UserAndEntity = user::Model;
224
225 let user_and_entity = self
226 .database
227 .transaction::<_, UserAndEntity, StorageError>(|database_transaction| {
228 Box::pin(async move {
229 #[cfg(feature = "entity")]
230 {
231 use sea_orm::ModelTrait;
232
233 let user_entity = user::Entity::find()
234 .filter(user::Column::Id.eq(Self::parse_uuid(&user.id)?))
235 .one(database_transaction)
236 .await
237 .map_err(|err| StorageError::Engine(err.to_string()))?
238 .ok_or_else(|| {
239 StorageError::NotFound("User".to_owned(), user.id.clone())
240 })?;
241
242 let mut entity_active_model: entity::ActiveModel = user_entity
243 .find_related(entity::Entity)
244 .one(database_transaction)
245 .await
246 .map_err(|err| StorageError::Engine(err.to_string()))?
247 .ok_or_else(|| {
248 StorageError::NotFound(
249 "Entity".to_owned(),
250 user_entity.entity_id.to_string(),
251 )
252 })?
253 .into();
254
255 if let Some(Some(name)) = user.name {
256 entity_active_model.name = ActiveValue::Set(name);
257 }
258
259 let entity = entity_active_model
260 .update(database_transaction)
261 .await
262 .map_err(|err| StorageError::Engine(err.to_string()))?;
263
264 Ok((user_entity, entity))
265 }
266
267 #[cfg(not(feature = "entity"))]
268 {
269 let user_entity = user::Entity::find()
270 .filter(user::Column::Id.eq(Self::parse_uuid(&user.id)?))
271 .one(database_transaction)
272 .await
273 .map_err(|err| StorageError::Engine(err.to_string()))?
274 .ok_or_else(|| {
275 StorageError::NotFound("User".to_owned(), user.id.clone())
276 })?;
277
278 let mut user_active_model: user::ActiveModel = user_entity.into();
279
280 #[cfg(not(feature = "entity"))]
281 if let Some(Some(name)) = user.name {
282 user_active_model.name = ActiveValue::Set(name);
283 }
284
285 let user = user_active_model
286 .update(database_transaction)
287 .await
288 .map_err(|err| StorageError::Engine(err.to_string()))?;
289
290 Ok(user)
291 }
292 })
293 })
294 .await
295 .map_err(|err| match err {
296 TransactionError::Connection(err) => StorageError::Engine(err.to_string()),
297 TransactionError::Transaction(err) => err,
298 })?;
299
300 #[cfg(feature = "entity")]
301 {
302 let (user, entity) = user_and_entity;
303 Ok(User::new(self.database.clone(), user, entity))
304 }
305
306 #[cfg(not(feature = "entity"))]
307 {
308 let user = user_and_entity;
309 Ok(User::new(self.database.clone(), user))
310 }
311 }
312
313 async fn delete_user(&self, user_id: &str) -> Result<(), StorageError> {
314 user::Entity::delete_by_id(Self::parse_uuid(user_id)?)
315 .exec(&self.database)
316 .await
317 .map_err(|err| StorageError::Engine(err.to_string()))
318 .map(|_| ())
319 }
320}