scythe-codegen 0.6.6

Polyglot code generation backends for scythe
Documentation
use scythe_backend::manifest::BackendManifest;
use scythe_backend::naming::{
    enum_type_name, enum_variant_name, fn_name, row_struct_name, to_pascal_case,
};
use std::fmt::Write;

use scythe_core::analyzer::{AnalyzedQuery, CompositeInfo, EnumInfo};
use scythe_core::errors::{ErrorCode, ScytheError};
use scythe_core::parser::QueryCommand;

use crate::backend_trait::{CodegenBackend, RbsGenerationContext, ResolvedColumn, ResolvedParam};

const DEFAULT_MANIFEST_TOML: &str = include_str!("../../manifests/ruby-trilogy.toml");

pub struct RubyTrilogyBackend {
    manifest: BackendManifest,
}

impl RubyTrilogyBackend {
    pub fn new(engine: &str) -> Result<Self, ScytheError> {
        match engine {
            "mysql" | "mariadb" => {}
            _ => {
                return Err(ScytheError::new(
                    ErrorCode::InternalError,
                    format!("ruby-trilogy only supports MySQL, got engine '{}'", engine),
                ));
            }
        }
        let manifest = super::load_or_default_manifest(
            "backends/ruby-trilogy/manifest.toml",
            DEFAULT_MANIFEST_TOML,
        )?;
        Ok(Self { manifest })
    }
}

/// Map a neutral type to a Ruby type coercion method for trilogy.
fn ruby_coercion(neutral_type: &str) -> &'static str {
    match neutral_type {
        "int16" | "int32" | "int64" => ".to_i",
        "float32" | "float64" => ".to_f",
        "bool" => " == 1",
        _ => "",
    }
}

impl CodegenBackend for RubyTrilogyBackend {
    fn name(&self) -> &str {
        "ruby-trilogy"
    }

    fn manifest(&self) -> &scythe_backend::manifest::BackendManifest {
        &self.manifest
    }

    fn supported_engines(&self) -> &[&str] {
        &["mysql", "mariadb"]
    }

    fn generate_rbs_file(&self, context: &RbsGenerationContext) -> Option<String> {
        Some(super::ruby_rbs::generate_rbs_content(context, "Trilogy"))
    }

    fn file_header(&self) -> String {
        "# frozen_string_literal: true\n\n# Auto-generated by scythe. Do not edit.\n\nmodule Queries"
            .to_string()
    }

    fn file_footer(&self) -> String {
        "end".to_string()
    }

    fn generate_row_struct(
        &self,
        query_name: &str,
        columns: &[ResolvedColumn],
    ) -> Result<String, ScytheError> {
        let struct_name = row_struct_name(query_name, &self.manifest.naming);
        let fields = columns
            .iter()
            .map(|c| format!(":{}", c.field_name))
            .collect::<Vec<_>>()
            .join(", ");
        let mut out = String::new();
        let _ = writeln!(out, "  {} = Data.define({})", struct_name, fields);
        Ok(out)
    }

    fn generate_model_struct(
        &self,
        table_name: &str,
        columns: &[ResolvedColumn],
    ) -> Result<String, ScytheError> {
        let name = to_pascal_case(table_name);
        self.generate_row_struct(&name, columns)
    }

    fn generate_query_fn(
        &self,
        analyzed: &AnalyzedQuery,
        struct_name: &str,
        columns: &[ResolvedColumn],
        params: &[ResolvedParam],
    ) -> Result<String, ScytheError> {
        let func_name = fn_name(&analyzed.name, &self.manifest.naming);
        let sql = super::clean_sql_with_optional(
            &analyzed.sql,
            &analyzed.optional_params,
            &analyzed.params,
        );
        let mut out = String::new();

        let param_list = params
            .iter()
            .map(|p| p.field_name.clone())
            .collect::<Vec<_>>()
            .join(", ");
        let sep = if param_list.is_empty() { "" } else { ", " };

        let _param_array = if params.is_empty() {
            "[]".to_string()
        } else {
            format!(
                "[{}]",
                params
                    .iter()
                    .map(|p| p.field_name.clone())
                    .collect::<Vec<_>>()
                    .join(", ")
            )
        };

        match &analyzed.command {
            QueryCommand::One | QueryCommand::Opt => {
                let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
                // Trilogy does not support parameterized queries, so we use string interpolation
                // Replace ? placeholders with #{param_name} for Ruby string interpolation
                // String types need to be quoted for SQL
                let mut sql_interpolated = sql.clone();
                for param in params.iter() {
                    if let Some(pos) = sql_interpolated.find('?') {
                        // Quote strings for SQL, keep numbers unquoted
                        let _param_expr = if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            format!(
                                "\\\"#{{{}}}\\\"\"\n                             ",
                                param.field_name
                            )
                            .replace("\n                             ", "")
                            .replace("\\\"", "'")
                        } else {
                            format!("#{{{}}}", param.field_name)
                        };

                        if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            sql_interpolated.replace_range(
                                pos..pos + 1,
                                &format!("'#{{{}}}'", param.field_name),
                            );
                        } else {
                            sql_interpolated
                                .replace_range(pos..pos + 1, &format!("#{{{}}}", param.field_name));
                        }
                    }
                }

                if params.is_empty() {
                    let _ = writeln!(out, "    results = client.query(\"{}\")", sql);
                } else {
                    let _ = writeln!(out, "    results = client.query(\"{}\")", sql_interpolated);
                }
                let _ = writeln!(out, "    row = results.first");
                let _ = writeln!(out, "    return nil if row.nil?");

                let fields = columns
                    .iter()
                    .enumerate()
                    .map(|(i, c)| {
                        let coercion = ruby_coercion(&c.neutral_type);
                        if c.nullable {
                            format!("{}: row[{}]&.then {{ |v| v{} }}", c.field_name, i, coercion)
                        } else {
                            format!("{}: row[{}]{}", c.field_name, i, coercion)
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                let _ = writeln!(out, "    {}.new({})", struct_name, fields);
            }
            QueryCommand::Batch => {
                let batch_fn_name = format!("{}_batch", func_name);
                // Batch writes its own function signature (not the outer one)
                let _ = writeln!(out, "  def self.{}(client, items)", batch_fn_name);
                let _ = writeln!(out, "    items.each do |item|");
                if params.is_empty() {
                    let _ = writeln!(out, "      client.query(\"{}\")", sql);
                } else if params.len() > 1 {
                    // Multiple params: destructure from tuple/array
                    let mut sql_with_params = sql.clone();
                    for (i, param) in params.iter().enumerate() {
                        if let Some(pos) = sql_with_params.find('?') {
                            let item_expr = format!("item[{}]", i);
                            if param.neutral_type.starts_with("enum::")
                                || param.neutral_type == "string"
                            {
                                sql_with_params
                                    .replace_range(pos..pos + 1, &format!("'#{{{}}}'", item_expr));
                            } else {
                                sql_with_params
                                    .replace_range(pos..pos + 1, &format!("#{{{}}}", item_expr));
                            }
                        }
                    }
                    let _ = writeln!(out, "      client.query(\"{}\")", sql_with_params);
                } else {
                    // Single param
                    let mut sql_with_param = sql.clone();
                    if let Some(pos) = sql_with_param.find('?') {
                        let param = &params[0];
                        if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            sql_with_param.replace_range(pos..pos + 1, "'#{item}'");
                        } else {
                            sql_with_param.replace_range(pos..pos + 1, "#{item}");
                        }
                    }
                    let _ = writeln!(out, "      client.query(\"{}\")", sql_with_param);
                }
                let _ = writeln!(out, "    end");
                let _ = write!(out, "  end");
                return Ok(out);
            }
            QueryCommand::Many => {
                let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
                // Replace ? placeholders with #{param_name} for Ruby string interpolation
                // String types need to be quoted for SQL
                let mut sql_interpolated = sql.clone();
                for param in params.iter() {
                    if let Some(pos) = sql_interpolated.find('?') {
                        if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            sql_interpolated.replace_range(
                                pos..pos + 1,
                                &format!("'#{{{}}}'", param.field_name),
                            );
                        } else {
                            sql_interpolated
                                .replace_range(pos..pos + 1, &format!("#{{{}}}", param.field_name));
                        }
                    }
                }

                if params.is_empty() {
                    let _ = writeln!(out, "    results = client.query(\"{}\")", sql);
                } else {
                    let _ = writeln!(out, "    results = client.query(\"{}\")", sql_interpolated);
                }
                let _ = writeln!(out, "    results.map do |row|");
                let fields = columns
                    .iter()
                    .enumerate()
                    .map(|(i, c)| {
                        let coercion = ruby_coercion(&c.neutral_type);
                        if c.nullable {
                            format!("{}: row[{}]&.then {{ |v| v{} }}", c.field_name, i, coercion)
                        } else {
                            format!("{}: row[{}]{}", c.field_name, i, coercion)
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                let _ = writeln!(out, "      {}.new({})", struct_name, fields);
                let _ = writeln!(out, "    end");
            }
            QueryCommand::Exec => {
                let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
                // Replace ? placeholders with #{param_name} for Ruby string interpolation
                // String types need to be quoted for SQL
                let mut sql_interpolated = sql.clone();
                for param in params.iter() {
                    if let Some(pos) = sql_interpolated.find('?') {
                        if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            sql_interpolated.replace_range(
                                pos..pos + 1,
                                &format!("'#{{{}}}'", param.field_name),
                            );
                        } else {
                            sql_interpolated
                                .replace_range(pos..pos + 1, &format!("#{{{}}}", param.field_name));
                        }
                    }
                }

                if params.is_empty() {
                    let _ = writeln!(out, "    client.query(\"{}\")", sql);
                } else {
                    let _ = writeln!(out, "    client.query(\"{}\")", sql_interpolated);
                }
                let _ = writeln!(out, "    nil");
            }
            QueryCommand::ExecResult | QueryCommand::ExecRows => {
                let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
                // Replace ? placeholders with #{param_name} for Ruby string interpolation
                // String types need to be quoted for SQL
                let mut sql_interpolated = sql.clone();
                for param in params.iter() {
                    if let Some(pos) = sql_interpolated.find('?') {
                        if param.neutral_type.starts_with("enum::")
                            || param.neutral_type == "string"
                        {
                            sql_interpolated.replace_range(
                                pos..pos + 1,
                                &format!("'#{{{}}}'", param.field_name),
                            );
                        } else {
                            sql_interpolated
                                .replace_range(pos..pos + 1, &format!("#{{{}}}", param.field_name));
                        }
                    }
                }

                if params.is_empty() {
                    let _ = writeln!(out, "    client.query(\"{}\")", sql);
                } else {
                    let _ = writeln!(out, "    client.query(\"{}\")", sql_interpolated);
                }
                let _ = writeln!(out, "    client.affected_rows");
            }
            QueryCommand::Grouped => unreachable!("handled as Many in codegen"),
        }

        let _ = write!(out, "  end");
        Ok(out)
    }

    fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
        let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
        let mut out = String::new();
        let _ = writeln!(out, "  module {}", type_name);
        for value in &enum_info.values {
            let variant = enum_variant_name(value, &self.manifest.naming);
            let _ = writeln!(out, "    {} = \"{}\"", variant, value);
        }
        let all_values = enum_info
            .values
            .iter()
            .map(|v| enum_variant_name(v, &self.manifest.naming))
            .collect::<Vec<_>>()
            .join(", ");
        let _ = writeln!(out, "    ALL = [{}].freeze", all_values);
        let _ = write!(out, "  end");
        Ok(out)
    }

    fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
        let name = to_pascal_case(&composite.sql_name);
        let mut out = String::new();
        if composite.fields.is_empty() {
            let _ = writeln!(out, "  {} = Data.define()", name);
        } else {
            let fields = composite
                .fields
                .iter()
                .map(|f| format!(":{}", f.name))
                .collect::<Vec<_>>()
                .join(", ");
            let _ = writeln!(out, "  {} = Data.define({})", name, fields);
        }
        Ok(out)
    }
}