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::Rng;
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 )]
141 async fn lookup(&mut self, id: Ulid) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
142 let res = sqlx::query_as!(
143 LinkLookup,
144 r#"
145 SELECT
146 upstream_oauth_link_id,
147 upstream_oauth_provider_id,
148 user_id,
149 subject,
150 human_account_name,
151 created_at
152 FROM upstream_oauth_links
153 WHERE upstream_oauth_link_id = $1
154 "#,
155 Uuid::from(id),
156 )
157 .traced()
158 .fetch_optional(&mut *self.conn)
159 .await?
160 .map(Into::into);
161
162 Ok(res)
163 }
164
165 #[tracing::instrument(
166 name = "db.upstream_oauth_link.find_by_subject",
167 skip_all,
168 fields(
169 db.query.text,
170 upstream_oauth_link.subject = subject,
171 %upstream_oauth_provider.id,
172 upstream_oauth_provider.issuer = upstream_oauth_provider.issuer,
173 %upstream_oauth_provider.client_id,
174 ),
175 )]
176 async fn find_by_subject(
177 &mut self,
178 upstream_oauth_provider: &UpstreamOAuthProvider,
179 subject: &str,
180 ) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
181 let res = sqlx::query_as!(
182 LinkLookup,
183 r#"
184 SELECT
185 upstream_oauth_link_id,
186 upstream_oauth_provider_id,
187 user_id,
188 subject,
189 human_account_name,
190 created_at
191 FROM upstream_oauth_links
192 WHERE upstream_oauth_provider_id = $1
193 AND subject = $2
194 "#,
195 Uuid::from(upstream_oauth_provider.id),
196 subject,
197 )
198 .traced()
199 .fetch_optional(&mut *self.conn)
200 .await?
201 .map(Into::into);
202
203 Ok(res)
204 }
205
206 #[tracing::instrument(
207 name = "db.upstream_oauth_link.add",
208 skip_all,
209 fields(
210 db.query.text,
211 upstream_oauth_link.id,
212 upstream_oauth_link.subject = subject,
213 upstream_oauth_link.human_account_name = human_account_name,
214 %upstream_oauth_provider.id,
215 upstream_oauth_provider.issuer = upstream_oauth_provider.issuer,
216 %upstream_oauth_provider.client_id,
217 ),
218 )]
219 async fn add(
220 &mut self,
221 rng: &mut (dyn Rng + Send),
222 clock: &dyn Clock,
223 upstream_oauth_provider: &UpstreamOAuthProvider,
224 subject: String,
225 human_account_name: Option<String>,
226 ) -> Result<UpstreamOAuthLink, Self::Error> {
227 let created_at = clock.now();
228 let id = Ulid::from_datetime_with_rng(created_at, rng);
229 tracing::Span::current().record("upstream_oauth_link.id", tracing::field::display(id));
230
231 sqlx::query!(
232 r#"
233 INSERT INTO upstream_oauth_links (
234 upstream_oauth_link_id,
235 upstream_oauth_provider_id,
236 user_id,
237 subject,
238 human_account_name,
239 created_at
240 ) VALUES ($1, $2, NULL, $3, $4, $5)
241 "#,
242 Uuid::from(id),
243 Uuid::from(upstream_oauth_provider.id),
244 &subject,
245 human_account_name.as_deref(),
246 created_at,
247 )
248 .traced()
249 .execute(&mut *self.conn)
250 .await?;
251
252 Ok(UpstreamOAuthLink {
253 id,
254 provider_id: upstream_oauth_provider.id,
255 user_id: None,
256 subject,
257 human_account_name,
258 created_at,
259 })
260 }
261
262 #[tracing::instrument(
263 name = "db.upstream_oauth_link.associate_to_user",
264 skip_all,
265 fields(
266 db.query.text,
267 %upstream_oauth_link.id,
268 %upstream_oauth_link.subject,
269 %user.id,
270 %user.username,
271 ),
272 )]
273 async fn associate_to_user(
274 &mut self,
275 upstream_oauth_link: &UpstreamOAuthLink,
276 user: &User,
277 ) -> Result<(), Self::Error> {
278 sqlx::query!(
279 r#"
280 UPDATE upstream_oauth_links
281 SET user_id = $1
282 WHERE upstream_oauth_link_id = $2
283 "#,
284 Uuid::from(user.id),
285 Uuid::from(upstream_oauth_link.id),
286 )
287 .traced()
288 .execute(&mut *self.conn)
289 .await?;
290
291 Ok(())
292 }
293
294 #[tracing::instrument(
295 name = "db.upstream_oauth_link.list",
296 skip_all,
297 fields(
298 db.query.text,
299 ),
300 )]
301 async fn list(
302 &mut self,
303 filter: UpstreamOAuthLinkFilter<'_>,
304 pagination: Pagination,
305 ) -> Result<Page<UpstreamOAuthLink>, DatabaseError> {
306 let (sql, arguments) = Query::select()
307 .expr_as(
308 Expr::col((
309 UpstreamOAuthLinks::Table,
310 UpstreamOAuthLinks::UpstreamOAuthLinkId,
311 )),
312 LinkLookupIden::UpstreamOauthLinkId,
313 )
314 .expr_as(
315 Expr::col((
316 UpstreamOAuthLinks::Table,
317 UpstreamOAuthLinks::UpstreamOAuthProviderId,
318 )),
319 LinkLookupIden::UpstreamOauthProviderId,
320 )
321 .expr_as(
322 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::UserId)),
323 LinkLookupIden::UserId,
324 )
325 .expr_as(
326 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::Subject)),
327 LinkLookupIden::Subject,
328 )
329 .expr_as(
330 Expr::col((
331 UpstreamOAuthLinks::Table,
332 UpstreamOAuthLinks::HumanAccountName,
333 )),
334 LinkLookupIden::HumanAccountName,
335 )
336 .expr_as(
337 Expr::col((UpstreamOAuthLinks::Table, UpstreamOAuthLinks::CreatedAt)),
338 LinkLookupIden::CreatedAt,
339 )
340 .from(UpstreamOAuthLinks::Table)
341 .apply_filter(filter)
342 .generate_pagination(
343 (
344 UpstreamOAuthLinks::Table,
345 UpstreamOAuthLinks::UpstreamOAuthLinkId,
346 ),
347 pagination,
348 )
349 .build_sqlx(PostgresQueryBuilder);
350
351 let edges: Vec<LinkLookup> = sqlx::query_as_with(&sql, arguments)
352 .traced()
353 .fetch_all(&mut *self.conn)
354 .await?;
355
356 let page = pagination.process(edges).map(UpstreamOAuthLink::from);
357
358 Ok(page)
359 }
360
361 #[tracing::instrument(
362 name = "db.upstream_oauth_link.count",
363 skip_all,
364 fields(
365 db.query.text,
366 ),
367 )]
368 async fn count(&mut self, filter: UpstreamOAuthLinkFilter<'_>) -> Result<usize, Self::Error> {
369 let (sql, arguments) = Query::select()
370 .expr(
371 Expr::col((
372 UpstreamOAuthLinks::Table,
373 UpstreamOAuthLinks::UpstreamOAuthLinkId,
374 ))
375 .count(),
376 )
377 .from(UpstreamOAuthLinks::Table)
378 .apply_filter(filter)
379 .build_sqlx(PostgresQueryBuilder);
380
381 let count: i64 = sqlx::query_scalar_with(&sql, arguments)
382 .traced()
383 .fetch_one(&mut *self.conn)
384 .await?;
385
386 count
387 .try_into()
388 .map_err(DatabaseError::to_invalid_operation)
389 }
390
391 #[tracing::instrument(
392 name = "db.upstream_oauth_link.remove",
393 skip_all,
394 fields(
395 db.query.text,
396 upstream_oauth_link.id,
397 upstream_oauth_link.provider_id,
398 %upstream_oauth_link.subject,
399 ),
400 )]
401 async fn remove(
402 &mut self,
403 clock: &dyn Clock,
404 upstream_oauth_link: UpstreamOAuthLink,
405 ) -> Result<(), Self::Error> {
406 let span = tracing::info_span!(
409 "db.upstream_oauth_link.remove.unlink",
410 { DB_QUERY_TEXT } = tracing::field::Empty
411 );
412 sqlx::query!(
413 r#"
414 UPDATE upstream_oauth_authorization_sessions SET
415 upstream_oauth_link_id = NULL,
416 unlinked_at = $2
417 WHERE upstream_oauth_link_id = $1
418 "#,
419 Uuid::from(upstream_oauth_link.id),
420 clock.now()
421 )
422 .record(&span)
423 .execute(&mut *self.conn)
424 .instrument(span)
425 .await?;
426
427 let span = tracing::info_span!(
429 "db.upstream_oauth_link.remove.delete",
430 { DB_QUERY_TEXT } = tracing::field::Empty
431 );
432 let res = sqlx::query!(
433 r#"
434 DELETE FROM upstream_oauth_links
435 WHERE upstream_oauth_link_id = $1
436 "#,
437 Uuid::from(upstream_oauth_link.id),
438 )
439 .record(&span)
440 .execute(&mut *self.conn)
441 .instrument(span)
442 .await?;
443
444 DatabaseError::ensure_affected_rows(&res, 1)?;
445
446 Ok(())
447 }
448
449 #[tracing::instrument(
450 name = "db.upstream_oauth_link.cleanup_orphaned",
451 skip_all,
452 fields(
453 db.query.text,
454 since = since.map(tracing::field::display),
455 until = %until,
456 limit = limit,
457 ),
458 )]
459 async fn cleanup_orphaned(
460 &mut self,
461 since: Option<Ulid>,
462 until: Ulid,
463 limit: usize,
464 ) -> Result<(usize, Option<Ulid>), Self::Error> {
465 let res = sqlx::query_scalar!(
469 r#"
470 WITH
471 to_delete AS (
472 SELECT upstream_oauth_link_id
473 FROM upstream_oauth_links
474 WHERE user_id IS NULL
475 AND ($1::uuid IS NULL OR upstream_oauth_link_id > $1)
476 AND upstream_oauth_link_id <= $2
477 ORDER BY upstream_oauth_link_id
478 LIMIT $3
479 ),
480 deleted_sessions AS (
481 DELETE FROM upstream_oauth_authorization_sessions
482 USING to_delete
483 WHERE upstream_oauth_authorization_sessions.upstream_oauth_link_id = to_delete.upstream_oauth_link_id
484 )
485 DELETE FROM upstream_oauth_links
486 USING to_delete
487 WHERE upstream_oauth_links.upstream_oauth_link_id = to_delete.upstream_oauth_link_id
488 RETURNING upstream_oauth_links.upstream_oauth_link_id
489 "#,
490 since.map(Uuid::from),
491 Uuid::from(until),
492 i64::try_from(limit).unwrap_or(i64::MAX)
493 )
494 .traced()
495 .fetch_all(&mut *self.conn)
496 .await?;
497
498 let count = res.len();
499 let max_id = res.into_iter().max();
500
501 Ok((count, max_id.map(Ulid::from)))
502 }
503}