Skip to main content

scythe_codegen/backends/
ruby_mysql2.rs

1use std::fmt::Write;
2use std::path::Path;
3
4use scythe_backend::manifest::{BackendManifest, load_manifest};
5use scythe_backend::naming::{
6    enum_type_name, enum_variant_name, fn_name, row_struct_name, to_pascal_case,
7};
8
9use scythe_core::analyzer::{AnalyzedQuery, CompositeInfo, EnumInfo};
10use scythe_core::errors::{ErrorCode, ScytheError};
11use scythe_core::parser::QueryCommand;
12
13use crate::backend_trait::{CodegenBackend, RbsGenerationContext, ResolvedColumn, ResolvedParam};
14
15const DEFAULT_MANIFEST_TOML: &str = include_str!("../../manifests/ruby-mysql2.toml");
16
17pub struct RubyMysql2Backend {
18    manifest: BackendManifest,
19}
20
21impl RubyMysql2Backend {
22    pub fn new(engine: &str) -> Result<Self, ScytheError> {
23        match engine {
24            "mysql" | "mariadb" => {}
25            _ => {
26                return Err(ScytheError::new(
27                    ErrorCode::InternalError,
28                    format!("ruby-mysql2 only supports MySQL, got engine '{}'", engine),
29                ));
30            }
31        }
32        let manifest_path = Path::new("backends/ruby-mysql2/manifest.toml");
33        let manifest = if manifest_path.exists() {
34            load_manifest(manifest_path)
35                .map_err(|e| ScytheError::new(ErrorCode::InternalError, format!("manifest: {e}")))?
36        } else {
37            toml::from_str(DEFAULT_MANIFEST_TOML)
38                .map_err(|e| ScytheError::new(ErrorCode::InternalError, format!("manifest: {e}")))?
39        };
40        Ok(Self { manifest })
41    }
42}
43
44/// Map a neutral type to a Ruby type coercion method for mysql2.
45fn ruby_coercion(neutral_type: &str) -> &'static str {
46    match neutral_type {
47        "int16" | "int32" | "int64" => ".to_i",
48        "float32" | "float64" => ".to_f",
49        "bool" => " == 1",
50        _ => "",
51    }
52}
53
54impl CodegenBackend for RubyMysql2Backend {
55    fn name(&self) -> &str {
56        "ruby-mysql2"
57    }
58
59    fn manifest(&self) -> &scythe_backend::manifest::BackendManifest {
60        &self.manifest
61    }
62
63    fn supported_engines(&self) -> &[&str] {
64        &["mysql"]
65    }
66
67    fn generate_rbs_file(&self, context: &RbsGenerationContext) -> Option<String> {
68        Some(super::ruby_rbs::generate_rbs_content(
69            context,
70            "Mysql2::Client",
71        ))
72    }
73
74    fn file_header(&self) -> String {
75        "# frozen_string_literal: true\n\n# Auto-generated by scythe. Do not edit.\n\nmodule Queries"
76            .to_string()
77    }
78
79    fn file_footer(&self) -> String {
80        "end".to_string()
81    }
82
83    fn generate_row_struct(
84        &self,
85        query_name: &str,
86        columns: &[ResolvedColumn],
87    ) -> Result<String, ScytheError> {
88        let struct_name = row_struct_name(query_name, &self.manifest.naming);
89        let fields = columns
90            .iter()
91            .map(|c| format!(":{}", c.field_name))
92            .collect::<Vec<_>>()
93            .join(", ");
94        let mut out = String::new();
95        let _ = writeln!(out, "  {} = Data.define({})", struct_name, fields);
96        Ok(out)
97    }
98
99    fn generate_model_struct(
100        &self,
101        table_name: &str,
102        columns: &[ResolvedColumn],
103    ) -> Result<String, ScytheError> {
104        let name = to_pascal_case(table_name);
105        self.generate_row_struct(&name, columns)
106    }
107
108    fn generate_query_fn(
109        &self,
110        analyzed: &AnalyzedQuery,
111        struct_name: &str,
112        columns: &[ResolvedColumn],
113        params: &[ResolvedParam],
114    ) -> Result<String, ScytheError> {
115        let func_name = fn_name(&analyzed.name, &self.manifest.naming);
116        let sql = super::clean_sql_with_optional(
117            &analyzed.sql,
118            &analyzed.optional_params,
119            &analyzed.params,
120        );
121        let mut out = String::new();
122
123        let param_list = params
124            .iter()
125            .map(|p| p.field_name.clone())
126            .collect::<Vec<_>>()
127            .join(", ");
128        let sep = if param_list.is_empty() { "" } else { ", " };
129
130        let _ = writeln!(out, "  def self.{}(client{}{})", func_name, sep, param_list);
131
132        let param_array = if params.is_empty() {
133            "[]".to_string()
134        } else {
135            format!(
136                "[{}]",
137                params
138                    .iter()
139                    .map(|p| p.field_name.clone())
140                    .collect::<Vec<_>>()
141                    .join(", ")
142            )
143        };
144
145        match &analyzed.command {
146            QueryCommand::One => {
147                let _ = writeln!(out, "    stmt = client.prepare(\"{}\")", sql);
148                let _ = writeln!(
149                    out,
150                    "    results = stmt.execute({})",
151                    param_array.trim_start_matches('[').trim_end_matches(']')
152                );
153                let _ = writeln!(out, "    row = results.first");
154                let _ = writeln!(out, "    return nil if row.nil?");
155
156                let fields = columns
157                    .iter()
158                    .map(|c| {
159                        let coercion = ruby_coercion(&c.neutral_type);
160                        if c.nullable {
161                            format!(
162                                "{}: row[\"{}\"]&.then {{ |v| v{} }}",
163                                c.field_name, c.name, coercion
164                            )
165                        } else {
166                            format!("{}: row[\"{}\"]{}", c.field_name, c.name, coercion)
167                        }
168                    })
169                    .collect::<Vec<_>>()
170                    .join(", ");
171                let _ = writeln!(out, "    {}.new({})", struct_name, fields);
172            }
173            QueryCommand::Batch => {
174                let batch_fn_name = format!("{}_batch", func_name);
175                let _ = writeln!(out, "  def self.{}(client, items)", batch_fn_name);
176                let _ = writeln!(out, "    stmt = client.prepare(\"{}\")", sql);
177                let _ = writeln!(out, "    items.each do |item|");
178                if params.len() > 1 {
179                    let _ = writeln!(out, "      stmt.execute(*item)");
180                } else if params.len() == 1 {
181                    let _ = writeln!(out, "      stmt.execute(item)");
182                } else {
183                    let _ = writeln!(out, "      stmt.execute");
184                }
185                let _ = writeln!(out, "    end");
186                let _ = write!(out, "  end");
187                return Ok(out);
188            }
189            QueryCommand::Many => {
190                let _ = writeln!(out, "    stmt = client.prepare(\"{}\")", sql);
191                let _ = writeln!(
192                    out,
193                    "    results = stmt.execute({})",
194                    param_array.trim_start_matches('[').trim_end_matches(']')
195                );
196                let _ = writeln!(out, "    results.map do |row|");
197                let fields = columns
198                    .iter()
199                    .map(|c| {
200                        let coercion = ruby_coercion(&c.neutral_type);
201                        if c.nullable {
202                            format!(
203                                "{}: row[\"{}\"]&.then {{ |v| v{} }}",
204                                c.field_name, c.name, coercion
205                            )
206                        } else {
207                            format!("{}: row[\"{}\"]{}", c.field_name, c.name, coercion)
208                        }
209                    })
210                    .collect::<Vec<_>>()
211                    .join(", ");
212                let _ = writeln!(out, "      {}.new({})", struct_name, fields);
213                let _ = writeln!(out, "    end");
214            }
215            QueryCommand::Exec => {
216                let _ = writeln!(out, "    stmt = client.prepare(\"{}\")", sql);
217                let _ = writeln!(
218                    out,
219                    "    stmt.execute({})",
220                    param_array.trim_start_matches('[').trim_end_matches(']')
221                );
222                let _ = writeln!(out, "    nil");
223            }
224            QueryCommand::ExecResult | QueryCommand::ExecRows => {
225                let _ = writeln!(out, "    stmt = client.prepare(\"{}\")", sql);
226                let _ = writeln!(
227                    out,
228                    "    stmt.execute({})",
229                    param_array.trim_start_matches('[').trim_end_matches(']')
230                );
231                let _ = writeln!(out, "    client.affected_rows");
232            }
233            QueryCommand::Grouped => unreachable!("handled as Many in codegen"),
234        }
235
236        let _ = write!(out, "  end");
237        Ok(out)
238    }
239
240    fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
241        let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
242        let mut out = String::new();
243        let _ = writeln!(out, "  module {}", type_name);
244        for value in &enum_info.values {
245            let variant = enum_variant_name(value, &self.manifest.naming);
246            let _ = writeln!(out, "    {} = \"{}\"", variant, value);
247        }
248        let all_values = enum_info
249            .values
250            .iter()
251            .map(|v| enum_variant_name(v, &self.manifest.naming))
252            .collect::<Vec<_>>()
253            .join(", ");
254        let _ = writeln!(out, "    ALL = [{}].freeze", all_values);
255        let _ = write!(out, "  end");
256        Ok(out)
257    }
258
259    fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
260        let name = to_pascal_case(&composite.sql_name);
261        let mut out = String::new();
262        if composite.fields.is_empty() {
263            let _ = writeln!(out, "  {} = Data.define()", name);
264        } else {
265            let fields = composite
266                .fields
267                .iter()
268                .map(|f| format!(":{}", f.name))
269                .collect::<Vec<_>>()
270                .join(", ");
271            let _ = writeln!(out, "  {} = Data.define({})", name, fields);
272        }
273        Ok(out)
274    }
275}