From 8402a75a7df1e5577dbc34700c58490a15812ab1 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 29 Aug 2023 14:41:08 +0200
Subject: [PATCH] storage: Look up compat sessions by device_id

---
 ...3233e70b44277e7fe2ce4317c9d821512d86b.json | 53 +++++++++++++++++++
 crates/storage-pg/src/compat/mod.rs           | 13 +++++
 crates/storage-pg/src/compat/session.rs       | 40 ++++++++++++++
 crates/storage/src/compat/session.rs          | 24 +++++++++
 4 files changed, 130 insertions(+)
 create mode 100644 crates/storage-pg/.sqlx/query-89e8bb889b8e604d1b4669e904d3233e70b44277e7fe2ce4317c9d821512d86b.json

diff --git a/crates/storage-pg/.sqlx/query-89e8bb889b8e604d1b4669e904d3233e70b44277e7fe2ce4317c9d821512d86b.json b/crates/storage-pg/.sqlx/query-89e8bb889b8e604d1b4669e904d3233e70b44277e7fe2ce4317c9d821512d86b.json
new file mode 100644
index 00000000..2aff21fa
--- /dev/null
+++ b/crates/storage-pg/.sqlx/query-89e8bb889b8e604d1b4669e904d3233e70b44277e7fe2ce4317c9d821512d86b.json
@@ -0,0 +1,53 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                SELECT compat_session_id\n                     , device_id\n                     , user_id\n                     , created_at\n                     , finished_at\n                     , is_synapse_admin\n                FROM compat_sessions\n                WHERE user_id = $1\n                  AND device_id = $2\n            ",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "compat_session_id",
+        "type_info": "Uuid"
+      },
+      {
+        "ordinal": 1,
+        "name": "device_id",
+        "type_info": "Text"
+      },
+      {
+        "ordinal": 2,
+        "name": "user_id",
+        "type_info": "Uuid"
+      },
+      {
+        "ordinal": 3,
+        "name": "created_at",
+        "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 4,
+        "name": "finished_at",
+        "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 5,
+        "name": "is_synapse_admin",
+        "type_info": "Bool"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Text"
+      ]
+    },
+    "nullable": [
+      false,
+      false,
+      false,
+      false,
+      true,
+      false
+    ]
+  },
+  "hash": "89e8bb889b8e604d1b4669e904d3233e70b44277e7fe2ce4317c9d821512d86b"
+}
diff --git a/crates/storage-pg/src/compat/mod.rs b/crates/storage-pg/src/compat/mod.rs
index 12da6159..96d97955 100644
--- a/crates/storage-pg/src/compat/mod.rs
+++ b/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());
diff --git a/crates/storage-pg/src/compat/session.rs b/crates/storage-pg/src/compat/session.rs
index cd1f8a93..cbfeb710 100644
--- a/crates/storage-pg/src/compat/session.rs
+++ b/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,
diff --git a/crates/storage/src/compat/session.rs b/crates/storage/src/compat/session.rs
index dd2b39cd..ec28ebcf 100644
--- a/crates/storage/src/compat/session.rs
+++ b/crates/storage/src/compat/session.rs
@@ -148,6 +148,24 @@ pub trait CompatSessionRepository: Send + Sync {
     /// Returns [`Self::Error`] if the underlying repository fails
     async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatSession>, Self::Error>;
 
+    /// Find a compatibility session by its device ID
+    ///
+    /// Returns the compat session if it exists, `None` otherwise
+    ///
+    /// # Parameters
+    ///
+    /// * `user`: The user to lookup the compat session for
+    /// * `device`: The device ID of the compat session to lookup
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn find_by_device(
+        &mut self,
+        user: &User,
+        device: &Device,
+    ) -> Result<Option<CompatSession>, Self::Error>;
+
     /// Start a new compat session
     ///
     /// Returns the newly created compat session
@@ -223,6 +241,12 @@ pub trait CompatSessionRepository: Send + Sync {
 repository_impl!(CompatSessionRepository:
     async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatSession>, Self::Error>;
 
+    async fn find_by_device(
+        &mut self,
+        user: &User,
+        device: &Device,
+    ) -> Result<Option<CompatSession>, Self::Error>;
+
     async fn add(
         &mut self,
         rng: &mut (dyn RngCore + Send),