From 006636e81f984ab670f9527b22a9b415c830385b Mon Sep 17 00:00:00 2001 From: Jacob Magnusson Date: Wed, 29 Mar 2023 09:32:39 +0200 Subject: [PATCH 1/2] feat: Support custom PostgreSQL Docker images Useful when in need of extra extensions such as PostGIS or TimescaleDB. Also add support for changing port and container name. --- benches/codegen.rs | 11 +-- benches/execution/main.rs | 11 +-- crates/cornucopia/src/cli.rs | 23 ++++--- crates/cornucopia/src/codegen.rs | 4 +- crates/cornucopia/src/conn.rs | 6 +- crates/cornucopia/src/container.rs | 106 ++++++++++++++++++----------- crates/cornucopia/src/lib.rs | 9 +-- test_integration/src/main.rs | 20 +++--- 8 files changed, 116 insertions(+), 74 deletions(-) diff --git a/benches/codegen.rs b/benches/codegen.rs index 4eabc30f..b0cfaabd 100644 --- a/benches/codegen.rs +++ b/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); diff --git a/benches/execution/main.rs b/benches/execution/main.rs index 64d07a86..5b233677 100644 --- a/benches/execution/main.rs +++ b/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); diff --git a/crates/cornucopia/src/cli.rs b/crates/cornucopia/src/cli.rs index 73692260..a1077c75 100644 --- a/crates/cornucopia/src/cli.rs +++ b/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, + /// 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); } } diff --git a/crates/cornucopia/src/codegen.rs b/crates/cornucopia/src/codegen.rs index af6c4b8a..d4e2165f 100644 --- a/crates/cornucopia/src/codegen.rs +++ b/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::>(); 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: &mut W, module: &PreparedModule, query: &PreparedQu .map(|idx| param_field[*idx].param_ergo_ty(traits, ctx)) .collect(); let params_name = order.iter().map(|idx| ¶m_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; diff --git a/crates/cornucopia/src/conn.rs b/crates/cornucopia/src/conn.rs index 83a54e0e..81ea5270 100644 --- a/crates/cornucopia/src/conn.rs +++ b/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 { } /// Create a non-TLS connection to the container managed by Cornucopia. -pub fn cornucopia_conn() -> Result { +pub fn cornucopia_conn(opts: &ContainerOpts) -> Result { Ok(Config::new() .user("postgres") .password("postgres") .host("127.0.0.1") - .port(5435) + .port(opts.port) .dbname("postgres") .connect(NoTls)?) } diff --git a/crates/cornucopia/src/container.rs b/crates/cornucopia/src/container.rs index 5ef00506..ac903f82 100644 --- a/crates/cornucopia/src/container.rs +++ b/crates/cornucopia/src/container.rs @@ -2,44 +2,73 @@ 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 { +fn is_postgres_healthy(opts: &ContainerOpts) -> Result { Ok(cmd( - podman, - &["exec", "cornucopia_postgres", "pg_isready"], + opts, + &["exec", &opts.container_name, "pg_isready"], "check container health", ) .is_ok()) @@ -69,21 +98,21 @@ fn healthcheck(podman: bool, max_retries: u64, ms_per_retry: u64) -> Result<(), } /// 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,7 +125,7 @@ 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, )) } } @@ -104,6 +133,7 @@ fn cmd(podman: bool, args: &[&'static str], action: &'static str) -> Result<(), pub(crate) mod error { use std::fmt::Debug; + use super::ContainerOpts; use miette::Diagnostic; use thiserror::Error as ThisError; @@ -116,15 +146,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), } } } diff --git a/crates/cornucopia/src/lib.rs b/crates/cornucopia/src/lib.rs index b879ee28..5d35eb2f 100644 --- a/crates/cornucopia/src/lib.rs +++ b/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,7 +76,7 @@ pub fn generate_managed>( queries_path: P, schema_files: &[P], destination: Option

, - podman: bool, + container_opts: &ContainerOpts, settings: CodegenSettings, ) -> Result { // Read @@ -83,12 +84,12 @@ pub fn generate_managed>( .into_iter() .map(parse_query_module) .collect::>()?; - 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)?; diff --git a/test_integration/src/main.rs b/test_integration/src/main.rs index cf21be31..247f816e 100644 --- a/test_integration/src/main.rs +++ b/test_integration/src/main.rs @@ -2,7 +2,7 @@ use std::{fmt::Display, process::ExitCode}; use crate::{codegen::run_codegen_test, errors::run_errors_test}; use clap::Parser; -use cornucopia::container; +use cornucopia::container::{self, ContainerOpts}; mod codegen; mod errors; @@ -19,9 +19,8 @@ struct Args { /// Update the project's generated code #[clap(long)] apply_codegen: bool, - /// Use `podman` instead of `docker` - #[clap(short, long)] - podman: bool, + #[command(flatten)] + container_opts: ContainerOpts, } /// Print error to stderr @@ -37,18 +36,18 @@ fn test( Args { apply_errors, apply_codegen, - podman, + container_opts, }: Args, ) -> bool { // Start by removing previous container if it was left open - container::cleanup(podman).ok(); - container::setup(podman).unwrap(); + container::cleanup(&container_opts).ok(); + container::setup(&container_opts).unwrap(); let successful = std::panic::catch_unwind(|| { - let mut client = cornucopia::conn::cornucopia_conn().unwrap(); + let mut client = cornucopia::conn::cornucopia_conn(&container_opts).unwrap(); display(run_errors_test(&mut client, apply_errors)).unwrap() && display(run_codegen_test(&mut client, apply_codegen)).unwrap() }); - container::cleanup(podman).unwrap(); + container::cleanup(&container_opts).unwrap(); successful.unwrap() } @@ -64,6 +63,7 @@ fn main() -> ExitCode { #[cfg(test)] mod test { + use crate::container::ContainerOpts; use crate::test; #[test] @@ -71,7 +71,7 @@ mod test { assert!(test(crate::Args { apply_errors: false, apply_codegen: false, - podman: false + container_opts: ContainerOpts::default(), })) } } From 9b645f2d19a9d34d241bace0ceb337f084826b65 Mon Sep 17 00:00:00 2001 From: Jacob Magnusson Date: Wed, 29 Mar 2023 09:37:08 +0200 Subject: [PATCH 2/2] feat: More reliable container healthcheck The postgres Docker image restarts the database server on first launch. The pg_isready command that is used for the health check returns true before the restart, which I assume is why the 250 ms additional delay exists. When using a custom image which installs extra extensions, this 250 ms delay might not be enough for the container to actually become ready. This is now resolved by replacing the 250 ms delay with a check to ensure that pg_isready responds with OK 5 times in a row. --- crates/cornucopia/src/container.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/crates/cornucopia/src/container.rs b/crates/cornucopia/src/container.rs index ac903f82..bdd06cb2 100644 --- a/crates/cornucopia/src/container.rs +++ b/crates/cornucopia/src/container.rs @@ -75,14 +75,20 @@ fn is_postgres_healthy(opts: &ContainerOpts) -> Result { } /// 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)); @@ -92,8 +98,6 @@ 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(()) }