Skip to content

perf(supervisor): fix extension filtering + batch post_start into one psql #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jaredLunde merged 1 commit into from
Jun 26, 2026
Merged

perf(supervisor): fix extension filtering + batch post_start into one psql #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 49 additions & 17 deletions src/pg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,25 +149,57 @@ pub async fn psql_env(sql: &str, env: &[(&str, &str)]) -> Result<(), PgError> {
Err(PgError::NonZeroExit { code, stderr })
}

/// Set the superuser password using dollar-quoting to avoid SQL injection.
///
/// Dollar-quoting avoids single-quote injection. The password is MMDS-controlled
/// but correct quoting is still the right habit.
pub async fn set_superuser_password(pw: &str) -> Result<(), PgError> {
set_role_password("postgres", pw).await
/// Build the dollar-quoted `ALTER ROLE … WITH PASSWORD` statement (no trailing
/// semicolon, so it composes into a larger script). Dollar-quoting guards against
/// single-quote injection; the `$_beyond_$` tag is unlikely to appear in a
/// password (if it does, this breaks — a documented unsupported MMDS edge case).
pub fn alter_role_password_sql(role: &str, pw: &str) -> String {
format!("ALTER ROLE {role} WITH PASSWORD $_beyond_${pw}$_beyond_$")
}

/// Set a role's password using dollar-quoting (same injection guard as
/// `set_superuser_password`). `role` is a fixed internal identifier (never user
/// input); `pw` is MMDS-controlled and dollar-quoted.
pub async fn set_role_password(role: &str, pw: &str) -> Result<(), PgError> {
// Dollar-quote tag chosen to be unlikely to appear in a password. If the
// password itself contains `$_beyond_$`, this breaks — documented as an
// unsupported edge case in the MMDS field.
psql(&format!(
"ALTER ROLE {role} WITH PASSWORD $_beyond_${pw}$_beyond_$"
))
.await
/// Run a multi-statement SQL script in a SINGLE `psql` invocation, aborting on
/// the first error (`ON_ERROR_STOP=1`). Collapses the per-statement process
/// spawn + reconnect overhead of calling [`psql`] once per statement — used by
/// post-start setup, which is a fixed batch of idempotent statements.
pub async fn psql_script(sql: &str) -> Result<(), PgError> {
use tokio::io::AsyncWriteExt as _;
let mut cmd = Command::new("psql");
cmd.args([
"-U",
POSTGRES_USER,
"-h",
PG_SOCKET_DIR,
"-p",
&PG_PORT.to_string(),
"-d",
"postgres",
"-v",
"ON_ERROR_STOP=1",
"-f",
"-",
])
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::piped());
// Same peer-auth requirement as `psql`: drop to the postgres OS user.
drop_to_postgres_user(&mut cmd);

let mut child = cmd.spawn()?;
{
let mut stdin = child
.stdin
.take()
.ok_or_else(|| PgError::Io(std::io::Error::other("psql stdin unavailable")))?;
stdin.write_all(sql.as_bytes()).await?;
stdin.shutdown().await?;
}
let out = child.wait_with_output().await?;
if out.status.success() {
return Ok(());
}
let code = out.status.code().ok_or(PgError::Signal)?;
let stderr = String::from_utf8_lossy(&out.stderr).into_owned();
Err(PgError::NonZeroExit { code, stderr })
}

// ---------------------------------------------------------------------------
Expand Down
24 changes: 24 additions & 0 deletions src/sql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,27 @@ BEGIN
END IF;
END
$$";

/// pgbouncer auth: the `pgbouncer` login role + the SECURITY DEFINER lookup
/// function PgBouncer's `auth_query` calls, plus the schema USAGE the role needs
/// to call it. Idempotent (CREATE IF NOT EXISTS / OR REPLACE), and every
/// statement is `;`-terminated so it composes into the post-start batch script.
///
/// USAGE on the schema is REQUIRED — without it `auth_query` fails with
/// "permission denied for schema pgbouncer" for every client and no one can
/// connect through the pooler (EXECUTE on the function alone is not enough).
pub const PGBOUNCER_AUTH_SQL: &str = "DO $$
BEGIN
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'pgbouncer') THEN
CREATE ROLE pgbouncer LOGIN PASSWORD NULL;
END IF;
END
$$;
CREATE SCHEMA IF NOT EXISTS pgbouncer;
CREATE OR REPLACE FUNCTION pgbouncer.get_auth(p_user text)
RETURNS TABLE(username text, password text)
SECURITY DEFINER LANGUAGE sql AS $$
SELECT usename::text, passwd::text FROM pg_shadow WHERE usename = p_user
$$;
GRANT USAGE ON SCHEMA pgbouncer TO pgbouncer;
GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer;";
244 changes: 143 additions & 101 deletions src/supervisor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1752,70 +1752,37 @@ async fn wait_for_sync_standby(name: &str, timeout: std::time::Duration) {
// ---------------------------------------------------------------------------

async fn post_start(cfg: &MmdsConfig) -> Result<(), Box<dyn std::error::Error>> {
pg::set_superuser_password(&cfg.postgres_password)
.await
.map_err(|e| format!("failed to set postgres password: {e}"))?;

setup_pgbouncer_auth()
.await
.map_err(|e| format!("failed to set up pgbouncer auth: {e}"))?;

// Required extensions — fail the supervisor if CREATE EXTENSION fails for
// an extension whose shared object IS installed; the process manager will
// restart and retry rather than running in a degraded state.
//
// But the standalone postgres primitive ships without the auth/queue
// milestone (beyond_auth/beyond_queue) and pgdg's pg_cron (version pin
// drift). When an extension's .so isn't installed, treat it as a warning,
// not a fatal — the same self-adapting posture as the
// shared_preload_libraries filter (see config::beyond_conf). With the
// extensions present, the behavior is unchanged (still hard-required).
// Resolve the extensions to (idempotently) create: required + optional,
// filtered to those whose control file is present. A REQUIRED extension
// missing its control file is downgraded to a warning (the standalone image
// may ship without the auth/queue milestone) and dropped from the batch.
// Same filter the build-time bake (template.rs) uses ⇒ every CREATE EXTENSION
// is a no-op against the pre-baked cluster.
let mut extensions: Vec<&str> = Vec::new();
for ext in REQUIRED_EXTENSIONS {
if !extension_installed(ext) {
if extension_installed(ext) {
extensions.push(ext);
} else {
warn!(
"required extension {ext} not installed (no {ext}.so in {EXTENSION_PKGLIBDIR}); \
skipping (auth/queue milestone not in this image)"
"required extension {ext} not installed (no {ext}.control in \
{EXTENSION_CONTROL_DIR}); skipping (auth/queue milestone not in this image)"
);
continue;
}
pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}"))
.await
.map_err(|e| format!("required extension {ext} failed: {e}"))?;
}

for ext in OPTIONAL_EXTENSIONS {
if let Err(e) = pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}")).await {
warn!("optional extension {ext} failed: {e}");
}
}

// Create the `replicator` role when anything needs to stream WAL from us: a
// remote replica (replication_password set), a WAL sink, or CDC. Remote
// consumers authenticate via scram (pg_hba `host replication replicator`),
// so when a replication password is configured we set it — without it the
// role is PASSWORD NULL and scram can't succeed.
if cfg.replication_password.is_some() || cfg.wal_sink.is_some() || cfg.cdc_enabled {
pg::psql(crate::sql::REPLICATOR_ROLE_SQL)
.await
.map_err(|e| format!("failed to create replicator role: {e}"))?;
if let Some(pw) = &cfg.replication_password {
pg::set_role_password("replicator", pw)
.await
.map_err(|e| format!("failed to set replicator password: {e}"))?;
}
}
extensions.extend(OPTIONAL_EXTENSIONS.iter().copied().filter(|e| extension_installed(e)));

if cfg.cdc_enabled {
warn!("CDC enabled: replication slot 'cdc' will accumulate WAL until a consumer connects");
pg::psql(crate::sql::CDC_SLOT_SQL)
.await
.map_err(|e| format!("failed to create CDC replication slot: {e}"))?;

pg::psql(crate::sql::CDC_PUBLICATION_SQL)
.await
.map_err(|e| format!("failed to create CDC publication: {e}"))?;
}

// ONE psql invocation for the whole batch (ON_ERROR_STOP) instead of ~15
// process spawns + reconnects. Every statement is idempotent and, on a
// well-formed image, a no-op or fast role op — so a failure is a real problem
// and rightly aborts post_start (→ restart).
pg::psql_script(&build_post_start_script(cfg, &extensions))
.await
.map_err(|e| format!("post-start setup failed: {e}"))?;

boot::run_hook_scripts("/etc/postgresql/18/hooks/post-start.d")
.await
.map_err(|e| format!("post-start hook failed: {e}"))?;
Expand All @@ -1824,39 +1791,49 @@ async fn post_start(cfg: &MmdsConfig) -> Result<(), Box<dyn std::error::Error>>
Ok(())
}

async fn setup_pgbouncer_auth() -> Result<(), crate::pg::PgError> {
// Create the pgbouncer role and the SECURITY DEFINER auth lookup function.
// Idempotent: CREATE IF NOT EXISTS / CREATE OR REPLACE.
pg::psql(
"DO $$
BEGIN
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'pgbouncer') THEN
CREATE ROLE pgbouncer LOGIN PASSWORD NULL;
END IF;
END
$$",
)
.await?;

pg::psql("CREATE SCHEMA IF NOT EXISTS pgbouncer").await?;

pg::psql(
"CREATE OR REPLACE FUNCTION pgbouncer.get_auth(p_user text)
RETURNS TABLE(username text, password text)
SECURITY DEFINER LANGUAGE sql AS $$
SELECT usename::text, passwd::text FROM pg_shadow WHERE usename = p_user
$$",
)
.await?;

// USAGE on the schema is REQUIRED for the pgbouncer role to call get_auth at
// all — without it auth_query fails with "permission denied for schema
// pgbouncer" for every client, and no one can connect through the pooler.
// (EXECUTE on the function is not enough; the caller also needs schema USAGE.)
pg::psql("GRANT USAGE ON SCHEMA pgbouncer TO pgbouncer").await?;
pg::psql("GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer").await?;
/// Build the single post-start SQL script (pure — no I/O, so it's unit-testable).
/// Order: superuser password → pgbouncer auth → `CREATE EXTENSION IF NOT EXISTS`
/// per `extensions` → (when WAL consumers exist) replicator role + password →
/// (when CDC enabled) slot + publication. Every statement is `;`-terminated so
/// the whole thing runs as one `psql -f -` with `ON_ERROR_STOP=1`.
fn build_post_start_script(cfg: &MmdsConfig, extensions: &[&str]) -> String {
let mut s = String::new();

Ok(())
// 1. Reset the superuser password from the per-instance MMDS secret.
s.push_str(&pg::alter_role_password_sql("postgres", &cfg.postgres_password));
s.push_str(";\n");

// 2. pgbouncer auth (role + SECURITY DEFINER lookup function + grants).
s.push_str(crate::sql::PGBOUNCER_AUTH_SQL);
s.push('\n');

// 3. Extensions (already filtered to installable by the caller).
for ext in extensions {
s.push_str(&format!("CREATE EXTENSION IF NOT EXISTS {ext};\n"));
}

// 4. replicator role (+ password) when a replica / WAL sink / CDC will stream
// from us. Remote consumers scram-auth (pg_hba `host replication
// replicator`); without a password the role is PASSWORD NULL and scram
// can't succeed.
if cfg.replication_password.is_some() || cfg.wal_sink.is_some() || cfg.cdc_enabled {
s.push_str(crate::sql::REPLICATOR_ROLE_SQL);
s.push_str(";\n");
if let Some(pw) = &cfg.replication_password {
s.push_str(&pg::alter_role_password_sql("replicator", pw));
s.push_str(";\n");
}
}

// 5. CDC slot + publication.
if cfg.cdc_enabled {
s.push_str(crate::sql::CDC_SLOT_SQL);
s.push_str(";\n");
s.push_str(crate::sql::CDC_PUBLICATION_SQL);
s.push_str(";\n");
}

s
}

// beyond_auth is NOT auto-created here: the beyond-auth service installs its
Expand All @@ -1865,32 +1842,43 @@ async fn setup_pgbouncer_auth() -> Result<(), crate::pg::PgError> {
// an extension would collide with that migration ("function already exists").
pub(crate) const REQUIRED_EXTENSIONS: &[&str] = &["beyond_queue", "pg_cron"];

/// Directory holding PostgreSQL extension shared objects (PG18 Debian layout,
/// `pg_config --pkglibdir`). Mirrors `config`'s PKGLIBDIR.
pub(crate) const EXTENSION_PKGLIBDIR: &str = "/usr/lib/postgresql/18/lib";

/// True iff the extension's shared object is present in the image. An extension
/// listed in [`REQUIRED_EXTENSIONS`] but with no installed `.so` (e.g. the
/// future auth/queue milestone, or a dropped pgdg `pg_cron`) is downgraded from
/// fatal to a warning so the standalone primitive still boots.
/// Directory holding PostgreSQL extension control files (PG18 Debian layout,
/// `pg_config --sharedir` + `/extension`). `CREATE EXTENSION <name>` resolves
/// `{name}.control` here.
pub(crate) const EXTENSION_CONTROL_DIR: &str = "/usr/share/postgresql/18/extension";

/// True iff `CREATE EXTENSION {ext}` can find it — i.e. its control file is
/// present. This is the CORRECT installability test, not the bare `.so`:
/// preload-only modules (`auto_explain`) ship a `.so` but NO control file, and
/// `postgis` ships `postgis-3.so` but `postgis.control`. Filtering on the control
/// file (vs the lib) is what lets the build-time bake (`template.rs`) and runtime
/// `post_start` agree on the SAME installable set, so every runtime CREATE
/// EXTENSION is a no-op against the pre-baked cluster. A REQUIRED extension whose
/// control file is absent is downgraded from fatal to a warning so the standalone
/// primitive still boots.
pub(crate) fn extension_installed(ext: &str) -> bool {
std::path::Path::new(EXTENSION_PKGLIBDIR)
.join(format!("{ext}.so"))
std::path::Path::new(EXTENSION_CONTROL_DIR)
.join(format!("{ext}.control"))
.exists()
}

/// Extensions auto-created on boot when present (no-ops against the pre-baked
/// cluster). Every entry must be the extension's REAL name (the `.control`
/// basename) — e.g. pgvector's extension is `vector`. Entries that aren't built
/// for this image (no control file) are filtered out by [`extension_installed`];
/// they're omitted here too so the list reflects what actually ships. Preload-only
/// modules (`auto_explain`, also `pg_stat_statements`) are loaded via
/// `shared_preload_libraries`, not `CREATE EXTENSION` — `pg_stat_statements` is
/// kept because it ALSO ships a control file (it exposes a view), but it's the
/// only such dual case.
pub(crate) const OPTIONAL_EXTENSIONS: &[&str] = &[
"pg_stat_statements",
"auto_explain",
"pg_trgm",
"pgvector",
"pgvectorscale",
"vector", // pgvector — the extension is named `vector`
"pg_partman",
"pg_jsonschema",
"hypopg",
"pg_repack",
"postgis",
"pg_search",
];

// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -1951,6 +1939,60 @@ fn parse_memtotal_kib(content: &str) -> Option<u64> {
mod tests {
use super::*;

// --- post_start batch script (pure) ---

fn pg_cfg(replication_password: Option<&str>, wal_sink: Option<&str>, cdc: bool) -> MmdsConfig {
MmdsConfig {
pg_tier: crate::mmds::PgTier::Single,
ephemeral: false,
postgres_password: "s3cr3t".into(),
postgres_database: "postgres".into(),
wal_sink: wal_sink.map(str::to_owned),
cdc_enabled: cdc,
recovery_target_time: None,
primary_conninfo: None,
replication_password: replication_password.map(str::to_owned),
ram_bytes: 4 * 1024 * 1024 * 1024,
vcpus: 2,
}
}

#[test]
fn post_start_script_is_one_terminated_batch() {
let s = build_post_start_script(&pg_cfg(None, None, false), &["vector", "postgis"]);
// superuser password (dollar-quoted) + pgbouncer auth, each terminated.
assert!(s.contains("ALTER ROLE postgres WITH PASSWORD $_beyond_$s3cr3t$_beyond_$;"));
assert!(s.contains("CREATE SCHEMA IF NOT EXISTS pgbouncer;"));
assert!(s.contains("GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer;"));
// each filtered extension, semicolon-terminated.
assert!(s.contains("CREATE EXTENSION IF NOT EXISTS vector;"));
assert!(s.contains("CREATE EXTENSION IF NOT EXISTS postgis;"));
// nothing replication/CDC when not configured.
assert!(!s.contains("replicator"));
assert!(!s.contains("PUBLICATION cdc"));
// every statement is terminated: no trailing un-`;`d fragment.
assert!(s.trim_end().ends_with(';'));
}

#[test]
fn post_start_script_adds_replicator_and_cdc_when_enabled() {
// WAL sink ⇒ replicator role (no password set); CDC ⇒ slot + publication.
let s = build_post_start_script(&pg_cfg(Some("rpw"), Some("http://sink:9000"), true), &[]);
assert!(s.contains("CREATE ROLE replicator LOGIN REPLICATION PASSWORD NULL"));
assert!(s.contains("ALTER ROLE replicator WITH PASSWORD $_beyond_$rpw$_beyond_$;"));
assert!(s.contains("pg_create_logical_replication_slot('cdc'"));
assert!(s.contains("CREATE PUBLICATION cdc"));
assert!(s.trim_end().ends_with(';'));
}

#[test]
fn post_start_script_replicator_without_password_when_only_sink() {
// wal_sink alone creates the role but sets NO password (PASSWORD NULL).
let s = build_post_start_script(&pg_cfg(None, Some("http://sink:9000"), false), &[]);
assert!(s.contains("CREATE ROLE replicator"));
assert!(!s.contains("ALTER ROLE replicator WITH PASSWORD"));
}

// --- PgBouncer scaler decision (pure) ---

const COOL: Duration = Duration::from_secs(3600); // past every cooldown
Expand Down
Loading