feat: Support custom PostgreSQL Docker images

benches/codegen.rs

@@ -1,10 +1,11 @@
-use cornucopia::{conn::cornucopia_conn, CodegenSettings};
+use cornucopia::{conn::cornucopia_conn, container::ContainerOpts, CodegenSettings};
 use criterion::Criterion;
 
 fn bench(c: &mut Criterion) {
-    cornucopia::container::cleanup(false).ok();
-    cornucopia::container::setup(false).unwrap();
-    let client = &mut cornucopia_conn().unwrap();
+    let container_opts = ContainerOpts::default();
+    cornucopia::container::cleanup(&container_opts).ok();
+    cornucopia::container::setup(&container_opts).unwrap();
+    let client = &mut cornucopia_conn(&container_opts).unwrap();
 
     cornucopia::load_schema(client, &["../codegen_test/schema.sql"]).unwrap();
     c.bench_function("codegen_sync", |b| {
@@ -37,7 +38,7 @@ fn bench(c: &mut Criterion) {
             .unwrap()
         })
     });
-    cornucopia::container::cleanup(false).unwrap();
+    cornucopia::container::cleanup(&container_opts).unwrap();
 }
 criterion::criterion_group!(benches, bench);
 criterion::criterion_main!(benches);

benches/execution/main.rs

@@ -1,6 +1,6 @@
 use std::fmt::Write;
 
-use cornucopia::conn::cornucopia_conn;
+use cornucopia::{conn::cornucopia_conn, container::ContainerOpts};
 use criterion::{BenchmarkId, Criterion};
 use diesel::{Connection, PgConnection};
 use postgres::{fallible_iterator::FallibleIterator, Client, NoTls};
@@ -126,9 +126,10 @@ fn prepare_full(client: &mut Client) {
 }
 
 fn bench(c: &mut Criterion) {
-    cornucopia::container::cleanup(false).ok();
-    cornucopia::container::setup(false).unwrap();
-    let client = &mut cornucopia_conn().unwrap();
+    let container_opts = ContainerOpts::default();
+    cornucopia::container::cleanup(&container_opts).ok();
+    cornucopia::container::setup(&container_opts).unwrap();
+    let client = &mut cornucopia_conn(&container_opts).unwrap();
     let rt: &'static Runtime = Box::leak(Box::new(Runtime::new().unwrap()));
     let async_client = &mut rt.block_on(async {
         let (client, conn) = tokio_postgres::connect(
@@ -237,7 +238,7 @@ fn bench(c: &mut Criterion) {
         group.finish();
     }
 
-    cornucopia::container::cleanup(false).unwrap();
+    cornucopia::container::cleanup(&container_opts).unwrap();
 }
 criterion::criterion_group!(benches, bench);
 criterion::criterion_main!(benches);

crates/cornucopia/src/cli.rs

@@ -2,15 +2,17 @@ use std::path::PathBuf;
 
 use clap::{Parser, Subcommand};
 
-use crate::{conn, container, error::Error, generate_live, generate_managed, CodegenSettings};
+use crate::{
+    conn,
+    container::{self, ContainerOpts},
+    error::Error,
+    generate_live, generate_managed, CodegenSettings,
+};
 
 /// Command line interface to interact with Cornucopia SQL.
 #[derive(Parser, Debug)]
 #[clap(version)]
 struct Args {
-    /// Use `podman` instead of `docker`
-    #[clap(short, long)]
-    podman: bool,
     /// Folder containing the queries
     #[clap(short, long, default_value = "queries/")]
     queries_path: PathBuf,
@@ -41,13 +43,15 @@ enum Action {
     Schema {
         /// SQL files containing the database schema
         schema_files: Vec<PathBuf>,
+        /// Container options
+        #[command(flatten)]
+        container_opts: ContainerOpts,
     },
 }
 
 // Main entrypoint of the CLI. Parses the args and calls the appropriate routines.
 pub fn run() -> Result<(), Error> {
     let Args {
-        podman,
         queries_path,
         destination,
         action,
@@ -67,16 +71,19 @@ pub fn run() -> Result<(), Error> {
             let mut client = conn::from_url(&url)?;
             generate_live(&mut client, &queries_path, Some(&destination), settings)?;
         }
-        Action::Schema { schema_files } => {
+        Action::Schema {
+            schema_files,
+            container_opts,
+        } => {
             // Run the generate command. If the command is unsuccessful, cleanup Cornucopia's container
             if let Err(e) = generate_managed(
                 queries_path,
                 &schema_files,
                 Some(destination),
-                podman,
+                &container_opts,
                 settings,
             ) {
-                container::cleanup(podman).ok();
+                container::cleanup(&container_opts).ok();
                 return Err(e);
             }
         }

crates/cornucopia/src/codegen.rs

@@ -323,7 +323,7 @@ fn gen_params_struct(w: &mut impl Write, params: &PreparedItem, ctx: &GenCtx) {
             .map(|p| p.param_ergo_ty(traits, ctx))
             .collect::<Vec<_>>();
         let fields_name = fields.iter().map(|p| &p.ident.rs);
-        let traits_idx = (1..=traits.len()).into_iter().map(idx_char);
+        let traits_idx = (1..=traits.len()).map(idx_char);
         code!(w =>
             #[derive($copy Debug)]
             pub struct $name<$lifetime $($traits_idx: $traits,)> {
@@ -507,7 +507,7 @@ fn gen_query_fn<W: Write>(w: &mut W, module: &PreparedModule, query: &PreparedQu
         .map(|idx| param_field[*idx].param_ergo_ty(traits, ctx))
         .collect();
     let params_name = order.iter().map(|idx| &param_field[*idx].ident.rs);
-    let traits_idx = (1..=traits.len()).into_iter().map(idx_char);
+    let traits_idx = (1..=traits.len()).map(idx_char);
     let lazy_impl = |w: &mut W| {
         if let Some((idx, index)) = row {
             let item = module.rows.get_index(*idx).unwrap().1;

crates/cornucopia/src/conn.rs

@@ -1,5 +1,7 @@
 use postgres::{Client, Config, NoTls};
 
+use crate::container::ContainerOpts;
+
 use self::error::Error;
 
 /// Creates a non-TLS connection from a URL.
@@ -8,12 +10,12 @@ pub(crate) fn from_url(url: &str) -> Result<Client, Error> {
 }
 
 /// Create a non-TLS connection to the container managed by Cornucopia.
-pub fn cornucopia_conn() -> Result<Client, Error> {
+pub fn cornucopia_conn(opts: &ContainerOpts) -> Result<Client, Error> {
     Ok(Config::new()
         .user("postgres")
         .password("postgres")
         .host("127.0.0.1")
-        .port(5435)
+        .port(opts.port)
         .dbname("postgres")
         .connect(NoTls)?)
 }

crates/cornucopia/src/container.rs

@@ -2,58 +2,93 @@ use std::process::{Command, Stdio};
 
 use self::error::Error;
 
+static DEFAULT_CONTAINER_NAME: &str = "cornucopia_postgres";
+static DEFAULT_DOCKER_IMAGE: &str = "docker.io/library/postgres:latest";
+static DEFAULT_PORT: u16 = 5435;
+
+/// Container related options
+#[derive(Clone, Debug, clap::Parser)]
+pub struct ContainerOpts {
+    /// Run using podman instead of docker
+    #[clap(short = 'p', long = "podman")]
+    pub podman: bool,
+    /// Port under which to expose the cornucopia specific PostgreSQL container
+    #[clap(long = "port", default_value_t = DEFAULT_PORT)]
+    pub port: u16,
+    /// Name of the container to start
+    #[clap(long = "container-name", default_value = DEFAULT_CONTAINER_NAME)]
+    pub container_name: String,
+    #[clap(long = "docker-image", default_value = DEFAULT_DOCKER_IMAGE)]
+    pub docker_image: String,
+}
+
+impl Default for ContainerOpts {
+    fn default() -> Self {
+        Self {
+            podman: false,
+            port: DEFAULT_PORT,
+            container_name: DEFAULT_CONTAINER_NAME.into(),
+            docker_image: DEFAULT_DOCKER_IMAGE.into(),
+            // startup_delay: STARTUP_DELAY,
+        }
+    }
+}
+
 /// Starts Cornucopia's database container and wait until it reports healthy.
-pub fn setup(podman: bool) -> Result<(), Error> {
-    spawn_container(podman)?;
-    healthcheck(podman, 120, 50)?;
+pub fn setup(opts: &ContainerOpts) -> Result<(), Error> {
+    spawn_container(opts)?;
+    healthcheck(opts, 120, 50)?;
     Ok(())
 }
 
 /// Stop and remove a container and its volume.
-pub fn cleanup(podman: bool) -> Result<(), Error> {
-    stop_container(podman)?;
-    remove_container(podman)?;
+pub fn cleanup(opts: &ContainerOpts) -> Result<(), Error> {
+    stop_container(opts)?;
+    remove_container(opts)?;
     Ok(())
 }
 
 /// Starts Cornucopia's database container.
-fn spawn_container(podman: bool) -> Result<(), Error> {
-    cmd(
-        podman,
-        &[
-            "run",
-            "-d",
-            "--name",
-            "cornucopia_postgres",
-            "-p",
-            "5435:5432",
-            "-e",
-            "POSTGRES_PASSWORD=postgres",
-            "docker.io/library/postgres:latest",
-        ],
-        "spawn container",
-    )
+fn spawn_container(opts: &ContainerOpts) -> Result<(), Error> {
+    let args = [
+        "run",
+        "-d",
+        "--name",
+        &opts.container_name,
+        "-p",
+        &format!("{}:5432", opts.port),
+        "-e",
+        "POSTGRES_PASSWORD=postgres",
+        &opts.docker_image,
+    ];
+    cmd(opts, &args, "spawn container")
 }
 
 /// Checks if Cornucopia's container reports healthy
-fn is_postgres_healthy(podman: bool) -> Result<bool, Error> {
+fn is_postgres_healthy(opts: &ContainerOpts) -> Result<bool, Error> {
     Ok(cmd(
-        podman,
-        &["exec", "cornucopia_postgres", "pg_isready"],
+        opts,
+        &["exec", &opts.container_name, "pg_isready"],
         "check container health",
     )
     .is_ok())
 }
 
 /// This function controls how the healthcheck retries are handled.
-fn healthcheck(podman: bool, max_retries: u64, ms_per_retry: u64) -> Result<(), Error> {
+fn healthcheck(opts: &ContainerOpts, max_retries: u64, ms_per_retry: u64) -> Result<(), Error> {
     let slow_threshold = 10 + max_retries / 10;
     let mut nb_retries = 0;
-    while !is_postgres_healthy(podman)? {
+    let mut consecutive_successes = 0;
+    while consecutive_successes < 5 {
+        if is_postgres_healthy(opts)? {
+            consecutive_successes += 1;
+        } else {
+            consecutive_successes = 0;
+        }
         if nb_retries >= max_retries {
             return Err(Error::new(
                 String::from("Cornucopia reached the max number of connection retries"),
-                podman,
+                opts,
             ));
         };
         std::thread::sleep(std::time::Duration::from_millis(ms_per_retry));
@@ -63,27 +98,25 @@ fn healthcheck(podman: bool, max_retries: u64, ms_per_retry: u64) -> Result<(),
             println!("Container startup slower than expected ({nb_retries} retries out of {max_retries})");
         }
     }
-    // Just for extra safety...
-    std::thread::sleep(std::time::Duration::from_millis(250));
     Ok(())
 }
 
 /// Stops Cornucopia's container.
-fn stop_container(podman: bool) -> Result<(), Error> {
-    cmd(podman, &["stop", "cornucopia_postgres"], "stop container")
+fn stop_container(opts: &ContainerOpts) -> Result<(), Error> {
+    cmd(opts, &["stop", &opts.container_name], "stop container")
 }
 
 /// Removes Cornucopia's container and its volume.
-fn remove_container(podman: bool) -> Result<(), Error> {
+fn remove_container(opts: &ContainerOpts) -> Result<(), Error> {
     cmd(
-        podman,
-        &["rm", "-v", "cornucopia_postgres"],
+        opts,
+        &["rm", "-v", &opts.container_name],
         "remove container",
     )
 }
 
-fn cmd(podman: bool, args: &[&'static str], action: &'static str) -> Result<(), Error> {
-    let command = if podman { "podman" } else { "docker" };
+fn cmd(opts: &ContainerOpts, args: &[&str], action: &'static str) -> Result<(), Error> {
+    let command = if opts.podman { "podman" } else { "docker" };
     let output = Command::new(command)
         .args(args)
         .stderr(Stdio::piped())
@@ -96,14 +129,15 @@ fn cmd(podman: bool, args: &[&'static str], action: &'static str) -> Result<(),
         let err = String::from_utf8_lossy(&output.stderr);
         Err(Error::new(
             format!("`{command}` couldn't {action}: {err}"),
-            podman,
+            opts,
         ))
     }
 }
 
 pub(crate) mod error {
     use std::fmt::Debug;
 
+    use super::ContainerOpts;
     use miette::Diagnostic;
     use thiserror::Error as ThisError;
 
@@ -116,15 +150,15 @@ pub(crate) mod error {
     }
 
     impl Error {
-        pub fn new(msg: String, podman: bool) -> Self {
-            let help = if podman {
-                "Make sure that port 5435 is usable and that no container named `cornucopia_postgres` already exists."
+        pub fn new(msg: String, opts: &ContainerOpts) -> Self {
+            let help = if opts.podman {
+                format!("Make sure that port {} is usable and that no container named `{}` already exists.", opts.port, opts.container_name)
             } else {
-                "First, check that the docker daemon is up-and-running. Then, make sure that port 5435 is usable and that no container named `cornucopia_postgres` already exists."
+                format!("First, check that the docker daemon is up-and-running. Then, make sure that port {} is usable and that no container named `{}` already exists.", opts.port, opts.container_name)
             };
             Error {
                 msg,
-                help: Some(String::from(help)),
+                help: Some(help),
             }
         }
     }

crates/cornucopia/src/lib.rs

@@ -16,6 +16,7 @@ pub mod container;
 
 use std::path::Path;
 
+use container::ContainerOpts;
 use postgres::Client;
 
 use codegen::generate as generate_internal;
@@ -75,20 +76,20 @@ pub fn generate_managed<P: AsRef<Path>>(
     queries_path: P,
     schema_files: &[P],
     destination: Option<P>,
-    podman: bool,
+    container_opts: &ContainerOpts,
     settings: CodegenSettings,
 ) -> Result<String, Error> {
     // Read
     let modules = read_query_modules(queries_path.as_ref())?
         .into_iter()
         .map(parse_query_module)
         .collect::<Result<_, parser::error::Error>>()?;
-    container::setup(podman)?;
-    let mut client = conn::cornucopia_conn()?;
+    container::setup(container_opts)?;
+    let mut client = conn::cornucopia_conn(container_opts)?;
     load_schema(&mut client, schema_files)?;
     let prepared_modules = prepare(&mut client, modules)?;
     let generated_code = generate_internal(prepared_modules, settings);
-    container::cleanup(podman)?;
+    container::cleanup(container_opts)?;
 
     if let Some(destination) = destination {
         write_generated_code(destination.as_ref(), &generated_code)?;