nahpu_db 0.1.1

A utility crate for handling database schema and models, auto-generating Rust structs from the NAHPU SQLite Drift schema.
//! # Build Script for dwc-converter
//!
//! This build script fetches a remote `.drift` file containing SQL table definitions,
//! parses it, and generates Rust structs for each table. The generated code is written
//! to `OUT_DIR/generated_models.rs` and is intended to be included in the main crate.
//!
//! ## Steps performed:
//! 1. Download the `.drift` file from a specified URL.
//! 2. Parse the SQL statements using `sqlparser`.
//! 3. Map SQL types to Rust types, handling nullability.
//! 4. Generate Rust structs with Serde support for each table.
//! 5. Write the generated code to the build output directory.
use heck::{ToPascalCase, ToSnakeCase};
use reqwest::blocking::get;
use sqlparser::ast::{ColumnOption, DataType, Statement};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

const NAHPU_TABLES_URL: &str =
    "https://raw.githubusercontent.com/nahpu/nahpu/main/lib/services/database/tables.drift";

const GENERATED_FILE: &str = "nahpu_sqlite.rs";

/// Maps SQL data types from the .drift file to corresponding Rust types.
///
/// # Arguments
/// * `sql_type` - The SQL data type parsed from the drift file.
/// * `is_nullable` - Whether the column is nullable.
///
/// # Returns
/// A `String` representing the Rust type, wrapped in `Option<>` if nullable.
fn sql_type_to_rust(sql_type: &DataType, is_nullable: bool) -> String {
    let base_type = match sql_type {
        DataType::Text | DataType::Varchar(_) | DataType::String(_) => "String".to_string(),
        DataType::Int(_) | DataType::Integer(_) => "i32".to_string(),
        DataType::Real | DataType::Float(_) | DataType::Double(_) => "f64".to_string(),
        DataType::Boolean => "bool".to_string(),
        _ => {
            // Default to String for any unknown or custom SQL types.
            eprintln!(
                "Warning: Unknown SQL type {:?}, defaulting to String.",
                sql_type
            );
            "String".to_string()
        }
    };

    if is_nullable {
        format!("Option<{}>", base_type)
    } else {
        base_type
    }
}

/// The drift definition is fetched from a NAHPU Github URL.
/// The TABLE definitions contains the schema and sql types.
/// ```sql
/// CREATE TABLE specimenPart (
///     id INT UNIQUE PRIMARY KEY AUTOINCREMENT, -- internal id
///     specimenUuid TEXT,
///     personnelId TEXT,
///     tissueID TEXT,
///     barcodeID TEXT,
///     type TEXT,
///     count TEXT,
///     treatment TEXT,
///     additionalTreatment TEXT,
///     dateTaken TEXT,
///     timeTaken TEXT,
///     pmi TEXT,
///     museumPermanent TEXT,
///     museumLoan TEXT,
///     remark TEXT,
///     FOREIGN KEY(specimenUuid) REFERENCES specimen(uuid),
///     FOREIGN KEY(personnelId) REFERENCES personnel(uuid)
/// );
///
/// listProject: SELECT uuid,name,created,lastAccessed FROM project;
/// ```
/// We need to convert the SQL types to Rust types and generate and ignore any non-table statements.
fn clean_drift_content(content: &str) -> String {
    // Include all CREATE TABLE statements and ignore others.
    content
        .split(';') // Split into individual statements
        .filter(|stmt| stmt.trim().to_uppercase().starts_with("CREATE TABLE"))
        .map(|stmt| format!("{};", stmt.trim())) // Add the semicolon back and ensure it's clean
        .collect()
}

use std::env;

fn write_rust_file(content: &str) {
    let out_dir = env::var("OUT_DIR").expect("OUT_DIR environment variable is not set");
    let dest_path = Path::new(&out_dir).join(GENERATED_FILE);
    let file = File::create(&dest_path).expect("Unable to create file");
    let mut writer = BufWriter::new(file);
    // Write auto-generated indicator and source link at the top of the file.
    let header = format!(
        r#"// This file is auto-generated by build.rs. Do not edit directly.
// It is derived from the NAHPU Drift schema (CREATE TABLE statements only).
// Source: {}
// Regenerate by running `cargo build`."#,
        NAHPU_TABLES_URL
    );
    writeln!(writer, "{}\n\n{}", header, content).expect("Unable to write data to file");
    println!("Generated Rust code written to {:?}", dest_path);
}

fn create_rust_code(drift_content: &str) -> String {
    let dialect = GenericDialect {}; // A generic SQL dialect for parsing.
    let ast =
        Parser::parse_sql(&dialect, drift_content).expect("Failed to parse drift file as SQL");

    let mut rust_code = String::new();

    rust_code.push_str("use serde::{Deserialize, Serialize};\n\n");

    for statement in ast {
        if let Statement::CreateTable(create_table) = statement {
            let table_name = match create_table.name.0.last().unwrap() {
                sqlparser::ast::ObjectNamePart::Identifier(ident) => ident.value.clone(),
                _ => panic!("Expected identifier"),
            };
            let struct_name = table_name.to_pascal_case();

            // Add Serde attributes for JSON serialization/deserialization.
            rust_code.push_str("#[derive(Serialize, Deserialize, Debug, Clone)]\n");
            rust_code.push_str("#[serde(rename_all = \"camelCase\")]\n");
            rust_code.push_str(&format!("pub struct {} {{\n", struct_name));

            for col in create_table.columns {
                let col_name = col.name.value.clone();
                let field_name_snake = col_name.to_snake_case();

                // A column is considered nullable if it doesn't have "NOT NULL".
                let is_nullable = !col
                    .options
                    .iter()
                    .any(|opt| matches!(opt.option, ColumnOption::NotNull));
                let rust_type = sql_type_to_rust(&col.data_type, is_nullable);

                // Handle Rust keywords by renaming the field and adding a serde attribute.
                if field_name_snake == "type" {
                    // Add serde rename attribute to map the JSON 'type' field
                    // to the Rust 'type_' field, overriding the struct-level rename_all.
                    rust_code.push_str(&format!("    #[serde(rename = \"type\")]\n"));
                    rust_code.push_str(&format!("    pub type_: {},\n", rust_type));
                } else {
                    rust_code.push_str(&format!("    pub {}: {},\n", field_name_snake, rust_type));
                }
            }

            rust_code.push_str("}\n\n");
        }
    }
    rust_code
}

/// Main entry point for the build script.
///
/// Downloads the drift file, parses it, generates Rust structs, and writes them to OUT_DIR.
fn main() {
    let drift_file_url =
        "https://raw.githubusercontent.com/nahpu/nahpu/main/lib/services/database/tables.drift";
    println!("cargo:rerun-if-changed=build.rs"); // Re-run if build script itself changes.
    println!("Fetching drift file from: {}", drift_file_url);
    // Fetch the drift file content from the URL.
    let drift_content = match get(drift_file_url) {
        Ok(response) => response.text().expect("Failed to read response text"),
        Err(e) => {
            println!(
                "cargo:warning=Failed to fetch drift file: {}. Skipping generation.",
                e
            );
            return;
        }
    };

    let cleaned_content = clean_drift_content(&drift_content);
    // Close the 'models' module.
    let rust_code = create_rust_code(&cleaned_content);
    write_rust_file(&rust_code);
}