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,
-        CompatSsoLoginState, Device,
+        CompatAccessToken, CompatRefreshToken, CompatRefreshTokenState, CompatSession,
+        CompatSessionState, CompatSsoLogin, CompatSsoLoginState, Device,
     },
     oauth2::{
         AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, Client,