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

scythe_codegen/backends/
go_pgx.rs
