rok-sqlc 0.2.0

Generate SQLx Rust database modules from SQL schema files
Documentation
//! The root `db/mod.rs`, with its dependency hints, and the `db/schema.rs`
//! snapshot.

use super::{HEADER, QueriesModule, queries, sql, write};
use crate::schema::{Table, rust_type};

/// What `db/mod.rs` declares besides the resource modules.
#[derive(Debug, Clone, Copy)]
pub struct DbMod<'a> {
    /// Whether `types.rs` is generated.
    pub has_types: bool,
    /// Whether `session.rs` is generated.
    pub has_session: bool,
    pub queries: QueriesModule<'a>,
    /// Whether `derive = ["serde"]` is on, which adds `nullable.rs`.
    pub serde: bool,
}

/// Renders `db/mod.rs`: the dependency hints, a module per table of
/// `tables` (generated or kept) and per shared module of `db_mod`, and
/// `UpdateOutcome` when a table's `update` uses it.
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');
    // Sorted the way rustfmt sorts `mod` declarations.
    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
}

/// The result of every versioned `update` (a table with a version column),
/// defined once in `db/mod.rs`.
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,
}
";

/// The `Cargo.toml` lines the generated code needs. With `serde`, the crates
/// whose types appear in serialized items need their `serde` feature.
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::");
    // `PgTimeTz` is encoded and decoded through SQLx's `chrono` feature, and
    // its default parameters are `chrono` types.
    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::*;

    /// `db/mod.rs` for `tables` with no `types.rs` or `session.rs`.
    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"));
    }
}