Use the list function with a fixed amount of results instead of a custom function

Signed-off-by: MTRNord <[email protected]>

Author: MTRNord

crates/handlers/src/graphql/model/upstream_oauth.rs

@@ -7,7 +7,11 @@
 use anyhow::Context as _;
 use async_graphql::{Context, Object, ID};
 use chrono::{DateTime, Utc};
-use mas_storage::{upstream_oauth2::UpstreamOAuthProviderRepository, user::UserRepository};
+use mas_storage::{
+    upstream_oauth2::{UpstreamOAuthLinkFilter, UpstreamOAuthProviderRepository},
+    user::UserRepository,
+    Pagination,
+};
 
 use super::{NodeType, User};
 use crate::graphql::state::ContextExt;
@@ -57,20 +61,28 @@ impl UpstreamOAuth2Provider {
         ctx: &Context<'_>,
     ) -> Result<Vec<UpstreamOAuth2Link>, async_graphql::Error> {
         let state = ctx.state();
-        let user_id = ctx
+        let user = ctx
             .requester()
             .user()
-            .ok_or_else(|| async_graphql::Error::new("User ID not found in the request context"))?
-            .id;
+            .ok_or_else(|| async_graphql::Error::new("User ID not found in the request context"))?;
 
         let mut repo = state.repository().await?;
+        let filter = UpstreamOAuthLinkFilter::new()
+            .for_provider(&self.provider)
+            .for_user(&user);
         let links = repo
             .upstream_oauth_link()
-            .find_by_user_id(&self.provider, user_id)
+            // Hardcoded limit of 100 links. We do not expect reasonably more links
+            // See also https://github.com/element-hq/matrix-authentication-service/pull/3245#discussion_r1776850096
+            .list(filter, Pagination::first(100))
             .await?;
         repo.cancel().await?;
 
-        Ok(links.into_iter().map(UpstreamOAuth2Link::new).collect())
+        Ok(links
+            .edges
+            .into_iter()
+            .map(UpstreamOAuth2Link::new)
+            .collect())
     }
 }
 

crates/storage-pg/src/upstream_oauth2/link.rs

@@ -179,47 +179,6 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
         Ok(res)
     }
 
-    #[tracing::instrument(
-        name = "db.upstream_oauth_link.find_by_user_id",
-        skip_all,
-        fields(
-            db.query.text,
-            upstream_oauth_link.user_id = user_id.0,
-            %upstream_oauth_provider.id,
-            %upstream_oauth_provider.issuer,
-            %upstream_oauth_provider.client_id,
-        ),
-        err,
-    )]
-    async fn find_by_user_id(
-        &mut self,
-        upstream_oauth_provider: &UpstreamOAuthProvider,
-        user_id: Ulid,
-    ) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
-        let res = sqlx::query_as!(
-            LinkLookup,
-            r#"
-                SELECT
-                    upstream_oauth_link_id,
-                    upstream_oauth_provider_id,
-                    user_id,
-                    subject,
-                    created_at
-                FROM upstream_oauth_links
-                WHERE upstream_oauth_provider_id = $1
-                  AND user_id = $2
-            "#,
-            Uuid::from(upstream_oauth_provider.id),
-            Uuid::from(user_id),
-        )
-        .traced()
-        .fetch_optional(&mut *self.conn)
-        .await?
-        .map(Into::into);
-
-        Ok(res)
-    }
-
     #[tracing::instrument(
         name = "db.upstream_oauth_link.add",
         skip_all,

crates/storage/src/upstream_oauth2/link.rs

@@ -117,25 +117,6 @@ pub trait UpstreamOAuthLinkRepository: Send + Sync {
         subject: &str,
     ) -> Result<Option<UpstreamOAuthLink>, Self::Error>;
 
-    /// Find an upstream OAuth link for a provider by the associated user id
-    ///
-    /// Returns `None` if no matching upstream OAuth link was found
-    ///
-    /// # Parameters
-    ///
-    /// * `upstream_oauth_provider`: The upstream OAuth provider on which to
-    ///  find the link
-    /// * `user_id`: The user id of the upstream OAuth link to find
-    ///
-    /// # Errors
-    ///
-    /// Returns [`Self::Error`] if the underlying repository fails
-    async fn find_by_user_id(
-        &mut self,
-        upstream_oauth_provider: &UpstreamOAuthProvider,
-        user_id: Ulid,
-    ) -> Result<Option<UpstreamOAuthLink>, Self::Error>;
-
     /// Add a new upstream OAuth link
     ///
     /// Returns the newly created upstream OAuth link
@@ -214,12 +195,6 @@ repository_impl!(UpstreamOAuthLinkRepository:
         subject: &str,
     ) -> Result<Option<UpstreamOAuthLink>, Self::Error>;
 
-    async fn find_by_user_id(
-        &mut self,
-        upstream_oauth_provider: &UpstreamOAuthProvider,
-        user_id: Ulid,
-    ) -> Result<Option<UpstreamOAuthLink>, Self::Error>;
-
     async fn add(
         &mut self,
         rng: &mut (dyn RngCore + Send),