storage: Look up compat sessions by device_id

crates/storage-pg/src/compat/mod.rs

@@ -129,6 +129,19 @@ mod tests {
         assert!(session_lookup.is_valid());
         assert!(!session_lookup.is_finished());
 
+        // Look up the session by device
+        let session_lookup = repo
+            .compat_session()
+            .find_by_device(&user, &session.device)
+            .await
+            .unwrap()
+            .expect("compat session not found");
+        assert_eq!(session_lookup.id, session.id);
+        assert_eq!(session_lookup.user_id, user.id);
+        assert_eq!(session_lookup.device.as_str(), device_str);
+        assert!(session_lookup.is_valid());
+        assert!(!session_lookup.is_finished());
+
         // Finish the session
         let session = repo.compat_session().finish(&clock, session).await.unwrap();
         assert!(!session.is_valid());

crates/storage-pg/src/compat/session.rs

@@ -221,6 +221,46 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
         Ok(Some(res.try_into()?))
     }
 
+    #[tracing::instrument(
+        name = "db.compat_session.find_by_device",
+        skip_all,
+        fields(
+            db.statement,
+            %user.id,
+            %user.username,
+            compat_session.device.id = device.as_str(),
+        ),
+    )]
+    async fn find_by_device(
+        &mut self,
+        user: &User,
+        device: &Device,
+    ) -> Result<Option<CompatSession>, Self::Error> {
+        let res = sqlx::query_as!(
+            CompatSessionLookup,
+            r#"
+                SELECT compat_session_id
+                     , device_id
+                     , user_id
+                     , created_at
+                     , finished_at
+                     , is_synapse_admin
+                FROM compat_sessions
+                WHERE user_id = $1
+                  AND device_id = $2
+            "#,
+            Uuid::from(user.id),
+            device.as_str(),
+        )
+        .traced()
+        .fetch_optional(&mut *self.conn)
+        .await?;
+
+        let Some(res) = res else { return Ok(None) };
+
+        Ok(Some(res.try_into()?))
+    }
+
     #[tracing::instrument(
         name = "db.compat_session.add",
         skip_all,