Mark sessions as finished in bulk

crates/cli/src/commands/manage.rs

@@ -26,11 +26,12 @@ use mas_handlers::HttpClientFactory;
 use mas_matrix::HomeserverConnection;
 use mas_matrix_synapse::SynapseConnection;
 use mas_storage::{
-    compat::{CompatAccessTokenRepository, CompatSessionRepository},
+    compat::{CompatAccessTokenRepository, CompatSessionFilter, CompatSessionRepository},
     job::{
         DeactivateUserJob, JobRepositoryExt, ProvisionUserJob, ReactivateUserJob, SyncDevicesJob,
     },
-    user::{UserEmailRepository, UserPasswordRepository, UserRepository},
+    oauth2::OAuth2SessionFilter,
+    user::{BrowserSessionFilter, UserEmailRepository, UserPasswordRepository, UserRepository},
     Clock, RepositoryAccess, SystemClock,
 };
 use mas_storage_pg::{DatabaseError, PgRepository};
@@ -348,83 +349,43 @@ impl Options {
                     .await?
                     .context("User not found")?;
 
-                let compat_sessions_ids: Vec<Uuid> = sqlx::query_scalar(
-                    r"
-                        SELECT compat_session_id FROM compat_sessions
-                        WHERE user_id = $1 AND finished_at IS NULL
-                    ",
-                )
-                .bind(Uuid::from(user.id))
-                .fetch_all(&mut **repo)
-                .await?;
-
-                for id in compat_sessions_ids {
-                    let id = id.into();
-                    let compat_session = repo
-                        .compat_session()
-                        .lookup(id)
-                        .await?
-                        .context("Session not found")?;
-                    info!(%compat_session.id, %compat_session.device, "Killing compat session");
-
-                    if dry_run {
-                        continue;
-                    }
-                }
-
-                let oauth2_sessions_ids: Vec<Uuid> = sqlx::query_scalar(
-                    r"
-                        SELECT oauth2_sessions.oauth2_session_id
-                        FROM oauth2_sessions
-                        INNER JOIN user_sessions USING (user_session_id)
-                        WHERE user_sessions.user_id = $1 AND oauth2_sessions.finished_at IS NULL
-                    ",
-                )
-                .bind(Uuid::from(user.id))
-                .fetch_all(&mut **repo)
-                .await?;
-
-                for id in oauth2_sessions_ids {
-                    let id = id.into();
-                    let oauth2_session = repo
-                        .oauth2_session()
-                        .lookup(id)
-                        .await?
-                        .context("Session not found")?;
-                    info!(%oauth2_session.id, %oauth2_session.scope, "Killing oauth2 session");
+                let filter = CompatSessionFilter::new().for_user(&user).active_only();
+                let affected = if dry_run {
+                    repo.compat_session().count(filter).await?
+                } else {
+                    repo.compat_session().finish_bulk(&clock, filter).await?
+                };
 
-                    if dry_run {
-                        continue;
-                    }
-                    repo.oauth2_session().finish(&clock, oauth2_session).await?;
+                match affected {
+                    0 => info!("No active compatibility sessions to end"),
+                    1 => info!("Ended 1 active compatibility session"),
+                    _ => info!("Ended {affected} active compatibility sessions"),
                 }
 
-                let user_sessions_ids: Vec<Uuid> = sqlx::query_scalar(
-                    r"
-                        SELECT user_session_id FROM user_sessions
-                        WHERE user_id = $1 AND finished_at IS NULL
-                    ",
-                )
-                .bind(Uuid::from(user.id))
-                .fetch_all(&mut **repo)
-                .await?;
+                let filter = OAuth2SessionFilter::new().for_user(&user).active_only();
+                let affected = if dry_run {
+                    repo.oauth2_session().count(filter).await?
+                } else {
+                    repo.oauth2_session().finish_bulk(&clock, filter).await?
+                };
 
-                for id in user_sessions_ids {
-                    let id = id.into();
-                    let browser_session = repo
-                        .browser_session()
-                        .lookup(id)
-                        .await?
-                        .context("Session not found")?;
-                    info!(%browser_session.id, "Killing browser session");
+                match affected {
+                    0 => info!("No active compatibility sessions to end"),
+                    1 => info!("Ended 1 active OAuth 2.0 session"),
+                    _ => info!("Ended {affected} active OAuth 2.0 sessions"),
+                };
 
-                    if dry_run {
-                        continue;
-                    }
+                let filter = BrowserSessionFilter::new().for_user(&user).active_only();
+                let affected = if dry_run {
+                    repo.browser_session().count(filter).await?
+                } else {
+                    repo.browser_session().finish_bulk(&clock, filter).await?
+                };
 
-                    repo.browser_session()
-                        .finish(&clock, browser_session)
-                        .await?;
+                match affected {
+                    0 => info!("No active browser sessions to end"),
+                    1 => info!("Ended 1 active browser session"),
+                    _ => info!("Ended {affected} active browser sessions"),
                 }
 
                 // Schedule a job to sync the devices of the user with the homeserver

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

@@ -288,6 +288,16 @@ mod tests {
                 .unwrap(),
             1
         );
+
+        // Check that we can batch finish sessions
+        let affected = repo
+            .compat_session()
+            .finish_bulk(&clock, all.sso_login_only().active_only())
+            .await
+            .unwrap();
+        assert_eq!(affected, 1);
+        assert_eq!(repo.compat_session().count(finished).await.unwrap(), 2);
+        assert_eq!(repo.compat_session().count(active).await.unwrap(), 0);
     }
 
     #[sqlx::test(migrator = "crate::MIGRATOR")]

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

@@ -271,7 +271,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
 
         sqlx::query!(
             r#"
-                INSERT INTO compat_sessions 
+                INSERT INTO compat_sessions
                     (compat_session_id, user_id, device_id,
                      user_session_id, created_at, is_synapse_admin)
                 VALUES ($1, $2, $3, $4, $5, $6)
@@ -341,6 +341,64 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
         Ok(compat_session)
     }
 
+    #[tracing::instrument(
+        name = "db.compat_session.finish_bulk",
+        skip_all,
+        fields(db.statement),
+        err,
+    )]
+    async fn finish_bulk(
+        &mut self,
+        clock: &dyn Clock,
+        filter: CompatSessionFilter<'_>,
+    ) -> Result<usize, Self::Error> {
+        let finished_at = clock.now();
+        let (sql, arguments) = Query::update()
+            .table(CompatSessions::Table)
+            .value(CompatSessions::FinishedAt, finished_at)
+            .and_where_option(filter.user().map(|user| {
+                Expr::col((CompatSessions::Table, CompatSessions::UserId)).eq(Uuid::from(user.id))
+            }))
+            .and_where_option(filter.state().map(|state| {
+                if state.is_active() {
+                    Expr::col((CompatSessions::Table, CompatSessions::FinishedAt)).is_null()
+                } else {
+                    Expr::col((CompatSessions::Table, CompatSessions::FinishedAt)).is_not_null()
+                }
+            }))
+            .and_where_option(filter.auth_type().map(|auth_type| {
+                // This builds either a:
+                // `WHERE compat_session_id = ANY(...)`
+                // or a `WHERE compat_session_id <> ALL(...)`
+                let compat_sso_logins = Query::select()
+                    .expr(Expr::col((
+                        CompatSsoLogins::Table,
+                        CompatSsoLogins::CompatSessionId,
+                    )))
+                    .from(CompatSsoLogins::Table)
+                    .take();
+
+                if auth_type.is_sso_login() {
+                    Expr::col((CompatSessions::Table, CompatSessions::CompatSessionId))
+                        .eq(Expr::any(compat_sso_logins))
+                } else {
+                    Expr::col((CompatSessions::Table, CompatSessions::CompatSessionId))
+                        .ne(Expr::all(compat_sso_logins))
+                }
+            }))
+            .and_where_option(filter.device().map(|device| {
+                Expr::col((CompatSessions::Table, CompatSessions::DeviceId)).eq(device.as_str())
+            }))
+            .build_sqlx(PostgresQueryBuilder);
+
+        let res = sqlx::query_with(&sql, arguments)
+            .traced()
+            .execute(&mut *self.conn)
+            .await?;
+
+        Ok(res.rows_affected().try_into().unwrap_or(usize::MAX))
+    }
+
     #[tracing::instrument(
         name = "db.compat_session.list",
         skip_all,
@@ -545,7 +603,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
                   , last_active_ip = COALESCE(t.last_active_ip, compat_sessions.last_active_ip)
                 FROM (
                     SELECT *
-                    FROM UNNEST($1::uuid[], $2::timestamptz[], $3::inet[]) 
+                    FROM UNNEST($1::uuid[], $2::timestamptz[], $3::inet[])
                         AS t(compat_session_id, last_active_at, last_active_ip)
                 ) AS t
                 WHERE compat_sessions.compat_session_id = t.compat_session_id

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

@@ -709,6 +709,37 @@ mod tests {
         assert_eq!(list.edges.len(), 1);
         assert_eq!(list.edges[0], session11);
         assert_eq!(repo.oauth2_session().count(filter).await.unwrap(), 1);
+
+        // Finish all sessions of a client in batch
+        let affected = repo
+            .oauth2_session()
+            .finish_bulk(
+                &clock,
+                OAuth2SessionFilter::new()
+                    .for_client(&client1)
+                    .active_only(),
+            )
+            .await
+            .unwrap();
+        assert_eq!(affected, 1);
+
+        // We should have 3 finished sessions
+        assert_eq!(
+            repo.oauth2_session()
+                .count(OAuth2SessionFilter::new().finished_only())
+                .await
+                .unwrap(),
+            3
+        );
+
+        // We should have 1 active sessions
+        assert_eq!(
+            repo.oauth2_session()
+                .count(OAuth2SessionFilter::new().active_only())
+                .await
+                .unwrap(),
+            1
+        );
     }
 
     /// Test the [`OAuth2DeviceCodeGrantRepository`] implementation

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

@@ -206,6 +206,51 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> {
         })
     }
 
+    #[tracing::instrument(
+        name = "db.oauth2_session.finish_bulk",
+        skip_all,
+        fields(
+            db.statement,
+        ),
+        err,
+    )]
+    async fn finish_bulk(
+        &mut self,
+        clock: &dyn Clock,
+        filter: OAuth2SessionFilter<'_>,
+    ) -> Result<usize, Self::Error> {
+        let finished_at = clock.now();
+        let (sql, arguments) = Query::update()
+            .table(OAuth2Sessions::Table)
+            .value(OAuth2Sessions::FinishedAt, finished_at)
+            .and_where_option(filter.user().map(|user| {
+                Expr::col((OAuth2Sessions::Table, OAuth2Sessions::UserId)).eq(Uuid::from(user.id))
+            }))
+            .and_where_option(filter.client().map(|client| {
+                Expr::col((OAuth2Sessions::Table, OAuth2Sessions::OAuth2ClientId))
+                    .eq(Uuid::from(client.id))
+            }))
+            .and_where_option(filter.state().map(|state| {
+                if state.is_active() {
+                    Expr::col((OAuth2Sessions::Table, OAuth2Sessions::FinishedAt)).is_null()
+                } else {
+                    Expr::col((OAuth2Sessions::Table, OAuth2Sessions::FinishedAt)).is_not_null()
+                }
+            }))
+            .and_where_option(filter.scope().map(|scope| {
+                let scope: Vec<String> = scope.iter().map(|s| s.as_str().to_owned()).collect();
+                Expr::col((OAuth2Sessions::Table, OAuth2Sessions::ScopeList)).contains(scope)
+            }))
+            .build_sqlx(PostgresQueryBuilder);
+
+        let res = sqlx::query_with(&sql, arguments)
+            .traced()
+            .execute(&mut *self.conn)
+            .await?;
+
+        Ok(res.rows_affected().try_into().unwrap_or(usize::MAX))
+    }
+
     #[tracing::instrument(
         name = "db.oauth2_session.finish",
         skip_all,