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

systemprompt_sync/database/
mod.rs
