Link removal storage API

From #3245 with changes from review

Author: MTRNord

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

@@ -30,6 +30,12 @@ pub enum UpstreamOAuthAuthorizationSessionState {
         extra_callback_parameters: Option<serde_json::Value>,
         userinfo: Option<serde_json::Value>,
     },
+    Unlinked {
+        completed_at: DateTime<Utc>,
+        consumed_at: Option<DateTime<Utc>>,
+        unlinked_at: DateTime<Utc>,
+        id_token: Option<String>,
+    },
 }
 
 impl UpstreamOAuthAuthorizationSessionState {
@@ -57,7 +63,9 @@ impl UpstreamOAuthAuthorizationSessionState {
                 extra_callback_parameters,
                 userinfo,
             }),
-            Self::Completed { .. } | Self::Consumed { .. } => Err(InvalidTransitionError),
+            Self::Completed { .. } | Self::Consumed { .. } | Self::Unlinked { .. } => {
+                Err(InvalidTransitionError)
+            }
         }
     }
 
@@ -85,7 +93,9 @@ impl UpstreamOAuthAuthorizationSessionState {
                 extra_callback_parameters,
                 userinfo,
             }),
-            Self::Pending | Self::Consumed { .. } => Err(InvalidTransitionError),
+            Self::Pending | Self::Consumed { .. } | Self::Unlinked { .. } => {
+                Err(InvalidTransitionError)
+            }
         }
     }
 
@@ -98,7 +108,7 @@ impl UpstreamOAuthAuthorizationSessionState {
     #[must_use]
     pub fn link_id(&self) -> Option<Ulid> {
         match self {
-            Self::Pending => None,
+            Self::Pending | Self::Unlinked { .. } => None,
             Self::Completed { link_id, .. } | Self::Consumed { link_id, .. } => Some(*link_id),
         }
     }
@@ -114,9 +124,9 @@ impl UpstreamOAuthAuthorizationSessionState {
     pub fn completed_at(&self) -> Option<DateTime<Utc>> {
         match self {
             Self::Pending => None,
-            Self::Completed { completed_at, .. } | Self::Consumed { completed_at, .. } => {
-                Some(*completed_at)
-            }
+            Self::Completed { completed_at, .. }
+            | Self::Consumed { completed_at, .. }
+            | Self::Unlinked { completed_at, .. } => Some(*completed_at),
         }
     }
 
@@ -130,9 +140,9 @@ impl UpstreamOAuthAuthorizationSessionState {
     pub fn id_token(&self) -> Option<&str> {
         match self {
             Self::Pending => None,
-            Self::Completed { id_token, .. } | Self::Consumed { id_token, .. } => {
-                id_token.as_deref()
-            }
+            Self::Completed { id_token, .. }
+            | Self::Consumed { id_token, .. }
+            | Self::Unlinked { id_token, .. } => id_token.as_deref(),
         }
     }
 
@@ -145,7 +155,7 @@ impl UpstreamOAuthAuthorizationSessionState {
     #[must_use]
     pub fn extra_callback_parameters(&self) -> Option<&serde_json::Value> {
         match self {
-            Self::Pending => None,
+            Self::Pending | Self::Unlinked { .. } => None,
             Self::Completed {
                 extra_callback_parameters,
                 ..
@@ -160,7 +170,7 @@ impl UpstreamOAuthAuthorizationSessionState {
     #[must_use]
     pub fn userinfo(&self) -> Option<&serde_json::Value> {
         match self {
-            Self::Pending => None,
+            Self::Pending | Self::Unlinked { .. } => None,
             Self::Completed { userinfo, .. } | Self::Consumed { userinfo, .. } => userinfo.as_ref(),
         }
     }
@@ -177,6 +187,22 @@ impl UpstreamOAuthAuthorizationSessionState {
         match self {
             Self::Pending | Self::Completed { .. } => None,
             Self::Consumed { consumed_at, .. } => Some(*consumed_at),
+            Self::Unlinked { consumed_at, .. } => *consumed_at,
+        }
+    }
+
+    /// Get the time at which the upstream OAuth 2.0 authorization session was
+    /// unlinked.
+    ///
+    /// Returns `None` if the upstream OAuth 2.0 authorization session state is
+    /// not [`Unlinked`].
+    ///
+    /// [`Unlinked`]: UpstreamOAuthAuthorizationSessionState::Unlinked
+    #[must_use]
+    pub fn unlinked_at(&self) -> Option<DateTime<Utc>> {
+        match self {
+            Self::Pending | Self::Completed { .. } | Self::Consumed { .. } => None,
+            Self::Unlinked { unlinked_at, .. } => Some(*unlinked_at),
         }
     }
 
@@ -206,6 +232,15 @@ impl UpstreamOAuthAuthorizationSessionState {
     pub fn is_consumed(&self) -> bool {
         matches!(self, Self::Consumed { .. })
     }
+
+    /// Returns `true` if the upstream OAuth 2.0 authorization session state is
+    /// [`Unlinked`].
+    ///
+    /// [`Unlinked`]: UpstreamOAuthAuthorizationSessionState::Unlinked
+    #[must_use]
+    pub fn is_unlinked(&self) -> bool {
+        matches!(self, Self::Unlinked { .. })
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]

crates/storage-pg/.sqlx/query-ea30b3809fd7c1d4e9983909c0219f343953a89f2a43f6b8c4ab4fbea7645ccc.json renamed to crates/storage-pg/.sqlx/query-37a124678323380357fa9d1375fd125fb35476ac3008e5adbd04a761d5edcd42.json

@@ -0,0 +1,7 @@
+-- Copyright 2025 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+ALTER TABLE upstream_oauth_authorization_sessions
+  ADD COLUMN unlinked_at TIMESTAMP WITH TIME ZONE;

crates/storage-pg/.sqlx/query-3ed73cfce8ef6a1108f454e18b1668f64b76975dba07e67d04ed7a52e2e8107f.json

@@ -374,4 +374,53 @@ impl UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'_> {
             .try_into()
             .map_err(DatabaseError::to_invalid_operation)
     }
+
+    #[tracing::instrument(
+        name = "db.upstream_oauth_link.remove",
+        skip_all,
+        fields(
+            db.query.text,
+            upstream_oauth_link.id,
+            upstream_oauth_link.provider_id,
+            %upstream_oauth_link.subject,
+        ),
+        err,
+    )]
+    async fn remove(
+        &mut self,
+        clock: &dyn Clock,
+        upstream_oauth_link: UpstreamOAuthLink,
+    ) -> Result<(), Self::Error> {
+        // Unlink the authorization sessions first, as they have a foreign key
+        // constraint on the links.
+        sqlx::query!(
+            r#"
+                UPDATE upstream_oauth_authorization_sessions SET
+                    upstream_oauth_link_id = NULL,
+                    unlinked_at = $2
+                WHERE upstream_oauth_link_id = $1
+            "#,
+            Uuid::from(upstream_oauth_link.id),
+            clock.now()
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        // Then delete the link itself
+        let res = sqlx::query!(
+            r#"
+                DELETE FROM upstream_oauth_links
+                WHERE upstream_oauth_link_id = $1
+            "#,
+            Uuid::from(upstream_oauth_link.id),
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        DatabaseError::ensure_affected_rows(&res, 1)?;
+
+        Ok(())
+    }
 }

crates/storage-pg/.sqlx/query-cc60ad934d347fb4546205d1fe07e9d2f127cb15b1bb650d1ea3805a4c55b196.json

@@ -45,6 +45,7 @@ struct SessionLookup {
     completed_at: Option<DateTime<Utc>>,
     consumed_at: Option<DateTime<Utc>>,
     extra_callback_parameters: Option<serde_json::Value>,
+    unlinked_at: Option<DateTime<Utc>>,
 }
 
 impl TryFrom<SessionLookup> for UpstreamOAuthAuthorizationSession {
@@ -59,15 +60,19 @@ impl TryFrom<SessionLookup> for UpstreamOAuthAuthorizationSession {
             value.userinfo,
             value.completed_at,
             value.consumed_at,
+            value.unlinked_at,
         ) {
-            (None, None, None, None, None, None) => UpstreamOAuthAuthorizationSessionState::Pending,
+            (None, None, None, None, None, None, None) => {
+                UpstreamOAuthAuthorizationSessionState::Pending
+            }
             (
                 Some(link_id),
                 id_token,
                 extra_callback_parameters,
                 userinfo,
                 Some(completed_at),
                 None,
+                None,
             ) => UpstreamOAuthAuthorizationSessionState::Completed {
                 completed_at,
                 link_id: link_id.into(),
@@ -82,6 +87,7 @@ impl TryFrom<SessionLookup> for UpstreamOAuthAuthorizationSession {
                 userinfo,
                 Some(completed_at),
                 Some(consumed_at),
+                None,
             ) => UpstreamOAuthAuthorizationSessionState::Consumed {
                 completed_at,
                 link_id: link_id.into(),
@@ -90,6 +96,14 @@ impl TryFrom<SessionLookup> for UpstreamOAuthAuthorizationSession {
                 userinfo,
                 consumed_at,
             },
+            (_, id_token, _, _, Some(completed_at), consumed_at, Some(unlinked_at)) => {
+                UpstreamOAuthAuthorizationSessionState::Unlinked {
+                    completed_at,
+                    id_token,
+                    consumed_at,
+                    unlinked_at,
+                }
+            }
             _ => {
                 return Err(DatabaseInconsistencyError::on(
                     "upstream_oauth_authorization_sessions",
@@ -142,7 +156,8 @@ impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> {
                     userinfo,
                     created_at,
                     completed_at,
-                    consumed_at
+                    consumed_at,
+                    unlinked_at
                 FROM upstream_oauth_authorization_sessions
                 WHERE upstream_oauth_authorization_session_id = $1
             "#,

crates/storage-pg/.sqlx/query-fcd8b4b9e003d1540357c6bf1ff9c715560d011d4c01112703a9c046170c84f1.json

@@ -200,6 +200,22 @@ pub trait UpstreamOAuthLinkRepository: Send + Sync {
     ///
     /// Returns [`Self::Error`] if the underlying repository fails
     async fn count(&mut self, filter: UpstreamOAuthLinkFilter<'_>) -> Result<usize, Self::Error>;
+
+    /// Delete a [`UpstreamOAuthLink`]
+    ///
+    /// # Parameters
+    ///
+    /// * `clock`: The clock used to generate timestamps
+    /// * `upstream_oauth_link`: The [`UpstreamOAuthLink`] to delete
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn remove(
+        &mut self,
+        clock: &dyn Clock,
+        upstream_oauth_link: UpstreamOAuthLink,
+    ) -> Result<(), Self::Error>;
 }
 
 repository_impl!(UpstreamOAuthLinkRepository:
@@ -233,4 +249,6 @@ repository_impl!(UpstreamOAuthLinkRepository:
     ) -> Result<Page<UpstreamOAuthLink>, Self::Error>;
 
     async fn count(&mut self, filter: UpstreamOAuthLinkFilter<'_>) -> Result<usize, Self::Error>;
+
+    async fn remove(&mut self, clock: &dyn Clock, upstream_oauth_link: UpstreamOAuthLink) -> Result<(), Self::Error>;
 );