nautilus_codegen/

generator.rs

//! Code generator for Nautilus models, delegates, and builders.

use heck::ToSnakeCase;
use nautilus_schema::ir::{ModelIr, ResolvedFieldType, SchemaIr};
use serde::Serialize;
use std::collections::{HashMap, HashSet};
use tera::{Context, Tera};

use crate::type_helpers::field_to_rust_type;

pub static TEMPLATES: std::sync::LazyLock<Tera> = std::sync::LazyLock::new(|| {
    let mut tera = Tera::default();
    tera.add_raw_templates(vec![
        (
            "columns_struct.tera",
            include_str!("../templates/rust/columns_struct.tera"),
        ),
        (
            "column_impl.tera",
            include_str!("../templates/rust/column_impl.tera"),
        ),
        ("create.tera", include_str!("../templates/rust/create.tera")),
        (
            "create_many.tera",
            include_str!("../templates/rust/create_many.tera"),
        ),
        (
            "delegate.tera",
            include_str!("../templates/rust/delegate.tera"),
        ),
        ("delete.tera", include_str!("../templates/rust/delete.tera")),
        ("enum.tera", include_str!("../templates/rust/enum.tera")),
        (
            "find_many.tera",
            include_str!("../templates/rust/find_many.tera"),
        ),
        (
            "from_row_impl.tera",
            include_str!("../templates/rust/from_row_impl.tera"),
        ),
        (
            "model_file.tera",
            include_str!("../templates/rust/model_file.tera"),
        ),
        (
            "model_struct.tera",
            include_str!("../templates/rust/model_struct.tera"),
        ),
        ("update.tera", include_str!("../templates/rust/update.tera")),
        (
            "composite_type.tera",
            include_str!("../templates/rust/composite_type.tera"),
        ),
    ])
    .expect("embedded Rust templates must parse");
    tera
});

fn render(template: &str, ctx: &Context) -> String {
    TEMPLATES
        .render(template, ctx)
        .unwrap_or_else(|e| panic!("template rendering failed for '{}': {:?}", template, e))
}

/// Template context for a single model field in the Rust codegen backend.
///
/// This struct is intentionally separate from [`PythonFieldContext`] in
/// `python/generator.rs`: the two backends expose different template
/// variables (Rust needs `rust_type` / `column_type`; Python needs
/// `python_type` / `base_type` / `is_enum` / `has_default` / `default`) and
/// are expected to evolve independently.
#[derive(Debug, Clone, Serialize)]
struct FieldContext {
    name: String,
    db_name: String,
    rust_type: String,
    column_type: String,
    is_array: bool,
    index: usize,
    is_pk: bool,
    /// `true` when the field maps to an `Option<T>` Rust type
    /// (i.e. the schema field is not required and is not a relation).
    is_optional: bool,
    /// `true` when the field has `@updatedAt` — auto-defaults to `now()` if not provided.
    is_updated_at: bool,
    /// `true` when the field is a `@computed` generated column (read-only from client side).
    is_computed: bool,
}

/// Serialisable (logical_name, db_name) pair for primary-key fields.
/// Used in templates to generate cursor predicate slices.
#[derive(Debug, Clone, Serialize)]
struct PkFieldContext {
    /// Snake-case logical name — used as the cursor map key in generated code.
    name: String,
    /// Database column name — used to build the `table__db_col` column reference.
    db_name: String,
}

#[derive(Debug, Clone, Serialize)]
struct RelationContext {
    field_name: String,
    target_model: String,
    target_table: String,
    is_array: bool,
    fields: Vec<String>,
    references: Vec<String>,
    fields_db: Vec<String>,
    references_db: Vec<String>,
    target_scalar_fields: Vec<FieldContext>,
}

/// Generate complete code for a model (struct, impls, delegate, builders).
///
/// `is_async` determines whether the generated delegate methods and internal
/// builders use `async fn`/`.await` (`true`) or blocking sync wrappers (`false`).
pub fn generate_model(model: &ModelIr, ir: &SchemaIr, is_async: bool) -> String {
    let mut context = Context::new();

    context.insert("model_name", &model.logical_name);
    context.insert("table_name", &model.db_name);
    context.insert("delegate_name", &format!("{}Delegate", model.logical_name));
    context.insert("columns_name", &format!("{}Columns", model.logical_name));
    context.insert("find_many_name", &format!("{}FindMany", model.logical_name));
    context.insert("create_name", &format!("{}Create", model.logical_name));
    context.insert(
        "create_many_name",
        &format!("{}CreateMany", model.logical_name),
    );
    context.insert("entry_name", &format!("{}CreateEntry", model.logical_name));
    context.insert("update_name", &format!("{}Update", model.logical_name));
    context.insert("delete_name", &format!("{}Delete", model.logical_name));

    let pk_field_names = model.primary_key.fields();
    context.insert("primary_key_fields", &pk_field_names);

    let pk_fields_with_db: Vec<PkFieldContext> = pk_field_names
        .iter()
        .filter_map(|logical| {
            model
                .scalar_fields()
                .find(|f| f.logical_name.as_str() == *logical)
                .map(|f| PkFieldContext {
                    name: f.logical_name.to_snake_case(),
                    db_name: f.db_name.clone(),
                })
        })
        .collect();
    context.insert("pk_fields_with_db", &pk_fields_with_db);

    // --- Single-pass scalar field context building ---
    let mut enum_imports = HashSet::new();
    let mut composite_type_imports = HashSet::new();

    let mut scalar_fields: Vec<FieldContext> = Vec::new();
    let mut create_fields: Vec<FieldContext> = Vec::new();
    let mut updated_at_fields: Vec<FieldContext> = Vec::new();

    for (idx, field) in model.scalar_fields().enumerate() {
        // --- Import tracking ---
        match &field.field_type {
            ResolvedFieldType::Enum { enum_name } => {
                if ir.enums.contains_key(enum_name) {
                    enum_imports.insert(enum_name.clone());
                }
            }
            ResolvedFieldType::CompositeType { type_name } => {
                if ir.composite_types.contains_key(type_name) {
                    composite_type_imports.insert(type_name.clone());
                }
            }
            _ => {}
        }

        // --- Computed values ---
        let column_type = match &field.field_type {
            ResolvedFieldType::Scalar(scalar) => scalar.rust_type().to_string(),
            ResolvedFieldType::Enum { enum_name } => enum_name.clone(),
            _ => String::new(),
        };
        let is_pk = pk_field_names.contains(&field.logical_name.as_str());

        let field_ctx = FieldContext {
            name: field.logical_name.to_snake_case(),
            db_name: field.db_name.clone(),
            rust_type: field_to_rust_type(field),
            column_type,
            is_array: field.is_array,
            index: idx,
            is_pk,
            is_optional: !field.is_required && !field.is_array,
            is_updated_at: field.is_updated_at,
            is_computed: field.computed.is_some(),
        };

        create_fields.push(field_ctx.clone());

        // Updated-at fields
        if field.is_updated_at {
            updated_at_fields.push(field_ctx.clone());
        }

        scalar_fields.push(field_ctx);
    }

    // --- Relation imports ---
    let mut relation_imports = HashSet::new();
    for field in model.relation_fields() {
        if let ResolvedFieldType::Relation(rel) = &field.field_type {
            relation_imports.insert(rel.target_model.clone());
        }
    }

    context.insert("has_enums", &!enum_imports.is_empty());
    context.insert(
        "enum_imports",
        &enum_imports.into_iter().collect::<Vec<_>>(),
    );
    context.insert("has_relations", &!relation_imports.is_empty());
    context.insert(
        "relation_imports",
        &relation_imports.into_iter().collect::<Vec<_>>(),
    );
    context.insert("has_composite_types", &!composite_type_imports.is_empty());
    context.insert(
        "composite_type_imports",
        &composite_type_imports.into_iter().collect::<Vec<_>>(),
    );

    let relation_fields: Vec<FieldContext> = model
        .relation_fields()
        .map(|field| FieldContext {
            name: field.logical_name.to_snake_case(),
            db_name: field.db_name.clone(),
            rust_type: field_to_rust_type(field),
            column_type: String::new(),
            is_array: field.is_array,
            index: 0,
            is_pk: false,
            is_optional: true,
            is_updated_at: false,
            is_computed: false,
        })
        .collect();

    let relations: Vec<RelationContext> = model
        .relation_fields()
        .filter_map(|field| {
            if let ResolvedFieldType::Relation(rel) = &field.field_type {
                if let Some(target_model) = ir.models.get(&rel.target_model) {
                    let target_pk_names = target_model.primary_key.fields();
                    let target_scalar_fields: Vec<FieldContext> = target_model
                        .scalar_fields()
                        .enumerate()
                        .map(|(idx, f)| {
                            let column_type = match &f.field_type {
                                ResolvedFieldType::Scalar(scalar) => scalar.rust_type().to_string(),
                                ResolvedFieldType::Enum { enum_name } => enum_name.clone(),
                                _ => String::new(),
                            };
                            let f_is_pk = target_pk_names.contains(&f.logical_name.as_str());
                            FieldContext {
                                name: f.logical_name.to_snake_case(),
                                db_name: f.db_name.clone(),
                                rust_type: field_to_rust_type(f),
                                column_type,
                                is_array: f.is_array,
                                index: idx,
                                is_pk: f_is_pk,
                                is_optional: !f.is_required && !f.is_array,
                                is_updated_at: f.is_updated_at,
                                is_computed: f.computed.is_some(),
                            }
                        })
                        .collect();

                    let (fields, references) = if rel.fields.is_empty() {
                        let inverse = target_model.relation_fields().find(|f| {
                            if let ResolvedFieldType::Relation(inv_rel) = &f.field_type {
                                inv_rel.target_model == model.logical_name
                            } else {
                                false
                            }
                        });

                        if let Some(inverse_field) = inverse {
                            if let ResolvedFieldType::Relation(inv_rel) = &inverse_field.field_type
                            {
                                // Swap fields and references for the many-side
                                (inv_rel.references.clone(), inv_rel.fields.clone())
                            } else {
                                (vec![], vec![])
                            }
                        } else {
                            (vec![], vec![])
                        }
                    } else {
                        (rel.fields.clone(), rel.references.clone())
                    };

                    let fields_db: Vec<String> = fields
                        .iter()
                        .filter_map(|logical_name| {
                            model
                                .fields
                                .iter()
                                .find(|f| &f.logical_name == logical_name)
                                .map(|f| f.db_name.clone())
                        })
                        .collect();

                    let references_db: Vec<String> = references
                        .iter()
                        .filter_map(|logical_name| {
                            target_model
                                .fields
                                .iter()
                                .find(|f| &f.logical_name == logical_name)
                                .map(|f| f.db_name.clone())
                        })
                        .collect();

                    Some(RelationContext {
                        field_name: field.logical_name.to_snake_case(),
                        target_model: rel.target_model.clone(),
                        target_table: target_model.db_name.clone(),
                        is_array: field.is_array,
                        fields,
                        references,
                        fields_db,
                        references_db,
                        target_scalar_fields,
                    })
                } else {
                    None
                }
            } else {
                None
            }
        })
        .collect();

    context.insert("scalar_fields", &scalar_fields);
    context.insert("relation_fields", &relation_fields);
    context.insert("relations", &relations);
    context.insert("create_fields", &create_fields);
    context.insert("updated_at_fields", &updated_at_fields);
    context.insert("all_scalar_fields", &scalar_fields);
    context.insert("is_async", &is_async);

    render("model_file.tera", &context)
}

/// Generate all models from a schema IR.
///
/// `is_async` is forwarded to every [`generate_model`] call.
pub fn generate_all_models(ir: &SchemaIr, is_async: bool) -> HashMap<String, String> {
    let mut generated = HashMap::new();

    for (model_name, model_ir) in &ir.models {
        let code = generate_model(model_ir, ir, is_async);
        generated.insert(model_name.clone(), code);
    }

    generated
}