use super::{HEADER, QueriesModule, queries, sql, write};
use crate::schema::{Table, rust_type};
#[derive(Debug, Clone, Copy)]
pub struct DbMod<'a> {
pub has_types: bool,
pub has_session: bool,
pub queries: QueriesModule<'a>,
pub serde: bool,
}
pub fn render_db_mod(tables: &[Table], db_mod: &DbMod<'_>) -> String {
let DbMod {
has_types,
has_session,
queries,
serde,
} = *db_mod;
let mut types: Vec<String> = tables
.iter()
.flat_map(|table| &table.columns)
.map(|column| column.rust_type.clone())
.collect();
types.extend(queries::mentioned_types(queries.analyzed()));
let mut output =
format!("{HEADER}\n//! Typed database access for the schema, generated by rok-sqlc.\n\n");
output.push_str(
"// The generated code requires these dependencies in the consuming crate\n\
// (SQLx 0.9 needs Rust 1.94 or newer):\n",
);
output.push_str(&dependency_hints(&types, serde));
output.push('\n');
let mut modules: Vec<&str> = tables
.iter()
.map(|table| table.module_name.as_str())
.collect();
modules.push("schema");
modules.push("search");
if serde {
modules.push("nullable");
}
if has_types {
modules.push("types");
}
if has_session {
modules.push("session");
}
if queries.is_declared() {
modules.push("queries");
}
modules.sort_unstable();
for module in modules {
output.push_str(&format!("pub mod {module};\n"));
}
if tables.iter().any(write::has_versioned_update) {
output.push_str(UPDATE_OUTCOME);
}
output
}
const UPDATE_OUTCOME: &str = "
/// What an `update` on a table with a version column did.
#[must_use]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UpdateOutcome<T> {
/// The row was at the expected version and is now updated, as stored.
Updated(T),
/// No row has the given primary key.
NotFound,
/// The row is at another version, so nothing changed.
VersionConflict,
}
";
fn dependency_hints(types: &[String], serde: bool) -> String {
let uses = |needle: &str| types.iter().any(|rust_type| rust_type.contains(needle));
let uses_uuid = uses("uuid::");
let uses_chrono = uses("chrono::") || uses(rust_type::TIME_TZ);
let uses_serde_json = uses("serde_json::");
let uses_rust_decimal = uses("rust_decimal::");
let mut features = String::from("\"runtime-tokio\", \"postgres\", \"derive\"");
if uses_uuid {
features.push_str(", \"uuid\"");
}
if uses_chrono {
features.push_str(", \"chrono\"");
}
if uses_serde_json {
features.push_str(", \"json\"");
}
if uses_rust_decimal {
features.push_str(", \"rust_decimal\"");
}
let crate_hint = |name: &str, version: &str| {
if serde {
format!("// {name} = {{ version = \"{version}\", features = [\"serde\"] }}\n")
} else {
format!("// {name} = \"{version}\"\n")
}
};
let mut hints = format!("// sqlx = {{ version = \"0.9\", features = [{features}] }}\n");
if serde {
hints.push_str("// serde = { version = \"1\", features = [\"derive\"] }\n");
}
if uses_uuid {
hints.push_str(&crate_hint("uuid", "1"));
}
if uses_chrono {
hints.push_str(&crate_hint("chrono", "0.4"));
}
if uses_serde_json {
hints.push_str("// serde_json = \"1\"\n");
}
if uses_rust_decimal {
hints.push_str(&crate_hint("rust_decimal", "1"));
}
hints
}
pub fn render_schema_snapshot(schema_sql: &str) -> String {
let hashes = sql::raw_hashes(schema_sql);
format!(
"{HEADER}\n//! The schema the code was generated from.\n\n/// Verbatim snapshot of the input schema; never executed.\npub const SCHEMA_SQL: &str = r{hashes}\"{schema_sql}\"{hashes};\n"
)
}
#[cfg(test)]
mod tests {
use super::super::test_support::{column, table};
use super::*;
fn db_mod(tables: &[Table], queries: QueriesModule<'_>, serde: bool) -> String {
render_db_mod(
tables,
&DbMod {
has_types: false,
has_session: false,
queries,
serde,
},
)
}
#[test]
fn db_mod_lists_modules_and_conditional_dependencies() {
let tables = vec![table(vec![column("id", "uuid::Uuid")], &["id"])];
let rendered = db_mod(&tables, QueriesModule::None, false);
assert!(rendered.contains("pub mod users;"));
assert!(rendered.contains("pub mod schema;"));
assert!(rendered.contains("\"uuid\""));
assert!(!rendered.contains("\"chrono\""));
assert!(rendered.contains("uuid = \"1\""));
assert!(rendered.contains(
"// sqlx = { version = \"0.9\", features = [\"runtime-tokio\", \"postgres\", \"derive\", \"uuid\"] }"
));
assert!(rendered.contains("(SQLx 0.9 needs Rust 1.94 or newer)"));
}
#[test]
fn db_mod_declares_queries_when_generated_or_kept() {
let tables = vec![table(vec![column("id", "i64")], &["id"])];
assert!(!db_mod(&tables, QueriesModule::None, false).contains("pub mod queries;"));
assert!(
db_mod(&tables, QueriesModule::Keep, false)
.contains("pub mod queries;\npub mod schema;\npub mod search;\npub mod users;")
);
assert!(db_mod(&tables, QueriesModule::Generate(&[]), false).contains("pub mod queries;"));
}
#[test]
fn db_mod_declares_the_nullable_helper_with_serde() {
let tables = vec![table(vec![column("id", "i64")], &["id"])];
assert!(!db_mod(&tables, QueriesModule::None, false).contains("pub mod nullable;"));
assert!(
db_mod(&tables, QueriesModule::None, true)
.contains("pub mod nullable;\npub mod schema;\npub mod search;\npub mod users;")
);
}
#[test]
fn db_mod_hints_rust_decimal_when_used() {
let tables = vec![table(vec![column("id", "rust_decimal::Decimal")], &["id"])];
let rendered = db_mod(&tables, QueriesModule::None, false);
assert!(rendered.contains("\"rust_decimal\""));
assert!(rendered.contains("rust_decimal = \"1\""));
assert!(!rendered.contains("\"uuid\""));
}
#[test]
fn snapshot_escapes_raw_string_delimiters() {
let schema = "CREATE TABLE tricky (x TEXT, marker TEXT); -- contains \"# inside\n";
let rendered = render_schema_snapshot(schema);
assert!(rendered.contains("r##\""));
assert!(rendered.ends_with("\"##;\n"));
assert!(rendered.contains(schema));
}
#[test]
fn snapshot_uses_single_hash_when_safe() {
let rendered = render_schema_snapshot("CREATE TABLE plain (x TEXT);\n");
assert!(rendered.contains("r#\""));
assert!(rendered.ends_with("\"#;\n"));
}
}