Add support for type mapping

Similar to Kanel's customTypeMap.

One use case for this is allowing one to treat PG extension types as
some other types that they behave as (e.g., map public.ulid to String).

Fixes #246

Author: duarten

benches/codegen.rs

@@ -17,6 +17,7 @@ fn bench(c: &mut Criterion) {
                     gen_sync: true,
                     gen_async: false,
                     derive_ser: true,
+                    config: Default::default(),
                 },
             )
             .unwrap()
@@ -32,6 +33,7 @@ fn bench(c: &mut Criterion) {
                     gen_sync: true,
                     gen_async: false,
                     derive_ser: true,
+                    config: Default::default(),
                 },
             )
             .unwrap()

benches/execution/diesel_benches.rs

@@ -164,10 +164,7 @@ pub fn bench_insert(b: &mut Bencher, conn: &mut PgConnection, size: usize) {
     };
     let insert = &insert;
 
-    b.iter(|| {
-        let insert = insert;
-        insert(conn)
-    })
+    b.iter(|| insert(conn))
 }
 
 pub fn loading_associations_sequentially(b: &mut Bencher, conn: &mut PgConnection) {

crates/cornucopia/Cargo.toml

@@ -33,3 +33,7 @@ heck = "0.4.0"
 
 # Order-preserving map to work around borrowing issues
 indexmap = "2.0.2"
+
+# Config handling
+serde = { version = "1.0.203", features = ["derive"] }
+toml = "0.8.14"

crates/cornucopia/src/cli.rs

@@ -1,8 +1,12 @@
-use std::path::PathBuf;
+use miette::Diagnostic;
+use std::{fs, path::PathBuf};
+use thiserror::Error as ThisError;
 
 use clap::{Parser, Subcommand};
 
-use crate::{conn, container, error::Error, generate_live, generate_managed, CodegenSettings};
+use crate::{
+    config::Config, conn, container, error::Error, generate_live, generate_managed, CodegenSettings,
+};
 
 /// Command line interface to interact with Cornucopia SQL.
 #[derive(Parser, Debug)]
@@ -28,6 +32,13 @@ struct Args {
     /// Derive serde's `Serialize` trait for generated types.
     #[clap(long)]
     serialize: bool,
+    /// The location of the configuration file.
+    #[clap(short, long, default_value = default_config_path())]
+    config: PathBuf,
+}
+
+const fn default_config_path() -> &'static str {
+    "cornucopia.toml"
 }
 
 #[derive(Debug, Subcommand)]
@@ -44,8 +55,26 @@ enum Action {
     },
 }
 
+/// Enumeration of the errors reported by the CLI.
+#[derive(ThisError, Debug, Diagnostic)]
+pub enum CliError {
+    /// An error occurred while loading the configuration file.
+    #[error("Could not load config `{path}`: ({err})")]
+    MissingConfig { path: String, err: std::io::Error },
+    /// An error occurred while parsing the configuration file.
+    #[error("Could not parse config `{path}`: ({err})")]
+    ConfigContents {
+        path: String,
+        err: Box<dyn std::error::Error + Send + Sync>,
+    },
+    /// An error occurred while running the CLI.
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    Internal(#[from] Error),
+}
+
 // Main entrypoint of the CLI. Parses the args and calls the appropriate routines.
-pub fn run() -> Result<(), Error> {
+pub fn run() -> Result<(), CliError> {
     let Args {
         podman,
         queries_path,
@@ -54,17 +83,40 @@ pub fn run() -> Result<(), Error> {
         sync,
         r#async,
         serialize,
+        config,
     } = Args::parse();
 
+    let config = match fs::read_to_string(config.as_path()) {
+        Ok(contents) => match toml::from_str(&contents) {
+            Ok(config) => config,
+            Err(err) => {
+                return Err(CliError::ConfigContents {
+                    path: config.to_string_lossy().into_owned(),
+                    err: err.into(),
+                });
+            }
+        },
+        Err(err) => {
+            if config.as_path().as_os_str() != default_config_path() {
+                return Err(CliError::MissingConfig {
+                    path: config.to_string_lossy().into_owned(),
+                    err,
+                });
+            } else {
+                Config::default()
+            }
+        }
+    };
     let settings = CodegenSettings {
         gen_async: r#async || !sync,
         gen_sync: sync,
         derive_ser: serialize,
+        config,
     };
 
     match action {
         Action::Live { url } => {
-            let mut client = conn::from_url(&url)?;
+            let mut client = conn::from_url(&url).map_err(|e| CliError::Internal(e.into()))?;
             generate_live(&mut client, &queries_path, Some(&destination), settings)?;
         }
         Action::Schema { schema_files } => {
@@ -77,7 +129,7 @@ pub fn run() -> Result<(), Error> {
                 settings,
             ) {
                 container::cleanup(podman).ok();
-                return Err(e);
+                return Err(CliError::Internal(e));
             }
         }
     };

crates/cornucopia/src/config.rs

@@ -0,0 +1,12 @@
+//! Configuration for Cornucopia.
+
+use std::collections::HashMap;
+
+use serde::Deserialize;
+
+/// Configuration for Cornucopia.
+#[derive(Clone, Deserialize, Default, Debug)]
+pub struct Config {
+    /// Contains a map of what given type should map to.
+    pub custom_type_map: HashMap<String, String>,
+}

crates/cornucopia/src/lib.rs

@@ -1,5 +1,6 @@
 mod cli;
 mod codegen;
+mod config;
 mod error;
 mod load_schema;
 mod parser;
@@ -16,6 +17,7 @@ pub mod container;
 
 use std::path::Path;
 
+use config::Config;
 use postgres::Client;
 
 use codegen::generate as generate_internal;
@@ -31,11 +33,12 @@ pub use error::Error;
 pub use load_schema::load_schema;
 
 /// Struct containing the settings for code generation.
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 pub struct CodegenSettings {
     pub gen_async: bool,
     pub gen_sync: bool,
     pub derive_ser: bool,
+    pub config: Config,
 }
 
 /// Generates Rust queries from PostgreSQL queries located at `queries_path`,
@@ -54,7 +57,7 @@ pub fn generate_live<P: AsRef<Path>>(
         .map(parse_query_module)
         .collect::<Result<_, parser::error::Error>>()?;
     // Generate
-    let prepared_modules = prepare(client, modules)?;
+    let prepared_modules = prepare(client, modules, settings.clone())?;
     let generated_code = generate_internal(prepared_modules, settings);
     // Write
     if let Some(d) = destination {
@@ -86,7 +89,7 @@ pub fn generate_managed<P: AsRef<Path>>(
     container::setup(podman)?;
     let mut client = conn::cornucopia_conn()?;
     load_schema(&mut client, schema_files)?;
-    let prepared_modules = prepare(&mut client, modules)?;
+    let prepared_modules = prepare(&mut client, modules, settings.clone())?;
     let generated_code = generate_internal(prepared_modules, settings);
     container::cleanup(podman)?;
 

crates/cornucopia/src/prepare_queries.rs

@@ -9,10 +9,9 @@ use crate::{
     codegen::GenCtx,
     parser::{Module, NullableIdent, Query, Span, TypeAnnotation},
     read_queries::ModuleInfo,
-    type_registrar::CornucopiaType,
-    type_registrar::TypeRegistrar,
+    type_registrar::{CornucopiaType, TypeRegistrar},
     utils::KEYWORD,
-    validation,
+    validation, CodegenSettings,
 };
 
 use self::error::Error;
@@ -226,8 +225,12 @@ impl PreparedModule {
 }
 
 /// Prepares all modules
-pub(crate) fn prepare(client: &mut Client, modules: Vec<Module>) -> Result<Preparation, Error> {
-    let mut registrar = TypeRegistrar::default();
+pub(crate) fn prepare(
+    client: &mut Client,
+    modules: Vec<Module>,
+    settings: CodegenSettings,
+) -> Result<Preparation, Error> {
+    let mut registrar = TypeRegistrar::new(settings.config.custom_type_map);
     let mut tmp = Preparation {
         modules: Vec::new(),
         types: IndexMap::new(),
@@ -244,16 +247,12 @@ pub(crate) fn prepare(client: &mut Client, modules: Vec<Module>) -> Result<Prepa
     }
 
     // Prepare types grouped by schema
-    for ((schema, name), ty) in &registrar.types {
-        if let Some(ty) = prepare_type(&registrar, name, ty, &declared) {
-            match tmp.types.entry(schema.clone()) {
-                Entry::Occupied(mut entry) => {
-                    entry.get_mut().push(ty);
-                }
-                Entry::Vacant(entry) => {
-                    entry.insert(vec![ty]);
-                }
-            }
+    for (schema_key, ty) in registrar.types() {
+        if let Some(ty) = prepare_type(&registrar, schema_key.name, ty, &declared) {
+            tmp.types
+                .entry(schema_key.schema.to_owned())
+                .or_default()
+                .push(ty);
         }
     }
     Ok(tmp)
@@ -301,7 +300,9 @@ fn prepare_type(
                     })
                     .collect(),
             ),
-            _ => unreachable!(),
+            _ => {
+                return None;
+            }
         };
         Some(PreparedType {
             name: name.to_string(),