diff --git a/crates/storage-pg/.sqlx/query-d0c02576b1550fe2eb877d24f7cdfc819307ee0c47af9fbbf1a3b484290b321d.json b/crates/storage-pg/.sqlx/query-047990a99794b565c2cad396946299db5b617f52f6c24bcca0a24c0c185c4478.json similarity index 61% rename from crates/storage-pg/.sqlx/query-d0c02576b1550fe2eb877d24f7cdfc819307ee0c47af9fbbf1a3b484290b321d.json rename to crates/storage-pg/.sqlx/query-047990a99794b565c2cad396946299db5b617f52f6c24bcca0a24c0c185c4478.json index 550b6ece..8eb80f53 100644 --- a/crates/storage-pg/.sqlx/query-d0c02576b1550fe2eb877d24f7cdfc819307ee0c47af9fbbf1a3b484290b321d.json +++ b/crates/storage-pg/.sqlx/query-047990a99794b565c2cad396946299db5b617f52f6c24bcca0a24c0c185c4478.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n UPDATE oauth2_sessions\n SET last_active_at = GREATEST(t.last_active_at, oauth2_sessions.last_active_at)\n , last_active_ip = COALESCE(t.last_active_ip, oauth2_sessions.last_active_ip)\n FROM (\n SELECT *\n FROM UNNEST($1::uuid[], $2::timestamptz[], $3::inet[]) \n AS t(oauth2_session_id, last_active_at, last_active_ip)\n ) AS t\n WHERE oauth2_sessions.oauth2_session_id = t.oauth2_session_id\n ", + "query": "\n UPDATE oauth2_sessions\n SET last_active_at = GREATEST(t.last_active_at, oauth2_sessions.last_active_at)\n , last_active_ip = COALESCE(t.last_active_ip, oauth2_sessions.last_active_ip)\n FROM (\n SELECT *\n FROM UNNEST($1::uuid[], $2::timestamptz[], $3::inet[])\n AS t(oauth2_session_id, last_active_at, last_active_ip)\n ) AS t\n WHERE oauth2_sessions.oauth2_session_id = t.oauth2_session_id\n ", "describe": { "columns": [], "parameters": { @@ -12,5 +12,5 @@ }, "nullable": [] }, - "hash": "d0c02576b1550fe2eb877d24f7cdfc819307ee0c47af9fbbf1a3b484290b321d" + "hash": "047990a99794b565c2cad396946299db5b617f52f6c24bcca0a24c0c185c4478" } diff --git a/crates/storage-pg/src/compat/session.rs b/crates/storage-pg/src/compat/session.rs index b8580108..e372585c 100644 --- a/crates/storage-pg/src/compat/session.rs +++ b/crates/storage-pg/src/compat/session.rs @@ -33,6 +33,7 @@ use url::Url; use uuid::Uuid; use crate::{ + filter::{Filter, StatementExt, StatementWithJoinsExt}, iden::{CompatSessions, CompatSsoLogins}, pagination::QueryBuilderExt, tracing::ExecuteExt, @@ -203,6 +204,57 @@ impl TryFrom for (CompatSession, Option { + fn generate_condition(&self, has_joins: bool) -> impl sea_query::IntoCondition { + sea_query::Condition::all() + .add_option(self.user().map(|user| { + Expr::col((CompatSessions::Table, CompatSessions::UserId)).eq(Uuid::from(user.id)) + })) + .add_option(self.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() + } + })) + .add_option(self.auth_type().map(|auth_type| { + // In in the SELECT to list sessions, we can rely on the JOINed table, whereas + // in other queries we need to do a subquery + if has_joins { + if auth_type.is_sso_login() { + Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)) + .is_not_null() + } else { + Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)) + .is_null() + } + } else { + // 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)) + } + } + })) + .add_option(self.device().map(|device| { + Expr::col((CompatSessions::Table, CompatSessions::DeviceId)).eq(device.as_str()) + })) + } +} + #[async_trait] impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { type Error = DatabaseError; @@ -356,39 +408,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { 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()) - })) + .apply_filter(filter) .build_sqlx(PostgresQueryBuilder); let res = sqlx::query_with(&sql, arguments) @@ -483,27 +503,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { Expr::col((CompatSessions::Table, CompatSessions::CompatSessionId)) .equals((CompatSsoLogins::Table, CompatSsoLogins::CompatSessionId)), ) - .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| { - if auth_type.is_sso_login() { - Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)) - .is_not_null() - } else { - Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSsoLoginId)).is_null() - } - })) - .and_where_option(filter.device().map(|device| { - Expr::col((CompatSessions::Table, CompatSessions::DeviceId)).eq(device.as_str()) - })) + .apply_filter_with_joins(filter) .generate_pagination( (CompatSessions::Table, CompatSessions::CompatSessionId), pagination, @@ -532,36 +532,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { let (sql, arguments) = sea_query::Query::select() .expr(Expr::col((CompatSessions::Table, CompatSessions::CompatSessionId)).count()) .from(CompatSessions::Table) - .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| { - // Check if it is an SSO login by checking if there is a SSO login for the - // session. - let exists = Expr::exists( - Query::select() - .expr(Expr::cust("1")) - .from(CompatSsoLogins::Table) - .and_where( - Expr::col((CompatSsoLogins::Table, CompatSsoLogins::CompatSessionId)) - .equals((CompatSessions::Table, CompatSessions::CompatSessionId)), - ) - .take(), - ); - - if auth_type.is_sso_login() { - exists - } else { - exists.not() - } - })) + .apply_filter(filter) .build_sqlx(PostgresQueryBuilder); let count: i64 = sqlx::query_scalar_with(&sql, arguments)