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::Rng;
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 )]
160 async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatSsoLogin>, Self::Error> {
161 let res = sqlx::query_as!(
162 CompatSsoLoginLookup,
163 r#"
164 SELECT compat_sso_login_id
165 , login_token
166 , redirect_uri
167 , created_at
168 , fulfilled_at
169 , exchanged_at
170 , compat_session_id
171 , user_session_id
172
173 FROM compat_sso_logins
174 WHERE compat_sso_login_id = $1
175 "#,
176 Uuid::from(id),
177 )
178 .traced()
179 .fetch_optional(&mut *self.conn)
180 .await?;
181
182 let Some(res) = res else { return Ok(None) };
183
184 Ok(Some(res.try_into()?))
185 }
186
187 #[tracing::instrument(
188 name = "db.compat_sso_login.find_for_session",
189 skip_all,
190 fields(
191 db.query.text,
192 %compat_session.id,
193 ),
194 )]
195 async fn find_for_session(
196 &mut self,
197 compat_session: &CompatSession,
198 ) -> Result<Option<CompatSsoLogin>, Self::Error> {
199 let res = sqlx::query_as!(
200 CompatSsoLoginLookup,
201 r#"
202 SELECT compat_sso_login_id
203 , login_token
204 , redirect_uri
205 , created_at
206 , fulfilled_at
207 , exchanged_at
208 , compat_session_id
209 , user_session_id
210
211 FROM compat_sso_logins
212 WHERE compat_session_id = $1
213 "#,
214 Uuid::from(compat_session.id),
215 )
216 .traced()
217 .fetch_optional(&mut *self.conn)
218 .await?;
219
220 let Some(res) = res else { return Ok(None) };
221
222 Ok(Some(res.try_into()?))
223 }
224
225 #[tracing::instrument(
226 name = "db.compat_sso_login.find_by_token",
227 skip_all,
228 fields(
229 db.query.text,
230 ),
231 )]
232 async fn find_by_token(
233 &mut self,
234 login_token: &str,
235 ) -> Result<Option<CompatSsoLogin>, Self::Error> {
236 let res = sqlx::query_as!(
237 CompatSsoLoginLookup,
238 r#"
239 SELECT compat_sso_login_id
240 , login_token
241 , redirect_uri
242 , created_at
243 , fulfilled_at
244 , exchanged_at
245 , compat_session_id
246 , user_session_id
247
248 FROM compat_sso_logins
249 WHERE login_token = $1
250 "#,
251 login_token,
252 )
253 .traced()
254 .fetch_optional(&mut *self.conn)
255 .await?;
256
257 let Some(res) = res else { return Ok(None) };
258
259 Ok(Some(res.try_into()?))
260 }
261
262 #[tracing::instrument(
263 name = "db.compat_sso_login.add",
264 skip_all,
265 fields(
266 db.query.text,
267 compat_sso_login.id,
268 compat_sso_login.redirect_uri = %redirect_uri,
269 ),
270 )]
271 async fn add(
272 &mut self,
273 rng: &mut (dyn Rng + Send),
274 clock: &dyn Clock,
275 login_token: String,
276 redirect_uri: Url,
277 ) -> Result<CompatSsoLogin, Self::Error> {
278 let created_at = clock.now();
279 let id = Ulid::from_datetime_with_rng(created_at, rng);
280 tracing::Span::current().record("compat_sso_login.id", tracing::field::display(id));
281
282 sqlx::query!(
283 r#"
284 INSERT INTO compat_sso_logins
285 (compat_sso_login_id, login_token, redirect_uri, created_at)
286 VALUES ($1, $2, $3, $4)
287 "#,
288 Uuid::from(id),
289 &login_token,
290 redirect_uri.as_str(),
291 created_at,
292 )
293 .traced()
294 .execute(&mut *self.conn)
295 .await?;
296
297 Ok(CompatSsoLogin {
298 id,
299 login_token,
300 redirect_uri,
301 created_at,
302 state: CompatSsoLoginState::default(),
303 })
304 }
305
306 #[tracing::instrument(
307 name = "db.compat_sso_login.fulfill",
308 skip_all,
309 fields(
310 db.query.text,
311 %compat_sso_login.id,
312 %browser_session.id,
313 user.id = %browser_session.user.id,
314 ),
315 )]
316 async fn fulfill(
317 &mut self,
318 clock: &dyn Clock,
319 compat_sso_login: CompatSsoLogin,
320 browser_session: &BrowserSession,
321 ) -> Result<CompatSsoLogin, Self::Error> {
322 let fulfilled_at = clock.now();
323 let compat_sso_login = compat_sso_login
324 .fulfill(fulfilled_at, browser_session)
325 .map_err(DatabaseError::to_invalid_operation)?;
326
327 let res = sqlx::query!(
328 r#"
329 UPDATE compat_sso_logins
330 SET
331 user_session_id = $2,
332 fulfilled_at = $3
333 WHERE
334 compat_sso_login_id = $1
335 "#,
336 Uuid::from(compat_sso_login.id),
337 Uuid::from(browser_session.id),
338 fulfilled_at,
339 )
340 .traced()
341 .execute(&mut *self.conn)
342 .await?;
343
344 DatabaseError::ensure_affected_rows(&res, 1)?;
345
346 Ok(compat_sso_login)
347 }
348
349 #[tracing::instrument(
350 name = "db.compat_sso_login.exchange",
351 skip_all,
352 fields(
353 db.query.text,
354 %compat_sso_login.id,
355 %compat_session.id,
356 compat_session.device.id = compat_session.device.as_ref().map(mas_data_model::Device::as_str),
357 ),
358 )]
359 async fn exchange(
360 &mut self,
361 clock: &dyn Clock,
362 compat_sso_login: CompatSsoLogin,
363 compat_session: &CompatSession,
364 ) -> Result<CompatSsoLogin, Self::Error> {
365 let exchanged_at = clock.now();
366 let compat_sso_login = compat_sso_login
367 .exchange(exchanged_at, compat_session)
368 .map_err(DatabaseError::to_invalid_operation)?;
369
370 let res = sqlx::query!(
371 r#"
372 UPDATE compat_sso_logins
373 SET
374 exchanged_at = $2,
375 compat_session_id = $3
376 WHERE
377 compat_sso_login_id = $1
378 "#,
379 Uuid::from(compat_sso_login.id),
380 exchanged_at,
381 Uuid::from(compat_session.id),
382 )
383 .traced()
384 .execute(&mut *self.conn)
385 .await?;
386
387 DatabaseError::ensure_affected_rows(&res, 1)?;
388
389 Ok(compat_sso_login)
390 }
391
392 #[tracing::instrument(
393 name = "db.compat_sso_login.list",
394 skip_all,
395 fields(
396 db.query.text,
397 ),
398 )]
399 async fn list(
400 &mut self,
401 filter: CompatSsoLoginFilter<'_>,
402 pagination: Pagination,
403 ) -> Result<Page<CompatSsoLogin>, Self::Error> {
404 let (sql, arguments) = Query::select()
405 .expr_as(
406 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)),
407 CompatSsoLoginLookupIden::CompatSsoLoginId,
408 )
409 .expr_as(
410 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSessionId)),
411 CompatSsoLoginLookupIden::CompatSessionId,
412 )
413 .expr_as(
414 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::UserSessionId)),
415 CompatSsoLoginLookupIden::UserSessionId,
416 )
417 .expr_as(
418 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::LoginToken)),
419 CompatSsoLoginLookupIden::LoginToken,
420 )
421 .expr_as(
422 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::RedirectUri)),
423 CompatSsoLoginLookupIden::RedirectUri,
424 )
425 .expr_as(
426 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CreatedAt)),
427 CompatSsoLoginLookupIden::CreatedAt,
428 )
429 .expr_as(
430 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::FulfilledAt)),
431 CompatSsoLoginLookupIden::FulfilledAt,
432 )
433 .expr_as(
434 Expr::col((CompatSsoLogins::Table, CompatSsoLogins::ExchangedAt)),
435 CompatSsoLoginLookupIden::ExchangedAt,
436 )
437 .from(CompatSsoLogins::Table)
438 .apply_filter(filter)
439 .generate_pagination(
440 (CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId),
441 pagination,
442 )
443 .build_sqlx(PostgresQueryBuilder);
444
445 let edges: Vec<CompatSsoLoginLookup> = sqlx::query_as_with(&sql, arguments)
446 .traced()
447 .fetch_all(&mut *self.conn)
448 .await?;
449
450 let page = pagination.process(edges).try_map(TryFrom::try_from)?;
451
452 Ok(page)
453 }
454
455 #[tracing::instrument(
456 name = "db.compat_sso_login.count",
457 skip_all,
458 fields(
459 db.query.text,
460 ),
461 )]
462 async fn count(&mut self, filter: CompatSsoLoginFilter<'_>) -> Result<usize, Self::Error> {
463 let (sql, arguments) = Query::select()
464 .expr(Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)).count())
465 .from(CompatSsoLogins::Table)
466 .apply_filter(filter)
467 .build_sqlx(PostgresQueryBuilder);
468
469 let count: i64 = sqlx::query_scalar_with(&sql, arguments)
470 .traced()
471 .fetch_one(&mut *self.conn)
472 .await?;
473
474 count
475 .try_into()
476 .map_err(DatabaseError::to_invalid_operation)
477 }
478}