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};
8use scythe_backend::types::resolve_type;
9
10use scythe_core::analyzer::{AnalyzedQuery, CompositeInfo, EnumInfo};
11use scythe_core::errors::{ErrorCode, ScytheError};
12use scythe_core::parser::QueryCommand;
13
14use crate::backend_trait::{CodegenBackend, ResolvedColumn, ResolvedParam};
15
16pub struct GoDatabaseSqlBackend {
17 manifest: BackendManifest,
18 engine: String,
19}
20
21impl GoDatabaseSqlBackend {
22 pub fn new(engine: &str) -> Result<Self, ScytheError> {
23 let manifest_toml = match engine {
24 "mysql" | "mariadb" => include_str!("../../manifests/go-database-sql.mysql.toml"),
25 "sqlite" | "sqlite3" => include_str!("../../manifests/go-database-sql.sqlite.toml"),
26 _ => {
27 return Err(ScytheError::new(
28 ErrorCode::InternalError,
29 format!(
30 "go-database-sql supports MySQL and SQLite, got engine '{}'",
31 engine
32 ),
33 ));
34 }
35 };
36 let manifest_path = Path::new("backends/go-database-sql/manifest.toml");
37 let manifest = if manifest_path.exists() {
38 load_manifest(manifest_path)
39 .map_err(|e| ScytheError::new(ErrorCode::InternalError, format!("manifest: {e}")))?
40 } else {
41 toml::from_str(manifest_toml)
42 .map_err(|e| ScytheError::new(ErrorCode::InternalError, format!("manifest: {e}")))?
43 };
44 Ok(Self {
45 manifest,
46 engine: engine.to_string(),
47 })
48 }
49}
50
51impl CodegenBackend for GoDatabaseSqlBackend {
52 fn name(&self) -> &str {
53 "go-database-sql"
54 }
55
56 fn manifest(&self) -> &scythe_backend::manifest::BackendManifest {
57 &self.manifest
58 }
59
60 fn supported_engines(&self) -> &[&str] {
61 &["mysql", "sqlite"]
62 }
63
64 fn file_header(&self) -> String {
65 let uses_time = matches!(self.engine.as_str(), "mysql" | "mariadb");
66 let mut header =
67 String::from("package queries\n\nimport (\n\t\"context\"\n\t\"database/sql\"");
68 if uses_time {
69 header.push_str("\n\t\"time\"");
70 }
71 header.push_str("\n)\n");
72 header
73 }
74
75 fn generate_row_struct(
76 &self,
77 query_name: &str,
78 columns: &[ResolvedColumn],
79 ) -> Result<String, ScytheError> {
80 let struct_name = row_struct_name(query_name, &self.manifest.naming);
81 let mut out = String::new();
82 let _ = writeln!(out, "type {} struct {{", struct_name);
83 for col in columns {
84 let field = to_pascal_case(&col.field_name);
85 let json_tag = &col.field_name;
86 let _ = writeln!(out, "\t{} {} `json:\"{}\"`", field, col.full_type, json_tag);
87 }
88 let _ = write!(out, "}}");
89 Ok(out)
90 }
91
92 fn generate_model_struct(
93 &self,
94 table_name: &str,
95 columns: &[ResolvedColumn],
96 ) -> Result<String, ScytheError> {
97 let name = to_pascal_case(table_name);
98 self.generate_row_struct(&name, columns)
99 }
100
101 fn generate_query_fn(
102 &self,
103 analyzed: &AnalyzedQuery,
104 struct_name: &str,
105 columns: &[ResolvedColumn],
106 params: &[ResolvedParam],
107 ) -> Result<String, ScytheError> {
108 let func_name = fn_name(&analyzed.name, &self.manifest.naming);
109 let sql = super::clean_sql_oneline_with_optional(
110 &analyzed.sql,
111 &analyzed.optional_params,
112 &analyzed.params,
113 );
114
115 let param_list = params
116 .iter()
117 .map(|p| {
118 let field = to_pascal_case(&p.field_name);
119 format!("{} {}", field, p.full_type)
120 })
121 .collect::<Vec<_>>()
122 .join(", ");
123 let sep = if param_list.is_empty() { "" } else { ", " };
124
125 let args = params
126 .iter()
127 .map(|p| to_pascal_case(&p.field_name).into_owned())
128 .collect::<Vec<_>>();
129
130 let mut out = String::new();
131
132 match &analyzed.command {
133 QueryCommand::Exec => {
134 let _ = writeln!(
135 out,
136 "func {}(ctx context.Context, db *sql.DB{}{}) error {{",
137 func_name, sep, param_list
138 );
139 let args_str = if args.is_empty() {
140 String::new()
141 } else {
142 format!(", {}", args.join(", "))
143 };
144 let _ = writeln!(
145 out,
146 "\t_, err := db.ExecContext(ctx, \"{}\"{})",
147 sql, args_str
148 );
149 let _ = writeln!(out, "\treturn err");
150 let _ = write!(out, "}}");
151 }
152 QueryCommand::ExecResult | QueryCommand::ExecRows => {
153 let _ = writeln!(
154 out,
155 "func {}(ctx context.Context, db *sql.DB{}{}) (int64, error) {{",
156 func_name, sep, param_list
157 );
158 let args_str = if args.is_empty() {
159 String::new()
160 } else {
161 format!(", {}", args.join(", "))
162 };
163 let _ = writeln!(
164 out,
165 "\tresult, err := db.ExecContext(ctx, \"{}\"{})",
166 sql, args_str
167 );
168 let _ = writeln!(out, "\tif err != nil {{");
169 let _ = writeln!(out, "\t\treturn 0, err");
170 let _ = writeln!(out, "\t}}");
171 let _ = writeln!(out, "\treturn result.RowsAffected()");
172 let _ = write!(out, "}}");
173 }
174 QueryCommand::One => {
175 let _ = writeln!(
176 out,
177 "func {}(ctx context.Context, db *sql.DB{}{}) ({}, error) {{",
178 func_name, sep, param_list, struct_name
179 );
180 let args_str = if args.is_empty() {
181 String::new()
182 } else {
183 format!(", {}", args.join(", "))
184 };
185 let _ = writeln!(
186 out,
187 "\trow := db.QueryRowContext(ctx, \"{}\"{})",
188 sql, args_str
189 );
190 let _ = writeln!(out, "\tvar r {}", struct_name);
191 let scan_fields: Vec<String> = columns
192 .iter()
193 .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
194 .collect();
195 let _ = writeln!(out, "\terr := row.Scan({})", scan_fields.join(", "));
196 let _ = writeln!(out, "\treturn r, err");
197 let _ = write!(out, "}}");
198 }
199 QueryCommand::Batch => {
200 let batch_fn_name = format!("{}Batch", func_name);
201 if params.len() > 1 {
202 let params_struct_name = format!("{}BatchParams", func_name);
203 let _ = writeln!(out, "type {} struct {{", params_struct_name);
204 for p in params {
205 let field = to_pascal_case(&p.field_name);
206 let _ = writeln!(out, "\t{} {}", field, p.full_type);
207 }
208 let _ = writeln!(out, "}}");
209 let _ = writeln!(out);
210 let _ = writeln!(
211 out,
212 "func {}(ctx context.Context, db *sql.DB, items []{}) error {{",
213 batch_fn_name, params_struct_name
214 );
215 } else if params.len() == 1 {
216 let _ = writeln!(
217 out,
218 "func {}(ctx context.Context, db *sql.DB, items []{}) error {{",
219 batch_fn_name, params[0].full_type
220 );
221 } else {
222 let _ = writeln!(
223 out,
224 "func {}(ctx context.Context, db *sql.DB, count int) error {{",
225 batch_fn_name
226 );
227 }
228 let _ = writeln!(out, "\ttx, err := db.BeginTx(ctx, nil)");
229 let _ = writeln!(out, "\tif err != nil {{");
230 let _ = writeln!(out, "\t\treturn err");
231 let _ = writeln!(out, "\t}}");
232 let _ = writeln!(out, "\tdefer tx.Rollback()");
233 if params.is_empty() {
234 let _ = writeln!(out, "\tfor i := 0; i < count; i++ {{");
235 let _ = writeln!(out, "\t\t_, err := tx.ExecContext(ctx, \"{}\")", sql);
236 } else {
237 let _ = writeln!(out, "\tfor _, item := range items {{");
238 if params.len() > 1 {
239 let item_args: Vec<String> = params
240 .iter()
241 .map(|p| format!("item.{}", to_pascal_case(&p.field_name)))
242 .collect();
243 let _ = writeln!(
244 out,
245 "\t\t_, err := tx.ExecContext(ctx, \"{}\", {})",
246 sql,
247 item_args.join(", ")
248 );
249 } else {
250 let _ =
251 writeln!(out, "\t\t_, err := tx.ExecContext(ctx, \"{}\", item)", sql);
252 }
253 }
254 let _ = writeln!(out, "\t\tif err != nil {{");
255 let _ = writeln!(out, "\t\t\treturn err");
256 let _ = writeln!(out, "\t\t}}");
257 let _ = writeln!(out, "\t}}");
258 let _ = writeln!(out, "\treturn tx.Commit()");
259 let _ = write!(out, "}}");
260 }
261 QueryCommand::Many => {
262 let _ = writeln!(
263 out,
264 "func {}(ctx context.Context, db *sql.DB{}{}) ([]{}, error) {{",
265 func_name, sep, param_list, struct_name
266 );
267 let args_str = if args.is_empty() {
268 String::new()
269 } else {
270 format!(", {}", args.join(", "))
271 };
272 let _ = writeln!(
273 out,
274 "\trows, err := db.QueryContext(ctx, \"{}\"{})",
275 sql, args_str
276 );
277 let _ = writeln!(out, "\tif err != nil {{");
278 let _ = writeln!(out, "\t\treturn nil, err");
279 let _ = writeln!(out, "\t}}");
280 let _ = writeln!(out, "\tdefer rows.Close()");
281 let _ = writeln!(out, "\tvar result []{}", struct_name);
282 let _ = writeln!(out, "\tfor rows.Next() {{");
283 let _ = writeln!(out, "\t\tvar r {}", struct_name);
284 let scan_fields: Vec<String> = columns
285 .iter()
286 .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
287 .collect();
288 let _ = writeln!(
289 out,
290 "\t\tif err := rows.Scan({}); err != nil {{",
291 scan_fields.join(", ")
292 );
293 let _ = writeln!(out, "\t\t\treturn nil, err");
294 let _ = writeln!(out, "\t\t}}");
295 let _ = writeln!(out, "\t\tresult = append(result, r)");
296 let _ = writeln!(out, "\t}}");
297 let _ = writeln!(out, "\treturn result, rows.Err()");
298 let _ = write!(out, "}}");
299 }
300 }
301
302 Ok(out)
303 }
304
305 fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
306 let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
307 let mut out = String::new();
308 let _ = writeln!(out, "type {} string", type_name);
309 let _ = writeln!(out);
310 let _ = writeln!(out, "const (");
311 for value in &enum_info.values {
312 let variant = enum_variant_name(value, &self.manifest.naming);
313 let _ = writeln!(
314 out,
315 "\t{}{} {} = \"{}\"",
316 type_name, variant, type_name, value
317 );
318 }
319 let _ = write!(out, ")");
320 Ok(out)
321 }
322
323 fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
324 let name = to_pascal_case(&composite.sql_name);
325 let mut out = String::new();
326 let _ = writeln!(out, "type {} struct {{", name);
327 if !composite.fields.is_empty() {
328 for field in &composite.fields {
329 let field_name = to_pascal_case(&field.name);
330 let go_type = resolve_type(&field.neutral_type, &self.manifest, false)
331 .map(|t| t.into_owned())
332 .unwrap_or_else(|_| "interface{}".to_string());
333 let json_tag = &field.name;
334 let _ = writeln!(out, "\t{} {} `json:\"{}\"`", field_name, go_type, json_tag);
335 }
336 }
337 let _ = write!(out, "}}");
338 Ok(out)
339 }
340}