This commit is contained in:
Kelorastus
2026-09-22 13:37:29 +02:00
parent d5da361952
commit 4a92cc0b63
5 changed files with 612 additions and 21 deletions
+398 -5
View File
@@ -19,6 +19,8 @@ pub const PG_DATABASE: &str = "odian";
#[derive(Clone, Debug)]
pub struct DbFile {
pub section: String,
/// `schema.table` for a data dump, as the snapshot recorded it.
pub table: Option<String>,
/// Relative to `data\`, e.g. `db/20260915T082529Z/data/public.db_ground_poi.dump`.
pub path: String,
pub size: u64,
@@ -33,6 +35,8 @@ pub struct Pg {
data: PathBuf,
logs: PathBuf,
password: String,
/// Always PG_PORT for ODIAN; tests move it so a PC running ODIAN can still run them.
port: u16,
}
impl Pg {
@@ -42,9 +46,16 @@ impl Pg {
data: layout.pg_data(),
logs: layout.state_dir().join("pg"),
password: password.to_string(),
port: PG_PORT,
}
}
#[cfg(test)]
fn on_port(mut self, port: u16) -> Pg {
self.port = port;
self
}
fn tool(&self, name: &str) -> Command {
let mut cmd = command(&self.bin.join(format!("{name}.exe")), &self.logs, &[self.bin.clone()]);
cmd.env("PGPASSWORD", &self.password).env("PGCONNECT_TIMEOUT", "5").env("PGAPPNAME", "odian-launcher");
@@ -108,7 +119,7 @@ impl Pg {
return Ok(());
}
// Another install's Postgres on the shared port would leave only "exit code 1" in the log.
if crate::stack::port_taken(PG_PORT) {
if crate::stack::port_taken(self.port) {
return Err("Another ODIAN install is already running on this PC: it holds the database port. Stop it in the launcher window that started it, then try again."
.to_string());
}
@@ -120,7 +131,7 @@ impl Pg {
.arg("-l")
.arg(self.logs.join("postgres.log"))
.args(["-w", "-t", "120", "-o"])
.arg(format!("-p {PG_PORT} -c listen_addresses=127.0.0.1"));
.arg(format!("-p {} -c listen_addresses=127.0.0.1", self.port));
// The postmaster inherits pg_ctl's handles and outlives it, so a pipe would never
// reach EOF: its output goes to a file.
log_output(&mut cmd, &self.logs.join("pg_ctl.log"))?;
@@ -148,7 +159,7 @@ impl Pg {
}
fn psql(&self, database: &str, sql: &str) -> Result<String, String> {
let port = PG_PORT.to_string();
let port = self.port.to_string();
let mut cmd = self.tool("psql");
cmd
.args(["-X", "-A", "-t", "-v", "ON_ERROR_STOP=1", "-h", "127.0.0.1", "-p", port.as_str()])
@@ -160,7 +171,7 @@ impl Pg {
/// (unaligned, tuples only). For sync, whose JSON payloads are too large for a command line.
pub fn script(&self, sql: &str) -> Result<String, String> {
use std::io::Write;
let port = PG_PORT.to_string();
let port = self.port.to_string();
let mut cmd = self.tool("psql");
cmd
.env("PGCLIENTENCODING", "UTF8")
@@ -190,6 +201,7 @@ impl Pg {
&self,
data_root: &Path,
files: &[DbFile],
extensions: &[String],
workers: usize,
progress: RestoreProgress<'_>,
cancelled: &(dyn Fn() -> bool + Sync),
@@ -197,6 +209,7 @@ impl Pg {
self.ensure_logs()?;
self.psql("postgres", &format!("DROP DATABASE IF EXISTS {PG_DATABASE}"))?;
self.psql("postgres", &format!("CREATE DATABASE {PG_DATABASE} TEMPLATE template0 ENCODING 'UTF8'"))?;
self.create_extensions(extensions)?;
let total_bytes: u64 = files.iter().map(|f| f.size).sum();
let done_bytes = AtomicU64::new(0);
@@ -252,7 +265,7 @@ impl Pg {
fn restore_one(&self, data_root: &Path, file: &DbFile, pre_data: bool) -> Result<(), String> {
let dump = file.path.split('/').fold(data_root.to_path_buf(), |path, part| path.join(part));
let port = PG_PORT.to_string();
let port = self.port.to_string();
let mut cmd = self.tool("pg_restore");
cmd
.args(["-h", "127.0.0.1", "-p", port.as_str(), "-U", PG_USER, "-d", PG_DATABASE])
@@ -281,4 +294,384 @@ impl Pg {
std::fs::write(&path, kept).map_err(|e| format!("cannot write {}: {e}", path.display()))?;
Ok(path)
}
/// Create the snapshot's extensions (`postgis`). A PostGIS table needs its types before
/// its CREATE TABLE, and pg_dump -n never emits CREATE EXTENSION. The names come from
/// the signed manifest; checked anyway, since they are spliced into SQL.
fn create_extensions(&self, extensions: &[String]) -> Result<(), String> {
for name in extensions {
if !is_plain_name(name) {
return Err(format!("the snapshot names an extension this PC will not create: {name:?}"));
}
self.psql(PG_DATABASE, &format!("CREATE EXTENSION IF NOT EXISTS {name}"))?;
}
Ok(())
}
/// The snapshot's tables this database lacks, as `schema.table`.
///
/// A PC keeps its database when it updates -- that is where the operator's work lives --
/// so a table a later snapshot adds (0.59.0: the PostGIS overlays) would never reach it.
/// Its dump is on disk regardless: an update downloads the whole snapshot.
pub fn missing_tables(&self, files: &[DbFile]) -> Result<Vec<String>, String> {
let tables = candidate_tables(files);
if tables.is_empty() {
return Ok(Vec::new());
}
let list = tables.iter().map(|t| format!("'{t}'")).collect::<Vec<_>>().join(",");
let out = self.script(&format!(
"SELECT t FROM unnest(ARRAY[{list}]::text[]) AS t WHERE to_regclass(t) IS NULL ORDER BY t;"
))?;
Ok(out.lines().map(str::trim).filter(|l| !l.is_empty()).map(str::to_string).collect())
}
/// Drop the tables an add was in the middle of when it was cut short (power loss, the
/// launcher killed): they would read as present and never be filled. The marker names
/// them; add_tables writes it before it creates anything and removes it when done.
pub fn drop_interrupted_add(&self) -> Result<(), String> {
let marker = self.logs.join(ADDING_MARKER);
let Ok(text) = std::fs::read_to_string(&marker) else { return Ok(()) };
let tables: Vec<&str> = text.lines().map(str::trim).filter(|t| is_qualified_plain(t)).collect();
if !tables.is_empty() {
self.script(&format!("DROP TABLE IF EXISTS {} CASCADE;", tables.join(", ")))?;
}
std::fs::remove_file(&marker).map_err(|e| format!("cannot remove {}: {e}", marker.display()))
}
/// Add `tables` to this database from the snapshot on disk, touching nothing else: the
/// snapshot's extensions, then those tables' own schema statements, their rows, their
/// sequence positions, and last their indexes and constraints.
pub fn add_tables(
&self,
data_root: &Path,
files: &[DbFile],
extensions: &[String],
tables: &[String],
progress: &dyn Fn(String),
) -> Result<(), String> {
self.ensure_logs()?;
self.create_extensions(extensions)?;
let marker = self.logs.join(ADDING_MARKER);
std::fs::write(&marker, tables.join("\n")).map_err(|e| format!("cannot write {}: {e}", marker.display()))?;
let result = self.add_tables_once(data_root, files, tables, progress);
if result.is_err() {
// A half-added table would read as present from now on and never be filled: drop
// what this run created, so the next launch starts again from nothing. Should even
// this fail, the marker stays and the next launch drops them first.
let list = tables.join(", ");
if self.script(&format!("DROP TABLE IF EXISTS {list} CASCADE;")).is_err() {
return result;
}
}
let _ = std::fs::remove_file(&marker);
result
}
fn add_tables_once(&self, data_root: &Path, files: &[DbFile], tables: &[String], progress: &dyn Fn(String)) -> Result<(), String> {
let section = |name: &str| files.iter().find(|f| f.section == name);
let pre = section("pre-data").ok_or("the snapshot has no schema dump")?;
let (preamble, blocks) = sql_blocks(&self.render(data_root, pre)?);
let schema = schema_blocks(&blocks, tables);
let sequences: Vec<String> = schema.iter().filter(|b| b.kind == "SEQUENCE").map(|b| b.qualified()).collect();
progress(format!("creating {} table(s)", tables.len()));
self.script(&join_blocks(&preamble, &schema))?;
for file in files.iter().filter(|f| f.section == "data") {
if let Some(table) = table_of(file).filter(|t| tables.contains(t)) {
progress(format!("loading {table}"));
self.restore_one(data_root, file, false)?;
}
}
if let Some(file) = section("sequences") {
let (preamble, blocks) = sql_blocks(&self.render(data_root, file)?);
let positions: Vec<&SqlBlock> = blocks.iter().filter(|b| sequences.contains(&b.qualified())).collect();
if !positions.is_empty() {
self.script(&join_blocks(&preamble, &positions))?;
}
}
if let Some(file) = section("post-data") {
progress("building their indexes".to_string());
let (preamble, blocks) = sql_blocks(&self.render(data_root, file)?);
let indexes = index_blocks(&blocks, tables);
if !indexes.is_empty() {
self.script(&format!("SET maintenance_work_mem = '512MB';\n{}", join_blocks(&preamble, &indexes)))?;
}
}
Ok(())
}
/// A dump rendered as the SQL pg_restore would run.
fn render(&self, data_root: &Path, file: &DbFile) -> Result<String, String> {
let dump = file.path.split('/').fold(data_root.to_path_buf(), |path, part| path.join(part));
let mut cmd = self.tool("pg_restore");
cmd.args(["--no-owner", "--no-privileges", "-f", "-"]).arg(&dump);
Self::run(cmd, &format!("pg_restore {}", file.path))
}
}
/// The file naming the tables an add has started on; see `drop_interrupted_add`.
const ADDING_MARKER: &str = "adding-tables.txt";
/// Tables that describe the database they live in, not data: never copied into another
/// one. `odian_schema_migrations` lists the migrations applied HERE -- the live database's
/// copy would make this PC's Engine API skip migrations its schema never had.
const OWN_DATABASE_ONLY: [&str; 1] = ["public.odian_schema_migrations"];
/// The snapshot's data tables an existing database may lack and be given.
fn candidate_tables(files: &[DbFile]) -> Vec<String> {
files
.iter()
.filter(|f| f.section == "data")
.filter_map(table_of)
.filter(|t| !OWN_DATABASE_ONLY.contains(&t.as_str()))
.collect()
}
/// `schema.table`, both plain names.
fn is_qualified_plain(name: &str) -> bool {
name.split_once('.').is_some_and(|(schema, table)| is_plain_name(schema) && is_plain_name(table))
}
/// Lower-case letters, digits and `_`: safe to splice into SQL unquoted.
fn is_plain_name(name: &str) -> bool {
!name.is_empty() && name.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
}
/// A data dump's table as `schema.table`, when it is a plain name.
fn table_of(file: &DbFile) -> Option<String> {
let name = file.table.clone().or_else(|| {
let stem = file.path.rsplit('/').next()?.strip_suffix(".dump")?;
Some(stem.to_string())
})?;
let (schema, table) = name.split_once('.')?;
(is_plain_name(schema) && is_plain_name(table)).then_some(name)
}
/// One object from pg_restore's SQL: the `-- Name: x; Type: T; Schema: s` header pg_dump
/// writes before it, and the statements under it.
#[derive(Debug)]
struct SqlBlock {
name: String,
kind: String,
schema: String,
sql: String,
}
impl SqlBlock {
fn qualified(&self) -> String {
format!("{}.{}", self.schema, self.name)
}
}
/// Split pg_restore's SQL into its session settings and one block per object. Lines that
/// are psql meta-commands (`\restrict` and its pair) guard a whole dump; a partial replay
/// must leave them out or psql refuses to run it.
fn sql_blocks(text: &str) -> (String, Vec<SqlBlock>) {
let mut preamble = String::new();
let mut blocks: Vec<SqlBlock> = Vec::new();
for line in text.lines() {
if line.starts_with('\\') {
continue;
}
if let Some(header) = line.strip_prefix("-- Name: ") {
let field = |key: &str| header.split("; ").find_map(|p| p.strip_prefix(key)).unwrap_or_default().to_string();
let name = header.split("; ").next().unwrap_or_default().to_string();
blocks.push(SqlBlock { name, kind: field("Type: "), schema: field("Schema: "), sql: String::new() });
} else if let Some(block) = blocks.last_mut() {
block.sql.push_str(line);
block.sql.push('\n');
} else if !line.starts_with("--") {
preamble.push_str(line);
preamble.push('\n');
}
}
(preamble, blocks)
}
/// Does `sql` name `qualified` (`schema.table`) as a whole identifier, not as a prefix?
fn references(sql: &str, qualified: &str) -> bool {
let ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
sql.match_indices(qualified).any(|(at, found)| {
let before = sql[..at].chars().next_back();
let after = sql[at + found.len()..].chars().next();
!before.is_some_and(|c| ident(c) || c == '.' || c == '"') && !after.is_some_and(ident)
})
}
/// The schema statements that create `tables`: the tables themselves, their identity or
/// serial sequences, defaults and comments. Functions, views and the like are left alone:
/// they may already exist here, and CREATE without OR REPLACE would stop the whole script.
fn schema_blocks<'a>(blocks: &'a [SqlBlock], tables: &[String]) -> Vec<&'a SqlBlock> {
const TABLE_SCOPED: [&str; 5] = ["TABLE", "SEQUENCE", "SEQUENCE OWNED BY", "DEFAULT", "COMMENT"];
let mut keep: Vec<bool> = blocks
.iter()
.map(|b| TABLE_SCOPED.contains(&b.kind.as_str()) && tables.iter().any(|t| references(&b.sql, t)))
.collect();
// A serial column's CREATE SEQUENCE names no table; its OWNED BY and DEFAULT do.
let kept_sql: Vec<&str> = blocks.iter().zip(&keep).filter(|(_, k)| **k).map(|(b, _)| b.sql.as_str()).collect();
for (i, block) in blocks.iter().enumerate() {
if !keep[i] && block.kind == "SEQUENCE" && kept_sql.iter().any(|sql| references(sql, &block.qualified())) {
keep[i] = true;
}
}
blocks.iter().zip(keep).filter(|(_, k)| *k).map(|(b, _)| b).collect()
}
/// The indexes, constraints and triggers of `tables` -- never the sync marking triggers,
/// which call sync.mark() from a schema the launcher loads itself.
fn index_blocks<'a>(blocks: &'a [SqlBlock], tables: &[String]) -> Vec<&'a SqlBlock> {
const TABLE_SCOPED: [&str; 5] = ["INDEX", "CONSTRAINT", "FK CONSTRAINT", "TRIGGER", "INDEX ATTACH"];
blocks
.iter()
.filter(|b| TABLE_SCOPED.contains(&b.kind.as_str()))
.filter(|b| !(b.kind == "TRIGGER" && b.name.starts_with("sync_mark_")))
.filter(|b| tables.iter().any(|t| references(&b.sql, t)))
.collect()
}
fn join_blocks(preamble: &str, blocks: &[&SqlBlock]) -> String {
blocks.iter().fold(preamble.to_string(), |mut sql, block| {
sql.push_str(&block.sql);
sql
})
}
#[cfg(test)]
mod tests {
use super::*;
/// Abridged pg_restore output, shaped exactly like the 0.59.0 snapshot's.
const PRE: &str = "--\n-- PostgreSQL database dump\n--\n\n\\restrict KEY\n\nSET client_encoding = 'UTF8';\n\
SELECT pg_catalog.set_config('search_path', '', false);\n\n\
--\n-- Name: odb_note; Type: TABLE; Schema: public; Owner: -\n--\n\nCREATE TABLE public.odb_note (\n id integer\n);\n\n\
--\n-- Name: postgis_hospital; Type: TABLE; Schema: public; Owner: -\n--\n\nCREATE TABLE public.postgis_hospital (\n id integer NOT NULL,\n geom public.geometry\n);\n\n\
--\n-- Name: postgis_hospital_id_seq; Type: SEQUENCE; Schema: public; Owner: -\n--\n\n\
ALTER TABLE public.postgis_hospital ALTER COLUMN id ADD GENERATED ALWAYS AS IDENTITY (\n SEQUENCE NAME public.postgis_hospital_id_seq\n);\n\n\
--\n-- Name: postgis_hospital_x; Type: TABLE; Schema: public; Owner: -\n--\n\nCREATE TABLE public.postgis_hospital_x (\n id integer\n);\n\n\
--\n-- Name: hospital_count; Type: FUNCTION; Schema: public; Owner: -\n--\n\n\
CREATE FUNCTION public.hospital_count() RETURNS bigint AS $$ SELECT count(*) FROM public.postgis_hospital $$ LANGUAGE sql;\n\n\
\\unrestrict KEY\n";
#[test]
fn blocks_keep_the_session_settings_and_drop_psql_guards() {
let (preamble, blocks) = sql_blocks(PRE);
assert!(preamble.contains("set_config('search_path', '', false)"));
assert!(!preamble.contains("restrict") && blocks.iter().all(|b| !b.sql.contains("restrict")));
assert_eq!(blocks.len(), 5);
assert_eq!((blocks[1].name.as_str(), blocks[1].kind.as_str(), blocks[1].qualified()),
("postgis_hospital", "TABLE", "public.postgis_hospital".to_string()));
}
#[test]
fn a_table_is_referenced_whole_never_as_a_prefix() {
assert!(references("CREATE INDEX i ON public.postgis_hospital USING gist (geom);", "public.postgis_hospital"));
assert!(references("ALTER SEQUENCE s OWNED BY public.postgis_hospital.id;", "public.postgis_hospital"));
assert!(!references("CREATE TABLE public.postgis_hospital_x (", "public.postgis_hospital"));
assert!(!references("CREATE TABLE xpublic.postgis_hospital (", "public.postgis_hospital"));
}
#[test]
fn only_the_missing_tables_own_statements_are_replayed() {
let (_, blocks) = sql_blocks(PRE);
let chosen: Vec<String> = schema_blocks(&blocks, &["public.postgis_hospital".to_string()])
.iter()
.map(|b| b.qualified())
.collect();
// Its table and identity sequence; not a look-alike table, not an existing one, and not a
// function that merely queries it (CREATE FUNCTION would fail on a PC that has it).
assert_eq!(chosen, ["public.postgis_hospital", "public.postgis_hospital_id_seq"]);
}
#[test]
fn indexes_follow_their_table_but_sync_triggers_never_do() {
let post = "--\n-- Name: idx_postgis_hospital_geom; Type: INDEX; Schema: public; Owner: -\n--\n\n\
CREATE INDEX idx_postgis_hospital_geom ON public.postgis_hospital USING gist (geom);\n\n\
--\n-- Name: idx_note; Type: INDEX; Schema: public; Owner: -\n--\n\nCREATE INDEX idx_note ON public.odb_note USING btree (id);\n\n\
--\n-- Name: sync_mark_ins; Type: TRIGGER; Schema: public; Owner: -\n--\n\n\
CREATE TRIGGER sync_mark_ins AFTER INSERT ON public.postgis_hospital FOR EACH STATEMENT EXECUTE FUNCTION sync.mark();\n";
let (_, blocks) = sql_blocks(post);
let chosen: Vec<&str> = index_blocks(&blocks, &["public.postgis_hospital".to_string()])
.iter()
.map(|b| b.name.as_str())
.collect();
assert_eq!(chosen, ["idx_postgis_hospital_geom"]);
}
/// The real thing on real snapshots: a database restored from the snapshot before PostGIS
/// (as every PC that installed then has it), then given the missing tables of a later one.
/// ODIAN_ADD_TABLES_ROOT: a layout with versions/test/runtime/pgsql (carrying PostGIS) and
/// both snapshots under data/db/. ODIAN_ADD_TABLES_OLD / _NEW: their snapshot.json files.
#[test]
#[ignore = "needs real snapshots and a PostGIS runtime: see the doc comment"]
fn adds_the_new_tables_to_a_database_restored_before_them() {
let env = |key: &str| std::env::var(key).unwrap_or_else(|_| panic!("set {key}"));
let root = PathBuf::from(env("ODIAN_ADD_TABLES_ROOT"));
let snapshot_files = |json: &str| -> Vec<DbFile> {
let doc: serde_json::Value = serde_json::from_slice(&std::fs::read(json).unwrap()).unwrap();
doc["files"]
.as_array()
.unwrap()
.iter()
.map(|f| DbFile {
section: f["section"].as_str().unwrap().into(),
table: f["table"].as_str().map(str::to_string),
path: f["path"].as_str().unwrap().into(),
size: f["size"].as_u64().unwrap(),
})
.collect()
};
let old = snapshot_files(&env("ODIAN_ADD_TABLES_OLD"));
let new = snapshot_files(&env("ODIAN_ADD_TABLES_NEW"));
let data_root = root.join("data");
let pg = Pg::new(&Layout::new(&root, "test"), "pw").on_port(55466);
if !pg.is_initialized() {
pg.init().unwrap();
}
pg.start().unwrap();
let never = || false;
let quiet = |_: u64, _: u64, _: usize, _: usize, _: &str| {};
pg.restore(&data_root, &old, &[], 4, &quiet, &never).unwrap();
let pois_before = pg.script("SELECT count(*) FROM public.db_ground_poi;").unwrap();
let missing = pg.missing_tables(&new).unwrap();
eprintln!("missing before: {missing:?}");
let added = pg.add_tables(&data_root, &new, &["postgis".to_string()], &missing, &|step| eprintln!(" {step}"));
let hospitals = pg.script("SELECT count(*) FROM public.postgis_hospital;");
let spatial = pg.script(
"SELECT count(*) FROM pg_indexes WHERE tablename = 'postgis_hospital' AND indexdef LIKE '%USING gist%';",
);
let pois_after = pg.script("SELECT count(*) FROM public.db_ground_poi;");
let still_missing = pg.missing_tables(&new);
let _ = pg.stop();
added.unwrap();
assert!(missing.contains(&"public.postgis_hospital".to_string()), "{missing:?}");
assert!(hospitals.unwrap().trim().parse::<u64>().unwrap() > 100_000, "the rows arrived");
assert_eq!(spatial.unwrap().trim(), "1", "with their spatial index");
assert_eq!(pois_before, pois_after.unwrap(), "an existing table is not touched");
assert!(still_missing.unwrap().is_empty(), "nothing is left to add");
}
#[test]
fn the_migrations_ledger_never_travels_between_databases() {
let file = |table: &str| DbFile { section: "data".into(), table: Some(table.into()), path: "x".into(), size: 0 };
let files = [file("public.postgis_oil"), file("public.odian_schema_migrations"), file("public.osm_roads_import_runs")];
assert_eq!(candidate_tables(&files), ["public.postgis_oil", "public.osm_roads_import_runs"]);
}
#[test]
fn a_dump_names_its_table_even_without_the_manifest_field() {
let file = |table: Option<&str>, path: &str| DbFile {
section: "data".into(),
table: table.map(str::to_string),
path: path.into(),
size: 0,
};
assert_eq!(table_of(&file(Some("public.postgis_oil"), "x")), Some("public.postgis_oil".into()));
assert_eq!(table_of(&file(None, "db/s/data/public.postgis_gas.dump")), Some("public.postgis_gas".into()));
assert_eq!(table_of(&file(None, "db/s/data/public.Weird Name.dump")), None, "never spliced into SQL");
}
}