Optimising codegen using an asynchronous client and pipelining

Author: Virgiel

Cargo.lock

@@ -6,7 +6,7 @@ fn bench(c: &mut Criterion) {
     cornucopia::container::setup(false).unwrap();
     let client = &mut cornucopia_conn().unwrap();
 
-    cornucopia::load_schema(client, &["../codegen_test/schema.sql"]).unwrap();
+    cornucopia::load_schema(client, &["../test_codegen/schema.sql"]).unwrap();
     c.bench_function("codegen_sync", |b| {
         b.iter(|| {
             cornucopia::generate_live(

benches/codegen.rs

@@ -4,7 +4,6 @@ use cornucopia::conn::cornucopia_conn;
 use criterion::{BenchmarkId, Criterion};
 use diesel::{Connection, PgConnection};
 use postgres::{fallible_iterator::FallibleIterator, Client, NoTls};
-use tokio::runtime::Runtime;
 
 const QUERY_SIZE: &[usize] = &[1, 100, 10_000];
 const INSERT_SIZE: &[usize] = &[1, 100, 1000];
@@ -128,23 +127,16 @@ 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 rt: &'static Runtime = Box::leak(Box::new(Runtime::new().unwrap()));
-    let async_client = &mut rt.block_on(async {
-        let (client, conn) = tokio_postgres::connect(
-            "postgresql://postgres:[email protected]:5435/postgres",
-            NoTls,
-        )
-        .await
-        .unwrap();
-        rt.spawn(conn);
-        client
-    });
+    let client = &mut postgres::Client::connect(
+        "postgresql://postgres:[email protected]:5435/postgres",
+        NoTls,
+    )
+    .unwrap();
+    let async_client = &mut cornucopia_conn().unwrap();
     let conn =
         &mut PgConnection::establish("postgresql://postgres:[email protected]:5435/postgres")
             .unwrap();
-
-    cornucopia::load_schema(client, &["usage/cornucopia_benches/schema.sql"]).unwrap();
+    cornucopia::load_schema(async_client, &["execution/cornucopia_benches/schema.sql"]).unwrap();
     {
         let mut group = c.benchmark_group("bench_trivial_query");
         for size in QUERY_SIZE {

benches/execution/main.rs

@@ -15,9 +15,13 @@ keywords = ["postgresql", "query", "generator", "sql", "tokio-postgres"]
 codegen_template = { path = "../codegen_template", version = "0.1.0" }
 
 # Postgres interaction
-postgres = "0.19.4"
+tokio-postgres = "0.7.8"
 postgres-types = "0.2.4"
 
+# Async
+tokio = { version = "1.28.1", features = ["rt-multi-thread"] }
+futures = "0.3.28"
+
 # Error handling and reporting
 thiserror = "1.0.38"
 miette = { version = "5.5.0", features = ["fancy"] }

crates/cornucopia/Cargo.toml

@@ -64,8 +64,8 @@ pub fn run() -> Result<(), Error> {
 
     match action {
         Action::Live { url } => {
-            let mut client = conn::from_url(&url)?;
-            generate_live(&mut client, &queries_path, Some(&destination), settings)?;
+            let client = conn::from_url(&url)?;
+            generate_live(&client, &queries_path, Some(&destination), settings)?;
         }
         Action::Schema { schema_files } => {
             // Run the generate command. If the command is unsuccessful, cleanup Cornucopia's container

crates/cornucopia/src/cli.rs

@@ -1,27 +1,42 @@
-use postgres::{Client, Config, NoTls};
+use tokio::runtime::Runtime;
+use tokio_postgres::{Client, Config, NoTls};
 
 use self::error::Error;
 
 /// Creates a non-TLS connection from a URL.
 pub(crate) fn from_url(url: &str) -> Result<Client, Error> {
-    Ok(Client::connect(url, NoTls)?)
+    connect(url.parse()?)
 }
 
 /// Create a non-TLS connection to the container managed by Cornucopia.
 pub fn cornucopia_conn() -> Result<Client, Error> {
-    Ok(Config::new()
-        .user("postgres")
-        .password("postgres")
-        .host("127.0.0.1")
-        .port(5435)
-        .dbname("postgres")
-        .connect(NoTls)?)
+    connect(
+        Config::new()
+            .user("postgres")
+            .password("postgres")
+            .host("127.0.0.1")
+            .port(5435)
+            .dbname("postgres")
+            .clone(),
+    )
+}
+
+fn connect(config: Config) -> Result<Client, Error> {
+    let rt: &'static Runtime = Box::leak(Box::new(
+        Runtime::new().expect("Failed to start async Runtime"),
+    ));
+    let client = rt.block_on(async {
+        let (client, conn) = config.connect(NoTls).await.unwrap();
+        rt.spawn(conn);
+        client
+    });
+    Ok(client)
 }
 
 pub(crate) mod error {
     use miette::Diagnostic;
 
     #[derive(Debug, thiserror::Error, Diagnostic)]
     #[error("Couldn't establish a connection with the database.")]
-    pub struct Error(#[from] pub postgres::Error);
+    pub struct Error(#[from] pub tokio_postgres::Error);
 }

crates/cornucopia/src/conn.rs

@@ -16,7 +16,7 @@ pub mod container;
 
 use std::path::Path;
 
-use postgres::Client;
+use tokio_postgres::Client;
 
 use codegen::generate as generate_internal;
 use error::WriteOutputError;
@@ -43,7 +43,7 @@ pub struct CodegenSettings {
 /// the generated code will be written at that path. Code generation settings are
 /// set using the `settings` parameter.
 pub fn generate_live<P: AsRef<Path>>(
-    client: &mut Client,
+    client: &Client,
     queries_path: P,
     destination: Option<P>,
     settings: CodegenSettings,
@@ -84,9 +84,9 @@ pub fn generate_managed<P: AsRef<Path>>(
         .map(parse_query_module)
         .collect::<Result<_, parser::error::Error>>()?;
     container::setup(podman)?;
-    let mut client = conn::cornucopia_conn()?;
-    load_schema(&mut client, schema_files)?;
-    let prepared_modules = prepare(&mut client, modules)?;
+    let client = conn::cornucopia_conn()?;
+    load_schema(&client, schema_files)?;
+    let prepared_modules = prepare(&client, modules)?;
     let generated_code = generate_internal(prepared_modules, settings);
     container::cleanup(podman)?;
 

crates/cornucopia/src/lib.rs

@@ -1,7 +1,7 @@
 use std::path::Path;
 
 use miette::NamedSource;
-use postgres::Client;
+use tokio_postgres::Client;
 
 use crate::utils::db_err;
 
@@ -10,14 +10,14 @@ use self::error::Error;
 /// Loads PostgreSQL schemas into a database.
 ///
 /// Takes a list of file paths as parameter and loads them in their given order.
-pub fn load_schema<P: AsRef<Path>>(client: &mut Client, paths: &[P]) -> Result<(), Error> {
+pub fn load_schema<P: AsRef<Path>>(client: &Client, paths: &[P]) -> Result<(), Error> {
     for path in paths {
         let path = path.as_ref();
         let sql = std::fs::read_to_string(path).map_err(|err| Error::Io {
             path: path.to_string_lossy().to_string(),
             err,
         })?;
-        client.batch_execute(&sql).map_err(|err| {
+        futures::executor::block_on(client.batch_execute(&sql)).map_err(|err| {
             let msg = format!("{err:#}");
             let src = NamedSource::new(path.to_string_lossy(), sql);
             if let Some((position, msg, help)) = db_err(&err) {