main.rsannotatedmain.rssource487 lines · 19.3 KB · raw
1//! `postjevsql-sidecar`: the one engine that sets up sidecar mode
2//! (CLAUDE.md *Deployment*, "Launching it").
3//!
4//! ```text
5//! postjevsql-sidecar [--sidecar <conninfo>] <config.toml> converge
6//! postjevsql-sidecar [--sidecar <conninfo>] <config.toml> sync [--apply]
7//! postjevsql-sidecar <config.toml> serve [-- <postgres options>…]
8//! ```
9//!
10//! `serve` is the container image's entrypoint (nix/sidecar-image.nix):
11//! it initialises `$PGDATA` if empty, starts Postgres, runs `converge`
12//! and `sync --apply`, stops it, and then execs `postgres` in its own
13//! place, so the server is PID 1 and gets the runtime's signals
14//! directly. A refusal stops the server and exits non-zero with the
15//! refusal. Postgres is started directly and stopped with `pg_ctl stop`,
16//! neither of which needs a shell.
17//!
18//! Every secret is resolved before anything connects, so a config that
19//! names a secret twice, or a required one nowhere, is refused first.
20//!
21//! The target's catalog is read through the foreign server itself: an
22//! `IMPORT FOREIGN SCHEMA` into a scratch schema, in a transaction that
23//! is rolled back. So the CLI holds no second connection to the target,
24//! and the types it compares are exactly what postgres_fdw imports.
25#![forbid(unsafe_code)]
26
27use std::collections::BTreeMap;
28use std::os::unix::process::CommandExt;
29use std::path::PathBuf;
30use std::process::ExitCode;
31use std::time::{Duration, Instant};
32
33use postjevsql_sidecar::config::Config;
34use postjevsql_sidecar::secret::Source;
35use postjevsql_sidecar::sync::{self, Column, Table, Touches, ident, lit};
36use tokio_postgres::{Client, GenericClient, NoTls};
37
38const USAGE: &str =
39    "usage: postjevsql-sidecar [--sidecar <conninfo>] <config.toml> converge | sync [--apply] | serve [-- <postgres options>…]";
40const DEFAULT_SIDECAR: &str = "host=/run/postgresql user=postgres dbname=postgres";
41
42type Res<T> = Result<T, String>;
43
44enum Command {
45    Converge,
46    Sync { apply: bool },
47    /// Options after `--` go to `postgres` as given.
48    Serve { postgres: Vec<String> },
49}
50
51struct Args {
52    sidecar: String,
53    config: String,
54    command: Command,
55}
56
57fn parse_args(args: &[String]) -> Res<Args> {
58    let mut sidecar = DEFAULT_SIDECAR.to_owned();
59    let mut rest = Vec::new();
60    let (args, postgres) = match args.iter().position(|a| a == "--") {
61        Some(i) => (&args[..i], Some(args[i + 1..].to_vec())),
62        None => (args, None),
63    };
64    let mut it = args.iter();
65    while let Some(a) = it.next() {
66        match a.as_str() {
67            "--sidecar" => sidecar = it.next().ok_or(USAGE)?.clone(),
68            _ => rest.push(a.as_str()),
69        }
70    }
71    let command = match rest.get(1..) {
72        Some(["converge"]) => Command::Converge,
73        Some(["sync"]) => Command::Sync { apply: false },
74        Some(["sync", "--apply"]) => Command::Sync { apply: true },
75        Some(["serve"]) => Command::Serve { postgres: postgres.clone().unwrap_or_default() },
76        _ => return Err(USAGE.into()),
77    };
78    if postgres.is_some() && !matches!(command, Command::Serve { .. }) {
79        return Err(USAGE.into());
80    }
81    Ok(Args { sidecar, config: rest[0].to_owned(), command })
82}
83
84/// Each mapping's password, `None` where no source is set.
85struct Secrets {
86    passwords: Vec<Option<String>>,
87}
88
89/// Resolves every secret the config names; the API key must be set.
90/// The API key's value is not kept: the sidecar's server reads it itself.
91fn resolve(config: &Config) -> Res<Secrets> {
92    config.jev.api_key().resolve(true).map_err(|e| e.to_string())?;
93    let passwords = config
94        .mappings
95        .iter()
96        .map(|m| m.password().resolve(false).map_err(|e| e.to_string()))
97        .collect::<Res<_>>()?;
98    Ok(Secrets { passwords })
99}
100
101#[tokio::main(flavor = "current_thread")]
102async fn main() -> ExitCode {
103    let args: Vec<String> = std::env::args().skip(1).collect();
104    match run(&args).await {
105        Ok(()) => ExitCode::SUCCESS,
106        Err(e) => {
107            eprintln!("postjevsql-sidecar: {e}");
108            ExitCode::FAILURE
109        }
110    }
111}
112
113async fn run(args: &[String]) -> Res<()> {
114    let args = parse_args(args)?;
115    let text = std::fs::read_to_string(&args.config).map_err(|e| format!("{}: {e}", args.config))?;
116    let config = Config::parse(&text).map_err(|e| format!("{}: {e}", args.config))?;
117    let secrets = resolve(&config)?;
118    if let Command::Serve { postgres } = &args.command {
119        return serve(&args.sidecar, &config, &secrets, postgres).await;
120    }
121    run_command(&args.sidecar, &config, &secrets, &args.command).await
122}
123
124async fn run_command(sidecar: &str, config: &Config, secrets: &Secrets, command: &Command) -> Res<()> {
125    let (mut client, connection) =
126        tokio_postgres::connect(sidecar, NoTls).await.map_err(|e| format!("sidecar: {e}"))?;
127    tokio::spawn(async move {
128        if let Err(e) = connection.await {
129            eprintln!("postjevsql-sidecar: sidecar connection: {e}");
130        }
131    });
132    match command {
133        Command::Converge => converge(&mut client, config, secrets).await,
134        Command::Sync { apply } => sync_tables(&mut client, config, *apply).await,
135        Command::Serve { .. } => unreachable!("serve is dispatched in run"),
136    }
137}
138
139/// How long a started server may take to accept connections.
140const READY_TIMEOUT: Duration = Duration::from_secs(120);
141
142/// The container entrypoint; see the module header. Only returns on a
143/// failure: on success the process becomes `postgres`.
144async fn serve(sidecar: &str, config: &Config, secrets: &Secrets, extra: &[String]) -> Res<()> {
145    let pgdata = PathBuf::from(std::env::var_os("PGDATA").ok_or("serve: PGDATA is not set")?);
146    if !pgdata.join("PG_VERSION").exists() {
147        let status = std::process::Command::new("initdb")
148            .arg("-D")
149            .arg(&pgdata)
150            .args(["-U", "postgres", "--auth-local=peer", "--auth-host=scram-sha-256", "--no-locale", "-E", "UTF8"])
151            .status()
152            .map_err(|e| format!("initdb: {e}"))?;
153        if !status.success() {
154            return Err(format!("initdb: {status}"));
155        }
156    }
157
158    // The server reads the API key itself: from the file the config
159    // names, or from $TYPESAFE_API_KEY, which it inherits. Never both,
160    // which `resolve` has already refused.
161    let api_key = config.jev.api_key();
162    let env = std::env::var(&api_key.env).ok();
163    let mut options: Vec<String> = Vec::new();
164    if let Some(Source::File(file)) = api_key.choose(env.as_deref()).map_err(|e| e.to_string())? {
165        options.extend(["-c".into(), format!("jev.api_key_file={}", file.display())]);
166    }
167    options.extend(extra.iter().cloned());
168    let postgres = |options: &[String]| {
169        let mut c = std::process::Command::new("postgres");
170        c.arg("-D").arg(&pgdata).args(["-c", "listen_addresses=*", "-c", "unix_socket_directories=/run/postgresql"]);
171        c.args(options);
172        c
173    };
174
175    let mut server = postgres(&options).spawn().map_err(|e| format!("postgres: {e}"))?;
176    let setup = async {
177        wait_ready(sidecar, &mut server).await?;
178        run_command(sidecar, config, secrets, &Command::Converge).await?;
179        run_command(sidecar, config, secrets, &Command::Sync { apply: true }).await
180    }
181    .await;
182    // Stopped either way: on success to exec it in our place, on a
183    // refusal so the container exits with the refusal.
184    let stop = std::process::Command::new("pg_ctl")
185        .arg("stop")
186        .arg("-D")
187        .arg(&pgdata)
188        .args(["-m", "fast", "-w"])
189        .status();
190    let _ = server.wait();
191    setup?;
192    match stop {
193        Ok(s) if s.success() => {}
194        Ok(s) => return Err(format!("pg_ctl stop: {s}")),
195        Err(e) => return Err(format!("pg_ctl stop: {e}")),
196    }
197    Err(format!("postgres: {}", postgres(&options).exec()))
198}
199
200/// Waits until the server accepts a connection, failing if it exits.
201async fn wait_ready(sidecar: &str, server: &mut std::process::Child) -> Res<()> {
202    let start = Instant::now();
203    loop {
204        if let Some(status) = server.try_wait().map_err(|e| format!("postgres: {e}"))? {
205            return Err(format!("postgres exited before it was ready: {status}"));
206        }
207        match tokio_postgres::connect(sidecar, NoTls).await {
208            Ok(_) => return Ok(()),
209            Err(e) if start.elapsed() > READY_TIMEOUT => {
210                return Err(format!("postgres not ready after {}s: {e}", READY_TIMEOUT.as_secs()));
211            }
212            Err(_) => tokio::time::sleep(Duration::from_millis(200)).await,
213        }
214    }
215}
216
217fn db(e: tokio_postgres::Error) -> String {
218    match e.as_db_error() {
219        Some(d) => format!("{}: {}", d.code().code(), d.message()),
220        None => e.to_string(),
221    }
222}
223
224/// The `ALTER … OPTIONS (…)` items that turn `have` into exactly `want`:
225/// undeclared options (a hand-added `extensions`) are dropped.
226fn option_diff(have: &BTreeMap<String, String>, want: &BTreeMap<&str, String>) -> Vec<String> {
227    let mut items = Vec::new();
228    for (k, v) in want {
229        match have.get(*k) {
230            Some(h) if h == v => {}
231            Some(_) => items.push(format!("SET {} {}", ident(k), lit(v))),
232            None => items.push(format!("ADD {} {}", ident(k), lit(v))),
233        }
234    }
235    for k in have.keys() {
236        if !want.contains_key(k.as_str()) {
237            items.push(format!("DROP {}", ident(k)));
238        }
239    }
240    items
241}
242
243fn options(raw: Option<Vec<String>>) -> BTreeMap<String, String> {
244    raw.unwrap_or_default()
245        .into_iter()
246        .filter_map(|kv| kv.split_once('=').map(|(k, v)| (k.to_owned(), v.to_owned())))
247        .collect()
248}
249
250async fn converge(client: &mut Client, config: &Config, secrets: &Secrets) -> Res<()> {
251    let version: i32 = client
252        .query_one("SELECT current_setting('server_version_num')::int", &[])
253        .await
254        .map_err(db)?
255        .get(0);
256    // `use_scram_passthrough` is PG18's; before it a mapping needs a password.
257    for (m, pw) in config.mappings.iter().zip(&secrets.passwords) {
258        if pw.is_none() && version < 180000 {
259            return Err(format!(
260                "{} is not set, and use_scram_passthrough needs PostgreSQL 18 on the sidecar (this is {version}): \
261                 name a file in the config or set ${}",
262                m.password().what,
263                m.password_env
264            ));
265        }
266    }
267
268    let t = &config.target;
269    let mut want: BTreeMap<&str, String> = BTreeMap::new();
270    want.insert("host", t.host.clone());
271    want.insert("port", t.port.to_string());
272    want.insert("dbname", t.dbname.clone());
273    want.insert("sslmode", t.sslmode.as_str().into());
274    want.insert("use_remote_estimate", t.use_remote_estimate.to_string());
275    want.insert("fetch_size", t.fetch_size.to_string());
276    if let Some(root) = &t.sslrootcert {
277        want.insert("sslrootcert", root.clone());
278    }
279
280    let tx = client.transaction().await.map_err(db)?;
281    // The mapping statements carry the password as a literal.
282    tx.batch_execute(
283        "SET LOCAL log_statement = 'none'; SET LOCAL log_min_error_statement = 'panic';
284         CREATE EXTENSION IF NOT EXISTS postgres_fdw;
285         CREATE EXTENSION IF NOT EXISTS postjevsql;",
286    )
287    .await
288    .map_err(db)?;
289    let server = ident(&t.server);
290    tx.batch_execute(&format!("CREATE SERVER IF NOT EXISTS {server} FOREIGN DATA WRAPPER postgres_fdw"))
291        .await
292        .map_err(db)?;
293    let have = options(
294        tx.query_one("SELECT srvoptions FROM pg_foreign_server WHERE srvname = $1", &[&t.server])
295            .await
296            .map_err(db)?
297            .get(0),
298    );
299    alter(&tx, &format!("ALTER SERVER {server}"), &option_diff(&have, &want)).await?;
300
301    for (m, pw) in config.mappings.iter().zip(&secrets.passwords) {
302        let mut want: BTreeMap<&str, String> = BTreeMap::new();
303        want.insert("user", m.remote_user.clone());
304        match pw {
305            Some(pw) => want.insert("password", pw.clone()),
306            None => want.insert("use_scram_passthrough", "true".into()),
307        };
308        let user = ident(&m.local_user);
309        tx.batch_execute(&format!("CREATE USER MAPPING IF NOT EXISTS FOR {user} SERVER {server}"))
310            .await
311            .map_err(db)?;
312        let have = options(
313            tx.query_one(
314                "SELECT umoptions FROM pg_user_mapping um JOIN pg_foreign_server s ON s.oid = um.umserver
315                 WHERE s.srvname = $1 AND um.umuser = $2::text::regrole",
316                &[&t.server, &m.local_user],
317            )
318            .await
319            .map_err(db)?
320            .get(0),
321        );
322        alter(&tx, &format!("ALTER USER MAPPING FOR {user} SERVER {server}"), &option_diff(&have, &want)).await?;
323    }
324    tx.commit().await.map_err(db)?;
325    println!("converged server {}", t.server);
326    Ok(())
327}
328
329async fn alter(tx: &impl GenericClient, head: &str, items: &[String]) -> Res<()> {
330    if items.is_empty() {
331        return Ok(());
332    }
333    tx.batch_execute(&format!("{head} OPTIONS ({})", items.join(", "))).await.map_err(db)
334}
335
336/// Tables, in attnum order, of the foreign tables on `server` in
337/// `schemas`, each column named by its `column_name` option. `rename`
338/// maps a scanned schema back to the one it stands for.
339async fn read_tables(
340    c: &impl GenericClient,
341    server: &str,
342    schemas: &[String],
343    rename: &dyn Fn(&str) -> String,
344) -> Res<Vec<Table>> {
345    let rows = c
346        .query(
347            "SELECT n.nspname::text, r.relname::text,
348                    coalesce((SELECT substr(o, 13) FROM unnest(a.attfdwoptions) o
349                              WHERE o LIKE 'column_name=%'), a.attname::text),
350                    format_type(a.atttypid, a.atttypmod), a.attnotnull
351             FROM pg_foreign_table f
352             JOIN pg_foreign_server s ON s.oid = f.ftserver
353             JOIN pg_class r ON r.oid = f.ftrelid
354             JOIN pg_namespace n ON n.oid = r.relnamespace
355             JOIN pg_attribute a ON a.attrelid = r.oid AND a.attnum > 0 AND NOT a.attisdropped
356             WHERE s.srvname = $1 AND n.nspname = ANY($2)
357             ORDER BY n.nspname, r.relname, a.attnum",
358            &[&server, &schemas],
359        )
360        .await
361        .map_err(db)?;
362    let mut tables: Vec<Table> = Vec::new();
363    for row in rows {
364        let (schema, name) = (rename(row.get(0)), row.get::<_, String>(1));
365        let column = Column { name: row.get(2), type_sql: row.get(3), not_null: row.get(4) };
366        match tables.last_mut() {
367            Some(t) if t.schema == schema && t.name == name => t.columns.push(column),
368            _ => tables.push(Table { schema, name, columns: vec![column] }),
369        }
370    }
371    Ok(tables)
372}
373
374/// The target's tables, imported into scratch schemas and rolled back.
375async fn target_tables(client: &mut Client, config: &Config) -> Res<Vec<Table>> {
376    let t = &config.target;
377    let scratch: Vec<String> = (0..t.schemas.len()).map(|i| format!("postjevsql_sync_{i}")).collect();
378    let tx = client.transaction().await.map_err(db)?;
379    let server = ident(&t.server);
380    // Import as the first mapping's user, whose credentials the scans use.
381    let role = config.mappings.first().map(|m| ident(&m.local_user));
382    for local in &scratch {
383        tx.batch_execute(&format!("CREATE SCHEMA {}", ident(local))).await.map_err(db)?;
384        if let Some(role) = &role {
385            tx.batch_execute(&format!(
386                "GRANT USAGE ON FOREIGN SERVER {server} TO {role}; GRANT CREATE ON SCHEMA {} TO {role}",
387                ident(local)
388            ))
389            .await
390            .map_err(db)?;
391        }
392    }
393    if let Some(role) = &role {
394        tx.batch_execute(&format!("SET LOCAL ROLE {role}")).await.map_err(db)?;
395    }
396    for (remote, local) in t.schemas.iter().zip(&scratch) {
397        tx.batch_execute(&format!(
398            "IMPORT FOREIGN SCHEMA {} FROM SERVER {server} INTO {}",
399            ident(remote),
400            ident(local)
401        ))
402        .await
403        .map_err(|e| format!("reading the target's schema {remote}: {}", db(e)))?;
404    }
405    let back = |s: &str| -> String {
406        scratch.iter().position(|x| x == s).map_or_else(|| s.to_owned(), |i| t.schemas[i].clone())
407    };
408    let tables = read_tables(&tx, &t.server, &scratch, &back).await?;
409    tx.rollback().await.map_err(db)?;
410    Ok(tables)
411}
412
413/// Everything on the sidecar that depends on what `touch` names:
414/// views and `BEGIN ATOMIC` functions (pg_depend), and grants on the
415/// table or its columns. A function with a string body is not tracked by
416/// Postgres, so it cannot be found here.
417async fn dependents(c: &impl GenericClient, server: &str, touch: &Touches<'_>) -> Res<Vec<String>> {
418    let rows = c
419        .query(
420            "WITH rel AS (
421               SELECT r.oid FROM pg_foreign_table f
422               JOIN pg_foreign_server s ON s.oid = f.ftserver
423               JOIN pg_class r ON r.oid = f.ftrelid
424               JOIN pg_namespace n ON n.oid = r.relnamespace
425               WHERE s.srvname = $1 AND n.nspname = $2 AND r.relname = $3),
426             col AS (
427               SELECT a.attrelid, a.attnum, a.attacl, a.attname FROM pg_attribute a JOIN rel ON a.attrelid = rel.oid
428               WHERE a.attnum > 0 AND NOT a.attisdropped
429                 AND ($4::text IS NULL OR coalesce((SELECT substr(o, 13) FROM unnest(a.attfdwoptions) o
430                                                    WHERE o LIKE 'column_name=%'), a.attname::text) = $4))
431             SELECT DISTINCT pg_describe_object(d.classid, d.objid, 0)
432             FROM pg_depend d JOIN col ON d.refclassid = 'pg_class'::regclass
433               AND d.refobjid = col.attrelid AND (d.refobjsubid = col.attnum OR ($4::text IS NULL AND d.refobjsubid = 0))
434             WHERE d.deptype = 'n'
435             UNION
436             SELECT DISTINCT 'grant on column ' || quote_ident(col.attname::text) || ' of ' || col.attrelid::regclass
437             FROM col WHERE col.attacl IS NOT NULL
438             UNION
439             SELECT 'grant on table ' || r.oid::regclass FROM rel JOIN pg_class r ON r.oid = rel.oid
440             WHERE $4::text IS NULL AND r.relacl IS NOT NULL
441             ORDER BY 1",
442            &[&server, &touch.schema, &touch.table, &touch.column],
443        )
444        .await
445        .map_err(db)?;
446    Ok(rows.iter().map(|r| r.get(0)).collect())
447}
448
449async fn sync_tables(client: &mut Client, config: &Config, apply: bool) -> Res<()> {
450    let t = &config.target;
451    let target = target_tables(client, config).await?;
452    let local = read_tables(&*client, &t.server, &t.schemas, &|s| s.to_owned()).await?;
453    let plan = sync::plan(&target, &local);
454    if plan.is_empty() {
455        println!("in sync: nothing to change");
456        return Ok(());
457    }
458    for change in &plan {
459        println!("{};", change.sql(&t.server));
460    }
461    if !apply {
462        return Ok(());
463    }
464
465    let tx = client.transaction().await.map_err(db)?;
466    let mut refusals = Vec::new();
467    for change in &plan {
468        let Some(touch) = change.touches() else { continue };
469        let deps = dependents(&tx, &t.server, &touch).await?;
470        if !deps.is_empty() {
471            refusals.push(format!("{change}: depended on by {}", deps.join(", ")));
472        }
473    }
474    if !refusals.is_empty() {
475        tx.rollback().await.map_err(db)?;
476        return Err(format!("refused, nothing changed:\n  {}", refusals.join("\n  ")));
477    }
478    for schema in &t.schemas {
479        tx.batch_execute(&format!("CREATE SCHEMA IF NOT EXISTS {}", ident(schema))).await.map_err(db)?;
480    }
481    for change in &plan {
482        tx.batch_execute(&change.sql(&t.server)).await.map_err(|e| format!("{change}: {}", db(e)))?;
483    }
484    tx.commit().await.map_err(db)?;
485    println!("applied {} changes", plan.len());
486    Ok(())
487}