Skip to main content

scythe_codegen/backends/
go_database_sql.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};
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(&analyzed.sql);
110
111        let param_list = params
112            .iter()
113            .map(|p| {
114                let field = to_pascal_case(&p.field_name);
115                format!("{} {}", field, p.full_type)
116            })
117            .collect::<Vec<_>>()
118            .join(", ");
119        let sep = if param_list.is_empty() { "" } else { ", " };
120
121        let args = params
122            .iter()
123            .map(|p| to_pascal_case(&p.field_name).into_owned())
124            .collect::<Vec<_>>();
125
126        let mut out = String::new();
127
128        match &analyzed.command {
129            QueryCommand::Exec => {
130                let _ = writeln!(
131                    out,
132                    "func {}(ctx context.Context, db *sql.DB{}{}) error {{",
133                    func_name, sep, param_list
134                );
135                let args_str = if args.is_empty() {
136                    String::new()
137                } else {
138                    format!(", {}", args.join(", "))
139                };
140                let _ = writeln!(
141                    out,
142                    "\t_, err := db.ExecContext(ctx, \"{}\"{})",
143                    sql, args_str
144                );
145                let _ = writeln!(out, "\treturn err");
146                let _ = write!(out, "}}");
147            }
148            QueryCommand::ExecResult | QueryCommand::ExecRows => {
149                let _ = writeln!(
150                    out,
151                    "func {}(ctx context.Context, db *sql.DB{}{}) (int64, error) {{",
152                    func_name, sep, param_list
153                );
154                let args_str = if args.is_empty() {
155                    String::new()
156                } else {
157                    format!(", {}", args.join(", "))
158                };
159                let _ = writeln!(
160                    out,
161                    "\tresult, err := db.ExecContext(ctx, \"{}\"{})",
162                    sql, args_str
163                );
164                let _ = writeln!(out, "\tif err != nil {{");
165                let _ = writeln!(out, "\t\treturn 0, err");
166                let _ = writeln!(out, "\t}}");
167                let _ = writeln!(out, "\treturn result.RowsAffected()");
168                let _ = write!(out, "}}");
169            }
170            QueryCommand::One => {
171                let _ = writeln!(
172                    out,
173                    "func {}(ctx context.Context, db *sql.DB{}{}) ({}, error) {{",
174                    func_name, sep, param_list, struct_name
175                );
176                let args_str = if args.is_empty() {
177                    String::new()
178                } else {
179                    format!(", {}", args.join(", "))
180                };
181                let _ = writeln!(
182                    out,
183                    "\trow := db.QueryRowContext(ctx, \"{}\"{})",
184                    sql, args_str
185                );
186                let _ = writeln!(out, "\tvar r {}", struct_name);
187                let scan_fields: Vec<String> = columns
188                    .iter()
189                    .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
190                    .collect();
191                let _ = writeln!(out, "\terr := row.Scan({})", scan_fields.join(", "));
192                let _ = writeln!(out, "\treturn r, err");
193                let _ = write!(out, "}}");
194            }
195            QueryCommand::Many | QueryCommand::Batch => {
196                let _ = writeln!(
197                    out,
198                    "func {}(ctx context.Context, db *sql.DB{}{}) ([]{}, error) {{",
199                    func_name, sep, param_list, struct_name
200                );
201                let args_str = if args.is_empty() {
202                    String::new()
203                } else {
204                    format!(", {}", args.join(", "))
205                };
206                let _ = writeln!(
207                    out,
208                    "\trows, err := db.QueryContext(ctx, \"{}\"{})",
209                    sql, args_str
210                );
211                let _ = writeln!(out, "\tif err != nil {{");
212                let _ = writeln!(out, "\t\treturn nil, err");
213                let _ = writeln!(out, "\t}}");
214                let _ = writeln!(out, "\tdefer rows.Close()");
215                let _ = writeln!(out, "\tvar result []{}", struct_name);
216                let _ = writeln!(out, "\tfor rows.Next() {{");
217                let _ = writeln!(out, "\t\tvar r {}", struct_name);
218                let scan_fields: Vec<String> = columns
219                    .iter()
220                    .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
221                    .collect();
222                let _ = writeln!(
223                    out,
224                    "\t\tif err := rows.Scan({}); err != nil {{",
225                    scan_fields.join(", ")
226                );
227                let _ = writeln!(out, "\t\t\treturn nil, err");
228                let _ = writeln!(out, "\t\t}}");
229                let _ = writeln!(out, "\t\tresult = append(result, r)");
230                let _ = writeln!(out, "\t}}");
231                let _ = writeln!(out, "\treturn result, rows.Err()");
232                let _ = write!(out, "}}");
233            }
234        }
235
236        Ok(out)
237    }
238
239    fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
240        let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
241        let mut out = String::new();
242        let _ = writeln!(out, "type {} string", type_name);
243        let _ = writeln!(out);
244        let _ = writeln!(out, "const (");
245        for value in &enum_info.values {
246            let variant = enum_variant_name(value, &self.manifest.naming);
247            let _ = writeln!(
248                out,
249                "\t{}{} {} = \"{}\"",
250                type_name, variant, type_name, value
251            );
252        }
253        let _ = write!(out, ")");
254        Ok(out)
255    }
256
257    fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
258        let name = to_pascal_case(&composite.sql_name);
259        let mut out = String::new();
260        let _ = writeln!(out, "type {} struct {{", name);
261        if !composite.fields.is_empty() {
262            for field in &composite.fields {
263                let field_name = to_pascal_case(&field.name);
264                let go_type = resolve_type(&field.neutral_type, &self.manifest, false)
265                    .map(|t| t.into_owned())
266                    .unwrap_or_else(|_| "interface{}".to_string());
267                let json_tag = &field.name;
268                let _ = writeln!(out, "\t{} {} `json:\"{}\"`", field_name, go_type, json_tag);
269            }
270        }
271        let _ = write!(out, "}}");
272        Ok(out)
273    }
274}