scythe-codegen 0.6.9

Polyglot code generation backends for scythe
Documentation
use std::fmt::Write;

use scythe_backend::manifest::BackendManifest;
use scythe_backend::naming::{
    enum_type_name, enum_variant_name, fn_name, row_struct_name, to_pascal_case,
};

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-tiny-tds.toml");

/// Check if a neutral type should not be escaped (numeric or boolean types).
/// Neutral types are core type identifiers like "int32", "bool", "decimal".
fn is_numeric_or_bool_type(neutral_type: &str) -> bool {
    // Remove container wrappers and check the base type
    let base_type = if neutral_type.contains('<') {
        // Handle Array<T> or similar generic types
        neutral_type.split('<').next().unwrap_or("")
    } else {
        neutral_type
    };

    matches!(
        base_type,
        "bool" | "int16" | "int32" | "int64" | "float32" | "float64" | "decimal"
    )
}

/// Check if a neutral type is specifically a boolean.
fn is_bool_type(neutral_type: &str) -> bool {
    let base_type = if neutral_type.contains('<') {
        neutral_type.split('<').next().unwrap_or("")
    } else {
        neutral_type
    };
    base_type == "bool"
}

/// Generate the inline SQL interpolation for a parameter value.
/// Returns the string that replaces `@pN` directly in the SQL body.
/// - Booleans: `#{var ? 1 : 0}`
/// - Other numeric/decimal types: `#{var}`
/// - Nullable strings: `#{var.nil? ? 'NULL' : "'#{client.escape(var)}'"}`
/// - Non-nullable strings: `'#{client.escape(var)}'`
/// - `var_access` is used to reference the variable (e.g., "id" or "item[:id]")
fn generate_inline_value(neutral_type: &str, var_access: &str, nullable: bool) -> String {
    if is_bool_type(neutral_type) {
        format!("#{{{} ? 1 : 0}}", var_access)
    } else if is_numeric_or_bool_type(neutral_type) {
        format!("#{{{}}}", var_access)
    } else if nullable {
        // Nullable string: emit SQL NULL literal or a quoted, escaped value.
        // Ruby evaluates the ternary at runtime: nil → NULL, otherwise → 'escaped_value'.
        // The inner double-quoted string `"'#{client.escape(...)}'"`  uses its own interpolation,
        // which Ruby handles correctly even when nested inside the outer string's #{} block.
        format!(
            "#{{ {}.nil? ? 'NULL' : \"'#{{client.escape({})}}'\"}}",
            var_access, var_access
        )
    } else {
        format!("'#{{client.escape({})}}'", var_access)
    }
}

/// Replace `@p1`, `@p2`, ... placeholders in the SQL string with inline Ruby interpolations.
fn inline_params(sql: &str, params: &[(String, String, bool)]) -> String {
    let mut result = sql.to_string();
    // Replace in reverse order so that @p10 is replaced before @p1
    for (index, (neutral_type, var_access, nullable)) in params.iter().enumerate().rev() {
        let placeholder = format!("@p{}", index + 1);
        let value = generate_inline_value(neutral_type, var_access, *nullable);
        result = result.replace(&placeholder, &value);
    }
    result
}

pub struct RubyTinyTdsBackend {
    manifest: BackendManifest,
}

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

/// Rewrite $1, $2, ... to @p1, @p2, ... for MSSQL.
impl CodegenBackend for RubyTinyTdsBackend {
    fn name(&self) -> &str {
        "ruby-tiny-tds"
    }

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

    fn supported_engines(&self) -> &[&str] {
        &["mssql"]
    }

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

    fn file_header(&self) -> String {
        "# frozen_string_literal: true\n\n# Auto-generated by scythe. Do not edit.\n\nrequire \"tiny_tds\"\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::rewrite_pg_placeholders(
            &super::clean_sql_with_optional(
                &analyzed.sql,
                &analyzed.optional_params,
                &analyzed.params,
            ),
            |n| format!("@p{n}"),
        );
        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 { ", " };

        // WARNING: TinyTDS does not support native parameterized queries.
        // We use Ruby string interpolation with client.escape() for SQL injection prevention.
        // client.escape() handles string values; callers must ensure numeric types are validated.
        // Consider sp_executesql for stronger parameterization if needed.
        if !matches!(analyzed.command, QueryCommand::Batch) {
            let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
        }

        match &analyzed.command {
            QueryCommand::One | QueryCommand::Opt => {
                if params.is_empty() {
                    let _ = writeln!(out, "    result = client.execute(\"{}\").first", sql);
                } else {
                    let params_info: Vec<(String, String, bool)> = params
                        .iter()
                        .map(|p| (p.neutral_type.clone(), p.field_name.clone(), p.nullable))
                        .collect();
                    let inlined_sql = inline_params(&sql, &params_info);
                    let _ = writeln!(out, "    sql = \"{}\"", inlined_sql);
                    let _ = writeln!(out, "    result = client.execute(sql).first");
                }
                let _ = writeln!(out, "    return nil if result.nil?");

                let fields = columns
                    .iter()
                    .map(|c| format!("{}: result[\"{}\"]", c.field_name, c.name))
                    .collect::<Vec<_>>()
                    .join(", ");
                let _ = writeln!(out, "    {}.new({})", struct_name, fields);
            }
            QueryCommand::Batch => {
                let batch_fn_name = format!("{}_batch", func_name);
                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.execute(\"{}\").do", sql);
                } else {
                    let params_info: Vec<(String, String, bool)> = params
                        .iter()
                        .map(|p| {
                            let var_access = if params.len() == 1 {
                                "item".to_string()
                            } else {
                                format!("item[:{}]", p.field_name)
                            };
                            (p.neutral_type.clone(), var_access, p.nullable)
                        })
                        .collect();
                    let inlined_sql = inline_params(&sql, &params_info);
                    let _ = writeln!(out, "      sql = \"{}\"", inlined_sql);
                    let _ = writeln!(out, "      client.execute(sql).do");
                }
                let _ = writeln!(out, "    end");
                let _ = write!(out, "  end");
                return Ok(out);
            }
            QueryCommand::Many => {
                if params.is_empty() {
                    let _ = writeln!(out, "    results = client.execute(\"{}\")", sql);
                } else {
                    let params_info: Vec<(String, String, bool)> = params
                        .iter()
                        .map(|p| (p.neutral_type.clone(), p.field_name.clone(), p.nullable))
                        .collect();
                    let inlined_sql = inline_params(&sql, &params_info);
                    let _ = writeln!(out, "    sql = \"{}\"", inlined_sql);
                    let _ = writeln!(out, "    results = client.execute(sql)");
                }
                let _ = writeln!(out, "    results.map do |row|");
                let fields = columns
                    .iter()
                    .map(|c| format!("{}: row[\"{}\"]", c.field_name, c.name))
                    .collect::<Vec<_>>()
                    .join(", ");
                let _ = writeln!(out, "      {}.new({})", struct_name, fields);
                let _ = writeln!(out, "    end");
            }
            QueryCommand::Exec => {
                if params.is_empty() {
                    let _ = writeln!(out, "    client.execute(\"{}\").do", sql);
                } else {
                    let params_info: Vec<(String, String, bool)> = params
                        .iter()
                        .map(|p| (p.neutral_type.clone(), p.field_name.clone(), p.nullable))
                        .collect();
                    let inlined_sql = inline_params(&sql, &params_info);
                    let _ = writeln!(out, "    sql = \"{}\"", inlined_sql);
                    let _ = writeln!(out, "    client.execute(sql).do");
                }
                let _ = writeln!(out, "    nil");
            }
            QueryCommand::ExecResult | QueryCommand::ExecRows => {
                if params.is_empty() {
                    let _ = writeln!(out, "    client.execute(\"{}\").affected_rows", sql);
                } else {
                    let params_info: Vec<(String, String, bool)> = params
                        .iter()
                        .map(|p| (p.neutral_type.clone(), p.field_name.clone(), p.nullable))
                        .collect();
                    let inlined_sql = inline_params(&sql, &params_info);
                    let _ = writeln!(out, "    sql = \"{}\"", inlined_sql);
                    let _ = writeln!(out, "    client.execute(sql).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)
    }
}