1use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use mas_data_model::{
11 BrowserSession, Clock, CompatSession, CompatSsoLogin, CompatSsoLoginState, UlidExt as _,
12};
13use mas_storage::{
14 Page, Pagination,
15 compat::{CompatSsoLoginFilter, CompatSsoLoginRepository},
16 pagination::Node,
17};
18use rand::RngCore;
19use sea_query::{Expr, ExprTrait, PostgresQueryBuilder, Query, enum_def};
20use sea_query_sqlx::SqlxBinder;
21use sqlx::PgConnection;
22use ulid::Ulid;
23use url::Url;
24use uuid::Uuid;
25
26use crate::{
27 DatabaseError, DatabaseInconsistencyError,
28 filter::{Filter, StatementExt},
29 iden::{CompatSsoLogins, UserSessions},
30 pagination::QueryBuilderExt,
31 tracing::ExecuteExt,
32};
33
34pub struct PgCompatSsoLoginRepository<'c> {
37 conn: &'c mut PgConnection,
38}
39
40impl<'c> PgCompatSsoLoginRepository<'c> {
41 pub fn new(conn: &'c mut PgConnection) -> Self {
44 Self { conn }
45 }
46}
47
48#[derive(sqlx::FromRow, Debug)]
49#[enum_def]
50struct CompatSsoLoginLookup {
51 compat_sso_login_id: Uuid,
52 login_token: String,
53 redirect_uri: String,
54 created_at: DateTime<Utc>,
55 fulfilled_at: Option<DateTime<Utc>>,
56 exchanged_at: Option<DateTime<Utc>>,
57 user_session_id: Option<Uuid>,
58 compat_session_id: Option<Uuid>,
59}
60
61impl Node<Ulid> for CompatSsoLoginLookup {
62 fn cursor(&self) -> Ulid {
63 self.compat_sso_login_id.into()
64 }
65}
66
67impl TryFrom<CompatSsoLoginLookup> for CompatSsoLogin {
68 type Error = DatabaseInconsistencyError;
69
70 fn try_from(res: CompatSsoLoginLookup) -> Result<Self, Self::Error> {
71 let id = res.compat_sso_login_id.into();
72 let redirect_uri = Url::parse(&res.redirect_uri).map_err(|e| {
73 DatabaseInconsistencyError::on("compat_sso_logins")
74 .column("redirect_uri")
75 .row(id)
76 .source(e)
77 })?;
78
79 let state = match (
80 res.fulfilled_at,
81 res.exchanged_at,
82 res.user_session_id,
83 res.compat_session_id,
84 ) {
85 (None, None, None, None) => CompatSsoLoginState::Pending,
86 (Some(fulfilled_at), None, Some(browser_session_id), None) => {
87 CompatSsoLoginState::Fulfilled {
88 fulfilled_at,
89 browser_session_id: browser_session_id.into(),
90 }
91 }
92 (Some(fulfilled_at), Some(exchanged_at), _, Some(compat_session_id)) => {
93 CompatSsoLoginState::Exchanged {
94 fulfilled_at,
95 exchanged_at,
96 compat_session_id: compat_session_id.into(),
97 }
98 }
99 _ => return Err(DatabaseInconsistencyError::on("compat_sso_logins").row(id)),
100 };
101
102 Ok(CompatSsoLogin {
103 id,
104 login_token: res.login_token,
105 redirect_uri,
106 created_at: res.created_at,
107 state,
108 })
109 }
110}
111
112impl Filter for CompatSsoLoginFilter<'_> {
113 fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
114 sea_query::Condition::all()
115 .add_option(self.user().map(|user| {
116 Expr::exists(
117 Query::select()
118 .expr(Expr::cust("1"))
119 .from(UserSessions::Table)
120 .and_where(
121 Expr::col((UserSessions::Table, UserSessions::UserId))
122 .eq(Uuid::from(user.id)),
123 )
124 .and_where(
125 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::UserSessionId))
126 .equals((UserSessions::Table, UserSessions::UserSessionId)),
127 )
128 .take(),
129 )
130 }))
131 .add_option(self.state().map(|state| {
132 if state.is_exchanged() {
133 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::ExchangedAt)).is_not_null()
134 } else if state.is_fulfilled() {
135 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::FulfilledAt))
136 .is_not_null()
137 .and(
138 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::ExchangedAt))
139 .is_null(),
140 )
141 } else {
142 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::FulfilledAt)).is_null()
143 }
144 }))
145 }
146}
147
148#[async_trait]
149impl CompatSsoLoginRepository for PgCompatSsoLoginRepository<'_> {
150 type Error = DatabaseError;
151
152 #[tracing::instrument(
153 name = "db.compat_sso_login.lookup",
154 skip_all,
155 fields(
156 db.query.text,
157 compat_sso_login.id = %id,
158 ),
159 err,
160 )]
161 async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatSsoLogin>, Self::Error> {
162 let res = sqlx::query_as!(
163 CompatSsoLoginLookup,
164 r#"
165 SELECT compat_sso_login_id
166 , login_token
167 , redirect_uri
168 , created_at
169 , fulfilled_at
170 , exchanged_at
171 , compat_session_id
172 , user_session_id
173
174 FROM compat_sso_logins
175 WHERE compat_sso_login_id = $1
176 "#,
177 Uuid::from(id),
178 )
179 .traced()
180 .fetch_optional(&mut *self.conn)
181 .await?;
182
183 let Some(res) = res else { return Ok(None) };
184
185 Ok(Some(res.try_into()?))
186 }
187
188 #[tracing::instrument(
189 name = "db.compat_sso_login.find_for_session",
190 skip_all,
191 fields(
192 db.query.text,
193 %compat_session.id,
194 ),
195 err,
196 )]
197 async fn find_for_session(
198 &mut self,
199 compat_session: &CompatSession,
200 ) -> Result<Option<CompatSsoLogin>, Self::Error> {
201 let res = sqlx::query_as!(
202 CompatSsoLoginLookup,
203 r#"
204 SELECT compat_sso_login_id
205 , login_token
206 , redirect_uri
207 , created_at
208 , fulfilled_at
209 , exchanged_at
210 , compat_session_id
211 , user_session_id
212
213 FROM compat_sso_logins
214 WHERE compat_session_id = $1
215 "#,
216 Uuid::from(compat_session.id),
217 )
218 .traced()
219 .fetch_optional(&mut *self.conn)
220 .await?;
221
222 let Some(res) = res else { return Ok(None) };
223
224 Ok(Some(res.try_into()?))
225 }
226
227 #[tracing::instrument(
228 name = "db.compat_sso_login.find_by_token",
229 skip_all,
230 fields(
231 db.query.text,
232 ),
233 err,
234 )]
235 async fn find_by_token(
236 &mut self,
237 login_token: &str,
238 ) -> Result<Option<CompatSsoLogin>, Self::Error> {
239 let res = sqlx::query_as!(
240 CompatSsoLoginLookup,
241 r#"
242 SELECT compat_sso_login_id
243 , login_token
244 , redirect_uri
245 , created_at
246 , fulfilled_at
247 , exchanged_at
248 , compat_session_id
249 , user_session_id
250
251 FROM compat_sso_logins
252 WHERE login_token = $1
253 "#,
254 login_token,
255 )
256 .traced()
257 .fetch_optional(&mut *self.conn)
258 .await?;
259
260 let Some(res) = res else { return Ok(None) };
261
262 Ok(Some(res.try_into()?))
263 }
264
265 #[tracing::instrument(
266 name = "db.compat_sso_login.add",
267 skip_all,
268 fields(
269 db.query.text,
270 compat_sso_login.id,
271 compat_sso_login.redirect_uri = %redirect_uri,
272 ),
273 err,
274 )]
275 async fn add(
276 &mut self,
277 rng: &mut (dyn RngCore + Send),
278 clock: &dyn Clock,
279 login_token: String,
280 redirect_uri: Url,
281 ) -> Result<CompatSsoLogin, Self::Error> {
282 let created_at = clock.now();
283 let id = Ulid::from_datetime_with_rng(created_at, rng);
284 tracing::Span::current().record("compat_sso_login.id", tracing::field::display(id));
285
286 sqlx::query!(
287 r#"
288 INSERT INTO compat_sso_logins
289 (compat_sso_login_id, login_token, redirect_uri, created_at)
290 VALUES ($1, $2, $3, $4)
291 "#,
292 Uuid::from(id),
293 &login_token,
294 redirect_uri.as_str(),
295 created_at,
296 )
297 .traced()
298 .execute(&mut *self.conn)
299 .await?;
300
301 Ok(CompatSsoLogin {
302 id,
303 login_token,
304 redirect_uri,
305 created_at,
306 state: CompatSsoLoginState::default(),
307 })
308 }
309
310 #[tracing::instrument(
311 name = "db.compat_sso_login.fulfill",
312 skip_all,
313 fields(
314 db.query.text,
315 %compat_sso_login.id,
316 %browser_session.id,
317 user.id = %browser_session.user.id,
318 ),
319 err,
320 )]
321 async fn fulfill(
322 &mut self,
323 clock: &dyn Clock,
324 compat_sso_login: CompatSsoLogin,
325 browser_session: &BrowserSession,
326 ) -> Result<CompatSsoLogin, Self::Error> {
327 let fulfilled_at = clock.now();
328 let compat_sso_login = compat_sso_login
329 .fulfill(fulfilled_at, browser_session)
330 .map_err(DatabaseError::to_invalid_operation)?;
331
332 let res = sqlx::query!(
333 r#"
334 UPDATE compat_sso_logins
335 SET
336 user_session_id = $2,
337 fulfilled_at = $3
338 WHERE
339 compat_sso_login_id = $1
340 "#,
341 Uuid::from(compat_sso_login.id),
342 Uuid::from(browser_session.id),
343 fulfilled_at,
344 )
345 .traced()
346 .execute(&mut *self.conn)
347 .await?;
348
349 DatabaseError::ensure_affected_rows(&res, 1)?;
350
351 Ok(compat_sso_login)
352 }
353
354 #[tracing::instrument(
355 name = "db.compat_sso_login.exchange",
356 skip_all,
357 fields(
358 db.query.text,
359 %compat_sso_login.id,
360 %compat_session.id,
361 compat_session.device.id = compat_session.device.as_ref().map(mas_data_model::Device::as_str),
362 ),
363 err,
364 )]
365 async fn exchange(
366 &mut self,
367 clock: &dyn Clock,
368 compat_sso_login: CompatSsoLogin,
369 compat_session: &CompatSession,
370 ) -> Result<CompatSsoLogin, Self::Error> {
371 let exchanged_at = clock.now();
372 let compat_sso_login = compat_sso_login
373 .exchange(exchanged_at, compat_session)
374 .map_err(DatabaseError::to_invalid_operation)?;
375
376 let res = sqlx::query!(
377 r#"
378 UPDATE compat_sso_logins
379 SET
380 exchanged_at = $2,
381 compat_session_id = $3
382 WHERE
383 compat_sso_login_id = $1
384 "#,
385 Uuid::from(compat_sso_login.id),
386 exchanged_at,
387 Uuid::from(compat_session.id),
388 )
389 .traced()
390 .execute(&mut *self.conn)
391 .await?;
392
393 DatabaseError::ensure_affected_rows(&res, 1)?;
394
395 Ok(compat_sso_login)
396 }
397
398 #[tracing::instrument(
399 name = "db.compat_sso_login.list",
400 skip_all,
401 fields(
402 db.query.text,
403 ),
404 err
405 )]
406 async fn list(
407 &mut self,
408 filter: CompatSsoLoginFilter<'_>,
409 pagination: Pagination,
410 ) -> Result<Page<CompatSsoLogin>, Self::Error> {
411 let (sql, arguments) = Query::select()
412 .expr_as(
413 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)),
414 CompatSsoLoginLookupIden::CompatSsoLoginId,
415 )
416 .expr_as(
417 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSessionId)),
418 CompatSsoLoginLookupIden::CompatSessionId,
419 )
420 .expr_as(
421 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::UserSessionId)),
422 CompatSsoLoginLookupIden::UserSessionId,
423 )
424 .expr_as(
425 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::LoginToken)),
426 CompatSsoLoginLookupIden::LoginToken,
427 )
428 .expr_as(
429 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::RedirectUri)),
430 CompatSsoLoginLookupIden::RedirectUri,
431 )
432 .expr_as(
433 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CreatedAt)),
434 CompatSsoLoginLookupIden::CreatedAt,
435 )
436 .expr_as(
437 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::FulfilledAt)),
438 CompatSsoLoginLookupIden::FulfilledAt,
439 )
440 .expr_as(
441 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::ExchangedAt)),
442 CompatSsoLoginLookupIden::ExchangedAt,
443 )
444 .from(CompatSsoLogins::Table)
445 .apply_filter(filter)
446 .generate_pagination(
447 (CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId),
448 pagination,
449 )
450 .build_sqlx(PostgresQueryBuilder);
451
452 let edges: Vec<CompatSsoLoginLookup> = sqlx::query_as_with(&sql, arguments)
453 .traced()
454 .fetch_all(&mut *self.conn)
455 .await?;
456
457 let page = pagination.process(edges).try_map(TryFrom::try_from)?;
458
459 Ok(page)
460 }
461
462 #[tracing::instrument(
463 name = "db.compat_sso_login.count",
464 skip_all,
465 fields(
466 db.query.text,
467 ),
468 err
469 )]
470 async fn count(&mut self, filter: CompatSsoLoginFilter<'_>) -> Result<usize, Self::Error> {
471 let (sql, arguments) = Query::select()
472 .expr(Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)).count())
473 .from(CompatSsoLogins::Table)
474 .apply_filter(filter)
475 .build_sqlx(PostgresQueryBuilder);
476
477 let count: i64 = sqlx::query_scalar_with(&sql, arguments)
478 .traced()
479 .fetch_one(&mut *self.conn)
480 .await?;
481
482 count
483 .try_into()
484 .map_err(DatabaseError::to_invalid_operation)
485 }
486}