From d2f1ab08fb779f1214ac288a8a8bea81653bb1bb Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Tue, 22 Apr 2025 14:42:52 +0200 Subject: [PATCH] syn2mas: implement WriteBatch for MasNewCompatRefreshToken --- ...b986b1b4864a778525d0b8b0ad6678aef3e9f.json | 18 ++++ ...e5d8cac3836701fc24922f4f0e8b98d330796.json | 18 ---- crates/syn2mas/src/mas_writer/mod.rs | 99 ++++++++++--------- 3 files changed, 71 insertions(+), 64 deletions(-) create mode 100644 crates/syn2mas/.sqlx/query-1d1004d0fb5939fbf30c1986b80b986b1b4864a778525d0b8b0ad6678aef3e9f.json delete mode 100644 crates/syn2mas/.sqlx/query-88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796.json diff --git a/crates/syn2mas/.sqlx/query-1d1004d0fb5939fbf30c1986b80b986b1b4864a778525d0b8b0ad6678aef3e9f.json b/crates/syn2mas/.sqlx/query-1d1004d0fb5939fbf30c1986b80b986b1b4864a778525d0b8b0ad6678aef3e9f.json new file mode 100644 index 000000000..c65dfb7a4 --- /dev/null +++ b/crates/syn2mas/.sqlx/query-1d1004d0fb5939fbf30c1986b80b986b1b4864a778525d0b8b0ad6678aef3e9f.json @@ -0,0 +1,18 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO syn2mas__compat_refresh_tokens (\n compat_refresh_token_id,\n compat_session_id,\n compat_access_token_id,\n refresh_token,\n created_at)\n SELECT * FROM UNNEST(\n $1::UUID[],\n $2::UUID[],\n $3::UUID[],\n $4::TEXT[],\n $5::TIMESTAMP WITH TIME ZONE[])\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "UuidArray", + "UuidArray", + "UuidArray", + "TextArray", + "TimestamptzArray" + ] + }, + "nullable": [] + }, + "hash": "1d1004d0fb5939fbf30c1986b80b986b1b4864a778525d0b8b0ad6678aef3e9f" +} diff --git a/crates/syn2mas/.sqlx/query-88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796.json b/crates/syn2mas/.sqlx/query-88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796.json deleted file mode 100644 index cb251624d..000000000 --- a/crates/syn2mas/.sqlx/query-88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n INSERT INTO syn2mas__compat_refresh_tokens (\n compat_refresh_token_id,\n compat_session_id,\n compat_access_token_id,\n refresh_token,\n created_at)\n SELECT * FROM UNNEST(\n $1::UUID[],\n $2::UUID[],\n $3::UUID[],\n $4::TEXT[],\n $5::TIMESTAMP WITH TIME ZONE[])\n ", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "UuidArray", - "UuidArray", - "UuidArray", - "TextArray", - "TimestamptzArray" - ] - }, - "nullable": [] - }, - "hash": "88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796" -} diff --git a/crates/syn2mas/src/mas_writer/mod.rs b/crates/syn2mas/src/mas_writer/mod.rs index 182a4795b..0012ef04a 100644 --- a/crates/syn2mas/src/mas_writer/mod.rs +++ b/crates/syn2mas/src/mas_writer/mod.rs @@ -668,6 +668,58 @@ pub struct MasNewCompatRefreshToken { pub created_at: DateTime, } +impl WriteBatch for MasNewCompatRefreshToken { + async fn write_batch(conn: &mut PgConnection, batch: Vec) -> Result<(), Error> { + let mut refresh_token_ids: Vec = Vec::with_capacity(batch.len()); + let mut session_ids: Vec = Vec::with_capacity(batch.len()); + let mut access_token_ids: Vec = Vec::with_capacity(batch.len()); + let mut refresh_tokens: Vec = Vec::with_capacity(batch.len()); + let mut created_ats: Vec> = Vec::with_capacity(batch.len()); + + for MasNewCompatRefreshToken { + refresh_token_id, + session_id, + access_token_id, + refresh_token, + created_at, + } in batch + { + refresh_token_ids.push(refresh_token_id); + session_ids.push(session_id); + access_token_ids.push(access_token_id); + refresh_tokens.push(refresh_token); + created_ats.push(created_at); + } + + sqlx::query!( + r#" + INSERT INTO syn2mas__compat_refresh_tokens ( + compat_refresh_token_id, + compat_session_id, + compat_access_token_id, + refresh_token, + created_at) + SELECT * FROM UNNEST( + $1::UUID[], + $2::UUID[], + $3::UUID[], + $4::TEXT[], + $5::TIMESTAMP WITH TIME ZONE[]) + "#, + &refresh_token_ids[..], + &session_ids[..], + &access_token_ids[..], + &refresh_tokens[..], + &created_ats[..], + ) + .execute(&mut *conn) + .await + .into_database("writing compat refresh tokens to MAS")?; + + Ok(()) + } +} + /// 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. @@ -1103,52 +1155,7 @@ impl MasWriter { self.writer_pool .spawn_with_connection(move |conn| { Box::pin(async move { - let mut refresh_token_ids: Vec = Vec::with_capacity(tokens.len()); - let mut session_ids: Vec = Vec::with_capacity(tokens.len()); - let mut access_token_ids: Vec = Vec::with_capacity(tokens.len()); - let mut refresh_tokens: Vec = Vec::with_capacity(tokens.len()); - let mut created_ats: Vec> = Vec::with_capacity(tokens.len()); - - for MasNewCompatRefreshToken { - refresh_token_id, - session_id, - access_token_id, - refresh_token, - created_at, - } in tokens - { - refresh_token_ids.push(refresh_token_id); - session_ids.push(session_id); - access_token_ids.push(access_token_id); - refresh_tokens.push(refresh_token); - created_ats.push(created_at); - } - - sqlx::query!( - r#" - INSERT INTO syn2mas__compat_refresh_tokens ( - compat_refresh_token_id, - compat_session_id, - compat_access_token_id, - refresh_token, - created_at) - SELECT * FROM UNNEST( - $1::UUID[], - $2::UUID[], - $3::UUID[], - $4::TEXT[], - $5::TIMESTAMP WITH TIME ZONE[]) - "#, - &refresh_token_ids[..], - &session_ids[..], - &access_token_ids[..], - &refresh_tokens[..], - &created_ats[..], - ) - .execute(&mut *conn) - .await - .into_database("writing compat refresh tokens to MAS")?; - + MasNewCompatRefreshToken::write_batch(conn, tokens).await?; Ok(()) }) })