Better OAuth 2.0 sessions pagination and filtering

crates/graphql/src/model/browser_sessions.rs

@@ -19,10 +19,6 @@ use mas_storage::{user::BrowserSessionRepository, RepositoryAccess};
 use super::{NodeType, User};
 use crate::state::ContextExt;
 
-/// A browser session represents a logged in user in a browser.
-#[derive(Description)]
-pub struct BrowserSession(pub mas_data_model::BrowserSession);
-
 /// The state of a browser session.
 #[derive(Enum, Copy, Clone, Eq, PartialEq)]
 pub enum BrowserSessionState {
@@ -33,6 +29,10 @@ pub enum BrowserSessionState {
     Finished,
 }
 
+/// A browser session represents a logged in user in a browser.
+#[derive(Description)]
+pub struct BrowserSession(pub mas_data_model::BrowserSession);
+
 impl From<mas_data_model::BrowserSession> for BrowserSession {
     fn from(v: mas_data_model::BrowserSession) -> Self {
         Self(v)

crates/graphql/src/model/oauth.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use anyhow::Context as _;
-use async_graphql::{Context, Description, Object, ID};
+use async_graphql::{Context, Description, Enum, Object, ID};
 use chrono::{DateTime, Utc};
 use mas_data_model::SessionState;
 use mas_storage::{oauth2::OAuth2ClientRepository, user::BrowserSessionRepository};
@@ -24,6 +24,16 @@ use url::Url;
 use super::{BrowserSession, NodeType, User};
 use crate::state::ContextExt;
 
+/// The state of an OAuth 2.0 session.
+#[derive(Enum, Copy, Clone, Eq, PartialEq)]
+pub enum OAuth2SessionState {
+    /// The session is active.
+    Active,
+
+    /// The session is no longer active.
+    Finished,
+}
+
 /// An OAuth 2.0 session represents a client session which used the OAuth APIs
 /// to login.
 #[derive(Description)]

crates/graphql/src/model/users.rs

@@ -19,7 +19,7 @@ use async_graphql::{
 use chrono::{DateTime, Utc};
 use mas_storage::{
     compat::{CompatSessionFilter, CompatSsoLoginFilter, CompatSsoLoginRepository},
-    oauth2::OAuth2SessionRepository,
+    oauth2::{OAuth2SessionFilter, OAuth2SessionRepository},
     upstream_oauth2::UpstreamOAuthLinkRepository,
     user::{BrowserSessionFilter, BrowserSessionRepository, UserEmailFilter, UserEmailRepository},
     Pagination, RepositoryAccess,
@@ -34,6 +34,7 @@ use crate::{
         browser_sessions::BrowserSessionState,
         compat_sessions::{CompatSessionState, CompatSessionType},
         matrix::MatrixUser,
+        oauth::OAuth2SessionState,
         CompatSession,
     },
     state::ContextExt,
@@ -365,17 +366,23 @@ impl User {
     }
 
     /// Get the list of OAuth 2.0 sessions, chronologically sorted
+    #[allow(clippy::too_many_arguments)]
     async fn oauth2_sessions(
         &self,
         ctx: &Context<'_>,
 
+        #[graphql(name = "state", desc = "List only sessions in the given state.")]
+        state_param: Option<OAuth2SessionState>,
+
+        #[graphql(desc = "List only sessions for the given client.")] client: Option<ID>,
+
         #[graphql(desc = "Returns the elements in the list that come after the cursor.")]
         after: Option<String>,
         #[graphql(desc = "Returns the elements in the list that come before the cursor.")]
         before: Option<String>,
         #[graphql(desc = "Returns the first *n* elements from the list.")] first: Option<i32>,
         #[graphql(desc = "Returns the last *n* elements from the list.")] last: Option<i32>,
-    ) -> Result<Connection<Cursor, OAuth2Session>, async_graphql::Error> {
+    ) -> Result<Connection<Cursor, OAuth2Session, PreloadedTotalCount>, async_graphql::Error> {
         let state = ctx.state();
         let mut repo = state.repository().await?;
 
@@ -393,14 +400,49 @@ impl User {
                     .transpose()?;
                 let pagination = Pagination::try_new(before_id, after_id, first, last)?;
 
-                let page = repo
-                    .oauth2_session()
-                    .list_paginated(&self.0, pagination)
-                    .await?;
+                let client = if let Some(id) = client {
+                    // Load the client if we're filtering by it
+                    let id = NodeType::OAuth2Client.extract_ulid(&id)?;
+                    let client = repo
+                        .oauth2_client()
+                        .lookup(id)
+                        .await?
+                        .ok_or(async_graphql::Error::new("Unknown client ID"))?;
+
+                    Some(client)
+                } else {
+                    None
+                };
+
+                let filter = OAuth2SessionFilter::new().for_user(&self.0);
+
+                let filter = match state_param {
+                    Some(OAuth2SessionState::Active) => filter.active_only(),
+                    Some(OAuth2SessionState::Finished) => filter.finished_only(),
+                    None => filter,
+                };
+
+                let filter = match client.as_ref() {
+                    Some(client) => filter.for_client(client),
+                    None => filter,
+                };
+
+                let page = repo.oauth2_session().list(filter, pagination).await?;
+
+                let count = if ctx.look_ahead().field("totalCount").exists() {
+                    Some(repo.oauth2_session().count(filter).await?)
+                } else {
+                    None
+                };
 
                 repo.cancel().await?;
 
-                let mut connection = Connection::new(page.has_previous_page, page.has_next_page);
+                let mut connection = Connection::with_additional_fields(
+                    page.has_previous_page,
+                    page.has_next_page,
+                    PreloadedTotalCount(count),
+                );
+
                 connection.edges.extend(page.edges.into_iter().map(|s| {
                     Edge::new(
                         OpaqueCursor(NodeCursor(NodeType::OAuth2Session, s.id)),