Skip to main content

cornucopia/
lib.rs

1mod cli;
2mod codegen;
3mod error;
4mod load_schema;
5mod parser;
6mod prepare_queries;
7mod read_queries;
8mod type_registrar;
9mod utils;
10mod validation;
11
12pub mod config;
13/// Helpers to establish connections to database instances.
14pub mod conn;
15/// High-level interfaces to work with Cornucopia's container manager.
16pub mod container;
17
18/// Re-export of cargo_toml for working with manifest configurations
19pub use cargo_toml;
20use std::path::Path;
21
22use tokio_postgres::Client;
23
24use config::Config;
25use parser::{Module, parse_query_module};
26use prepare_queries::prepare;
27use read_queries::read_query_modules;
28
29#[doc(hidden)]
30pub use cli::run;
31
32pub use error::Error;
33use error::Warning;
34
35pub use load_schema::load_schema;
36
37fn warn_if_no_queries(modules: &[Module]) {
38    if modules.iter().all(|m| m.queries.is_empty()) {
39        Warning::NoQueries.emit();
40    }
41}
42
43/// Generates Rust queries from PostgreSQL queries located at `queries_path`,
44/// using a live database managed by you. Code generation settings are
45/// set using the `config` parameter.
46pub fn gen_live(client: &Client, config: Config) -> Result<(), Error> {
47    // Read
48    let modules = read_query_modules(config.queries.as_ref(), &config)?
49        .into_iter()
50        .map(parse_query_module)
51        .collect::<Result<Vec<_>, parser::error::Error>>()
52        .map_err(Box::new)?;
53
54    warn_if_no_queries(&modules);
55
56    // Generate
57    let prepared_modules = prepare(client, modules, &config).map_err(Box::new)?;
58    let generated = codegen::generate(prepared_modules, &config);
59
60    // Write
61    generated.persist(config.destination, config.static_files)?;
62
63    Ok(())
64}
65
66/// Generates Rust queries from PostgreSQL queries located at `queries_path`, using
67/// a container managed by cornucopia. The database schema is created using `schema_files`.
68/// Code generation settings are set using the `config` parameter.
69///
70/// By default, the container manager is Docker, but Podman can be used by setting the
71/// `podman` parameter to `true`.
72pub fn gen_managed<P: AsRef<Path>>(schema_files: &[P], config: Config) -> Result<(), Error> {
73    // Read
74    let modules = read_query_modules(config.queries.as_ref(), &config)?
75        .into_iter()
76        .map(parse_query_module)
77        .collect::<Result<Vec<_>, parser::error::Error>>()
78        .map_err(Box::new)?;
79
80    warn_if_no_queries(&modules);
81
82    container::setup(
83        config.podman,
84        &config.container_image,
85        config.container_wait,
86    )?;
87
88    let client = conn::cornucopia_conn()?;
89    load_schema(&client, schema_files).map_err(Box::new)?;
90    let prepared_modules = prepare(&client, modules, &config).map_err(Box::new)?;
91    let generated = codegen::generate(prepared_modules, &config);
92    container::cleanup(config.podman)?;
93
94    // Write
95    generated.persist(config.destination, config.static_files)?;
96
97    Ok(())
98}
99
100/// Generates Rust queries from PostgreSQL queries located at `queries_path`, using
101/// a temporary database created on an existing PostgreSQL server. The database schema
102/// is created using `schema_files`. Code generation settings are set using the `config` parameter.
103///
104/// This function creates a temporary database on the specified server, loads the schema,
105/// generates the code, and optionally drops the temporary database based on the `keep_db` parameter.
106pub fn gen_fresh<P: AsRef<Path>>(
107    url: &str,
108    db_name: &str,
109    schema_files: &[P],
110    search_path: Option<&str>,
111    keep_db: bool,
112    config: Config,
113) -> Result<(), Error> {
114    let modules = read_query_modules(config.queries.as_ref(), &config)?
115        .into_iter()
116        .map(parse_query_module)
117        .collect::<Result<Vec<_>, parser::error::Error>>()
118        .map_err(Box::new)?;
119
120    warn_if_no_queries(&modules);
121
122    let server_client = conn::from_url(url)?;
123
124    let create_db_query = format!("CREATE DATABASE \"{db_name}\"");
125    futures::executor::block_on(server_client.execute(&create_db_query, &[]))
126        .map_err(conn::error::Error)?;
127
128    let db_url = if url.contains('?') {
129        format!("{url}&dbname={db_name}")
130    } else if url.ends_with('/') {
131        format!("{url}{db_name}?")
132    } else {
133        format!("{url}/{db_name}?")
134    };
135
136    let generation_result = (|| -> Result<(), Error> {
137        let db_client = conn::from_url(&db_url)?;
138
139        if let Some(search_path) = search_path {
140            conn::set_search_path(&db_client, search_path)?;
141        }
142
143        load_schema(&db_client, schema_files).map_err(Box::new)?;
144
145        let prepared_modules = prepare(&db_client, modules, &config).map_err(Box::new)?;
146        let generated = codegen::generate(prepared_modules, &config);
147
148        generated.persist(config.destination, config.static_files)?;
149
150        Ok(())
151    })();
152
153    if !keep_db {
154        let drop_db_query = format!("DROP DATABASE \"{db_name}\"");
155        if let Err(e) = futures::executor::block_on(server_client.execute(&drop_db_query, &[])) {
156            eprintln!("Warning: Failed to drop temporary database '{db_name}': {e}");
157        }
158    }
159
160    generation_result
161}