systemprompt_users/repository/user/
merge.rs

1//! Anonymous-to-identified user merge operations.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use sqlx::{Acquire, Postgres, Transaction};
7use systemprompt_identifiers::UserId;
8
9use crate::error::Result;
10use crate::repository::UserRepository;
11
12#[derive(Debug, Clone, Copy)]
13pub struct MergeResult {
14    pub sessions: u64,
15    pub tasks: u64,
16    pub total_rows: u64,
17}
18
19pub const MERGE_EXCLUDED_SECURITY_TABLES: &[&str] = &[
20    "oauth_auth_codes",
21    "oauth_refresh_tokens",
22    "oauth_clients",
23    "webauthn_credentials",
24    "webauthn_challenges",
25    "webauthn_setup_tokens",
26    "user_api_keys",
27    "user_device_certs",
28    "bridge_sessions",
29    "bridge_exchange_codes",
30    "federated_identities",
31];
32
33impl UserRepository {
34    pub async fn merge_users(&self, source_id: &UserId, target_id: &UserId) -> Result<MergeResult> {
35        let mut conn = self.write_pool.acquire().await?;
36        let mut tx = conn.begin().await?;
37
38        let source = source_id.as_str();
39        let target = target_id.as_str();
40
41        let sessions = transfer_sessions(&mut tx, source, target).await?;
42        let tasks = transfer_tasks(&mut tx, source, target).await?;
43        let mut total_rows = sessions + tasks;
44        total_rows += transfer_audit_rows(&mut tx, source, target).await?;
45        total_rows += transfer_content_rows(&mut tx, source, target).await?;
46
47        sqlx::query!(
48            "UPDATE fingerprint_reputation SET associated_user_ids = \
49             array_replace(associated_user_ids, $2, $1) WHERE $2 = ANY(associated_user_ids)",
50            target,
51            source
52        )
53        .execute(&mut *tx)
54        .await?;
55
56        sqlx::query!(
57            "DELETE FROM ai_quota_buckets WHERE subject_kind = 'user' AND subject_id = $1",
58            source
59        )
60        .execute(&mut *tx)
61        .await?;
62
63        sqlx::query!("DELETE FROM users WHERE id = $1", source)
64            .execute(&mut *tx)
65            .await?;
66
67        tx.commit().await?;
68        Ok(MergeResult {
69            sessions,
70            tasks,
71            total_rows,
72        })
73    }
74}
75
76async fn transfer_sessions(
77    tx: &mut Transaction<'_, Postgres>,
78    source: &str,
79    target: &str,
80) -> Result<u64> {
81    let result = sqlx::query!(
82        "UPDATE user_sessions SET user_id = $1 WHERE user_id = $2",
83        target,
84        source
85    )
86    .execute(&mut **tx)
87    .await?;
88    Ok(result.rows_affected())
89}
90
91async fn transfer_tasks(
92    tx: &mut Transaction<'_, Postgres>,
93    source: &str,
94    target: &str,
95) -> Result<u64> {
96    let result = sqlx::query!(
97        "UPDATE agent_tasks SET user_id = $1 WHERE user_id = $2",
98        target,
99        source
100    )
101    .execute(&mut **tx)
102    .await?;
103    Ok(result.rows_affected())
104}
105
106async fn transfer_audit_rows(
107    tx: &mut Transaction<'_, Postgres>,
108    source: &str,
109    target: &str,
110) -> Result<u64> {
111    let mut moved = 0;
112    moved += sqlx::query!(
113        "UPDATE task_messages SET user_id = $1 WHERE user_id = $2",
114        target,
115        source
116    )
117    .execute(&mut **tx)
118    .await?
119    .rows_affected();
120    moved += sqlx::query!(
121        "UPDATE user_contexts SET user_id = $1 WHERE user_id = $2",
122        target,
123        source
124    )
125    .execute(&mut **tx)
126    .await?
127    .rows_affected();
128    moved += sqlx::query!(
129        "UPDATE mcp_tool_executions SET user_id = $1 WHERE user_id = $2",
130        target,
131        source
132    )
133    .execute(&mut **tx)
134    .await?
135    .rows_affected();
136    moved += sqlx::query!(
137        "UPDATE mcp_artifacts SET user_id = $1 WHERE user_id = $2",
138        target,
139        source
140    )
141    .execute(&mut **tx)
142    .await?
143    .rows_affected();
144    moved += sqlx::query!(
145        "UPDATE mcp_sessions SET user_id = $1 WHERE user_id = $2",
146        target,
147        source
148    )
149    .execute(&mut **tx)
150    .await?
151    .rows_affected();
152    moved += sqlx::query!(
153        "UPDATE governance_decisions SET user_id = $1 WHERE user_id = $2",
154        target,
155        source
156    )
157    .execute(&mut **tx)
158    .await?
159    .rows_affected();
160    moved += sqlx::query!(
161        "UPDATE logs SET user_id = $1 WHERE user_id = $2",
162        target,
163        source
164    )
165    .execute(&mut **tx)
166    .await?
167    .rows_affected();
168    Ok(moved)
169}
170
171async fn transfer_content_rows(
172    tx: &mut Transaction<'_, Postgres>,
173    source: &str,
174    target: &str,
175) -> Result<u64> {
176    let mut moved = 0;
177    moved += sqlx::query!(
178        "UPDATE ai_requests SET user_id = $1 WHERE user_id = $2",
179        target,
180        source
181    )
182    .execute(&mut **tx)
183    .await?
184    .rows_affected();
185    moved += sqlx::query!(
186        "UPDATE engagement_events SET user_id = $1 WHERE user_id = $2",
187        target,
188        source
189    )
190    .execute(&mut **tx)
191    .await?
192    .rows_affected();
193    moved += sqlx::query!(
194        "UPDATE analytics_events SET user_id = $1 WHERE user_id = $2",
195        target,
196        source
197    )
198    .execute(&mut **tx)
199    .await?
200    .rows_affected();
201    moved += sqlx::query!(
202        "UPDATE event_outbox SET user_id = $1 WHERE user_id = $2",
203        target,
204        source
205    )
206    .execute(&mut **tx)
207    .await?
208    .rows_affected();
209    moved += sqlx::query!(
210        "UPDATE files SET user_id = $1 WHERE user_id = $2",
211        target,
212        source
213    )
214    .execute(&mut **tx)
215    .await?
216    .rows_affected();
217    moved += sqlx::query!(
218        "UPDATE link_clicks SET user_id = $1 WHERE user_id = $2",
219        target,
220        source
221    )
222    .execute(&mut **tx)
223    .await?
224    .rows_affected();
225    Ok(moved)
226}
systemprompt_users/repository/user/merge.rs

systemprompt_users/repository/user/
merge.rs
