rok-sqlc 0.3.0

//! The root `db/mod.rs`, with its dependency hints, and the `db/schema.rs`
//! snapshot.

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

/// 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>,
    /// What `derive` asks for: `serde` adds `nullable.rs`, and `validate`
    /// needs `validator`.
    pub derives: DeriveOptions,
    /// The frameworks `http/` has an extractor for; none leaves `http/` out.
    pub extractors: &'a [Extractor],
}

/// 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,
        derives,
        extractors,
    } = *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",
    );
    let ids: Vec<(IdGeneration, BaseType)> = tables
        .iter()
        .flat_map(|table| &table.columns)
        .filter_map(|column| Some((column.id_generation?, column.base_type())))
        .collect();
    output.push_str(&dependency_hints(&types, &ids, derives, extractors));
    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 derives.serde {
        modules.push("nullable");
    }
    if has_types {
        modules.push("types");
    }
    if !extractors.is_empty() {
        modules.push("http");
    }
    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.
/// With `validate`, `validator`; with `extractors`, `serde_json` (to read
/// serde's errors) and each framework, and for axum `serde_path_to_error`,
/// whose error wraps serde's inside axum's rejection.
fn dependency_hints(
    types: &[String],
    ids: &[(IdGeneration, BaseType)],
    derives: DeriveOptions,
    extractors: &[Extractor],
) -> String {
    let serde = derives.serde;
    let json_bodies = !extractors.is_empty();
    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 uses_ipnetwork = uses("ipnetwork::");
    let uses_mac_address = uses("mac_address::");

    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\"");
    }
    if uses_ipnetwork {
        features.push_str(", \"ipnetwork\"");
    }
    if uses_mac_address {
        features.push_str(", \"mac_address\"");
    }

    let uses_ids = |wanted: IdGeneration| ids.iter().any(|(ids, _)| *ids == wanted);
    let uuid_features: Vec<&str> = [(IdGeneration::UuidV4, "v4"), (IdGeneration::UuidV7, "v7")]
        .into_iter()
        .filter(|(ids, _)| uses_ids(*ids))
        .map(|(_, feature)| feature)
        .collect();
    let crate_hint = |name: &str, version: &str| crate_hint_with(name, version, serde, &[]);
    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_with("uuid", "1", serde, &uuid_features));
    }
    if uses_chrono {
        hints.push_str(&crate_hint("chrono", "0.4"));
    }
    if uses_serde_json || json_bodies {
        hints.push_str("//   serde_json = \"1\"\n");
    }
    if uses_rust_decimal {
        hints.push_str(&crate_hint("rust_decimal", "1"));
    }
    if uses_ipnetwork {
        hints.push_str(&crate_hint("ipnetwork", "0.21"));
    }
    if uses_mac_address {
        hints.push_str(&crate_hint("mac_address", "1"));
    }
    // ULIDs are made only in `generate()`, never serialized as such.
    if uses_ids(IdGeneration::Ulid) {
        let on_uuid = ids
            .iter()
            .any(|(ids, base)| *ids == IdGeneration::Ulid && *base == BaseType::Uuid);
        let features: &[&str] = if on_uuid { &["uuid"] } else { &[] };
        hints.push_str(&crate_hint_with("ulid", "1", false, features));
    }
    if derives.validate {
        hints.push_str("//   validator = { version = \"0.21\", features = [\"derive\"] }\n");
    }
    if extractors.contains(&Extractor::Axum) {
        hints.push_str("//   axum = \"0.8\"\n//   serde_path_to_error = \"0.1\"\n");
    }
    if extractors.contains(&Extractor::ActixWeb) {
        hints.push_str("//   actix-web = \"4\"\n");
    }
    hints
}

/// One crate's `Cargo.toml` line: its `features`, and `serde` too when the
/// serde derive is on.
fn crate_hint_with(name: &str, version: &str, serde: bool, features: &[&str]) -> String {
    let features: Vec<String> = features
        .iter()
        .copied()
        .chain(serde.then_some("serde"))
        .map(|feature| format!("\"{feature}\""))
        .collect();
    if features.is_empty() {
        format!("//   {name} = \"{version}\"\n")
    } else {
        format!(
            "//   {name} = {{ version = \"{version}\", features = [{}] }}\n",
            features.join(", ")
        )
    }
}

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,
                derives: DeriveOptions {
                    serde,
                    validate: false,
                },
                extractors: &[],
            },
        )
    }

    #[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_validator_with_validate() {
        let tables = vec![table(vec![column("id", "i64")], &["id"])];
        let hint = "//   serde = { version = \"1\", features = [\"derive\"] }\n//   validator = { version = \"0.21\", features = [\"derive\"] }\n";
        let render = |validate| {
            render_db_mod(
                &tables,
                &DbMod {
                    has_types: false,
                    has_session: false,
                    queries: QueriesModule::None,
                    derives: DeriveOptions {
                        serde: true,
                        validate,
                    },
                    extractors: &[],
                },
            )
        };
        assert!(render(true).contains(hint), "{}", render(true));
        assert!(!render(false).contains("validator"));
    }

    #[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"));
    }
}