Kelorastus
721 lines
31 KiB
Rust
721 lines
31 KiB
Rust
//! The PC's own Postgres cluster (`data\pg`): create it, start and stop it, load a snapshot.
|
|
//!
|
|
//! It listens on 127.0.0.1 only, with a password generated on this PC. The
|
|
//! snapshot is restored in the order proven on the Data PC's dumps: pre-data,
|
|
//! every table's data (in parallel), sequences, post-data.
|
|
|
|
use std::collections::{BTreeSet, VecDeque};
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::{Command, Stdio};
|
|
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
|
use std::sync::Mutex;
|
|
|
|
use crate::stack::{command, log_output, Layout, PG_PORT};
|
|
|
|
pub const PG_USER: &str = "odian";
|
|
pub const PG_DATABASE: &str = "odian";
|
|
|
|
/// One dump file from the manifest's `db.files`.
|
|
#[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,
|
|
}
|
|
|
|
/// Restore progress: (done bytes, total bytes, done files, total files, current file).
|
|
pub type RestoreProgress<'a> = &'a (dyn Fn(u64, u64, usize, usize, &str) + Sync);
|
|
|
|
#[derive(Clone)]
|
|
pub struct Pg {
|
|
bin: PathBuf,
|
|
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 {
|
|
pub fn new(layout: &Layout, password: &str) -> Pg {
|
|
Pg {
|
|
bin: layout.pg_bin(),
|
|
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");
|
|
cmd
|
|
}
|
|
|
|
fn run(mut cmd: Command, what: &str) -> Result<String, String> {
|
|
let out = cmd
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped())
|
|
.output()
|
|
.map_err(|e| format!("cannot run {what}: {e}"))?;
|
|
if out.status.success() {
|
|
return Ok(String::from_utf8_lossy(&out.stdout).into_owned());
|
|
}
|
|
let stderr = String::from_utf8_lossy(&out.stderr);
|
|
let tail: String = stderr.chars().rev().take(1500).collect::<Vec<_>>().into_iter().rev().collect();
|
|
Err(format!("{what} failed ({}): {}", out.status, tail.trim()))
|
|
}
|
|
|
|
fn ensure_logs(&self) -> Result<(), String> {
|
|
std::fs::create_dir_all(&self.logs).map_err(|e| format!("cannot create {}: {e}", self.logs.display()))
|
|
}
|
|
|
|
pub fn is_initialized(&self) -> bool {
|
|
self.data.join("PG_VERSION").is_file()
|
|
}
|
|
|
|
/// Create the cluster. A `data\pg` without PG_VERSION is a failed earlier attempt and is replaced.
|
|
pub fn init(&self) -> Result<(), String> {
|
|
self.ensure_logs()?;
|
|
if self.data.exists() && !self.is_initialized() {
|
|
std::fs::remove_dir_all(&self.data).map_err(|e| format!("cannot clear {}: {e}", self.data.display()))?;
|
|
}
|
|
let pwfile = self.logs.join("initdb.pw");
|
|
std::fs::write(&pwfile, &self.password).map_err(|e| format!("cannot write {}: {e}", pwfile.display()))?;
|
|
let mut cmd = self.tool("initdb");
|
|
cmd
|
|
.arg("-D")
|
|
.arg(&self.data)
|
|
.args(["-U", PG_USER, "-A", "scram-sha-256", "-E", "UTF8", "--locale=C"])
|
|
.args(["--locale-provider=builtin", "--builtin-locale=C.UTF-8"])
|
|
.arg(format!("--pwfile={}", pwfile.display()));
|
|
let result = Self::run(cmd, "initdb");
|
|
let _ = std::fs::remove_file(&pwfile);
|
|
result.map(|_| ())
|
|
}
|
|
|
|
pub fn running(&self) -> bool {
|
|
let mut cmd = self.tool("pg_ctl");
|
|
cmd.arg("status").arg("-D").arg(&self.data).stdout(Stdio::null()).stderr(Stdio::null());
|
|
cmd.status().map(|status| status.success()).unwrap_or(false)
|
|
}
|
|
|
|
pub fn start(&self) -> Result<(), String> {
|
|
if !self.is_initialized() {
|
|
return Err("The local database isn't set up. Reinstall from the DHP or a drive.".to_string());
|
|
}
|
|
self.ensure_logs()?;
|
|
if self.running() {
|
|
return Ok(());
|
|
}
|
|
// Another install's Postgres on the shared port would leave only "exit code 1" in the log.
|
|
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());
|
|
}
|
|
let mut cmd = self.tool("pg_ctl");
|
|
cmd
|
|
.arg("start")
|
|
.arg("-D")
|
|
.arg(&self.data)
|
|
.arg("-l")
|
|
.arg(self.logs.join("postgres.log"))
|
|
.args(["-w", "-t", "120", "-o"])
|
|
.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"))?;
|
|
let status = cmd.status().map_err(|e| format!("cannot run pg_ctl: {e}"))?;
|
|
if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!("Postgres did not start ({status}). See state\\pg\\postgres.log."))
|
|
}
|
|
}
|
|
|
|
pub fn stop(&self) -> Result<(), String> {
|
|
if !self.is_initialized() || !self.running() {
|
|
return Ok(());
|
|
}
|
|
let mut cmd = self.tool("pg_ctl");
|
|
cmd.arg("stop").arg("-D").arg(&self.data).args(["-m", "fast", "-w", "-t", "60"]);
|
|
log_output(&mut cmd, &self.logs.join("pg_ctl.log"))?;
|
|
let status = cmd.status().map_err(|e| format!("cannot run pg_ctl: {e}"))?;
|
|
if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!("Postgres did not stop cleanly ({status})."))
|
|
}
|
|
}
|
|
|
|
fn psql(&self, database: &str, sql: &str) -> Result<String, 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()])
|
|
.args(["-U", PG_USER, "-d", database, "-c", sql]);
|
|
Self::run(cmd, "psql")
|
|
}
|
|
|
|
/// Run a script on the `odian` database through psql's stdin and return what it printed
|
|
/// (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 = self.port.to_string();
|
|
let mut cmd = self.tool("psql");
|
|
cmd
|
|
.env("PGCLIENTENCODING", "UTF8")
|
|
.args(["-X", "-q", "-A", "-t", "-v", "ON_ERROR_STOP=1", "-h", "127.0.0.1", "-p", port.as_str()])
|
|
.args(["-U", PG_USER, "-d", PG_DATABASE, "-f", "-"])
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
let mut child = cmd.spawn().map_err(|e| format!("cannot run psql: {e}"))?;
|
|
let mut stdin = child.stdin.take().ok_or("cannot write to psql")?;
|
|
// Feed from a thread: psql filling its stdout pipe while we still write would deadlock.
|
|
let body = sql.as_bytes().to_vec();
|
|
let writer = std::thread::spawn(move || stdin.write_all(&body));
|
|
let out = child.wait_with_output().map_err(|e| format!("psql failed: {e}"))?;
|
|
let _ = writer.join();
|
|
if out.status.success() {
|
|
return Ok(String::from_utf8_lossy(&out.stdout).into_owned());
|
|
}
|
|
let stderr = String::from_utf8_lossy(&out.stderr);
|
|
let tail: String = stderr.chars().rev().take(1500).collect::<Vec<_>>().into_iter().rev().collect();
|
|
Err(format!("psql failed ({}): {}", out.status, tail.trim()))
|
|
}
|
|
|
|
/// Replace the `odian` database with a snapshot. Only for a PC whose restore never
|
|
/// completed: anything an operator did in the old database is dropped.
|
|
pub fn restore(
|
|
&self,
|
|
data_root: &Path,
|
|
files: &[DbFile],
|
|
extensions: &[String],
|
|
workers: usize,
|
|
progress: RestoreProgress<'_>,
|
|
cancelled: &(dyn Fn() -> bool + Sync),
|
|
) -> Result<(), String> {
|
|
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);
|
|
let done_files = AtomicUsize::new(0);
|
|
let finished = |file: &DbFile| {
|
|
let bytes = done_bytes.fetch_add(file.size, Ordering::Relaxed) + file.size;
|
|
let count = done_files.fetch_add(1, Ordering::Relaxed) + 1;
|
|
progress(bytes, total_bytes, count, files.len(), &file.path);
|
|
};
|
|
let section = |name: &'static str| files.iter().filter(move |f| f.section == name);
|
|
|
|
for file in section("pre-data") {
|
|
self.restore_one(data_root, file, true)?;
|
|
finished(file);
|
|
}
|
|
|
|
let queue = Mutex::new(section("data").collect::<VecDeque<_>>());
|
|
let failure: Mutex<Option<String>> = Mutex::new(None);
|
|
std::thread::scope(|scope| {
|
|
for _ in 0..workers.max(1) {
|
|
scope.spawn(|| loop {
|
|
if cancelled() || failure.lock().map(|f| f.is_some()).unwrap_or(true) {
|
|
return;
|
|
}
|
|
let Some(file) = queue.lock().ok().and_then(|mut q| q.pop_front()) else {
|
|
return;
|
|
};
|
|
match self.restore_one(data_root, file, false) {
|
|
Ok(()) => finished(file),
|
|
Err(message) => {
|
|
if let Ok(mut slot) = failure.lock() {
|
|
slot.get_or_insert(message);
|
|
}
|
|
return;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
});
|
|
if let Some(message) = failure.into_inner().unwrap_or(None) {
|
|
return Err(message);
|
|
}
|
|
if cancelled() {
|
|
return Err("cancelled".to_string());
|
|
}
|
|
|
|
for file in section("sequences").chain(section("post-data")) {
|
|
self.restore_one(data_root, file, false)?;
|
|
finished(file);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
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 = 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])
|
|
.args(["--no-owner", "--no-privileges", "--exit-on-error"]);
|
|
if pre_data {
|
|
// A fresh database already has schema public; --exit-on-error would stop on the duplicate.
|
|
let drop = |line: &str| line.contains(" SCHEMA - public ") || line.contains(" COMMENT - SCHEMA public ");
|
|
cmd.arg("-L").arg(self.toc_without(&dump, "pre-data.list", &drop)?);
|
|
} else if file.section == "post-data" {
|
|
// A DHP running sync has marking triggers on its tables; they call sync.mark(), which
|
|
// this database gets only when the launcher loads the version's sync schema.
|
|
let drop = |line: &str| line.contains(" TRIGGER ") && line.contains(" sync_mark_");
|
|
cmd.arg("-L").arg(self.toc_without(&dump, "post-data.list", &drop)?);
|
|
}
|
|
cmd.arg(&dump);
|
|
Self::run(cmd, &format!("pg_restore {}", file.path)).map(|_| ())
|
|
}
|
|
|
|
/// The dump's table of contents without the entries `drop` names, for `pg_restore -L`.
|
|
fn toc_without(&self, dump: &Path, list_name: &str, drop: &dyn Fn(&str) -> bool) -> Result<PathBuf, String> {
|
|
let mut cmd = self.tool("pg_restore");
|
|
cmd.arg("-l").arg(dump);
|
|
let listing = Self::run(cmd, "pg_restore -l")?;
|
|
let kept: String = listing.lines().filter(|line| !drop(line)).map(|line| format!("{line}\n")).collect();
|
|
let path = self.logs.join(list_name);
|
|
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 files of a snapshot a database lacking only `missing` needs: the schema and sequence
|
|
/// dumps (small, and adding tables reads them) and the data of the missing tables. The rest
|
|
/// is data the database already holds. A data dump with no plain table name is kept, the
|
|
/// safe side of not knowing.
|
|
pub fn files_needed(files: &[DbFile], missing: &[String]) -> BTreeSet<String> {
|
|
files
|
|
.iter()
|
|
.filter(|f| f.section != "data" || table_of(f).is_none_or(|table| missing.contains(&table)))
|
|
.map(|f| f.path.clone())
|
|
.collect()
|
|
}
|
|
|
|
/// 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_database_keeping_its_data_needs_only_the_schema_and_the_tables_it_lacks() {
|
|
let file = |section: &str, table: Option<&str>, path: &str| DbFile {
|
|
section: section.to_string(),
|
|
table: table.map(str::to_string),
|
|
path: path.to_string(),
|
|
size: 1,
|
|
};
|
|
let files = [
|
|
file("pre-data", None, "db/S/pre-data.dump"),
|
|
file("data", Some("public.odb_note"), "db/S/data/public.odb_note.dump"),
|
|
file("data", Some("public.postgis_hospital"), "db/S/data/public.postgis_hospital.dump"),
|
|
file("data", None, "db/S/data/Weird Name.dump"),
|
|
file("data", Some("public.odian_schema_migrations"), "db/S/data/public.odian_schema_migrations.dump"),
|
|
file("sequences", None, "db/S/sequences.dump"),
|
|
file("post-data", None, "db/S/post-data.dump"),
|
|
];
|
|
let needed = files_needed(&files, &["public.postgis_hospital".to_string()]);
|
|
assert_eq!(
|
|
needed.into_iter().collect::<Vec<_>>(),
|
|
[
|
|
"db/S/data/Weird Name.dump",
|
|
"db/S/data/public.postgis_hospital.dump",
|
|
"db/S/post-data.dump",
|
|
"db/S/pre-data.dump",
|
|
"db/S/sequences.dump",
|
|
]
|
|
);
|
|
assert_eq!(files_needed(&files, &[]).len(), 4, "nothing missing: the schema dumps and the unnamed one");
|
|
}
|
|
|
|
#[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");
|
|
}
|
|
}
|