systemprompt_sync/database/
mod.rs

1//! Database push / pull: serialise the user and context tables to
2//! JSON and round-trip them between a local Postgres and a cloud Postgres
3//! using compile-time-checked `sqlx` upserts.
4//!
5//! Copyright (c) systemprompt.io — Business Source License 1.1.
6//! See <https://systemprompt.io> for licensing details.
7
8mod upsert;
9
10use chrono::{DateTime, Utc};
11use serde::{Deserialize, Serialize};
12use sqlx::PgPool;
13use sqlx::prelude::FromRow;
14use systemprompt_identifiers::{ContextId, SessionId, UserId};
15use systemprompt_models::ContextKind;
16
17use crate::error::SyncResult;
18use crate::{SyncDirection, SyncOperationResult};
19
20use upsert::{upsert_context, upsert_user};
21
22#[derive(Clone, Debug, Serialize, Deserialize)]
23pub struct DatabaseExport {
24    pub users: Vec<UserExport>,
25    pub contexts: Vec<ContextExport>,
26    pub timestamp: DateTime<Utc>,
27}
28
29#[derive(Clone, Debug, Serialize, Deserialize, FromRow)]
30pub struct UserExport {
31    pub id: UserId,
32    pub name: String,
33    pub email: String,
34    pub full_name: Option<String>,
35    pub display_name: Option<String>,
36    pub status: String,
37    pub email_verified: bool,
38    pub roles: Vec<String>,
39    pub is_bot: bool,
40    pub is_scanner: bool,
41    pub avatar_url: Option<String>,
42    pub created_at: DateTime<Utc>,
43    pub updated_at: DateTime<Utc>,
44}
45
46#[derive(Clone, Debug, Serialize, Deserialize, FromRow)]
47pub struct ContextExport {
48    pub context_id: ContextId,
49    pub user_id: UserId,
50    pub session_id: Option<SessionId>,
51    pub name: String,
52    pub kind: ContextKind,
53    pub created_at: DateTime<Utc>,
54    pub updated_at: DateTime<Utc>,
55}
56
57#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
58pub struct ImportResult {
59    pub created: usize,
60    pub updated: usize,
61    pub skipped: usize,
62}
63
64#[derive(Debug)]
65pub struct DatabaseSyncService {
66    direction: SyncDirection,
67    dry_run: bool,
68    local_database_url: String,
69    cloud_database_url: String,
70}
71
72impl DatabaseSyncService {
73    pub fn new(
74        direction: SyncDirection,
75        dry_run: bool,
76        local_database_url: &str,
77        cloud_database_url: &str,
78    ) -> Self {
79        Self {
80            direction,
81            dry_run,
82            local_database_url: local_database_url.to_owned(),
83            cloud_database_url: cloud_database_url.to_owned(),
84        }
85    }
86
87    pub async fn sync(&self) -> SyncResult<SyncOperationResult> {
88        match self.direction {
89            SyncDirection::Push => self.push().await,
90            SyncDirection::Pull => self.pull().await,
91        }
92    }
93
94    async fn push(&self) -> SyncResult<SyncOperationResult> {
95        let export = export_from_database(&self.local_database_url).await?;
96        let count = export.users.len() + export.contexts.len();
97
98        if self.dry_run {
99            return Ok(SyncOperationResult::dry_run(
100                "database_push",
101                count,
102                serde_json::json!({
103                    "users": export.users.len(),
104                    "contexts": export.contexts.len(),
105                }),
106            ));
107        }
108
109        import_to_database(&self.cloud_database_url, &export).await?;
110        Ok(SyncOperationResult::success("database_push", count))
111    }
112
113    async fn pull(&self) -> SyncResult<SyncOperationResult> {
114        let export = export_from_database(&self.cloud_database_url).await?;
115        let count = export.users.len() + export.contexts.len();
116
117        if self.dry_run {
118            return Ok(SyncOperationResult::dry_run(
119                "database_pull",
120                count,
121                serde_json::json!({
122                    "users": export.users.len(),
123                    "contexts": export.contexts.len(),
124                }),
125            ));
126        }
127
128        import_to_database(&self.local_database_url, &export).await?;
129        Ok(SyncOperationResult::success("database_pull", count))
130    }
131}
132
133async fn export_from_database(database_url: &str) -> SyncResult<DatabaseExport> {
134    let pool = PgPool::connect(database_url).await?;
135
136    let users = sqlx::query_as!(
137        UserExport,
138        r#"SELECT id, name, email, full_name, display_name, status, email_verified,
139                  roles, is_bot, is_scanner, avatar_url, created_at, updated_at
140           FROM users"#
141    )
142    .fetch_all(&pool)
143    .await?;
144
145    let contexts = sqlx::query_as!(
146        ContextExport,
147        r#"SELECT context_id as "context_id!: ContextId",
148                  user_id as "user_id!: UserId",
149                  session_id as "session_id: SessionId",
150                  name, kind as "kind: ContextKind", created_at, updated_at
151           FROM user_contexts"#
152    )
153    .fetch_all(&pool)
154    .await?;
155
156    Ok(DatabaseExport {
157        users,
158        contexts,
159        timestamp: Utc::now(),
160    })
161}
162
163async fn import_to_database(
164    database_url: &str,
165    export: &DatabaseExport,
166) -> SyncResult<ImportResult> {
167    let pool = PgPool::connect(database_url).await?;
168    let mut created = 0;
169    let mut updated = 0;
170    let mut completed = 0usize;
171    let total = export.users.len() + export.contexts.len();
172
173    for user in &export.users {
174        match upsert_user(&pool, user).await {
175            Ok((c, u)) => {
176                created += c;
177                updated += u;
178                completed += 1;
179            },
180            Err(err) => {
181                return Err(crate::error::SyncError::PartialImport {
182                    completed,
183                    total,
184                    message: format!("user upsert failed: {err}"),
185                });
186            },
187        }
188    }
189
190    for context in &export.contexts {
191        match upsert_context(&pool, context).await {
192            Ok((c, u)) => {
193                created += c;
194                updated += u;
195                completed += 1;
196            },
197            Err(err) => {
198                return Err(crate::error::SyncError::PartialImport {
199                    completed,
200                    total,
201                    message: format!("context upsert failed: {err}"),
202                });
203            },
204        }
205    }
206
207    Ok(ImportResult {
208        created,
209        updated,
210        skipped: 0,
211    })
212}
systemprompt_sync/database/mod.rs

systemprompt_sync/database/
mod.rs
