Add MasWriter support for compat sessions

Author: reivilibre

crates/syn2mas/.sqlx/query-396c97dbfbc932c73301daa7376e36f4ef03e1224a0a89a921e5a3d398a5d35c.json

@@ -7,7 +7,7 @@
 //!
 //! This module is responsible for writing new records to MAS' database.
 
-use std::fmt::Display;
+use std::{fmt::Display, net::IpAddr};
 
 use chrono::{DateTime, Utc};
 use futures_util::{future::BoxFuture, FutureExt, TryStreamExt};
@@ -230,6 +230,18 @@ pub struct MasNewUpstreamOauthLink {
     pub created_at: DateTime<Utc>,
 }
 
+pub struct MasNewCompatSession {
+    pub session_id: Uuid,
+    pub user_id: Uuid,
+    pub device_id: String,
+    pub human_name: Option<String>,
+    pub created_at: DateTime<Utc>,
+    pub is_synapse_admin: bool,
+    pub last_active_at: Option<DateTime<Utc>>,
+    pub last_active_ip: Option<IpAddr>,
+    pub user_agent: Option<String>,
+}
+
 /// The 'version' of the password hashing scheme used for passwords when they
 /// are migrated from Synapse to MAS.
 /// This is version 1, as in the previous syn2mas script.
@@ -761,6 +773,85 @@ impl<'conn> MasWriter<'conn> {
             })
         }).boxed()
     }
+
+    #[tracing::instrument(skip_all, level = Level::DEBUG)]
+    pub fn write_compat_sessions(
+        &mut self,
+        sessions: Vec<MasNewCompatSession>,
+    ) -> BoxFuture<'_, Result<(), Error>> {
+        self.writer_pool
+            .spawn_with_connection(move |conn| {
+                Box::pin(async move {
+                    let mut session_ids: Vec<Uuid> = Vec::with_capacity(sessions.len());
+                    let mut user_ids: Vec<Uuid> = Vec::with_capacity(sessions.len());
+                    let mut device_ids: Vec<String> = Vec::with_capacity(sessions.len());
+                    let mut human_names: Vec<Option<String>> = Vec::with_capacity(sessions.len());
+                    let mut created_ats: Vec<DateTime<Utc>> = Vec::with_capacity(sessions.len());
+                    let mut is_synapse_admins: Vec<bool> = Vec::with_capacity(sessions.len());
+                    let mut last_active_ats: Vec<Option<DateTime<Utc>>> =
+                        Vec::with_capacity(sessions.len());
+                    let mut last_active_ips: Vec<Option<IpAddr>> =
+                        Vec::with_capacity(sessions.len());
+                    let mut user_agents: Vec<Option<String>> = Vec::with_capacity(sessions.len());
+
+                    for MasNewCompatSession {
+                        session_id,
+                        user_id,
+                        device_id,
+                        human_name,
+                        created_at,
+                        is_synapse_admin,
+                        last_active_at,
+                        last_active_ip,
+                        user_agent,
+                    } in sessions
+                    {
+                        session_ids.push(session_id);
+                        user_ids.push(user_id);
+                        device_ids.push(device_id);
+                        human_names.push(human_name);
+                        created_ats.push(created_at);
+                        is_synapse_admins.push(is_synapse_admin);
+                        last_active_ats.push(last_active_at);
+                        last_active_ips.push(last_active_ip);
+                        user_agents.push(user_agent);
+                    }
+
+                    sqlx::query!(
+                        r#"
+                        INSERT INTO syn2mas__compat_sessions (
+                          compat_session_id, user_id,
+                          device_id, human_name,
+                          created_at, is_synapse_admin,
+                          last_active_at, last_active_ip,
+                          user_agent)
+                        SELECT * FROM UNNEST(
+                          $1::UUID[], $2::UUID[],
+                          $3::TEXT[], $4::TEXT[],
+                          $5::TIMESTAMP WITH TIME ZONE[], $6::BOOLEAN[],
+                          $7::TIMESTAMP WITH TIME ZONE[], $8::INET[],
+                          $9::TEXT[])
+                        "#,
+                        &session_ids[..],
+                        &user_ids[..],
+                        &device_ids[..],
+                        &human_names[..] as &[Option<String>],
+                        &created_ats[..],
+                        &is_synapse_admins[..],
+                        // We need to override the typing for arrays of optionals (sqlx limitation)
+                        &last_active_ats[..] as &[Option<DateTime<Utc>>],
+                        &last_active_ips[..] as &[Option<IpAddr>],
+                        &user_agents[..] as &[Option<String>],
+                    )
+                    .execute(&mut *conn)
+                    .await
+                    .into_database("writing compat sessions to MAS")?;
+
+                    Ok(())
+                })
+            })
+            .boxed()
+    }
 }
 
 // How many entries to buffer at once, before writing a batch of rows to the
@@ -839,8 +930,8 @@ mod test {
 
     use crate::{
         mas_writer::{
-            MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser,
-            MasNewUserPassword,
+            MasNewCompatSession, MasNewEmailThreepid, MasNewUnsupportedThreepid,
+            MasNewUpstreamOauthLink, MasNewUser, MasNewUserPassword,
         },
         LockedMasDatabase, MasWriter,
     };
@@ -1105,4 +1196,41 @@ mod test {
 
         assert_db_snapshot!(&mut conn);
     }
+
+    /// Tests writing a single user, with a device (compat session).
+    #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
+    async fn test_write_user_with_device(pool: PgPool) {
+        let mut conn = pool.acquire().await.unwrap();
+        let mut writer = make_mas_writer(&pool, &mut conn).await;
+
+        writer
+            .write_users(vec![MasNewUser {
+                user_id: Uuid::from_u128(1u128),
+                username: "alice".to_owned(),
+                created_at: DateTime::default(),
+                locked_at: None,
+                can_request_admin: false,
+            }])
+            .await
+            .expect("failed to write user");
+
+        writer
+            .write_compat_sessions(vec![MasNewCompatSession {
+                user_id: Uuid::from_u128(1u128),
+                session_id: Uuid::from_u128(5u128),
+                created_at: DateTime::default(),
+                device_id: "ADEVICE".to_owned(),
+                human_name: Some("alice's pinephone".to_owned()),
+                is_synapse_admin: true,
+                last_active_at: Some(DateTime::default()),
+                last_active_ip: Some("203.0.113.1".parse().unwrap()),
+                user_agent: Some("Browser/5.0".to_owned()),
+            }])
+            .await
+            .expect("failed to write link");
+
+        writer.finish().await.expect("failed to finish MasWriter");
+
+        assert_db_snapshot!(&mut conn);
+    }
 }

crates/syn2mas/src/mas_writer/mod.rs

@@ -0,0 +1,23 @@
+---
+source: crates/syn2mas/src/mas_writer/mod.rs
+expression: db_snapshot
+---
+compat_sessions:
+  - compat_session_id: 00000000-0000-0000-0000-000000000005
+    created_at: "1970-01-01 00:00:00+00"
+    device_id: ADEVICE
+    finished_at: ~
+    human_name: "alice's pinephone"
+    is_synapse_admin: "true"
+    last_active_at: "1970-01-01 00:00:00+00"
+    last_active_ip: 203.0.113.1/32
+    user_agent: Browser/5.0
+    user_id: 00000000-0000-0000-0000-000000000001
+    user_session_id: ~
+users:
+  - can_request_admin: "false"
+    created_at: "1970-01-01 00:00:00+00"
+    locked_at: ~
+    primary_user_email_id: ~
+    user_id: 00000000-0000-0000-0000-000000000001
+    username: alice

crates/syn2mas/src/mas_writer/snapshots/syn2mas__mas_writer__test__write_user_with_device.snap

@@ -13,3 +13,6 @@ ALTER TABLE syn2mas__user_passwords RENAME TO user_passwords;
 ALTER TABLE syn2mas__user_emails RENAME TO user_emails;
 ALTER TABLE syn2mas__user_unsupported_third_party_ids RENAME TO user_unsupported_third_party_ids;
 ALTER TABLE syn2mas__upstream_oauth_links RENAME TO upstream_oauth_links;
+ALTER TABLE syn2mas__compat_sessions RENAME TO compat_sessions;
+ALTER TABLE syn2mas__compat_access_tokens RENAME TO compat_access_tokens;
+ALTER TABLE syn2mas__compat_refresh_tokens RENAME TO compat_refresh_tokens;

crates/syn2mas/src/mas_writer/syn2mas_revert_temporary_tables.sql

@@ -42,3 +42,6 @@ ALTER TABLE user_passwords RENAME TO syn2mas__user_passwords;
 ALTER TABLE user_emails RENAME TO syn2mas__user_emails;
 ALTER TABLE user_unsupported_third_party_ids RENAME TO syn2mas__user_unsupported_third_party_ids;
 ALTER TABLE upstream_oauth_links RENAME TO syn2mas__upstream_oauth_links;
+ALTER TABLE compat_sessions RENAME TO syn2mas__compat_sessions;
+ALTER TABLE compat_access_tokens RENAME TO syn2mas__compat_access_tokens;
+ALTER TABLE compat_refresh_tokens RENAME TO syn2mas__compat_refresh_tokens;