storage: do less joins in compat sessions

crates/data-model/src/compat/mod.rs

@@ -24,18 +24,83 @@ pub use self::{
     session::{CompatSession, CompatSessionState},
     sso_login::{CompatSsoLogin, CompatSsoLoginState},
 };
+use crate::InvalidTransitionError;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CompatAccessToken {
     pub id: Ulid,
+    pub session_id: Ulid,
     pub token: String,
     pub created_at: DateTime<Utc>,
     pub expires_at: Option<DateTime<Utc>>,
 }
 
+impl CompatAccessToken {
+    #[must_use]
+    pub fn is_valid(&self, now: DateTime<Utc>) -> bool {
+        if let Some(expires_at) = self.expires_at {
+            expires_at > now
+        } else {
+            true
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub enum CompatRefreshTokenState {
+    #[default]
+    Valid,
+    Consumed {
+        consumed_at: DateTime<Utc>,
+    },
+}
+
+impl CompatRefreshTokenState {
+    /// Returns `true` if the compat refresh token state is [`Valid`].
+    ///
+    /// [`Valid`]: CompatRefreshTokenState::Valid
+    #[must_use]
+    pub fn is_valid(&self) -> bool {
+        matches!(self, Self::Valid)
+    }
+
+    /// Returns `true` if the compat refresh token state is [`Consumed`].
+    ///
+    /// [`Consumed`]: CompatRefreshTokenState::Consumed
+    #[must_use]
+    pub fn is_consumed(&self) -> bool {
+        matches!(self, Self::Consumed { .. })
+    }
+
+    pub fn consume(self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        match self {
+            Self::Valid => Ok(Self::Consumed { consumed_at }),
+            Self::Consumed { .. } => Err(InvalidTransitionError),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CompatRefreshToken {
     pub id: Ulid,
+    pub state: CompatRefreshTokenState,
+    pub session_id: Ulid,
+    pub access_token_id: Ulid,
     pub token: String,
     pub created_at: DateTime<Utc>,
 }
+
+impl std::ops::Deref for CompatRefreshToken {
+    type Target = CompatRefreshTokenState;
+
+    fn deref(&self) -> &Self::Target {
+        &self.state
+    }
+}
+
+impl CompatRefreshToken {
+    pub fn consume(mut self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.consume(consumed_at)?;
+        Ok(self)
+    }
+}

crates/data-model/src/compat/session.rs

@@ -77,3 +77,10 @@ impl std::ops::Deref for CompatSession {
         &self.state
     }
 }
+
+impl CompatSession {
+    pub fn finish(mut self, finished_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.finish(finished_at)?;
+        Ok(self)
+    }
+}

crates/data-model/src/compat/sso_login.rs

@@ -25,12 +25,12 @@ pub enum CompatSsoLoginState {
     Pending,
     Fulfilled {
         fulfilled_at: DateTime<Utc>,
-        session: CompatSession,
+        session_id: Ulid,
     },
     Exchanged {
         fulfilled_at: DateTime<Utc>,
         exchanged_at: DateTime<Utc>,
-        session: CompatSession,
+        session_id: Ulid,
     },
 }
 
@@ -78,22 +78,24 @@ impl CompatSsoLoginState {
     }
 
     #[must_use]
-    pub fn session(&self) -> Option<&CompatSession> {
+    pub fn session_id(&self) -> Option<Ulid> {
         match self {
             Self::Pending => None,
-            Self::Fulfilled { session, .. } | Self::Exchanged { session, .. } => Some(session),
+            Self::Fulfilled { session_id, .. } | Self::Exchanged { session_id, .. } => {
+                Some(*session_id)
+            }
         }
     }
 
     pub fn fulfill(
         self,
         fulfilled_at: DateTime<Utc>,
-        session: CompatSession,
+        session: &CompatSession,
     ) -> Result<Self, InvalidTransitionError> {
         match self {
             Self::Pending => Ok(Self::Fulfilled {
                 fulfilled_at,
-                session,
+                session_id: session.id,
             }),
             Self::Fulfilled { .. } | Self::Exchanged { .. } => Err(InvalidTransitionError),
         }
@@ -103,11 +105,11 @@ impl CompatSsoLoginState {
         match self {
             Self::Fulfilled {
                 fulfilled_at,
-                session,
+                session_id,
             } => Ok(Self::Exchanged {
                 fulfilled_at,
                 exchanged_at,
-                session,
+                session_id,
             }),
             Self::Pending { .. } | Self::Exchanged { .. } => Err(InvalidTransitionError),
         }
@@ -135,7 +137,7 @@ impl CompatSsoLogin {
     pub fn fulfill(
         mut self,
         fulfilled_at: DateTime<Utc>,
-        session: CompatSession,
+        session: &CompatSession,
     ) -> Result<Self, InvalidTransitionError> {
         self.state = self.state.fulfill(fulfilled_at, session)?;
         Ok(self)

crates/data-model/src/lib.rs

@@ -37,8 +37,8 @@ pub struct InvalidTransitionError;
 
 pub use self::{
     compat::{
-        CompatAccessToken, CompatRefreshToken, CompatSession, CompatSessionState, CompatSsoLogin,
+        CompatAccessToken, CompatRefreshToken, CompatRefreshTokenState, CompatSession,
-        CompatSsoLoginState, Device,
+        CompatSessionState, CompatSsoLogin, CompatSsoLoginState, Device,
     },
     oauth2::{
         AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, Client,

crates/graphql/src/model/compat_sessions.rs

@@ -15,7 +15,7 @@
 use anyhow::Context as _;
 use async_graphql::{Context, Description, Object, ID};
 use chrono::{DateTime, Utc};
-use mas_storage::{user::UserRepository, Repository};
+use mas_storage::{compat::lookup_compat_session, user::UserRepository, Repository};
 use sqlx::PgPool;
 use url::Url;
 
@@ -94,7 +94,17 @@ impl CompatSsoLogin {
     }
 
     /// The compat session which was started by this login.
-    async fn session(&self) -> Option<CompatSession> {
+    async fn session(
-        self.0.session().cloned().map(CompatSession)
+        &self,
+        ctx: &Context<'_>,
+    ) -> Result<Option<CompatSession>, async_graphql::Error> {
+        let Some(session_id) = self.0.session_id() else { return Ok(None) };
+
+        let mut conn = ctx.data::<PgPool>()?.acquire().await?;
+        let session = lookup_compat_session(&mut conn, session_id)
+            .await?
+            .context("Could not load compat session")?;
+
+        Ok(Some(CompatSession(session)))
     }
 }

crates/handlers/src/compat/login.rs

@@ -19,7 +19,7 @@ use mas_data_model::{CompatSession, CompatSsoLoginState, Device, TokenType, User
 use mas_storage::{
     compat::{
         add_compat_access_token, add_compat_refresh_token, get_compat_sso_login_by_token,
-        mark_compat_sso_login_as_exchanged, start_compat_session,
+        lookup_compat_session, mark_compat_sso_login_as_exchanged, start_compat_session,
     },
     user::{UserPasswordRepository, UserRepository},
     Clock, Repository,
@@ -137,6 +137,9 @@ pub enum RouteError {
     #[error("user not found")]
     UserNotFound,
 
+    #[error("session not found")]
+    SessionNotFound,
+
     #[error("user has no password")]
     NoPassword,
 
@@ -156,7 +159,7 @@ impl_from_error_for_route!(mas_storage::DatabaseError);
 impl IntoResponse for RouteError {
     fn into_response(self) -> axum::response::Response {
         match self {
-            Self::Internal(_) => MatrixError {
+            Self::Internal(_) | Self::SessionNotFound => MatrixError {
                 errcode: "M_UNKNOWN",
                 error: "Internal server error",
                 status: StatusCode::INTERNAL_SERVER_ERROR,
@@ -268,7 +271,7 @@ async fn token_login(
         .ok_or(RouteError::InvalidLoginToken)?;
 
     let now = clock.now();
-    let user_id = match login.state {
+    let session_id = match login.state {
         CompatSsoLoginState::Pending => {
             tracing::error!(
                 compat_sso_login.id = %login.id,
@@ -277,21 +280,26 @@ async fn token_login(
             return Err(RouteError::InvalidLoginToken);
         }
         CompatSsoLoginState::Fulfilled {
-            fulfilled_at: fullfilled_at,
+            fulfilled_at,
-            ref session,
+            session_id,
             ..
         } => {
-            if now > fullfilled_at + Duration::seconds(30) {
+            if now > fulfilled_at + Duration::seconds(30) {
                 return Err(RouteError::LoginTookTooLong);
             }
 
-            session.user_id
+            session_id
         }
-        CompatSsoLoginState::Exchanged { exchanged_at, .. } => {
+        CompatSsoLoginState::Exchanged {
+            exchanged_at,
+            session_id,
+            ..
+        } => {
             if now > exchanged_at + Duration::seconds(30) {
                 // TODO: log that session out
                 tracing::error!(
                     compat_sso_login.id = %login.id,
+                    compat_session.id = %session_id,
                     "Login token exchanged a second time more than 30s after"
                 );
             }
@@ -300,18 +308,19 @@ async fn token_login(
         }
     };
 
+    let session = lookup_compat_session(&mut *txn, session_id)
+        .await?
+        .ok_or(RouteError::SessionNotFound)?;
+
     let user = txn
         .user()
-        .lookup(user_id)
+        .lookup(session.user_id)
         .await?
         .ok_or(RouteError::UserNotFound)?;
 
-    let login = mark_compat_sso_login_as_exchanged(&mut *txn, clock, login).await?;
+    mark_compat_sso_login_as_exchanged(&mut *txn, clock, login).await?;
 
-    match login.state {
+    Ok((session, user))
-        CompatSsoLoginState::Exchanged { session, .. } => Ok((session, user)),
-        _ => unreachable!(),
-    }
 }
 
 async fn user_password_login(

crates/handlers/src/compat/logout.rs

@@ -16,7 +16,10 @@ use axum::{extract::State, response::IntoResponse, Json, TypedHeader};
 use headers::{authorization::Bearer, Authorization};
 use hyper::StatusCode;
 use mas_data_model::TokenType;
-use mas_storage::{compat::compat_logout, Clock};
+use mas_storage::{
+    compat::{end_compat_session, find_compat_access_token, lookup_compat_session},
+    Clock,
+};
 use sqlx::PgPool;
 use thiserror::Error;
 
@@ -36,12 +39,10 @@ pub enum RouteError {
 
     #[error("Invalid access token")]
     InvalidAuthorization,
-
-    #[error("Logout failed")]
-    LogoutFailed,
 }
 
 impl_from_error_for_route!(sqlx::Error);
+impl_from_error_for_route!(mas_storage::DatabaseError);
 
 impl IntoResponse for RouteError {
     fn into_response(self) -> axum::response::Response {
@@ -56,7 +57,7 @@ impl IntoResponse for RouteError {
                 error: "Missing access token",
                 status: StatusCode::UNAUTHORIZED,
             },
-            Self::InvalidAuthorization | Self::LogoutFailed | Self::TokenFormat(_) => MatrixError {
+            Self::InvalidAuthorization | Self::TokenFormat(_) => MatrixError {
                 errcode: "M_UNKNOWN_TOKEN",
                 error: "Invalid access token",
                 status: StatusCode::UNAUTHORIZED,
@@ -71,7 +72,7 @@ pub(crate) async fn post(
     maybe_authorization: Option<TypedHeader<Authorization<Bearer>>>,
 ) -> Result<impl IntoResponse, RouteError> {
     let clock = Clock::default();
-    let mut conn = pool.acquire().await?;
+    let mut txn = pool.begin().await?;
 
     let TypedHeader(authorization) = maybe_authorization.ok_or(RouteError::MissingAuthorization)?;
 
@@ -82,9 +83,19 @@ pub(crate) async fn post(
         return Err(RouteError::InvalidAuthorization);
     }
 
-    if !compat_logout(&mut conn, &clock, token).await? {
+    let token = find_compat_access_token(&mut txn, token)
-        return Err(RouteError::LogoutFailed);
+        .await?
-    }
+        .filter(|t| t.is_valid(clock.now()))
+        .ok_or(RouteError::InvalidAuthorization)?;
+
+    let session = lookup_compat_session(&mut txn, token.session_id)
+        .await?
+        .filter(|s| s.is_valid())
+        .ok_or(RouteError::InvalidAuthorization)?;
+
+    end_compat_session(&mut txn, &clock, session).await?;
+
+    txn.commit().await?;
 
     Ok(Json(serde_json::json!({})))
 }

crates/handlers/src/compat/refresh.rs

@@ -18,7 +18,8 @@ use hyper::StatusCode;
 use mas_data_model::{TokenFormatError, TokenType};
 use mas_storage::compat::{
     add_compat_access_token, add_compat_refresh_token, consume_compat_refresh_token,
-    expire_compat_access_token, lookup_active_compat_refresh_token,
+    expire_compat_access_token, find_compat_refresh_token, lookup_compat_access_token,
+    lookup_compat_session,
 };
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, DurationMilliSeconds};
@@ -40,17 +41,26 @@ pub enum RouteError {
 
     #[error("invalid token")]
     InvalidToken,
+
+    #[error("refresh token already consumed")]
+    RefreshTokenConsumed,
+
+    #[error("invalid session")]
+    InvalidSession,
+
+    #[error("unknown session")]
+    UnknownSession,
 }
 
 impl IntoResponse for RouteError {
     fn into_response(self) -> axum::response::Response {
         match self {
-            Self::Internal(_) => MatrixError {
+            Self::Internal(_) | Self::UnknownSession => MatrixError {
                 errcode: "M_UNKNOWN",
                 error: "Internal error",
                 status: StatusCode::INTERNAL_SERVER_ERROR,
             },
-            Self::InvalidToken => MatrixError {
+            Self::InvalidToken | Self::InvalidSession | Self::RefreshTokenConsumed => MatrixError {
                 errcode: "M_UNKNOWN_TOKEN",
                 error: "Invalid refresh token",
                 status: StatusCode::UNAUTHORIZED,
@@ -91,10 +101,25 @@ pub(crate) async fn post(
         return Err(RouteError::InvalidToken);
     }
 
-    let (refresh_token, access_token, session) =
+    let refresh_token = find_compat_refresh_token(&mut txn, &input.refresh_token)
-        lookup_active_compat_refresh_token(&mut txn, &input.refresh_token)
+        .await?
-            .await?
+        .ok_or(RouteError::InvalidToken)?;
-            .ok_or(RouteError::InvalidToken)?;
+
+    if !refresh_token.is_valid() {
+        return Err(RouteError::RefreshTokenConsumed);
+    }
+
+    let session = lookup_compat_session(&mut txn, refresh_token.session_id)
+        .await?
+        .ok_or(RouteError::UnknownSession)?;
+
+    if !session.is_valid() {
+        return Err(RouteError::InvalidSession);
+    }
+
+    let access_token = lookup_compat_access_token(&mut txn, refresh_token.access_token_id)
+        .await?
+        .filter(|t| t.is_valid(clock.now()));
 
     let new_refresh_token_str = TokenType::CompatRefreshToken.generate(&mut rng);
     let new_access_token_str = TokenType::CompatAccessToken.generate(&mut rng);
@@ -120,7 +145,10 @@ pub(crate) async fn post(
     .await?;
 
     consume_compat_refresh_token(&mut txn, &clock, refresh_token).await?;
-    expire_compat_access_token(&mut txn, &clock, access_token).await?;
+
+    if let Some(access_token) = access_token {
+        expire_compat_access_token(&mut txn, &clock, access_token).await?;
+    }
 
     txn.commit().await?;
 

crates/handlers/src/oauth2/introspection.rs

@@ -22,7 +22,7 @@ use mas_data_model::{TokenFormatError, TokenType};
 use mas_iana::oauth::{OAuthClientAuthenticationMethod, OAuthTokenTypeHint};
 use mas_keystore::Encrypter;
 use mas_storage::{
-    compat::{lookup_active_compat_access_token, lookup_active_compat_refresh_token},
+    compat::{find_compat_access_token, find_compat_refresh_token, lookup_compat_session},
     oauth2::{
         access_token::lookup_active_access_token, refresh_token::lookup_active_refresh_token,
     },
@@ -194,6 +194,7 @@ pub(crate) async fn post(
                 jti: None,
             }
         }
+
         TokenType::RefreshToken => {
             let (token, session) = lookup_active_refresh_token(&mut conn, token)
                 .await?
@@ -221,9 +222,16 @@ pub(crate) async fn post(
                 jti: None,
             }
         }
+
         TokenType::CompatAccessToken => {
-            let (token, session) = lookup_active_compat_access_token(&mut conn, &clock, token)
+            let token = find_compat_access_token(&mut conn, token)
                 .await?
+                .filter(|t| t.is_valid(clock.now()))
+                .ok_or(RouteError::UnknownToken)?;
+
+            let session = lookup_compat_session(&mut conn, token.session_id)
+                .await?
+                .filter(|s| s.is_valid())
                 .ok_or(RouteError::UnknownToken)?;
 
             let user = conn
@@ -251,11 +259,17 @@ pub(crate) async fn post(
                 jti: None,
             }
         }
+
         TokenType::CompatRefreshToken => {
-            let (refresh_token, _access_token, session) =
+            let refresh_token = find_compat_refresh_token(&mut conn, token)
-                lookup_active_compat_refresh_token(&mut conn, token)
+                .await?
-                    .await?
+                .filter(|t| t.is_valid())
-                    .ok_or(RouteError::UnknownToken)?;
+                .ok_or(RouteError::UnknownToken)?;
+
+            let session = lookup_compat_session(&mut conn, refresh_token.session_id)
+                .await?
+                .filter(|s| s.is_valid())
+                .ok_or(RouteError::UnknownToken)?;
 
             let user = conn
                 .user()

crates/storage/sqlx-data.json

@@ -1,91 +1,5 @@
 {
   "db": "PostgreSQL",
-  "021f845e564500457e2e0c8614beb1d9fd10b4b5f13515478f7ca25b5474d016": {
-    "describe": {
-      "columns": [
-        {
-          "name": "compat_refresh_token_id",
-          "ordinal": 0,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_refresh_token",
-          "ordinal": 1,
-          "type_info": "Text"
-        },
-        {
-          "name": "compat_refresh_token_created_at",
-          "ordinal": 2,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_access_token_id",
-          "ordinal": 3,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_access_token",
-          "ordinal": 4,
-          "type_info": "Text"
-        },
-        {
-          "name": "compat_access_token_created_at",
-          "ordinal": 5,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_access_token_expires_at",
-          "ordinal": 6,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_id",
-          "ordinal": 7,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_session_created_at",
-          "ordinal": 8,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_finished_at",
-          "ordinal": 9,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_device_id",
-          "ordinal": 10,
-          "type_info": "Text"
-        },
-        {
-          "name": "user_id",
-          "ordinal": 11,
-          "type_info": "Uuid"
-        }
-      ],
-      "nullable": [
-        false,
-        false,
-        false,
-        false,
-        false,
-        false,
-        true,
-        false,
-        false,
-        true,
-        false,
-        false
-      ],
-      "parameters": {
-        "Left": [
-          "Text"
-        ]
-      }
-    },
-    "query": "\n            SELECT cr.compat_refresh_token_id\n                 , cr.refresh_token   AS \"compat_refresh_token\"\n                 , cr.created_at      AS \"compat_refresh_token_created_at\"\n                 , ct.compat_access_token_id\n                 , ct.access_token    AS \"compat_access_token\"\n                 , ct.created_at      AS \"compat_access_token_created_at\"\n                 , ct.expires_at      AS \"compat_access_token_expires_at\"\n                 , cs.compat_session_id\n                 , cs.created_at      AS \"compat_session_created_at\"\n                 , cs.finished_at     AS \"compat_session_finished_at\"\n                 , cs.device_id       AS \"compat_session_device_id\"\n                 , cs.user_id\n\n            FROM compat_refresh_tokens cr\n            INNER JOIN compat_sessions cs\n              USING (compat_session_id)\n            INNER JOIN compat_access_tokens ct\n              USING (compat_access_token_id)\n\n            WHERE cr.refresh_token = $1\n              AND cr.consumed_at IS NULL\n              AND cs.finished_at IS NULL\n        "
-  },
   "08d7df347c806ef14b6d0fb031cab041d79ba48528420160e23286369db7af35": {
     "describe": {
       "columns": [
@@ -312,6 +226,50 @@
     },
     "query": "\n            INSERT INTO compat_refresh_tokens\n                (compat_refresh_token_id, compat_session_id,\n                 compat_access_token_id, refresh_token, created_at)\n            VALUES ($1, $2, $3, $4, $5)\n        "
   },
+  "3cf8e061206620071b39d0262cd165bb367b12b8e904180730d8acfa5af3d4b9": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_session_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "device_id",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "user_id",
+          "ordinal": 2,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "created_at",
+          "ordinal": 3,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "finished_at",
+          "ordinal": 4,
+          "type_info": "Timestamptz"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        false,
+        true
+      ],
+      "parameters": {
+        "Left": [
+          "Uuid"
+        ]
+      }
+    },
+    "query": "\n            SELECT compat_session_id\n                 , device_id\n                 , user_id\n                 , created_at\n                 , finished_at\n            FROM compat_sessions\n            WHERE compat_session_id = $1\n        "
+  },
   "3d66f3121b11ce923b9c60609b510a8ca899640e78cc8f5b03168622928ffe94": {
     "describe": {
       "columns": [],
@@ -507,7 +465,20 @@
     },
     "query": "\n            UPDATE oauth2_authorization_grants AS og\n            SET\n                oauth2_session_id = os.oauth2_session_id,\n                fulfilled_at = os.created_at\n            FROM oauth2_sessions os\n            WHERE\n                og.oauth2_authorization_grant_id = $1\n                AND os.oauth2_session_id = $2\n            RETURNING fulfilled_at AS \"fulfilled_at!: DateTime<Utc>\"\n        "
   },
-  "4f8b0cd13d9488c2dd0f183d090d3856da15dcdb57a8c113febbee665a2a3ac5": {
+  "4c4dbb846bb98d84f6b7f886f8af9833c7efe27b8b4f297077887232bef322ee": {
+    "describe": {
+      "columns": [],
+      "nullable": [],
+      "parameters": {
+        "Left": [
+          "Uuid",
+          "Timestamptz"
+        ]
+      }
+    },
+    "query": "\n            UPDATE compat_sessions cs\n            SET finished_at = $2\n            WHERE compat_session_id = $1\n        "
+  },
+  "4f080990eb6dd9f6128f3a1aee195b99d5f286fa0f6c27d744f73848343879d4": {
     "describe": {
       "columns": [
         {
@@ -541,29 +512,9 @@
           "type_info": "Timestamptz"
         },
         {
-          "name": "compat_session_id?",
+          "name": "compat_session_id",
           "ordinal": 6,
           "type_info": "Uuid"
-        },
-        {
-          "name": "compat_session_created_at?",
-          "ordinal": 7,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_finished_at?",
-          "ordinal": 8,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_device_id?",
-          "ordinal": 9,
-          "type_info": "Text"
-        },
-        {
-          "name": "user_id?",
-          "ordinal": 10,
-          "type_info": "Uuid"
         }
       ],
       "nullable": [
@@ -573,19 +524,15 @@
         false,
         true,
         true,
-        false,
+        true
-        false,
-        true,
-        false,
-        false
       ],
       "parameters": {
         "Left": [
-          "Text"
+          "Uuid"
         ]
       }
     },
-    "query": "\n            SELECT cl.compat_sso_login_id\n                 , cl.login_token     AS \"compat_sso_login_token\"\n                 , cl.redirect_uri    AS \"compat_sso_login_redirect_uri\"\n                 , cl.created_at      AS \"compat_sso_login_created_at\"\n                 , cl.fulfilled_at    AS \"compat_sso_login_fulfilled_at\"\n                 , cl.exchanged_at    AS \"compat_sso_login_exchanged_at\"\n                 , cs.compat_session_id AS \"compat_session_id?\"\n                 , cs.created_at      AS \"compat_session_created_at?\"\n                 , cs.finished_at     AS \"compat_session_finished_at?\"\n                 , cs.device_id       AS \"compat_session_device_id?\"\n                 , cs.user_id         AS \"user_id?\"\n            FROM compat_sso_logins cl\n            LEFT JOIN compat_sessions cs\n              USING (compat_session_id)\n            WHERE cl.login_token = $1\n        "
+    "query": "\n            SELECT cl.compat_sso_login_id\n                 , cl.login_token     AS \"compat_sso_login_token\"\n                 , cl.redirect_uri    AS \"compat_sso_login_redirect_uri\"\n                 , cl.created_at      AS \"compat_sso_login_created_at\"\n                 , cl.fulfilled_at    AS \"compat_sso_login_fulfilled_at\"\n                 , cl.exchanged_at    AS \"compat_sso_login_exchanged_at\"\n                 , cl.compat_session_id AS \"compat_session_id\"\n\n            FROM compat_sso_logins cl\n            WHERE cl.compat_sso_login_id = $1\n        "
   },
   "51158bfcaa1a8d8e051bffe7c5ba0369bf53fb162f7622626054e89e68fc07bd": {
     "describe": {
@@ -608,27 +555,6 @@
     },
     "query": "\n            SELECT scope_token\n            FROM oauth2_consents\n            WHERE user_id = $1 AND oauth2_client_id = $2\n        "
   },
-  "559a486756d08d101eb7188ef6637b9d24c024d056795b8121f7f04a7f9db6a3": {
-    "describe": {
-      "columns": [
-        {
-          "name": "compat_session_id",
-          "ordinal": 0,
-          "type_info": "Uuid"
-        }
-      ],
-      "nullable": [
-        false
-      ],
-      "parameters": {
-        "Left": [
-          "Text",
-          "Timestamptz"
-        ]
-      }
-    },
-    "query": "\n            UPDATE compat_sessions cs\n            SET finished_at = $2\n            FROM compat_access_tokens ca\n            WHERE ca.access_token = $1\n              AND ca.compat_session_id = cs.compat_session_id\n              AND cs.finished_at IS NULL\n            RETURNING cs.compat_session_id\n        "
-  },
   "583ae9a0db9cd55fa57a179339550f3dab1bfc76f35ad488e1560ea37f7ed029": {
     "describe": {
       "columns": [],
@@ -1296,86 +1222,6 @@
     },
     "query": "\n                UPDATE user_email_confirmation_codes\n                SET consumed_at = $2\n                WHERE user_email_confirmation_code_id = $1\n            "
   },
-  "92ef320b75ca479ed1a38f6d654fdb953431188a8654c806fd5f98444b00c012": {
-    "describe": {
-      "columns": [
-        {
-          "name": "compat_sso_login_id",
-          "ordinal": 0,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_sso_login_token",
-          "ordinal": 1,
-          "type_info": "Text"
-        },
-        {
-          "name": "compat_sso_login_redirect_uri",
-          "ordinal": 2,
-          "type_info": "Text"
-        },
-        {
-          "name": "compat_sso_login_created_at",
-          "ordinal": 3,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_sso_login_fulfilled_at",
-          "ordinal": 4,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_sso_login_exchanged_at",
-          "ordinal": 5,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_id?",
-          "ordinal": 6,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_session_created_at?",
-          "ordinal": 7,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_finished_at?",
-          "ordinal": 8,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_device_id?",
-          "ordinal": 9,
-          "type_info": "Text"
-        },
-        {
-          "name": "user_id?",
-          "ordinal": 10,
-          "type_info": "Uuid"
-        }
-      ],
-      "nullable": [
-        false,
-        false,
-        false,
-        false,
-        true,
-        true,
-        false,
-        false,
-        true,
-        false,
-        false
-      ],
-      "parameters": {
-        "Left": [
-          "Uuid"
-        ]
-      }
-    },
-    "query": "\n            SELECT cl.compat_sso_login_id\n                 , cl.login_token     AS \"compat_sso_login_token\"\n                 , cl.redirect_uri    AS \"compat_sso_login_redirect_uri\"\n                 , cl.created_at      AS \"compat_sso_login_created_at\"\n                 , cl.fulfilled_at    AS \"compat_sso_login_fulfilled_at\"\n                 , cl.exchanged_at    AS \"compat_sso_login_exchanged_at\"\n                 , cs.compat_session_id AS \"compat_session_id?\"\n                 , cs.created_at      AS \"compat_session_created_at?\"\n                 , cs.finished_at     AS \"compat_session_finished_at?\"\n                 , cs.device_id       AS \"compat_session_device_id?\"\n                 , cs.user_id         AS \"user_id?\"\n\n            FROM compat_sso_logins cl\n            LEFT JOIN compat_sessions cs\n              USING (compat_session_id)\n            WHERE cl.compat_sso_login_id = $1\n        "
-  },
   "94fd96446b237c87bd6bf741f3c42b37ee751b87b7fcc459602bdf8c46962443": {
     "describe": {
       "columns": [
@@ -1868,6 +1714,106 @@
     },
     "query": "\n                INSERT INTO user_sessions (user_session_id, user_id, created_at)\n                VALUES ($1, $2, $3)\n            "
   },
+  "c3e60701299be7728108b8967ec5396fb186adaac360d6a0152d25e4a4f46f46": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_access_token_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "access_token",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "created_at",
+          "ordinal": 2,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "expires_at",
+          "ordinal": 3,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_session_id",
+          "ordinal": 4,
+          "type_info": "Uuid"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        true,
+        false
+      ],
+      "parameters": {
+        "Left": [
+          "Uuid"
+        ]
+      }
+    },
+    "query": "\n            SELECT compat_access_token_id\n                 , access_token\n                 , created_at\n                 , expires_at\n                 , compat_session_id\n\n            FROM compat_access_tokens\n\n            WHERE compat_access_token_id = $1\n        "
+  },
+  "c78246fc8737491352f71ea9410e79df8de88596c8197405cda36eb8c8187810": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_sso_login_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "compat_sso_login_token",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "compat_sso_login_redirect_uri",
+          "ordinal": 2,
+          "type_info": "Text"
+        },
+        {
+          "name": "compat_sso_login_created_at",
+          "ordinal": 3,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_sso_login_fulfilled_at",
+          "ordinal": 4,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_sso_login_exchanged_at",
+          "ordinal": 5,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_session_id",
+          "ordinal": 6,
+          "type_info": "Uuid"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        false,
+        true,
+        true,
+        true
+      ],
+      "parameters": {
+        "Left": [
+          "Text"
+        ]
+      }
+    },
+    "query": "\n            SELECT cl.compat_sso_login_id\n                 , cl.login_token     AS \"compat_sso_login_token\"\n                 , cl.redirect_uri    AS \"compat_sso_login_redirect_uri\"\n                 , cl.created_at      AS \"compat_sso_login_created_at\"\n                 , cl.fulfilled_at    AS \"compat_sso_login_fulfilled_at\"\n                 , cl.exchanged_at    AS \"compat_sso_login_exchanged_at\"\n                 , cl.compat_session_id AS \"compat_session_id\"\n            FROM compat_sso_logins cl\n            WHERE cl.login_token = $1\n        "
+  },
   "c88376abdba124ff0487a9a69d2345c7d69d7394f355111ec369cfa6d45fb40f": {
     "describe": {
       "columns": [],
@@ -1894,6 +1840,56 @@
     },
     "query": "\n            INSERT INTO oauth2_authorization_grants (\n                 oauth2_authorization_grant_id,\n                 oauth2_client_id,\n                 redirect_uri,\n                 scope,\n                 state,\n                 nonce,\n                 max_age,\n                 response_mode,\n                 code_challenge,\n                 code_challenge_method,\n                 response_type_code,\n                 response_type_id_token,\n                 authorization_code,\n                 requires_consent,\n                 created_at\n            )\n            VALUES\n                ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)\n        "
   },
+  "ca63558e877bd115aa7ca24de0cc3f78a13cb55105758fe0bd930da513f75504": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_refresh_token_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "refresh_token",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "created_at",
+          "ordinal": 2,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "consumed_at",
+          "ordinal": 3,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_session_id",
+          "ordinal": 4,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "compat_access_token_id",
+          "ordinal": 5,
+          "type_info": "Uuid"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        true,
+        false,
+        false
+      ],
+      "parameters": {
+        "Left": [
+          "Text"
+        ]
+      }
+    },
+    "query": "\n            SELECT compat_refresh_token_id\n                 , refresh_token\n                 , created_at\n                 , consumed_at\n                 , compat_session_id\n                 , compat_access_token_id\n\n            FROM compat_refresh_tokens\n\n            WHERE refresh_token = $1\n        "
+  },
   "caf54e4659306a746747aa61906bdb2cb8da51176e90435aa8b9754ebf3e4d60": {
     "describe": {
       "columns": [],
@@ -1909,6 +1905,50 @@
     },
     "query": "\n            INSERT INTO compat_sessions (compat_session_id, user_id, device_id, created_at)\n            VALUES ($1, $2, $3, $4)\n        "
   },
+  "cf43b82bdf534400f900cff3c5356083db0f9e5407e288b64f43d7ac100de058": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_access_token_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        },
+        {
+          "name": "access_token",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "created_at",
+          "ordinal": 2,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "expires_at",
+          "ordinal": 3,
+          "type_info": "Timestamptz"
+        },
+        {
+          "name": "compat_session_id",
+          "ordinal": 4,
+          "type_info": "Uuid"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        true,
+        false
+      ],
+      "parameters": {
+        "Left": [
+          "Text"
+        ]
+      }
+    },
+    "query": "\n            SELECT compat_access_token_id\n                 , access_token\n                 , created_at\n                 , expires_at\n                 , compat_session_id\n\n            FROM compat_access_tokens\n\n            WHERE access_token = $1\n        "
+  },
   "d12a513b81b3ef658eae1f0a719933323f28c6ee260b52cafe337dd3d19e865c": {
     "describe": {
       "columns": [
@@ -1944,75 +1984,6 @@
     },
     "query": "\n            INSERT INTO compat_sso_logins\n                (compat_sso_login_id, login_token, redirect_uri, created_at)\n            VALUES ($1, $2, $3, $4)\n        "
   },
-  "d4c6c070a0cd889cef9e0cfd65c64522a03f0bae12ee7c6b74343ec8f38d24c1": {
-    "describe": {
-      "columns": [
-        {
-          "name": "compat_access_token_id",
-          "ordinal": 0,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_access_token",
-          "ordinal": 1,
-          "type_info": "Text"
-        },
-        {
-          "name": "compat_access_token_created_at",
-          "ordinal": 2,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_access_token_expires_at",
-          "ordinal": 3,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_id",
-          "ordinal": 4,
-          "type_info": "Uuid"
-        },
-        {
-          "name": "compat_session_created_at",
-          "ordinal": 5,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_finished_at",
-          "ordinal": 6,
-          "type_info": "Timestamptz"
-        },
-        {
-          "name": "compat_session_device_id",
-          "ordinal": 7,
-          "type_info": "Text"
-        },
-        {
-          "name": "user_id!",
-          "ordinal": 8,
-          "type_info": "Uuid"
-        }
-      ],
-      "nullable": [
-        false,
-        false,
-        false,
-        true,
-        false,
-        false,
-        true,
-        false,
-        false
-      ],
-      "parameters": {
-        "Left": [
-          "Text",
-          "Timestamptz"
-        ]
-      }
-    },
-    "query": "\n            SELECT ct.compat_access_token_id\n                 , ct.access_token    AS \"compat_access_token\"\n                 , ct.created_at      AS \"compat_access_token_created_at\"\n                 , ct.expires_at      AS \"compat_access_token_expires_at\"\n                 , cs.compat_session_id\n                 , cs.created_at      AS \"compat_session_created_at\"\n                 , cs.finished_at     AS \"compat_session_finished_at\"\n                 , cs.device_id       AS \"compat_session_device_id\"\n                 , cs.user_id         AS \"user_id!\"\n\n            FROM compat_access_tokens ct\n            INNER JOIN compat_sessions cs\n              USING (compat_session_id)\n\n            WHERE ct.access_token = $1\n              AND (ct.expires_at < $2 OR ct.expires_at IS NULL)\n              AND cs.finished_at IS NULL \n        "
-  },
   "d8677b3b6ee594c230fad98c1aa1c6e3d983375bf5b701c7b52468e7f906abf9": {
     "describe": {
       "columns": [],

crates/storage/src/compat.rs

@@ -14,8 +14,8 @@
 
 use chrono::{DateTime, Duration, Utc};
 use mas_data_model::{
-    CompatAccessToken, CompatRefreshToken, CompatSession, CompatSessionState, CompatSsoLogin,
+    CompatAccessToken, CompatRefreshToken, CompatRefreshTokenState, CompatSession,
-    CompatSsoLoginState, Device, User,
+    CompatSessionState, CompatSsoLogin, CompatSsoLoginState, Device, User,
 };
 use rand::Rng;
 use sqlx::{Acquire, PgExecutor, Postgres, QueryBuilder};
@@ -29,71 +29,47 @@ use crate::{
     Clock, DatabaseError, DatabaseInconsistencyError, LookupResultExt,
 };
 
-struct CompatAccessTokenLookup {
+struct CompatSessionLookup {
-    compat_access_token_id: Uuid,
-    compat_access_token: String,
-    compat_access_token_created_at: DateTime<Utc>,
-    compat_access_token_expires_at: Option<DateTime<Utc>>,
     compat_session_id: Uuid,
-    compat_session_created_at: DateTime<Utc>,
+    device_id: String,
-    compat_session_finished_at: Option<DateTime<Utc>>,
-    compat_session_device_id: String,
     user_id: Uuid,
+    created_at: DateTime<Utc>,
+    finished_at: Option<DateTime<Utc>>,
 }
 
 #[tracing::instrument(skip_all, err)]
-pub async fn lookup_active_compat_access_token(
+pub async fn lookup_compat_session(
     executor: impl PgExecutor<'_>,
-    clock: &Clock,
+    session_id: Ulid,
-    token: &str,
+) -> Result<Option<CompatSession>, DatabaseError> {
-) -> Result<Option<(CompatAccessToken, CompatSession)>, DatabaseError> {
     let res = sqlx::query_as!(
-        CompatAccessTokenLookup,
+        CompatSessionLookup,
         r#"
-            SELECT ct.compat_access_token_id
+            SELECT compat_session_id
-                 , ct.access_token    AS "compat_access_token"
+                 , device_id
-                 , ct.created_at      AS "compat_access_token_created_at"
+                 , user_id
-                 , ct.expires_at      AS "compat_access_token_expires_at"
+                 , created_at
-                 , cs.compat_session_id
+                 , finished_at
-                 , cs.created_at      AS "compat_session_created_at"
+            FROM compat_sessions
-                 , cs.finished_at     AS "compat_session_finished_at"
+            WHERE compat_session_id = $1
-                 , cs.device_id       AS "compat_session_device_id"
-                 , cs.user_id         AS "user_id!"
-
-            FROM compat_access_tokens ct
-            INNER JOIN compat_sessions cs
-              USING (compat_session_id)
-
-            WHERE ct.access_token = $1
-              AND (ct.expires_at < $2 OR ct.expires_at IS NULL)
-              AND cs.finished_at IS NULL 
         "#,
-        token,
+        Uuid::from(session_id),
-        clock.now(),
     )
     .fetch_one(executor)
-    .instrument(info_span!("Fetch compat access token"))
     .await
     .to_option()?;
 
     let Some(res) = res else { return Ok(None) };
 
-    let token = CompatAccessToken {
-        id: res.compat_access_token_id.into(),
-        token: res.compat_access_token,
-        created_at: res.compat_access_token_created_at,
-        expires_at: res.compat_access_token_expires_at,
-    };
-
     let id = res.compat_session_id.into();
-    let device = Device::try_from(res.compat_session_device_id).map_err(|e| {
+    let device = Device::try_from(res.device_id).map_err(|e| {
         DatabaseInconsistencyError::on("compat_sessions")
             .column("device_id")
             .row(id)
             .source(e)
     })?;
 
-    let state = match res.compat_session_finished_at {
+    let state = match res.finished_at {
         None => CompatSessionState::Valid,
         Some(finished_at) => CompatSessionState::Finished { finished_at },
     };
@@ -103,103 +79,148 @@ pub async fn lookup_active_compat_access_token(
         state,
         user_id: res.user_id.into(),
         device,
-        created_at: res.compat_session_created_at,
+        created_at: res.created_at,
     };
 
-    Ok(Some((token, session)))
+    Ok(Some(session))
 }
 
-pub struct CompatRefreshTokenLookup {
+struct CompatAccessTokenLookup {
-    compat_refresh_token_id: Uuid,
-    compat_refresh_token: String,
-    compat_refresh_token_created_at: DateTime<Utc>,
     compat_access_token_id: Uuid,
-    compat_access_token: String,
+    access_token: String,
-    compat_access_token_created_at: DateTime<Utc>,
+    created_at: DateTime<Utc>,
-    compat_access_token_expires_at: Option<DateTime<Utc>>,
+    expires_at: Option<DateTime<Utc>>,
     compat_session_id: Uuid,
-    compat_session_created_at: DateTime<Utc>,
+}
-    compat_session_finished_at: Option<DateTime<Utc>>,
+
-    compat_session_device_id: String,
+impl From<CompatAccessTokenLookup> for CompatAccessToken {
-    user_id: Uuid,
+    fn from(value: CompatAccessTokenLookup) -> Self {
+        Self {
+            id: value.compat_access_token_id.into(),
+            session_id: value.compat_session_id.into(),
+            token: value.access_token,
+            created_at: value.created_at,
+            expires_at: value.expires_at,
+        }
+    }
 }
 
 #[tracing::instrument(skip_all, err)]
-#[allow(clippy::type_complexity)]
+pub async fn find_compat_access_token(
-pub async fn lookup_active_compat_refresh_token(
     executor: impl PgExecutor<'_>,
     token: &str,
-) -> Result<Option<(CompatRefreshToken, CompatAccessToken, CompatSession)>, DatabaseError> {
+) -> Result<Option<CompatAccessToken>, DatabaseError> {
     let res = sqlx::query_as!(
-        CompatRefreshTokenLookup,
+        CompatAccessTokenLookup,
         r#"
-            SELECT cr.compat_refresh_token_id
+            SELECT compat_access_token_id
-                 , cr.refresh_token   AS "compat_refresh_token"
+                 , access_token
-                 , cr.created_at      AS "compat_refresh_token_created_at"
+                 , created_at
-                 , ct.compat_access_token_id
+                 , expires_at
-                 , ct.access_token    AS "compat_access_token"
+                 , compat_session_id
-                 , ct.created_at      AS "compat_access_token_created_at"
-                 , ct.expires_at      AS "compat_access_token_expires_at"
-                 , cs.compat_session_id
-                 , cs.created_at      AS "compat_session_created_at"
-                 , cs.finished_at     AS "compat_session_finished_at"
-                 , cs.device_id       AS "compat_session_device_id"
-                 , cs.user_id
 
-            FROM compat_refresh_tokens cr
+            FROM compat_access_tokens
-            INNER JOIN compat_sessions cs
-              USING (compat_session_id)
-            INNER JOIN compat_access_tokens ct
-              USING (compat_access_token_id)
 
-            WHERE cr.refresh_token = $1
+            WHERE access_token = $1
-              AND cr.consumed_at IS NULL
+        "#,
-              AND cs.finished_at IS NULL
+        token,
+    )
+    .fetch_one(executor)
+    .await
+    .to_option()?;
+
+    let Some(res) = res else { return Ok(None) };
+
+    Ok(Some(res.into()))
+}
+
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_access_token.id = %id,
+    ),
+    err,
+)]
+pub async fn lookup_compat_access_token(
+    executor: impl PgExecutor<'_>,
+    id: Ulid,
+) -> Result<Option<CompatAccessToken>, DatabaseError> {
+    let res = sqlx::query_as!(
+        CompatAccessTokenLookup,
+        r#"
+            SELECT compat_access_token_id
+                 , access_token
+                 , created_at
+                 , expires_at
+                 , compat_session_id
+
+            FROM compat_access_tokens
+
+            WHERE compat_access_token_id = $1
+        "#,
+        Uuid::from(id),
+    )
+    .fetch_one(executor)
+    .await
+    .to_option()?;
+
+    let Some(res) = res else { return Ok(None) };
+
+    Ok(Some(res.into()))
+}
+
+pub struct CompatRefreshTokenLookup {
+    compat_refresh_token_id: Uuid,
+    refresh_token: String,
+    created_at: DateTime<Utc>,
+    consumed_at: Option<DateTime<Utc>>,
+    compat_access_token_id: Uuid,
+    compat_session_id: Uuid,
+}
+
+#[tracing::instrument(skip_all, err)]
+#[allow(clippy::type_complexity)]
+pub async fn find_compat_refresh_token(
+    executor: impl PgExecutor<'_>,
+    token: &str,
+) -> Result<Option<CompatRefreshToken>, DatabaseError> {
+    let res = sqlx::query_as!(
+        CompatRefreshTokenLookup,
+        r#"
+            SELECT compat_refresh_token_id
+                 , refresh_token
+                 , created_at
+                 , consumed_at
+                 , compat_session_id
+                 , compat_access_token_id
+
+            FROM compat_refresh_tokens
+
+            WHERE refresh_token = $1
         "#,
         token,
     )
     .fetch_one(executor)
-    .instrument(info_span!("Fetch compat refresh token"))
     .await
     .to_option()?;
 
     let Some(res) = res else { return Ok(None); };
 
+    let state = match res.consumed_at {
+        None => CompatRefreshTokenState::Valid,
+        Some(consumed_at) => CompatRefreshTokenState::Consumed { consumed_at },
+    };
+
     let refresh_token = CompatRefreshToken {
         id: res.compat_refresh_token_id.into(),
-        token: res.compat_refresh_token,
-        created_at: res.compat_refresh_token_created_at,
-    };
-
-    let access_token = CompatAccessToken {
-        id: res.compat_access_token_id.into(),
-        token: res.compat_access_token,
-        created_at: res.compat_access_token_created_at,
-        expires_at: res.compat_access_token_expires_at,
-    };
-
-    let id = res.compat_session_id.into();
-    let device = Device::try_from(res.compat_session_device_id).map_err(|e| {
-        DatabaseInconsistencyError::on("compat_sessions")
-            .column("device_id")
-            .row(id)
-            .source(e)
-    })?;
-
-    let state = match res.compat_session_finished_at {
-        None => CompatSessionState::Valid,
-        Some(finished_at) => CompatSessionState::Finished { finished_at },
-    };
-
-    let session = CompatSession {
-        id,
         state,
-        user_id: res.user_id.into(),
+        session_id: res.compat_session_id.into(),
-        device,
+        access_token_id: res.compat_access_token_id.into(),
-        created_at: res.compat_session_created_at,
+        token: res.refresh_token,
+        created_at: res.created_at,
     };
 
-    Ok(Some((refresh_token, access_token, session)))
+    Ok(Some(refresh_token))
 }
 
 #[tracing::instrument(
@@ -244,6 +265,7 @@ pub async fn add_compat_access_token(
 
     Ok(CompatAccessToken {
         id,
+        session_id: session.id,
         token,
         created_at,
         expires_at,
@@ -320,6 +342,9 @@ pub async fn add_compat_refresh_token(
 
     Ok(CompatRefreshToken {
         id,
+        state: CompatRefreshTokenState::default(),
+        session_id: session.id,
+        access_token_id: access_token.id,
         token,
         created_at,
     })
@@ -327,42 +352,35 @@ pub async fn add_compat_refresh_token(
 
 #[tracing::instrument(
     skip_all,
-    fields(compat_session.id),
+    fields(%compat_session.id),
     err,
 )]
-pub async fn compat_logout(
+pub async fn end_compat_session(
     executor: impl PgExecutor<'_>,
     clock: &Clock,
-    token: &str,
+    compat_session: CompatSession,
-) -> Result<bool, sqlx::Error> {
+) -> Result<CompatSession, DatabaseError> {
     let finished_at = clock.now();
-    // TODO: this does not check for token expiration
+
-    let res = sqlx::query_scalar!(
+    let res = sqlx::query!(
         r#"
             UPDATE compat_sessions cs
             SET finished_at = $2
-            FROM compat_access_tokens ca
+            WHERE compat_session_id = $1
-            WHERE ca.access_token = $1
-              AND ca.compat_session_id = cs.compat_session_id
-              AND cs.finished_at IS NULL
-            RETURNING cs.compat_session_id
         "#,
-        token,
+        Uuid::from(compat_session.id),
         finished_at,
     )
-    .fetch_one(executor)
+    .execute(executor)
-    .await
+    .await?;
-    .to_option()?;
 
-    if let Some(compat_session_id) = res {
+    DatabaseError::ensure_affected_rows(&res, 1)?;
-        tracing::Span::current().record(
+
-            "compat_session.id",
+    let compat_session = compat_session
-            tracing::field::display(compat_session_id),
+        .finish(finished_at)
-        );
+        .map_err(DatabaseError::to_invalid_operation)?;
-        Ok(true)
+
-    } else {
+    Ok(compat_session)
-        Ok(false)
-    }
 }
 
 #[tracing::instrument(
@@ -445,10 +463,6 @@ struct CompatSsoLoginLookup {
     compat_sso_login_fulfilled_at: Option<DateTime<Utc>>,
     compat_sso_login_exchanged_at: Option<DateTime<Utc>>,
     compat_session_id: Option<Uuid>,
-    compat_session_created_at: Option<DateTime<Utc>>,
-    compat_session_finished_at: Option<DateTime<Utc>>,
-    compat_session_device_id: Option<String>,
-    user_id: Option<Uuid>,
 }
 
 impl TryFrom<CompatSsoLoginLookup> for CompatSsoLogin {
@@ -463,58 +477,21 @@ impl TryFrom<CompatSsoLoginLookup> for CompatSsoLogin {
                 .source(e)
         })?;
 
-        let session = match (
-            res.compat_session_id,
-            res.compat_session_device_id,
-            res.compat_session_created_at,
-            res.compat_session_finished_at,
-            res.user_id,
-        ) {
-            (Some(id), Some(device_id), Some(created_at), finished_at, Some(user_id)) => {
-                let id = id.into();
-                let device = Device::try_from(device_id).map_err(|e| {
-                    DatabaseInconsistencyError::on("compat_sessions")
-                        .column("device")
-                        .row(id)
-                        .source(e)
-                })?;
-
-                let state = match finished_at {
-                    None => CompatSessionState::Valid,
-                    Some(finished_at) => CompatSessionState::Finished { finished_at },
-                };
-
-                Some(CompatSession {
-                    id,
-                    state,
-                    user_id: user_id.into(),
-                    device,
-                    created_at,
-                })
-            }
-            (None, None, None, None, None) => None,
-            _ => {
-                return Err(DatabaseInconsistencyError::on("compat_sso_logins")
-                    .column("compat_session_id")
-                    .row(id))
-            }
-        };
-
         let state = match (
             res.compat_sso_login_fulfilled_at,
             res.compat_sso_login_exchanged_at,
-            session,
+            res.compat_session_id,
         ) {
             (None, None, None) => CompatSsoLoginState::Pending,
-            (Some(fulfilled_at), None, Some(session)) => CompatSsoLoginState::Fulfilled {
+            (Some(fulfilled_at), None, Some(session_id)) => CompatSsoLoginState::Fulfilled {
                 fulfilled_at,
-                session,
+                session_id: session_id.into(),
             },
-            (Some(fulfilled_at), Some(exchanged_at), Some(session)) => {
+            (Some(fulfilled_at), Some(exchanged_at), Some(session_id)) => {
                 CompatSsoLoginState::Exchanged {
                     fulfilled_at,
                     exchanged_at,
-                    session,
+                    session_id: session_id.into(),
                 }
             }
             _ => return Err(DatabaseInconsistencyError::on("compat_sso_logins").row(id)),
@@ -550,15 +527,9 @@ pub async fn get_compat_sso_login_by_id(
                  , cl.created_at      AS "compat_sso_login_created_at"
                  , cl.fulfilled_at    AS "compat_sso_login_fulfilled_at"
                  , cl.exchanged_at    AS "compat_sso_login_exchanged_at"
-                 , cs.compat_session_id AS "compat_session_id?"
+                 , cl.compat_session_id AS "compat_session_id"
-                 , cs.created_at      AS "compat_session_created_at?"
-                 , cs.finished_at     AS "compat_session_finished_at?"
-                 , cs.device_id       AS "compat_session_device_id?"
-                 , cs.user_id         AS "user_id?"
 
             FROM compat_sso_logins cl
-            LEFT JOIN compat_sessions cs
-              USING (compat_session_id)
             WHERE cl.compat_sso_login_id = $1
         "#,
         Uuid::from(id),
@@ -589,8 +560,6 @@ pub async fn get_paginated_user_compat_sso_logins(
     first: Option<usize>,
     last: Option<usize>,
 ) -> Result<(bool, bool, Vec<CompatSsoLogin>), DatabaseError> {
-    // TODO: this queries too much (like user info) which we probably don't need
-    // because we already have them
     let mut query = QueryBuilder::new(
         r#"
             SELECT cl.compat_sso_login_id
@@ -599,14 +568,8 @@ pub async fn get_paginated_user_compat_sso_logins(
                  , cl.created_at      AS "compat_sso_login_created_at"
                  , cl.fulfilled_at    AS "compat_sso_login_fulfilled_at"
                  , cl.exchanged_at    AS "compat_sso_login_exchanged_at"
-                 , cs.compat_session_id AS "compat_session_id"
+                 , cl.compat_session_id AS "compat_session_id"
-                 , cs.created_at      AS "compat_session_created_at"
-                 , cs.finished_at     AS "compat_session_finished_at"
-                 , cs.device_id       AS "compat_session_device_id"
-                 , cs.user_id
             FROM compat_sso_logins cl
-            LEFT JOIN compat_sessions cs
-              USING (compat_session_id)
         "#,
     );
 
@@ -645,14 +608,8 @@ pub async fn get_compat_sso_login_by_token(
                  , cl.created_at      AS "compat_sso_login_created_at"
                  , cl.fulfilled_at    AS "compat_sso_login_fulfilled_at"
                  , cl.exchanged_at    AS "compat_sso_login_exchanged_at"
-                 , cs.compat_session_id AS "compat_session_id?"
+                 , cl.compat_session_id AS "compat_session_id"
-                 , cs.created_at      AS "compat_session_created_at?"
-                 , cs.finished_at     AS "compat_session_finished_at?"
-                 , cs.device_id       AS "compat_session_device_id?"
-                 , cs.user_id         AS "user_id?"
             FROM compat_sso_logins cl
-            LEFT JOIN compat_sessions cs
-              USING (compat_session_id)
             WHERE cl.login_token = $1
         "#,
         token,
@@ -739,7 +696,7 @@ pub async fn fullfill_compat_sso_login(
 
     let fulfilled_at = clock.now();
     let compat_sso_login = compat_sso_login
-        .fulfill(fulfilled_at, session)
+        .fulfill(fulfilled_at, &session)
         .map_err(DatabaseError::to_invalid_operation)?;
     sqlx::query!(
         r#"