use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use sqlx::{PgPool, Row, postgres::PgRow};
use std::io::{BufRead, Write};
const TABLES: [&str; 11] = [
"users",
"departments",
"agents",
"sessions",
"workdays",
"pauses",
"tasks",
"tags",
"task_tags",
"reports",
"audit_log",
];
const DEFERRED_COLUMNS: [(&str, &str); 2] = [("users", "department_id"), ("departments", "manager_id")];
const SETTINGS: &str = "settings";
#[derive(Debug, Serialize, Deserialize)]
pub struct Header {
pub format: String,
pub schema_version: i64,
pub server_version: String,
pub taken_at: chrono::DateTime<chrono::Utc>,
}
const FORMAT: &str = "kasl-server-backup";
#[derive(Debug, Serialize, Deserialize)]
struct Chunk {
table: String,
rows: Vec<serde_json::Value>,
}
pub async fn dump(pool: &PgPool, schema_version: i64, out: &mut impl Write) -> Result<Summary> {
let header = Header {
format: FORMAT.to_string(),
schema_version,
server_version: env!("CARGO_PKG_VERSION").to_string(),
taken_at: chrono::Utc::now(),
};
writeln!(out, "{}", serde_json::to_string(&header)?)?;
let mut summary = Summary::default();
for table in TABLES.into_iter().chain([SETTINGS]) {
let rows: Vec<serde_json::Value> = sqlx::query(sqlx::AssertSqlSafe(format!("SELECT row_to_json(t) AS row FROM {table} AS t")))
.fetch_all(pool)
.await
.with_context(|| format!("failed to read {table}"))?
.into_iter()
.map(|row: PgRow| row.get::<serde_json::Value, _>("row"))
.collect();
summary.rows += rows.len();
summary.tables += 1;
writeln!(
out,
"{}",
serde_json::to_string(&Chunk {
table: table.to_string(),
rows,
})?
)?;
}
out.flush()?;
Ok(summary)
}
#[derive(Debug, Default, PartialEq, Eq)]
pub struct Summary {
pub tables: usize,
pub rows: usize,
}
pub async fn load(pool: &PgPool, schema_version: i64, input: impl BufRead) -> Result<Summary> {
let mut lines = input.lines();
let header: Header = {
let first = lines.next().context("the backup is empty")??;
serde_json::from_str(&first).context("the first line is not a kasl-server backup header")?
};
if header.format != FORMAT {
bail!("this is not a kasl-server backup (format: {})", header.format);
}
if header.schema_version > schema_version {
bail!(
"the backup is from a newer server (schema {} against this server's {}); upgrade kasl-server before restoring",
header.schema_version,
schema_version
);
}
ensure_empty(pool).await?;
let mut summary = Summary::default();
let mut tx = pool.begin().await?;
let mut deferred: Vec<(&'static str, &'static str, serde_json::Value, serde_json::Value)> = Vec::new();
for line in lines {
let line = line?;
if line.trim().is_empty() {
continue;
}
let chunk: Chunk = serde_json::from_str(&line).context("a line of the backup could not be read")?;
summary.tables += 1;
for row in &chunk.rows {
let held = insert(&mut tx, &chunk.table, row).await?;
deferred.extend(held);
summary.rows += 1;
}
}
for (table, column, id, value) in deferred {
sqlx::query(sqlx::AssertSqlSafe(format!(
"UPDATE {table} SET {column} = ($1::text)::uuid WHERE id = ($2::text)::uuid"
)))
.bind(value.as_str().unwrap_or_default())
.bind(id.as_str().unwrap_or_default())
.execute(&mut *tx)
.await
.with_context(|| format!("failed to restore {table}.{column}"))?;
}
tx.commit().await?;
Ok(summary)
}
fn known_table(name: &str) -> Result<&'static str> {
TABLES
.into_iter()
.chain([SETTINGS])
.find(|table| *table == name)
.with_context(|| format!("the backup names a table this server does not have: {name}"))
}
type Deferred = (&'static str, &'static str, serde_json::Value, serde_json::Value);
async fn insert(tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, table: &str, row: &serde_json::Value) -> Result<Vec<Deferred>> {
let table = known_table(table)?;
let mut held = Vec::new();
let mut row = row.clone();
for (owner, column) in DEFERRED_COLUMNS {
if owner != table {
continue;
}
let Some(object) = row.as_object_mut() else { continue };
match object.get(column) {
Some(serde_json::Value::Null) | None => {}
Some(value) => {
let value = value.clone();
let id = object.get("id").cloned().unwrap_or(serde_json::Value::Null);
object.insert(column.to_string(), serde_json::Value::Null);
held.push((owner, column, id, value));
}
}
}
let row = &row;
if table == SETTINGS {
sqlx::query(sqlx::AssertSqlSafe(format!(
"UPDATE {SETTINGS} SET privacy_level = (r).privacy_level, updated_at = (r).updated_at
FROM (SELECT json_populate_record(NULL::{SETTINGS}, $1::json) AS r) AS s
WHERE singleton"
)))
.bind(row)
.execute(&mut **tx)
.await
.with_context(|| format!("failed to restore {table}"))?;
return Ok(held);
}
sqlx::query(sqlx::AssertSqlSafe(format!(
"INSERT INTO {table} SELECT * FROM json_populate_record(NULL::{table}, $1::json)"
)))
.bind(row)
.execute(&mut **tx)
.await
.with_context(|| format!("failed to restore a row of {table}"))?;
Ok(held)
}
async fn ensure_empty(pool: &PgPool) -> Result<()> {
let users: i64 = sqlx::query_scalar("SELECT count(*) FROM users").fetch_one(pool).await?;
if users > 0 {
bail!("this database already holds {users} accounts; restore into an empty one");
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_backup_names_itself_and_its_schema() {
let header = Header {
format: FORMAT.to_string(),
schema_version: 20260826000001,
server_version: "0.14.0".to_string(),
taken_at: chrono::Utc::now(),
};
let json = serde_json::to_string(&header).unwrap();
let read: Header = serde_json::from_str(&json).unwrap();
assert_eq!(read.format, FORMAT);
assert_eq!(read.schema_version, 20260826000001);
}
#[test]
fn parents_come_before_their_children() {
let position = |table: &str| TABLES.iter().position(|t| *t == table).unwrap_or_else(|| panic!("{table} is not in TABLES"));
for (child, parent) in [
("agents", "users"),
("sessions", "users"),
("workdays", "users"),
("pauses", "workdays"),
("tasks", "users"),
("tags", "users"),
("task_tags", "tasks"),
("task_tags", "tags"),
("reports", "users"),
("departments", "users"),
] {
assert!(position(parent) < position(child), "{parent} must be restored before {child}");
}
}
#[test]
fn a_table_name_from_a_file_has_to_be_one_of_ours() {
assert_eq!(known_table("users").unwrap(), "users");
assert_eq!(known_table(SETTINGS).unwrap(), SETTINGS);
for hostile in ["users; DROP TABLE users", "pg_shadow", "_sqlx_migrations", "users --", ""] {
let error = known_table(hostile).unwrap_err();
assert!(error.to_string().contains("does not have"), "`{hostile}` must be refused by name, got: {error}");
}
}
#[test]
fn every_table_is_listed_once() {
let mut seen = TABLES.to_vec();
seen.sort_unstable();
let count = seen.len();
seen.dedup();
assert_eq!(seen.len(), count, "a table listed twice would be restored twice");
assert!(!TABLES.contains(&SETTINGS), "settings is updated, not inserted, so it must not be in TABLES");
assert!(
!TABLES.contains(&"_sqlx_migrations"),
"the schema is applied by the server, not carried in a backup"
);
}
}