storage: list and count methods for upstream oauth sessions

crates/storage-pg/src/iden.rs

@@ -140,6 +140,29 @@ pub enum UpstreamOAuthLinks {
     CreatedAt,
 }
 
+#[derive(sea_query::Iden)]
+#[iden = "upstream_oauth_authorization_sessions"]
+pub enum UpstreamOAuthAuthorizationSessions {
+    Table,
+    #[iden = "upstream_oauth_authorization_session_id"]
+    UpstreamOAuthAuthorizationSessionId,
+    #[iden = "upstream_oauth_provider_id"]
+    UpstreamOAuthProviderId,
+    #[iden = "upstream_oauth_link_id"]
+    UpstreamOAuthLinkId,
+    State,
+    CodeChallengeVerifier,
+    Nonce,
+    IdToken,
+    IdTokenClaims,
+    ExtraCallbackParameters,
+    Userinfo,
+    CreatedAt,
+    CompletedAt,
+    ConsumedAt,
+    UnlinkedAt,
+}
+
 #[derive(sea_query::Iden)]
 pub enum UserRegistrationTokens {
     Table,

crates/storage-pg/src/upstream_oauth2/mod.rs

@@ -29,7 +29,7 @@ mod tests {
         upstream_oauth2::{
             UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository, UpstreamOAuthProviderFilter,
             UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository,
-            UpstreamOAuthSessionRepository,
+            UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository,
         },
         user::UserRepository,
     };
@@ -262,6 +262,29 @@ mod tests {
             1
         );
 
+        // Test listing and counting sessions
+        let session_filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
+
+        // Count the sessions for the provider
+        let session_count = repo
+            .upstream_oauth_session()
+            .count(session_filter)
+            .await
+            .unwrap();
+        assert_eq!(session_count, 1);
+
+        // List the sessions for the provider
+        let session_page = repo
+            .upstream_oauth_session()
+            .list(session_filter, Pagination::first(10))
+            .await
+            .unwrap();
+
+        assert_eq!(session_page.edges.len(), 1);
+        assert_eq!(session_page.edges[0].id, session.id);
+        assert!(!session_page.has_next_page);
+        assert!(!session_page.has_previous_page);
+
         // Try deleting the provider
         repo.upstream_oauth_provider()
             .delete(provider)
@@ -423,4 +446,136 @@ mod tests {
                 .is_empty()
         );
     }
+
+    /// Test that the pagination works as expected in the upstream OAuth
+    /// session repository
+    #[sqlx::test(migrator = "crate::MIGRATOR")]
+    async fn test_session_repository_pagination(pool: PgPool) {
+        let scope = Scope::from_iter([OPENID]);
+
+        let mut rng = rand_chacha::ChaChaRng::seed_from_u64(42);
+        let clock = MockClock::default();
+        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
+
+        // Create a provider
+        let provider = repo
+            .upstream_oauth_provider()
+            .add(
+                &mut rng,
+                &clock,
+                UpstreamOAuthProviderParams {
+                    issuer: Some("https://example.com/".to_owned()),
+                    human_name: None,
+                    brand_name: None,
+                    scope,
+                    token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None,
+                    id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
+                    fetch_userinfo: false,
+                    userinfo_signed_response_alg: None,
+                    token_endpoint_signing_alg: None,
+                    client_id: "client-id".to_owned(),
+                    encrypted_client_secret: None,
+                    claims_imports: UpstreamOAuthProviderClaimsImports::default(),
+                    token_endpoint_override: None,
+                    authorization_endpoint_override: None,
+                    userinfo_endpoint_override: None,
+                    jwks_uri_override: None,
+                    discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc,
+                    pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto,
+                    response_mode: None,
+                    additional_authorization_parameters: Vec::new(),
+                    forward_login_hint: false,
+                    ui_order: 0,
+                },
+            )
+            .await
+            .unwrap();
+
+        let filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
+
+        // Count the number of sessions before we start
+        assert_eq!(
+            repo.upstream_oauth_session().count(filter).await.unwrap(),
+            0
+        );
+
+        let mut ids = Vec::with_capacity(20);
+        // Create 20 sessions
+        for idx in 0..20 {
+            let state = format!("state-{idx}");
+            let session = repo
+                .upstream_oauth_session()
+                .add(&mut rng, &clock, &provider, state, None, None)
+                .await
+                .unwrap();
+            ids.push(session.id);
+            clock.advance(Duration::microseconds(10 * 1000 * 1000));
+        }
+
+        // Now we have 20 sessions
+        assert_eq!(
+            repo.upstream_oauth_session().count(filter).await.unwrap(),
+            20
+        );
+
+        // Lookup the first 10 items
+        let page = repo
+            .upstream_oauth_session()
+            .list(filter, Pagination::first(10))
+            .await
+            .unwrap();
+
+        // It returned the first 10 items
+        assert!(page.has_next_page);
+        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect();
+        assert_eq!(&edge_ids, &ids[..10]);
+
+        // Lookup the next 10 items
+        let page = repo
+            .upstream_oauth_session()
+            .list(filter, Pagination::first(10).after(ids[9]))
+            .await
+            .unwrap();
+
+        // It returned the next 10 items
+        assert!(!page.has_next_page);
+        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect();
+        assert_eq!(&edge_ids, &ids[10..]);
+
+        // Lookup the last 10 items
+        let page = repo
+            .upstream_oauth_session()
+            .list(filter, Pagination::last(10))
+            .await
+            .unwrap();
+
+        // It returned the last 10 items
+        assert!(page.has_previous_page);
+        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect();
+        assert_eq!(&edge_ids, &ids[10..]);
+
+        // Lookup the previous 10 items
+        let page = repo
+            .upstream_oauth_session()
+            .list(filter, Pagination::last(10).before(ids[10]))
+            .await
+            .unwrap();
+
+        // It returned the previous 10 items
+        assert!(!page.has_previous_page);
+        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect();
+        assert_eq!(&edge_ids, &ids[..10]);
+
+        // Lookup 5 items between two IDs
+        let page = repo
+            .upstream_oauth_session()
+            .list(filter, Pagination::first(10).after(ids[5]).before(ids[11]))
+            .await
+            .unwrap();
+
+        // It returned the items in between
+        assert!(!page.has_next_page);
+        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect();
+        assert_eq!(&edge_ids, &ids[6..11]);
+    }
 }

crates/storage-pg/src/upstream_oauth2/session.rs

@@ -10,13 +10,36 @@ use mas_data_model::{
     UpstreamOAuthAuthorizationSession, UpstreamOAuthAuthorizationSessionState, UpstreamOAuthLink,
     UpstreamOAuthProvider,
 };
-use mas_storage::{Clock, upstream_oauth2::UpstreamOAuthSessionRepository};
+use mas_storage::{
+    Clock, Page, Pagination,
+    upstream_oauth2::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository},
+};
 use rand::RngCore;
+use sea_query::{Expr, PostgresQueryBuilder, Query, enum_def};
+use sea_query_binder::SqlxBinder;
 use sqlx::PgConnection;
 use ulid::Ulid;
 use uuid::Uuid;
 
-use crate::{DatabaseError, DatabaseInconsistencyError, tracing::ExecuteExt};
+use crate::{
+    DatabaseError, DatabaseInconsistencyError,
+    filter::{Filter, StatementExt},
+    iden::UpstreamOAuthAuthorizationSessions,
+    pagination::QueryBuilderExt,
+    tracing::ExecuteExt,
+};
+
+impl Filter for UpstreamOAuthSessionFilter<'_> {
+    fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
+        sea_query::Condition::all().add_option(self.provider().map(|provider| {
+            Expr::col((
+                UpstreamOAuthAuthorizationSessions::Table,
+                UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId,
+            ))
+            .eq(Uuid::from(provider.id))
+        }))
+    }
+}
 
 /// An implementation of [`UpstreamOAuthSessionRepository`] for a PostgreSQL
 /// connection
@@ -32,6 +55,8 @@ impl<'c> PgUpstreamOAuthSessionRepository<'c> {
     }
 }
 
+#[derive(sqlx::FromRow)]
+#[enum_def]
 struct SessionLookup {
     upstream_oauth_authorization_session_id: Uuid,
     upstream_oauth_provider_id: Uuid,
@@ -346,4 +371,173 @@ impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> {
 
         Ok(upstream_oauth_authorization_session)
     }
+
+    #[tracing::instrument(
+        name = "db.upstream_oauth_authorization_session.list",
+        skip_all,
+        fields(
+            db.query.text,
+        ),
+        err,
+    )]
+    async fn list(
+        &mut self,
+        filter: UpstreamOAuthSessionFilter<'_>,
+        pagination: Pagination,
+    ) -> Result<Page<UpstreamOAuthAuthorizationSession>, Self::Error> {
+        let (sql, arguments) = Query::select()
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
+                )),
+                SessionLookupIden::UpstreamOauthAuthorizationSessionId,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId,
+                )),
+                SessionLookupIden::UpstreamOauthProviderId,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthLinkId,
+                )),
+                SessionLookupIden::UpstreamOauthLinkId,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::State,
+                )),
+                SessionLookupIden::State,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::CodeChallengeVerifier,
+                )),
+                SessionLookupIden::CodeChallengeVerifier,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::Nonce,
+                )),
+                SessionLookupIden::Nonce,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::IdToken,
+                )),
+                SessionLookupIden::IdToken,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::IdTokenClaims,
+                )),
+                SessionLookupIden::IdTokenClaims,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::ExtraCallbackParameters,
+                )),
+                SessionLookupIden::ExtraCallbackParameters,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::Userinfo,
+                )),
+                SessionLookupIden::Userinfo,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::CreatedAt,
+                )),
+                SessionLookupIden::CreatedAt,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::CompletedAt,
+                )),
+                SessionLookupIden::CompletedAt,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::ConsumedAt,
+                )),
+                SessionLookupIden::ConsumedAt,
+            )
+            .expr_as(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UnlinkedAt,
+                )),
+                SessionLookupIden::UnlinkedAt,
+            )
+            .from(UpstreamOAuthAuthorizationSessions::Table)
+            .apply_filter(filter)
+            .generate_pagination(
+                (
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
+                ),
+                pagination,
+            )
+            .build_sqlx(PostgresQueryBuilder);
+
+        let edges: Vec<SessionLookup> = sqlx::query_as_with(&sql, arguments)
+            .traced()
+            .fetch_all(&mut *self.conn)
+            .await?;
+
+        let page = pagination
+            .process(edges)
+            .try_map(UpstreamOAuthAuthorizationSession::try_from)?;
+
+        Ok(page)
+    }
+
+    #[tracing::instrument(
+        name = "db.upstream_oauth_authorization_session.count",
+        skip_all,
+        fields(
+            db.query.text,
+        ),
+        err,
+    )]
+    async fn count(
+        &mut self,
+        filter: UpstreamOAuthSessionFilter<'_>,
+    ) -> Result<usize, Self::Error> {
+        let (sql, arguments) = Query::select()
+            .expr(
+                Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
+                ))
+                .count(),
+            )
+            .from(UpstreamOAuthAuthorizationSessions::Table)
+            .apply_filter(filter)
+            .build_sqlx(PostgresQueryBuilder);
+
+        let count: i64 = sqlx::query_scalar_with(&sql, arguments)
+            .traced()
+            .fetch_one(&mut *self.conn)
+            .await?;
+
+        count
+            .try_into()
+            .map_err(DatabaseError::to_invalid_operation)
+    }
 }

crates/storage/src/upstream_oauth2/mod.rs

@@ -16,5 +16,5 @@ pub use self::{
     provider::{
         UpstreamOAuthProviderFilter, UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository,
     },
-    session::UpstreamOAuthSessionRepository,
+    session::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository},
 };

crates/storage/src/upstream_oauth2/session.rs

@@ -9,7 +9,36 @@ use mas_data_model::{UpstreamOAuthAuthorizationSession, UpstreamOAuthLink, Upstr
 use rand_core::RngCore;
 use ulid::Ulid;
 
-use crate::{Clock, repository_impl};
+use crate::{Clock, Pagination, pagination::Page, repository_impl};
+
+/// Filter parameters for listing upstream OAuth sessions
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
+pub struct UpstreamOAuthSessionFilter<'a> {
+    provider: Option<&'a UpstreamOAuthProvider>,
+}
+
+impl<'a> UpstreamOAuthSessionFilter<'a> {
+    /// Create a new [`UpstreamOAuthSessionFilter`] with default values
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Set the upstream OAuth provider for which to list sessions
+    #[must_use]
+    pub fn for_provider(mut self, provider: &'a UpstreamOAuthProvider) -> Self {
+        self.provider = Some(provider);
+        self
+    }
+
+    /// Get the upstream OAuth provider filter
+    ///
+    /// Returns [`None`] if no filter was set
+    #[must_use]
+    pub fn provider(&self) -> Option<&UpstreamOAuthProvider> {
+        self.provider
+    }
+}
 
 /// An [`UpstreamOAuthSessionRepository`] helps interacting with
 /// [`UpstreamOAuthAuthorizationSession`] saved in the storage backend
@@ -112,6 +141,36 @@ pub trait UpstreamOAuthSessionRepository: Send + Sync {
         clock: &dyn Clock,
         upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession,
     ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error>;
+
+    /// List [`UpstreamOAuthAuthorizationSession`] with the given filter and
+    /// pagination
+    ///
+    /// # Parameters
+    ///
+    /// * `filter`: The filter to apply
+    /// * `pagination`: The pagination parameters
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn list(
+        &mut self,
+        filter: UpstreamOAuthSessionFilter<'_>,
+        pagination: Pagination,
+    ) -> Result<Page<UpstreamOAuthAuthorizationSession>, Self::Error>;
+
+    /// Count the number of [`UpstreamOAuthAuthorizationSession`] with the given
+    /// filter
+    ///
+    /// # Parameters
+    ///
+    /// * `filter`: The filter to apply
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn count(&mut self, filter: UpstreamOAuthSessionFilter<'_>)
+    -> Result<usize, Self::Error>;
 }
 
 repository_impl!(UpstreamOAuthSessionRepository:
@@ -146,4 +205,12 @@ repository_impl!(UpstreamOAuthSessionRepository:
         clock: &dyn Clock,
         upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession,
     ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error>;
+
+    async fn list(
+        &mut self,
+        filter: UpstreamOAuthSessionFilter<'_>,
+        pagination: Pagination,
+    ) -> Result<Page<UpstreamOAuthAuthorizationSession>, Self::Error>;
+
+    async fn count(&mut self, filter: UpstreamOAuthSessionFilter<'_>) -> Result<usize, Self::Error>;
 );