prax_typegen/

zod.rs

//! Zod schema generator.

use convert_case::{Case, Casing};
use prax_schema::{CompositeType, Enum, Field, FieldType, Model, Schema, TypeModifier, View};

use crate::mapping::TypeMapper;

/// Generates Zod schemas from a Prax schema.
pub struct ZodGenerator;

impl ZodGenerator {
    /// Generate all Zod schemas for the given schema.
    pub fn generate(schema: &Schema) -> String {
        let mut out = String::with_capacity(4096);
        out.push_str("// Auto-generated by prax-typegen. Do not edit.\n\n");
        out.push_str("import { z } from 'zod';\n\n");

        for (_, enum_def) in &schema.enums {
            Self::write_enum(&mut out, enum_def);
            out.push('\n');
        }

        for (_, composite) in &schema.types {
            Self::write_composite(&mut out, composite, schema);
            out.push('\n');
        }

        for (_, model) in &schema.models {
            Self::write_model(&mut out, model, schema);
            out.push('\n');
        }

        for (_, view) in &schema.views {
            Self::write_view(&mut out, view, schema);
            out.push('\n');
        }

        out
    }

    fn write_enum(out: &mut String, enum_def: &Enum) {
        let name = enum_def.name();
        let variants: Vec<String> = enum_def
            .variants
            .iter()
            .map(|v| format!("'{}'", v.db_value()))
            .collect();

        out.push_str(&format!(
            "export const {name}Schema = z.enum([{}]);\n",
            variants.join(", ")
        ));
        out.push_str(&format!(
            "export type {name} = z.infer<typeof {name}Schema>;\n"
        ));
    }

    fn write_model(out: &mut String, model: &Model, schema: &Schema) {
        let name = model.name();
        out.push_str(&format!("export const {name}Schema = z.object({{\n"));

        for (_, field) in &model.fields {
            if field.is_relation() {
                Self::write_relation_field(out, field);
            } else {
                Self::write_field(out, field, schema);
            }
        }

        out.push_str("});\n");
        out.push_str(&format!(
            "export type {name} = z.infer<typeof {name}Schema>;\n"
        ));

        Self::write_create_schema(out, model, schema);
        Self::write_update_schema(out, model, schema);
    }

    fn write_view(out: &mut String, view: &View, schema: &Schema) {
        let name = view.name();
        out.push_str(&format!("export const {name}Schema = z.object({{\n"));

        for (_, field) in &view.fields {
            Self::write_field(out, field, schema);
        }

        out.push_str("});\n");
        out.push_str(&format!(
            "export type {name} = z.infer<typeof {name}Schema>;\n"
        ));
    }

    fn write_composite(out: &mut String, composite: &CompositeType, schema: &Schema) {
        let name = composite.name();
        out.push_str(&format!("export const {name}Schema = z.object({{\n"));

        for (_, field) in &composite.fields {
            Self::write_field(out, field, schema);
        }

        out.push_str("});\n");
        out.push_str(&format!(
            "export type {name} = z.infer<typeof {name}Schema>;\n"
        ));
    }

    fn write_field(out: &mut String, field: &Field, schema: &Schema) {
        let name = field.name().to_case(Case::Camel);
        let zod_expr = resolve_zod_type(field, schema);
        out.push_str(&format!("  {name}: {zod_expr},\n"));
    }

    fn write_relation_field(out: &mut String, field: &Field) {
        let name = field.name().to_case(Case::Camel);
        let type_name = field.field_type.type_name();
        let schema_ref = format!("z.lazy(() => {type_name}Schema)");

        let expr = if field.modifier.is_list() {
            format!("{schema_ref}.array()")
        } else if field.modifier.is_optional() {
            format!("{schema_ref}.nullable()")
        } else {
            schema_ref
        };

        out.push_str(&format!("  {name}: {expr},\n"));
    }

    fn write_create_schema(out: &mut String, model: &Model, schema: &Schema) {
        let name = model.name();
        out.push_str(&format!(
            "\nexport const {name}CreateSchema = z.object({{\n"
        ));

        for (_, field) in &model.fields {
            if field.is_relation() {
                continue;
            }
            let attrs = field.extract_attributes();
            if attrs.is_auto || attrs.is_updated_at {
                continue;
            }

            let fname = field.name().to_case(Case::Camel);
            let base_zod = resolve_zod_base(field, schema);

            let is_nullable = field.modifier.is_optional();
            let is_optional = is_nullable || attrs.default.is_some() || attrs.is_id;

            let expr = match (is_nullable, is_optional) {
                (true, _) => format!("{base_zod}.nullable().optional()"),
                (false, true) => format!("{base_zod}.optional()"),
                _ => base_zod,
            };

            out.push_str(&format!("  {fname}: {expr},\n"));
        }

        out.push_str("});\n");
        out.push_str(&format!(
            "export type {name}CreateInput = z.infer<typeof {name}CreateSchema>;\n"
        ));
    }

    fn write_update_schema(out: &mut String, model: &Model, schema: &Schema) {
        let name = model.name();
        out.push_str(&format!(
            "\nexport const {name}UpdateSchema = z.object({{\n"
        ));

        for (_, field) in &model.fields {
            if field.is_relation() {
                continue;
            }
            let attrs = field.extract_attributes();
            if attrs.is_auto || attrs.is_updated_at {
                continue;
            }

            let fname = field.name().to_case(Case::Camel);
            let base_zod = resolve_zod_base(field, schema);
            out.push_str(&format!("  {fname}: {base_zod}.optional(),\n"));
        }

        out.push_str("});\n");
        out.push_str(&format!(
            "export type {name}UpdateInput = z.infer<typeof {name}UpdateSchema>;\n"
        ));
    }
}

/// Resolve the full Zod expression for a field (including nullable/array wrappers).
fn resolve_zod_type(field: &Field, schema: &Schema) -> String {
    let base = resolve_zod_base(field, schema);

    match field.modifier {
        TypeModifier::List => format!("{base}.array()"),
        TypeModifier::Optional => format!("{base}.nullable()"),
        TypeModifier::OptionalList => format!("{base}.array().nullable()"),
        TypeModifier::Required => base,
    }
}

/// Resolve the base Zod type (without modifier wrappers).
fn resolve_zod_base(field: &Field, _schema: &Schema) -> String {
    match &field.field_type {
        FieldType::Scalar(s) => TypeMapper::zod_type(s).to_string(),
        FieldType::Enum(name) => format!("{name}Schema"),
        FieldType::Model(name) => format!("z.lazy(() => {name}Schema)"),
        FieldType::Composite(name) => format!("{name}Schema"),
        FieldType::Unsupported(_) => "z.unknown()".to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use prax_schema::parse_schema;

    #[test]
    fn test_generate_zod_model() {
        let schema = parse_schema(
            r#"
            model User {
                id    Int    @id @auto
                email String @unique
                name  String?
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("import { z } from 'zod';"));
        assert!(output.contains("export const UserSchema = z.object({"));
        assert!(output.contains("id: z.number().int(),"));
        assert!(output.contains("email: z.string(),"));
        assert!(output.contains("name: z.string().nullable(),"));
        assert!(output.contains("export type User = z.infer<typeof UserSchema>;"));
    }

    #[test]
    fn test_generate_zod_enum() {
        let schema = parse_schema(
            r#"
            enum Role {
                User
                Admin
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("export const RoleSchema = z.enum(['User', 'Admin']);"));
        assert!(output.contains("export type Role = z.infer<typeof RoleSchema>;"));
    }

    #[test]
    fn test_generate_zod_create_schema() {
        let schema = parse_schema(
            r#"
            model Post {
                id      Int    @id @auto
                title   String
                content String?
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("export const PostCreateSchema = z.object({"));
        assert!(output.contains("title: z.string(),"));
        assert!(output.contains("content: z.string().nullable().optional(),"));
    }

    #[test]
    fn test_generate_zod_update_schema() {
        let schema = parse_schema(
            r#"
            model User {
                id   Int    @id @auto
                name String
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("export const UserUpdateSchema = z.object({"));
        assert!(output.contains("name: z.string().optional(),"));
    }

    #[test]
    fn test_generate_zod_relation_lazy() {
        let schema = parse_schema(
            r#"
            model User {
                id    Int    @id @auto
                posts Post[]
            }
            model Post {
                id       Int  @id @auto
                authorId Int
                author   User @relation(fields: [authorId], references: [id])
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("posts: z.lazy(() => PostSchema).array(),"));
        assert!(output.contains("author: z.lazy(() => UserSchema),"));
    }

    #[test]
    fn test_generate_zod_uuid() {
        let schema = parse_schema(
            r#"
            model Token {
                id    Uuid @id
                value String
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("id: z.string().uuid(),"));
    }

    #[test]
    fn test_generate_zod_list() {
        let schema = parse_schema(
            r#"
            model User {
                id   Int      @id @auto
                tags String[]
            }
            "#,
        )
        .unwrap();

        let output = ZodGenerator::generate(&schema);
        assert!(output.contains("tags: z.string().array(),"));
    }
}