1use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use mas_data_model::{Clock, UlidExt as _, UpstreamOAuthLink, UpstreamOAuthProvider, User};
11use mas_storage::{
12 Page, Pagination,
13 pagination::Node,
14 upstream_oauth2::{UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository},
15};
16use opentelemetry_semantic_conventions::trace::DB_QUERY_TEXT;
17use rand::RngCore;
18use sea_query::{
19 Expr, ExprTrait, PostgresQueryBuilder, Query, enum_def, extension::postgres::PgExpr as _,
20};
21use sea_query_sqlx::SqlxBinder;
22use sqlx::PgConnection;
23use tracing::Instrument;
24use ulid::Ulid;
25use uuid::Uuid;
26
27use crate::{
28 DatabaseError,
29 filter::{Filter, StatementExt},
30 iden::{UpstreamOAuthLinks, UpstreamOAuthProviders},
31 pagination::QueryBuilderExt,
32 tracing::ExecuteExt,
33};
34
35pub struct PgUpstreamOAuthLinkRepository<'c> {
38 conn: &'c mut PgConnection,
39}
40
41impl<'c> PgUpstreamOAuthLinkRepository<'c> {
42 pub fn new(conn: &'c mut PgConnection) -> Self {
45 Self { conn }
46 }
47}
48
49#[derive(sqlx::FromRow)]
50#[enum_def]
51struct LinkLookup {
52 upstream_oauth_link_id: Uuid,
53 upstream_oauth_provider_id: Uuid,
54 user_id: Option<Uuid>,
55 subject: String,
56 human_account_name: Option<String>,
57 created_at: DateTime<Utc>,
58}
59
60impl Node<Ulid> for LinkLookup {
61 fn cursor(&self) -> Ulid {
62 self.upstream_oauth_link_id.into()
63 }
64}
65
66impl From<LinkLookup> for UpstreamOAuthLink {
67 fn from(value: LinkLookup) -> Self {
68 UpstreamOAuthLink {
69 id: Ulid::from(value.upstream_oauth_link_id),
70 provider_id: Ulid::from(value.upstream_oauth_provider_id),
71 user_id: value.user_id.map(Ulid::from),
72 subject: value.subject,
73 human_account_name: value.human_account_name,
74 created_at: value.created_at,
75 }
76 }
77}
78
79impl Filter for UpstreamOAuthLinkFilter<'_> {
80 fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
81 sea_query::Condition::all()
82 .add_option(self.user().map(|user| {
83 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::UserId))
84 .eq(Uuid::from(user.id))
85 }))
86 .add_option(self.provider().map(|provider| {
87 Expr::col((
88 UpstreamOAuthLinks::Table,
89 UpstreamOAuthLinks::UpstreamOAuthProviderId,
90 ))
91 .eq(Uuid::from(provider.id))
92 }))
93 .add_option(self.provider_enabled().map(|enabled| {
94 Expr::col((
95 UpstreamOAuthLinks::Table,
96 UpstreamOAuthLinks::UpstreamOAuthProviderId,
97 ))
98 .eq(Expr::any(
99 Query::select()
100 .expr(Expr::col((
101 UpstreamOAuthProviders::Table,
102 UpstreamOAuthProviders::UpstreamOAuthProviderId,
103 )))
104 .from(UpstreamOAuthProviders::Table)
105 .and_where(
106 Expr::col((
107 UpstreamOAuthProviders::Table,
108 UpstreamOAuthProviders::DisabledAt,
109 ))
110 .is_null()
111 .eq(enabled),
112 )
113 .take(),
114 ))
115 }))
116 .add_option(self.subject().map(|subject| {
117 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::Subject)).eq(subject)
118 }))
119 .add_option(self.human_account_name().map(|human_account_name| {
120 Expr::col((
121 UpstreamOAuthLinks::Table,
122 UpstreamOAuthLinks::HumanAccountName,
123 ))
124 .ilike(format!("%{human_account_name}%"))
125 }))
126 }
127}
128
129#[async_trait]
130impl UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'_> {
131 type Error = DatabaseError;
132
133 #[tracing::instrument(
134 name = "db.upstream_oauth_link.lookup",
135 skip_all,
136 fields(
137 db.query.text,
138 upstream_oauth_link.id = %id,
139 ),
140 err,
141 )]
142 async fn lookup(&mut self, id: Ulid) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
143 let res = sqlx::query_as!(
144 LinkLookup,
145 r#"
146 SELECT
147 upstream_oauth_link_id,
148 upstream_oauth_provider_id,
149 user_id,
150 subject,
151 human_account_name,
152 created_at
153 FROM upstream_oauth_links
154 WHERE upstream_oauth_link_id = $1
155 "#,
156 Uuid::from(id),
157 )
158 .traced()
159 .fetch_optional(&mut *self.conn)
160 .await?
161 .map(Into::into);
162
163 Ok(res)
164 }
165
166 #[tracing::instrument(
167 name = "db.upstream_oauth_link.find_by_subject",
168 skip_all,
169 fields(
170 db.query.text,
171 upstream_oauth_link.subject = subject,
172 %upstream_oauth_provider.id,
173 upstream_oauth_provider.issuer = upstream_oauth_provider.issuer,
174 %upstream_oauth_provider.client_id,
175 ),
176 err,
177 )]
178 async fn find_by_subject(
179 &mut self,
180 upstream_oauth_provider: &UpstreamOAuthProvider,
181 subject: &str,
182 ) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
183 let res = sqlx::query_as!(
184 LinkLookup,
185 r#"
186 SELECT
187 upstream_oauth_link_id,
188 upstream_oauth_provider_id,
189 user_id,
190 subject,
191 human_account_name,
192 created_at
193 FROM upstream_oauth_links
194 WHERE upstream_oauth_provider_id = $1
195 AND subject = $2
196 "#,
197 Uuid::from(upstream_oauth_provider.id),
198 subject,
199 )
200 .traced()
201 .fetch_optional(&mut *self.conn)
202 .await?
203 .map(Into::into);
204
205 Ok(res)
206 }
207
208 #[tracing::instrument(
209 name = "db.upstream_oauth_link.add",
210 skip_all,
211 fields(
212 db.query.text,
213 upstream_oauth_link.id,
214 upstream_oauth_link.subject = subject,
215 upstream_oauth_link.human_account_name = human_account_name,
216 %upstream_oauth_provider.id,
217 upstream_oauth_provider.issuer = upstream_oauth_provider.issuer,
218 %upstream_oauth_provider.client_id,
219 ),
220 err,
221 )]
222 async fn add(
223 &mut self,
224 rng: &mut (dyn RngCore + Send),
225 clock: &dyn Clock,
226 upstream_oauth_provider: &UpstreamOAuthProvider,
227 subject: String,
228 human_account_name: Option<String>,
229 ) -> Result<UpstreamOAuthLink, Self::Error> {
230 let created_at = clock.now();
231 let id = Ulid::from_datetime_with_rng(created_at, rng);
232 tracing::Span::current().record("upstream_oauth_link.id", tracing::field::display(id));
233
234 sqlx::query!(
235 r#"
236 INSERT INTO upstream_oauth_links (
237 upstream_oauth_link_id,
238 upstream_oauth_provider_id,
239 user_id,
240 subject,
241 human_account_name,
242 created_at
243 ) VALUES ($1, $2, NULL, $3, $4, $5)
244 "#,
245 Uuid::from(id),
246 Uuid::from(upstream_oauth_provider.id),
247 &subject,
248 human_account_name.as_deref(),
249 created_at,
250 )
251 .traced()
252 .execute(&mut *self.conn)
253 .await?;
254
255 Ok(UpstreamOAuthLink {
256 id,
257 provider_id: upstream_oauth_provider.id,
258 user_id: None,
259 subject,
260 human_account_name,
261 created_at,
262 })
263 }
264
265 #[tracing::instrument(
266 name = "db.upstream_oauth_link.associate_to_user",
267 skip_all,
268 fields(
269 db.query.text,
270 %upstream_oauth_link.id,
271 %upstream_oauth_link.subject,
272 %user.id,
273 %user.username,
274 ),
275 err,
276 )]
277 async fn associate_to_user(
278 &mut self,
279 upstream_oauth_link: &UpstreamOAuthLink,
280 user: &User,
281 ) -> Result<(), Self::Error> {
282 sqlx::query!(
283 r#"
284 UPDATE upstream_oauth_links
285 SET user_id = $1
286 WHERE upstream_oauth_link_id = $2
287 "#,
288 Uuid::from(user.id),
289 Uuid::from(upstream_oauth_link.id),
290 )
291 .traced()
292 .execute(&mut *self.conn)
293 .await?;
294
295 Ok(())
296 }
297
298 #[tracing::instrument(
299 name = "db.upstream_oauth_link.list",
300 skip_all,
301 fields(
302 db.query.text,
303 ),
304 err,
305 )]
306 async fn list(
307 &mut self,
308 filter: UpstreamOAuthLinkFilter<'_>,
309 pagination: Pagination,
310 ) -> Result<Page<UpstreamOAuthLink>, DatabaseError> {
311 let (sql, arguments) = Query::select()
312 .expr_as(
313 Expr::col((
314 UpstreamOAuthLinks::Table,
315 UpstreamOAuthLinks::UpstreamOAuthLinkId,
316 )),
317 LinkLookupIden::UpstreamOauthLinkId,
318 )
319 .expr_as(
320 Expr::col((
321 UpstreamOAuthLinks::Table,
322 UpstreamOAuthLinks::UpstreamOAuthProviderId,
323 )),
324 LinkLookupIden::UpstreamOauthProviderId,
325 )
326 .expr_as(
327 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::UserId)),
328 LinkLookupIden::UserId,
329 )
330 .expr_as(
331 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::Subject)),
332 LinkLookupIden::Subject,
333 )
334 .expr_as(
335 Expr::col((
336 UpstreamOAuthLinks::Table,
337 UpstreamOAuthLinks::HumanAccountName,
338 )),
339 LinkLookupIden::HumanAccountName,
340 )
341 .expr_as(
342 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::CreatedAt)),
343 LinkLookupIden::CreatedAt,
344 )
345 .from(UpstreamOAuthLinks::Table)
346 .apply_filter(filter)
347 .generate_pagination(
348 (
349 UpstreamOAuthLinks::Table,
350 UpstreamOAuthLinks::UpstreamOAuthLinkId,
351 ),
352 pagination,
353 )
354 .build_sqlx(PostgresQueryBuilder);
355
356 let edges: Vec<LinkLookup> = sqlx::query_as_with(&sql, arguments)
357 .traced()
358 .fetch_all(&mut *self.conn)
359 .await?;
360
361 let page = pagination.process(edges).map(UpstreamOAuthLink::from);
362
363 Ok(page)
364 }
365
366 #[tracing::instrument(
367 name = "db.upstream_oauth_link.count",
368 skip_all,
369 fields(
370 db.query.text,
371 ),
372 err,
373 )]
374 async fn count(&mut self, filter: UpstreamOAuthLinkFilter<'_>) -> Result<usize, Self::Error> {
375 let (sql, arguments) = Query::select()
376 .expr(
377 Expr::col((
378 UpstreamOAuthLinks::Table,
379 UpstreamOAuthLinks::UpstreamOAuthLinkId,
380 ))
381 .count(),
382 )
383 .from(UpstreamOAuthLinks::Table)
384 .apply_filter(filter)
385 .build_sqlx(PostgresQueryBuilder);
386
387 let count: i64 = sqlx::query_scalar_with(&sql, arguments)
388 .traced()
389 .fetch_one(&mut *self.conn)
390 .await?;
391
392 count
393 .try_into()
394 .map_err(DatabaseError::to_invalid_operation)
395 }
396
397 #[tracing::instrument(
398 name = "db.upstream_oauth_link.remove",
399 skip_all,
400 fields(
401 db.query.text,
402 upstream_oauth_link.id,
403 upstream_oauth_link.provider_id,
404 %upstream_oauth_link.subject,
405 ),
406 err,
407 )]
408 async fn remove(
409 &mut self,
410 clock: &dyn Clock,
411 upstream_oauth_link: UpstreamOAuthLink,
412 ) -> Result<(), Self::Error> {
413 let span = tracing::info_span!(
416 "db.upstream_oauth_link.remove.unlink",
417 { DB_QUERY_TEXT } = tracing::field::Empty
418 );
419 sqlx::query!(
420 r#"
421 UPDATE upstream_oauth_authorization_sessions SET
422 upstream_oauth_link_id = NULL,
423 unlinked_at = $2
424 WHERE upstream_oauth_link_id = $1
425 "#,
426 Uuid::from(upstream_oauth_link.id),
427 clock.now()
428 )
429 .record(&span)
430 .execute(&mut *self.conn)
431 .instrument(span)
432 .await?;
433
434 let span = tracing::info_span!(
436 "db.upstream_oauth_link.remove.delete",
437 { DB_QUERY_TEXT } = tracing::field::Empty
438 );
439 let res = sqlx::query!(
440 r#"
441 DELETE FROM upstream_oauth_links
442 WHERE upstream_oauth_link_id = $1
443 "#,
444 Uuid::from(upstream_oauth_link.id),
445 )
446 .record(&span)
447 .execute(&mut *self.conn)
448 .instrument(span)
449 .await?;
450
451 DatabaseError::ensure_affected_rows(&res, 1)?;
452
453 Ok(())
454 }
455
456 #[tracing::instrument(
457 name = "db.upstream_oauth_link.cleanup_orphaned",
458 skip_all,
459 fields(
460 db.query.text,
461 since = since.map(tracing::field::display),
462 until = %until,
463 limit = limit,
464 ),
465 err,
466 )]
467 async fn cleanup_orphaned(
468 &mut self,
469 since: Option<Ulid>,
470 until: Ulid,
471 limit: usize,
472 ) -> Result<(usize, Option<Ulid>), Self::Error> {
473 let res = sqlx::query_scalar!(
477 r#"
478 WITH
479 to_delete AS (
480 SELECT upstream_oauth_link_id
481 FROM upstream_oauth_links
482 WHERE user_id IS NULL
483 AND ($1::uuid IS NULL OR upstream_oauth_link_id > $1)
484 AND upstream_oauth_link_id <= $2
485 ORDER BY upstream_oauth_link_id
486 LIMIT $3
487 ),
488 deleted_sessions AS (
489 DELETE FROM upstream_oauth_authorization_sessions
490 USING to_delete
491 WHERE upstream_oauth_authorization_sessions.upstream_oauth_link_id = to_delete.upstream_oauth_link_id
492 )
493 DELETE FROM upstream_oauth_links
494 USING to_delete
495 WHERE upstream_oauth_links.upstream_oauth_link_id = to_delete.upstream_oauth_link_id
496 RETURNING upstream_oauth_links.upstream_oauth_link_id
497 "#,
498 since.map(Uuid::from),
499 Uuid::from(until),
500 i64::try_from(limit).unwrap_or(i64::MAX)
501 )
502 .traced()
503 .fetch_all(&mut *self.conn)
504 .await?;
505
506 let count = res.len();
507 let max_id = res.into_iter().max();
508
509 Ok((count, max_id.map(Ulid::from)))
510 }
511}