systemprompt_users/repository/user/
merge.rs1use sqlx::{Acquire, Postgres, Transaction};
7use systemprompt_identifiers::{ContextId, SessionId, UserId};
8
9use crate::error::Result;
10use crate::repository::UserRepository;
11
12const MERGE_TOOL_NAME: &str = "users.merge";
13const MERGE_POLICY: &str = "account_merge";
14
15#[derive(Debug, Clone, Copy)]
16pub struct MergeResult {
17 pub sessions: u64,
18 pub tasks: u64,
19 pub total_rows: u64,
20}
21
22pub const MERGE_EXCLUDED_SECURITY_TABLES: &[&str] = &[
23 "oauth_auth_codes",
24 "oauth_refresh_tokens",
25 "oauth_clients",
26 "webauthn_credentials",
27 "webauthn_challenges",
28 "webauthn_setup_tokens",
29 "user_api_keys",
30 "user_device_certs",
31 "bridge_sessions",
32 "bridge_exchange_codes",
33 "federated_identities",
34];
35
36impl UserRepository {
37 pub async fn merge_users(&self, source_id: &UserId, target_id: &UserId) -> Result<MergeResult> {
38 let mut conn = self.write_pool.acquire().await?;
39 let mut tx = conn.begin().await?;
40 sqlx::query!("SELECT public.begin_user_privacy() AS prepared")
41 .fetch_one(&mut *tx)
42 .await?;
43
44 let source = source_id.as_str();
45 let target = target_id.as_str();
46
47 let sessions = transfer_sessions(&mut tx, source, target).await?;
48 let tasks = transfer_tasks(&mut tx, source, target).await?;
49 let mut total_rows = sessions + tasks;
50 total_rows += transfer_audit_rows(&mut tx, source, target).await?;
51 total_rows += transfer_content_rows(&mut tx, source, target).await?;
52 record_merge_attribution(&mut tx, source, target).await?;
53
54 sqlx::query!(
55 "UPDATE fingerprint_reputation SET associated_user_ids = \
56 array_replace(associated_user_ids, $2, $1) WHERE $2 = ANY(associated_user_ids)",
57 target,
58 source
59 )
60 .execute(&mut *tx)
61 .await?;
62
63 sqlx::query!(
64 "DELETE FROM ai_quota_buckets WHERE subject_kind = 'user' AND subject_id = $1",
65 source
66 )
67 .execute(&mut *tx)
68 .await?;
69
70 sqlx::query!("DELETE FROM users WHERE id = $1", source)
71 .execute(&mut *tx)
72 .await?;
73
74 sqlx::query!("SELECT public.finish_user_privacy() AS finished")
75 .fetch_one(&mut *tx)
76 .await?;
77 tx.commit().await?;
78 Ok(MergeResult {
79 sessions,
80 tasks,
81 total_rows,
82 })
83 }
84}
85
86async fn transfer_sessions(
87 tx: &mut Transaction<'_, Postgres>,
88 source: &str,
89 target: &str,
90) -> Result<u64> {
91 let result = sqlx::query!(
92 "UPDATE user_sessions SET user_id = $1 WHERE user_id = $2",
93 target,
94 source
95 )
96 .execute(&mut **tx)
97 .await?;
98 Ok(result.rows_affected())
99}
100
101async fn transfer_tasks(
102 tx: &mut Transaction<'_, Postgres>,
103 source: &str,
104 target: &str,
105) -> Result<u64> {
106 let result = sqlx::query!(
107 "UPDATE agent_tasks SET user_id = $1 WHERE user_id = $2",
108 target,
109 source
110 )
111 .execute(&mut **tx)
112 .await?;
113 Ok(result.rows_affected())
114}
115
116async fn record_merge_attribution(
121 tx: &mut Transaction<'_, Postgres>,
122 source: &str,
123 target: &str,
124) -> Result<()> {
125 let id = uuid::Uuid::new_v4().to_string();
126 let context_id = ContextId::derived_from_session(&SessionId::new(id.clone()));
127 sqlx::query!(
128 "INSERT INTO governance_decisions (id, user_id, session_id, tool_name, decision, policy, \
129 reason, actor_kind, actor_id, context_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, \
130 $10)",
131 id,
132 target,
133 id,
134 MERGE_TOOL_NAME,
135 "allow",
136 MERGE_POLICY,
137 format!("account merge: {source} merged into {target}"),
138 "system",
139 target,
140 context_id.as_str(),
141 )
142 .execute(&mut **tx)
143 .await?;
144 Ok(())
145}
146
147async fn transfer_audit_rows(
148 tx: &mut Transaction<'_, Postgres>,
149 source: &str,
150 target: &str,
151) -> Result<u64> {
152 let mut moved = 0;
153 moved += sqlx::query!(
154 "UPDATE task_messages SET user_id = $1 WHERE user_id = $2",
155 target,
156 source
157 )
158 .execute(&mut **tx)
159 .await?
160 .rows_affected();
161 moved += sqlx::query!(
162 "UPDATE user_contexts SET user_id = $1 WHERE user_id = $2",
163 target,
164 source
165 )
166 .execute(&mut **tx)
167 .await?
168 .rows_affected();
169 moved += sqlx::query!(
170 "UPDATE mcp_tool_executions SET user_id = $1 WHERE user_id = $2",
171 target,
172 source
173 )
174 .execute(&mut **tx)
175 .await?
176 .rows_affected();
177 moved += sqlx::query!(
178 "UPDATE mcp_artifacts 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 mcp_sessions 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 logs SET user_id = $1 WHERE user_id = $2",
195 target,
196 source
197 )
198 .execute(&mut **tx)
199 .await?
200 .rows_affected();
201 Ok(moved)
202}
203
204async fn transfer_content_rows(
205 tx: &mut Transaction<'_, Postgres>,
206 source: &str,
207 target: &str,
208) -> Result<u64> {
209 let mut moved = 0;
210 moved += sqlx::query!(
211 "UPDATE ai_requests SET user_id = $1 WHERE user_id = $2",
212 target,
213 source
214 )
215 .execute(&mut **tx)
216 .await?
217 .rows_affected();
218 moved += sqlx::query!(
219 "UPDATE engagement_events SET user_id = $1 WHERE user_id = $2",
220 target,
221 source
222 )
223 .execute(&mut **tx)
224 .await?
225 .rows_affected();
226 moved += sqlx::query!(
227 "UPDATE analytics_events SET user_id = $1 WHERE user_id = $2",
228 target,
229 source
230 )
231 .execute(&mut **tx)
232 .await?
233 .rows_affected();
234 moved += sqlx::query!(
235 "UPDATE event_outbox SET user_id = $1 WHERE user_id = $2",
236 target,
237 source
238 )
239 .execute(&mut **tx)
240 .await?
241 .rows_affected();
242 moved += sqlx::query!(
243 "UPDATE files SET user_id = $1 WHERE user_id = $2",
244 target,
245 source
246 )
247 .execute(&mut **tx)
248 .await?
249 .rows_affected();
250 moved += sqlx::query!(
251 "UPDATE link_clicks SET user_id = $1 WHERE user_id = $2",
252 target,
253 source
254 )
255 .execute(&mut **tx)
256 .await?
257 .rows_affected();
258 Ok(moved)
259}