From 8962f355ff1178c3dcb335d541058e58a4ee7be5 Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Mon, 29 Sep 2025 14:46:28 +0200 Subject: [PATCH] storage: make the edges in pages include cursors --- crates/storage-pg/src/app_session.rs | 36 +++++--- crates/storage-pg/src/compat/mod.rs | 23 ++--- crates/storage-pg/src/compat/session.rs | 13 +++ crates/storage-pg/src/compat/sso_login.rs | 7 ++ crates/storage-pg/src/oauth2/mod.rs | 48 +++++----- crates/storage-pg/src/oauth2/session.rs | 7 ++ crates/storage-pg/src/upstream_oauth2/link.rs | 7 ++ crates/storage-pg/src/upstream_oauth2/mod.rs | 40 +++++---- .../src/upstream_oauth2/provider.rs | 7 ++ .../storage-pg/src/upstream_oauth2/session.rs | 7 ++ crates/storage-pg/src/user/email.rs | 7 ++ crates/storage-pg/src/user/mod.rs | 8 ++ .../storage-pg/src/user/registration_token.rs | 30 ++++--- crates/storage-pg/src/user/session.rs | 7 ++ crates/storage-pg/src/user/tests.rs | 14 +-- crates/storage/src/pagination.rs | 88 ++++++++++++++----- crates/storage/src/queue/tasks.rs | 6 +- 17 files changed, 246 insertions(+), 109 deletions(-) diff --git a/crates/storage-pg/src/app_session.rs b/crates/storage-pg/src/app_session.rs index 62f3979a9..4e12810cc 100644 --- a/crates/storage-pg/src/app_session.rs +++ b/crates/storage-pg/src/app_session.rs @@ -55,7 +55,9 @@ mod priv_ { use std::net::IpAddr; use chrono::{DateTime, Utc}; + use mas_storage::pagination::Node; use sea_query::enum_def; + use ulid::Ulid; use uuid::Uuid; #[derive(sqlx::FromRow)] @@ -77,6 +79,12 @@ mod priv_ { pub(super) last_active_at: Option>, pub(super) last_active_ip: Option, } + + impl Node for AppSessionLookup { + fn cursor(&self) -> Ulid { + self.cursor.into() + } + } } use priv_::{AppSessionLookup, AppSessionLookupIden}; @@ -592,13 +600,13 @@ mod tests { let full_list = repo.app_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 1); assert_eq!( - full_list.edges[0], + full_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); let active_list = repo.app_session().list(active, pagination).await.unwrap(); assert_eq!(active_list.edges.len(), 1); assert_eq!( - active_list.edges[0], + active_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); let finished_list = repo.app_session().list(finished, pagination).await.unwrap(); @@ -618,7 +626,7 @@ mod tests { let full_list = repo.app_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 1); assert_eq!( - full_list.edges[0], + full_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); let active_list = repo.app_session().list(active, pagination).await.unwrap(); @@ -626,7 +634,7 @@ mod tests { let finished_list = repo.app_session().list(finished, pagination).await.unwrap(); assert_eq!(finished_list.edges.len(), 1); assert_eq!( - finished_list.edges[0], + finished_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); @@ -680,25 +688,25 @@ mod tests { let full_list = repo.app_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 2); assert_eq!( - full_list.edges[0], + full_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); assert_eq!( - full_list.edges[1], + full_list.edges[1].node, AppSession::OAuth2(Box::new(oauth_session.clone())) ); let active_list = repo.app_session().list(active, pagination).await.unwrap(); assert_eq!(active_list.edges.len(), 1); assert_eq!( - active_list.edges[0], + active_list.edges[0].node, AppSession::OAuth2(Box::new(oauth_session.clone())) ); let finished_list = repo.app_session().list(finished, pagination).await.unwrap(); assert_eq!(finished_list.edges.len(), 1); assert_eq!( - finished_list.edges[0], + finished_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); @@ -716,11 +724,11 @@ mod tests { let full_list = repo.app_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 2); assert_eq!( - full_list.edges[0], + full_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); assert_eq!( - full_list.edges[1], + full_list.edges[1].node, AppSession::OAuth2(Box::new(oauth_session.clone())) ); @@ -730,11 +738,11 @@ mod tests { let finished_list = repo.app_session().list(finished, pagination).await.unwrap(); assert_eq!(finished_list.edges.len(), 2); assert_eq!( - finished_list.edges[0], + finished_list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); assert_eq!( - full_list.edges[1], + full_list.edges[1].node, AppSession::OAuth2(Box::new(oauth_session.clone())) ); @@ -744,7 +752,7 @@ mod tests { let list = repo.app_session().list(filter, pagination).await.unwrap(); assert_eq!(list.edges.len(), 1); assert_eq!( - list.edges[0], + list.edges[0].node, AppSession::Compat(Box::new(compat_session.clone())) ); @@ -753,7 +761,7 @@ mod tests { let list = repo.app_session().list(filter, pagination).await.unwrap(); assert_eq!(list.edges.len(), 1); assert_eq!( - list.edges[0], + list.edges[0].node, AppSession::OAuth2(Box::new(oauth_session.clone())) ); diff --git a/crates/storage-pg/src/compat/mod.rs b/crates/storage-pg/src/compat/mod.rs index db190db71..d42c9b1af 100644 --- a/crates/storage-pg/src/compat/mod.rs +++ b/crates/storage-pg/src/compat/mod.rs @@ -92,14 +92,14 @@ mod tests { let full_list = repo.compat_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 1); - assert_eq!(full_list.edges[0].0.id, session.id); + assert_eq!(full_list.edges[0].node.0.id, session.id); let active_list = repo .compat_session() .list(active, pagination) .await .unwrap(); assert_eq!(active_list.edges.len(), 1); - assert_eq!(active_list.edges[0].0.id, session.id); + assert_eq!(active_list.edges[0].node.0.id, session.id); let finished_list = repo .compat_session() .list(finished, pagination) @@ -150,7 +150,7 @@ mod tests { .await .unwrap(); assert_eq!(list.edges.len(), 1); - let session_lookup = &list.edges[0].0; + let session_lookup = &list.edges[0].node.0; assert_eq!(session_lookup.id, session.id); assert_eq!(session_lookup.user_id, user.id); assert_eq!(session.device.as_ref().unwrap().as_str(), device_str); @@ -168,7 +168,7 @@ mod tests { let full_list = repo.compat_session().list(all, pagination).await.unwrap(); assert_eq!(full_list.edges.len(), 1); - assert_eq!(full_list.edges[0].0.id, session.id); + assert_eq!(full_list.edges[0].node.0.id, session.id); let active_list = repo .compat_session() .list(active, pagination) @@ -181,7 +181,7 @@ mod tests { .await .unwrap(); assert_eq!(finished_list.edges.len(), 1); - assert_eq!(finished_list.edges[0].0.id, session.id); + assert_eq!(finished_list.edges[0].node.0.id, session.id); // Reload the session and check again let session_lookup = repo @@ -260,14 +260,14 @@ mod tests { .await .unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0].0.id, sso_login_session.id); + assert_eq!(list.edges[0].node.0.id, sso_login_session.id); let list = repo .compat_session() .list(unknown, pagination) .await .unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0].0.id, unknown_session.id); + assert_eq!(list.edges[0].node.0.id, unknown_session.id); // Check that combining the two filters works // At this point, there is one active SSO login session and one finished unknown @@ -696,7 +696,8 @@ mod tests { // List all logins let logins = repo.compat_sso_login().list(all, pagination).await.unwrap(); assert!(!logins.has_next_page); - assert_eq!(logins.edges, vec![login.clone()]); + assert_eq!(logins.edges.len(), 1); + assert_eq!(logins.edges[0].node, login); // List the logins for the user let logins = repo @@ -705,7 +706,8 @@ mod tests { .await .unwrap(); assert!(!logins.has_next_page); - assert_eq!(logins.edges, vec![login.clone()]); + assert_eq!(logins.edges.len(), 1); + assert_eq!(logins.edges[0].node, login); // List only the pending logins for the user let logins = repo @@ -732,6 +734,7 @@ mod tests { .await .unwrap(); assert!(!logins.has_next_page); - assert_eq!(logins.edges, &[login]); + assert_eq!(logins.edges.len(), 1); + assert_eq!(logins.edges[0].node, login); } } diff --git a/crates/storage-pg/src/compat/session.rs b/crates/storage-pg/src/compat/session.rs index 4ba5ee726..0fb21c487 100644 --- a/crates/storage-pg/src/compat/session.rs +++ b/crates/storage-pg/src/compat/session.rs @@ -15,6 +15,7 @@ use mas_data_model::{ use mas_storage::{ Page, Pagination, compat::{CompatSessionFilter, CompatSessionRepository}, + pagination::Node, }; use rand::RngCore; use sea_query::{Expr, PostgresQueryBuilder, Query, enum_def}; @@ -59,6 +60,12 @@ struct CompatSessionLookup { last_active_ip: Option, } +impl Node for CompatSessionLookup { + fn cursor(&self) -> Ulid { + self.compat_session_id.into() + } +} + impl From for CompatSession { fn from(value: CompatSessionLookup) -> Self { let id = value.compat_session_id.into(); @@ -106,6 +113,12 @@ struct CompatSessionAndSsoLoginLookup { compat_sso_login_exchanged_at: Option>, } +impl Node for CompatSessionAndSsoLoginLookup { + fn cursor(&self) -> Ulid { + self.compat_session_id.into() + } +} + impl TryFrom for (CompatSession, Option) { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/compat/sso_login.rs b/crates/storage-pg/src/compat/sso_login.rs index eeadff164..43ad4bead 100644 --- a/crates/storage-pg/src/compat/sso_login.rs +++ b/crates/storage-pg/src/compat/sso_login.rs @@ -10,6 +10,7 @@ use mas_data_model::{BrowserSession, Clock, CompatSession, CompatSsoLogin, Compa use mas_storage::{ Page, Pagination, compat::{CompatSsoLoginFilter, CompatSsoLoginRepository}, + pagination::Node, }; use rand::RngCore; use sea_query::{Expr, PostgresQueryBuilder, Query, enum_def}; @@ -54,6 +55,12 @@ struct CompatSsoLoginLookup { compat_session_id: Option, } +impl Node for CompatSsoLoginLookup { + fn cursor(&self) -> Ulid { + self.compat_sso_login_id.into() + } +} + impl TryFrom for CompatSsoLogin { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/oauth2/mod.rs b/crates/storage-pg/src/oauth2/mod.rs index 6b1c81d46..bf741b5f2 100644 --- a/crates/storage-pg/src/oauth2/mod.rs +++ b/crates/storage-pg/src/oauth2/mod.rs @@ -511,10 +511,10 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 4); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session12); - assert_eq!(list.edges[2], session21); - assert_eq!(list.edges[3], session22); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session12); + assert_eq!(list.edges[2].node, session21); + assert_eq!(list.edges[3].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 4); @@ -527,8 +527,8 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 2); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session21); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session21); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 2); @@ -541,8 +541,8 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 2); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session12); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session12); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 2); @@ -557,7 +557,7 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session22); + assert_eq!(list.edges[0].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); @@ -570,8 +570,8 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 2); - assert_eq!(list.edges[0], session12); - assert_eq!(list.edges[1], session21); + assert_eq!(list.edges[0].node, session12); + assert_eq!(list.edges[1].node, session21); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 2); @@ -584,8 +584,8 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 2); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session22); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 2); @@ -598,7 +598,7 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session22); + assert_eq!(list.edges[0].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); @@ -613,7 +613,7 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session22); + assert_eq!(list.edges[0].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); @@ -626,7 +626,7 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session12); + assert_eq!(list.edges[0].node, session12); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); @@ -641,7 +641,7 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session21); + assert_eq!(list.edges[0].node, session21); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); @@ -655,10 +655,10 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 4); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session12); - assert_eq!(list.edges[2], session21); - assert_eq!(list.edges[3], session22); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session12); + assert_eq!(list.edges[2].node, session21); + assert_eq!(list.edges[3].node, session22); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 4); // We should get all sessions with the "openid" and "email" scope @@ -671,8 +671,8 @@ mod tests { .unwrap(); assert!(!list.has_next_page); assert_eq!(list.edges.len(), 2); - assert_eq!(list.edges[0], session11); - assert_eq!(list.edges[1], session12); + assert_eq!(list.edges[0].node, session11); + assert_eq!(list.edges[1].node, session12); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 2); // Try combining the scope filter with the user filter @@ -685,7 +685,7 @@ mod tests { .await .unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0], session11); + assert_eq!(list.edges[0].node, session11); assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1); // Finish all sessions of a client in batch diff --git a/crates/storage-pg/src/oauth2/session.rs b/crates/storage-pg/src/oauth2/session.rs index c04e03df0..072691a06 100644 --- a/crates/storage-pg/src/oauth2/session.rs +++ b/crates/storage-pg/src/oauth2/session.rs @@ -12,6 +12,7 @@ use mas_data_model::{BrowserSession, Client, Clock, Session, SessionState, User} use mas_storage::{ Page, Pagination, oauth2::{OAuth2SessionFilter, OAuth2SessionRepository}, + pagination::Node, }; use oauth2_types::scope::{Scope, ScopeToken}; use rand::RngCore; @@ -61,6 +62,12 @@ struct OAuthSessionLookup { human_name: Option, } +impl Node for OAuthSessionLookup { + fn cursor(&self) -> Ulid { + self.oauth2_session_id.into() + } +} + impl TryFrom for Session { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/upstream_oauth2/link.rs b/crates/storage-pg/src/upstream_oauth2/link.rs index 6f671655d..c43dd8a18 100644 --- a/crates/storage-pg/src/upstream_oauth2/link.rs +++ b/crates/storage-pg/src/upstream_oauth2/link.rs @@ -9,6 +9,7 @@ use chrono::{DateTime, Utc}; use mas_data_model::{Clock, UpstreamOAuthLink, UpstreamOAuthProvider, User}; use mas_storage::{ Page, Pagination, + pagination::Node, upstream_oauth2::{UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository}, }; use opentelemetry_semantic_conventions::trace::DB_QUERY_TEXT; @@ -53,6 +54,12 @@ struct LinkLookup { created_at: DateTime, } +impl Node for LinkLookup { + fn cursor(&self) -> Ulid { + self.upstream_oauth_link_id.into() + } +} + impl From for UpstreamOAuthLink { fn from(value: LinkLookup) -> Self { UpstreamOAuthLink { diff --git a/crates/storage-pg/src/upstream_oauth2/mod.rs b/crates/storage-pg/src/upstream_oauth2/mod.rs index 6381c9b7f..d98e840b6 100644 --- a/crates/storage-pg/src/upstream_oauth2/mod.rs +++ b/crates/storage-pg/src/upstream_oauth2/mod.rs @@ -206,8 +206,8 @@ mod tests { assert!(!links.has_previous_page); assert!(!links.has_next_page); assert_eq!(links.edges.len(), 1); - assert_eq!(links.edges[0].id, link.id); - assert_eq!(links.edges[0].user_id, Some(user.id)); + assert_eq!(links.edges[0].node.id, link.id); + assert_eq!(links.edges[0].node.user_id, Some(user.id)); assert_eq!(repo.upstream_oauth_link().count(filter).await.unwrap(), 1); @@ -282,7 +282,7 @@ mod tests { .unwrap(); assert_eq!(session_page.edges.len(), 1); - assert_eq!(session_page.edges[0].id, session.id); + assert_eq!(session_page.edges[0].node.id, session.id); assert!(!session_page.has_next_page); assert!(!session_page.has_previous_page); @@ -374,7 +374,7 @@ mod tests { // It returned the first 10 items assert!(page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|p| p.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect(); assert_eq!(&edge_ids, &ids[..10]); // Getting the same page with the "enabled only" filter should return the same @@ -396,7 +396,7 @@ mod tests { // It returned the next 10 items assert!(!page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|p| p.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect(); assert_eq!(&edge_ids, &ids[10..]); // Lookup the last 10 items @@ -408,7 +408,7 @@ mod tests { // It returned the last 10 items assert!(page.has_previous_page); - let edge_ids: Vec<_> = page.edges.iter().map(|p| p.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect(); assert_eq!(&edge_ids, &ids[10..]); // Lookup the previous 10 items @@ -420,7 +420,7 @@ mod tests { // It returned the previous 10 items assert!(!page.has_previous_page); - let edge_ids: Vec<_> = page.edges.iter().map(|p| p.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect(); assert_eq!(&edge_ids, &ids[..10]); // Lookup 10 items between two IDs @@ -432,7 +432,7 @@ mod tests { // It returned the items in between assert!(!page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|p| p.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect(); assert_eq!(&edge_ids, &ids[6..8]); // There should not be any disabled providers @@ -560,7 +560,7 @@ mod tests { // It returned the first 10 items assert!(page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect(); assert_eq!(&edge_ids, &ids[..10]); // Lookup the next 10 items @@ -572,7 +572,7 @@ mod tests { // It returned the next 10 items assert!(!page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect(); assert_eq!(&edge_ids, &ids[10..]); // Lookup the last 10 items @@ -584,7 +584,7 @@ mod tests { // It returned the last 10 items assert!(page.has_previous_page); - let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect(); assert_eq!(&edge_ids, &ids[10..]); // Lookup the previous 10 items @@ -596,7 +596,7 @@ mod tests { // It returned the previous 10 items assert!(!page.has_previous_page); - let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect(); assert_eq!(&edge_ids, &ids[..10]); // Lookup 5 items between two IDs @@ -608,7 +608,7 @@ mod tests { // It returned the items in between assert!(!page.has_next_page); - let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect(); assert_eq!(&edge_ids, &ids[6..11]); // Check the sub/sid filters @@ -638,11 +638,21 @@ mod tests { assert_eq!(page.edges.len(), 4); for edge in page.edges { assert_eq!( - edge.id_token_claims().unwrap().get("sub").unwrap().as_str(), + edge.node + .id_token_claims() + .unwrap() + .get("sub") + .unwrap() + .as_str(), Some("alice") ); assert_eq!( - edge.id_token_claims().unwrap().get("sid").unwrap().as_str(), + edge.node + .id_token_claims() + .unwrap() + .get("sid") + .unwrap() + .as_str(), Some("one") ); } diff --git a/crates/storage-pg/src/upstream_oauth2/provider.rs b/crates/storage-pg/src/upstream_oauth2/provider.rs index 583f8ec0c..caade738d 100644 --- a/crates/storage-pg/src/upstream_oauth2/provider.rs +++ b/crates/storage-pg/src/upstream_oauth2/provider.rs @@ -9,6 +9,7 @@ use chrono::{DateTime, Utc}; use mas_data_model::{Clock, UpstreamOAuthProvider, UpstreamOAuthProviderClaimsImports}; use mas_storage::{ Page, Pagination, + pagination::Node, upstream_oauth2::{ UpstreamOAuthProviderFilter, UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository, }, @@ -74,6 +75,12 @@ struct ProviderLookup { on_backchannel_logout: String, } +impl Node for ProviderLookup { + fn cursor(&self) -> Ulid { + self.upstream_oauth_provider_id.into() + } +} + impl TryFrom for UpstreamOAuthProvider { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/upstream_oauth2/session.rs b/crates/storage-pg/src/upstream_oauth2/session.rs index 8b37aae2e..b961c4f8c 100644 --- a/crates/storage-pg/src/upstream_oauth2/session.rs +++ b/crates/storage-pg/src/upstream_oauth2/session.rs @@ -12,6 +12,7 @@ use mas_data_model::{ }; use mas_storage::{ Page, Pagination, + pagination::Node, upstream_oauth2::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository}, }; use rand::RngCore; @@ -91,6 +92,12 @@ struct SessionLookup { unlinked_at: Option>, } +impl Node for SessionLookup { + fn cursor(&self) -> Ulid { + self.upstream_oauth_authorization_session_id.into() + } +} + impl TryFrom for UpstreamOAuthAuthorizationSession { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/user/email.rs b/crates/storage-pg/src/user/email.rs index 916874ae8..0f998e55f 100644 --- a/crates/storage-pg/src/user/email.rs +++ b/crates/storage-pg/src/user/email.rs @@ -12,6 +12,7 @@ use mas_data_model::{ }; use mas_storage::{ Page, Pagination, + pagination::Node, user::{UserEmailFilter, UserEmailRepository}, }; use rand::RngCore; @@ -51,6 +52,12 @@ struct UserEmailLookup { created_at: DateTime, } +impl Node for UserEmailLookup { + fn cursor(&self) -> Ulid { + self.user_email_id.into() + } +} + impl From for UserEmail { fn from(e: UserEmailLookup) -> UserEmail { UserEmail { diff --git a/crates/storage-pg/src/user/mod.rs b/crates/storage-pg/src/user/mod.rs index 0be594556..bbf02567e 100644 --- a/crates/storage-pg/src/user/mod.rs +++ b/crates/storage-pg/src/user/mod.rs @@ -61,7 +61,9 @@ mod priv_ { #![allow(missing_docs)] use chrono::{DateTime, Utc}; + use mas_storage::pagination::Node; use sea_query::enum_def; + use ulid::Ulid; use uuid::Uuid; #[derive(Debug, Clone, sqlx::FromRow)] @@ -75,6 +77,12 @@ mod priv_ { pub(super) can_request_admin: bool, pub(super) is_guest: bool, } + + impl Node for UserLookup { + fn cursor(&self) -> Ulid { + self.user_id.into() + } + } } use priv_::{UserLookup, UserLookupIden}; diff --git a/crates/storage-pg/src/user/registration_token.rs b/crates/storage-pg/src/user/registration_token.rs index f64c8136d..5c9231aa1 100644 --- a/crates/storage-pg/src/user/registration_token.rs +++ b/crates/storage-pg/src/user/registration_token.rs @@ -8,6 +8,7 @@ use chrono::{DateTime, Utc}; use mas_data_model::{Clock, UserRegistrationToken}; use mas_storage::{ Page, Pagination, + pagination::Node, user::{UserRegistrationTokenFilter, UserRegistrationTokenRepository}, }; use rand::RngCore; @@ -53,6 +54,12 @@ struct UserRegistrationTokenLookup { revoked_at: Option>, } +impl Node for UserRegistrationTokenLookup { + fn cursor(&self) -> Ulid { + self.user_registration_token_id.into() + } +} + impl Filter for UserRegistrationTokenFilter { fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition { sea_query::Condition::all() @@ -230,7 +237,7 @@ impl UserRegistrationTokenRepository for PgUserRegistrationTokenRepository<'_> { filter: UserRegistrationTokenFilter, pagination: Pagination, ) -> Result, Self::Error> { - let (sql, values) = Query::select() + let (sql, arguments) = Query::select() .expr_as( Expr::col(( UserRegistrationTokens::Table, @@ -295,15 +302,14 @@ impl UserRegistrationTokenRepository for PgUserRegistrationTokenRepository<'_> { ) .build_sqlx(PostgresQueryBuilder); - let tokens = sqlx::query_as_with::<_, UserRegistrationTokenLookup, _>(&sql, values) + let edges: Vec = sqlx::query_as_with(&sql, arguments) .traced() .fetch_all(&mut *self.conn) - .await? - .into_iter() - .map(TryInto::try_into) - .collect::, _>>()?; + .await?; - let page = pagination.process(tokens); + let page = pagination + .process(edges) + .try_map(UserRegistrationToken::try_from)?; Ok(page) } @@ -705,7 +711,7 @@ mod tests { .await .unwrap(); - assert!(page.edges.iter().any(|t| t.id == unrevoked_token.id)); + assert!(page.edges.iter().any(|t| t.node.id == unrevoked_token.id)); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -867,7 +873,7 @@ mod tests { .await .unwrap(); assert_eq!(page.edges.len(), 1); - assert_eq!(page.edges[0].id, token2.id); + assert_eq!(page.edges[0].node.id, token2.id); // Test unused filter let unused_filter = UserRegistrationTokenFilter::new(clock.now()).with_been_used(false); @@ -886,7 +892,7 @@ mod tests { .await .unwrap(); assert_eq!(page.edges.len(), 1); - assert_eq!(page.edges[0].id, token3.id); + assert_eq!(page.edges[0].node.id, token3.id); let not_expired_filter = UserRegistrationTokenFilter::new(clock.now()).with_expired(false); let page = repo @@ -904,7 +910,7 @@ mod tests { .await .unwrap(); assert_eq!(page.edges.len(), 1); - assert_eq!(page.edges[0].id, token4.id); + assert_eq!(page.edges[0].node.id, token4.id); let not_revoked_filter = UserRegistrationTokenFilter::new(clock.now()).with_revoked(false); let page = repo @@ -941,7 +947,7 @@ mod tests { .await .unwrap(); assert_eq!(page.edges.len(), 1); - assert_eq!(page.edges[0].id, token4.id); + assert_eq!(page.edges[0].node.id, token4.id); // Test pagination let page = repo diff --git a/crates/storage-pg/src/user/session.rs b/crates/storage-pg/src/user/session.rs index 54645b3cb..8644d7666 100644 --- a/crates/storage-pg/src/user/session.rs +++ b/crates/storage-pg/src/user/session.rs @@ -14,6 +14,7 @@ use mas_data_model::{ }; use mas_storage::{ Page, Pagination, + pagination::Node, user::{BrowserSessionFilter, BrowserSessionRepository}, }; use rand::RngCore; @@ -64,6 +65,12 @@ struct SessionLookup { user_is_guest: bool, } +impl Node for SessionLookup { + fn cursor(&self) -> Ulid { + self.user_id.into() + } +} + impl TryFrom for BrowserSession { type Error = DatabaseInconsistencyError; diff --git a/crates/storage-pg/src/user/tests.rs b/crates/storage-pg/src/user/tests.rs index 0ee978914..98489d68d 100644 --- a/crates/storage-pg/src/user/tests.rs +++ b/crates/storage-pg/src/user/tests.rs @@ -190,7 +190,7 @@ async fn test_user_repo(pool: PgPool) { // Check the list method let list = repo.user().list(all, Pagination::first(10)).await.unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0].id, user.id); + assert_eq!(list.edges[0].node.id, user.id); let list = repo .user() @@ -205,7 +205,7 @@ async fn test_user_repo(pool: PgPool) { .await .unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0].id, user.id); + assert_eq!(list.edges[0].node.id, user.id); let list = repo .user() @@ -227,7 +227,7 @@ async fn test_user_repo(pool: PgPool) { .await .unwrap(); assert_eq!(list.edges.len(), 1); - assert_eq!(list.edges[0].id, user.id); + assert_eq!(list.edges[0].node.id, user.id); repo.save().await.unwrap(); } @@ -348,7 +348,7 @@ async fn test_user_email_repo(pool: PgPool) { .unwrap(); assert!(!emails.has_next_page); assert_eq!(emails.edges.len(), 1); - assert_eq!(emails.edges[0], user_email); + assert_eq!(emails.edges[0].node, user_email); // Listing emails from the email address should work let emails = repo @@ -358,7 +358,7 @@ async fn test_user_email_repo(pool: PgPool) { .unwrap(); assert!(!emails.has_next_page); assert_eq!(emails.edges.len(), 1); - assert_eq!(emails.edges[0], user_email); + assert_eq!(emails.edges[0].node, user_email); // Filtering on another email should not return anything let emails = repo @@ -648,7 +648,7 @@ async fn test_user_session(pool: PgPool) { .unwrap(); assert!(!session_list.has_next_page); assert_eq!(session_list.edges.len(), 1); - assert_eq!(session_list.edges[0], session); + assert_eq!(session_list.edges[0].node, session); let session_lookup = repo .browser_session() @@ -809,7 +809,7 @@ async fn test_user_session(pool: PgPool) { .await .unwrap(); assert_eq!(page.edges.len(), 1); - assert_eq!(page.edges[0].id, session.id); + assert_eq!(page.edges[0].node.id, session.id); // Try counting assert_eq!(repo.browser_session().count(filter).await.unwrap(), 1); diff --git a/crates/storage/src/pagination.rs b/crates/storage/src/pagination.rs index 01b8ed197..ad632cb10 100644 --- a/crates/storage/src/pagination.rs +++ b/crates/storage/src/pagination.rs @@ -16,12 +16,12 @@ pub struct InvalidPagination; /// Pagination parameters #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Pagination { +pub struct Pagination { /// The cursor to start from - pub before: Option, + pub before: Option, /// The cursor to end at - pub after: Option, + pub after: Option, /// The maximum number of items to return pub count: usize, @@ -40,16 +40,22 @@ pub enum PaginationDirection { Backward, } -impl Pagination { +/// A node in a page, with a cursor +pub trait Node { + /// The cursor of that particular node + fn cursor(&self) -> C; +} + +impl Pagination { /// Creates a new [`Pagination`] from user-provided parameters. /// /// # Errors /// /// Either `first` or `last` must be provided, else this function will /// return an [`InvalidPagination`] error. - pub const fn try_new( - before: Option, - after: Option, + pub fn try_new( + before: Option, + after: Option, first: Option, last: Option, ) -> Result { @@ -91,49 +97,57 @@ impl Pagination { /// Get items before the given cursor #[must_use] - pub const fn before(mut self, id: Ulid) -> Self { - self.before = Some(id); + pub fn before(mut self, cursor: C) -> Self { + self.before = Some(cursor); self } /// Clear the before cursor #[must_use] - pub const fn clear_before(mut self) -> Self { + pub fn clear_before(mut self) -> Self { self.before = None; self } /// Get items after the given cursor #[must_use] - pub const fn after(mut self, id: Ulid) -> Self { - self.after = Some(id); + pub fn after(mut self, cursor: C) -> Self { + self.after = Some(cursor); self } /// Clear the after cursor #[must_use] - pub const fn clear_after(mut self) -> Self { + pub fn clear_after(mut self) -> Self { self.after = None; self } /// Process a page returned by a paginated query #[must_use] - pub fn process(&self, mut edges: Vec) -> Page { - let is_full = edges.len() == (self.count + 1); + pub fn process>(&self, mut nodes: Vec) -> Page { + let is_full = nodes.len() == (self.count + 1); if is_full { - edges.pop(); + nodes.pop(); } let (has_previous_page, has_next_page) = match self.direction { PaginationDirection::Forward => (false, is_full), PaginationDirection::Backward => { // 6. If the last argument is provided, I reverse the order of the results - edges.reverse(); + nodes.reverse(); (is_full, false) } }; + let edges = nodes + .into_iter() + .map(|node| Edge { + cursor: node.cursor(), + node, + }) + .collect(); + Page { has_next_page, has_previous_page, @@ -142,9 +156,18 @@ impl Pagination { } } +/// An edge in a paginated result +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Edge { + /// The cursor of the edge + pub cursor: C, + /// The node of the edge + pub node: T, +} + /// A page of results returned by a paginated query #[derive(Debug, Clone, PartialEq, Eq)] -pub struct Page { +pub struct Page { /// When paginating forwards, this is true if there are more items after pub has_next_page: bool, @@ -152,21 +175,28 @@ pub struct Page { pub has_previous_page: bool, /// The items in the page - pub edges: Vec, + pub edges: Vec>, } -impl Page { +impl Page { /// Map the items in this page with the given function /// /// # Parameters /// /// * `f`: The function to map the items with #[must_use] - pub fn map(self, f: F) -> Page + pub fn map(self, mut f: F) -> Page where F: FnMut(T) -> T2, { - let edges = self.edges.into_iter().map(f).collect(); + let edges = self + .edges + .into_iter() + .map(|edge| Edge { + cursor: edge.cursor, + node: f(edge.node), + }) + .collect(); Page { has_next_page: self.has_next_page, has_previous_page: self.has_previous_page, @@ -183,11 +213,21 @@ impl Page { /// # Errors /// /// Returns the first error encountered while mapping the items - pub fn try_map(self, f: F) -> Result, E> + pub fn try_map(self, mut f: F) -> Result, E> where F: FnMut(T) -> Result, { - let edges: Result, E> = self.edges.into_iter().map(f).collect(); + let edges: Result>, E> = self + .edges + .into_iter() + .map(|edge| { + Ok(Edge { + cursor: edge.cursor, + node: f(edge.node)?, + }) + }) + .collect(); + Ok(Page { has_next_page: self.has_next_page, has_previous_page: self.has_previous_page, diff --git a/crates/storage/src/queue/tasks.rs b/crates/storage/src/queue/tasks.rs index eb16f6e29..3558314cf 100644 --- a/crates/storage/src/queue/tasks.rs +++ b/crates/storage/src/queue/tasks.rs @@ -384,7 +384,7 @@ impl ExpireInactiveOAuthSessionsJob { let last_edge = page.edges.last()?; Some(Self { threshold: self.threshold, - after: Some(last_edge.id), + after: Some(last_edge.cursor), }) } } @@ -441,7 +441,7 @@ impl ExpireInactiveCompatSessionsJob { let last_edge = page.edges.last()?; Some(Self { threshold: self.threshold, - after: Some(last_edge.id), + after: Some(last_edge.cursor), }) } } @@ -498,7 +498,7 @@ impl ExpireInactiveUserSessionsJob { let last_edge = page.edges.last()?; Some(Self { threshold: self.threshold, - after: Some(last_edge.id), + after: Some(last_edge.cursor), }) } }