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        "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"
56            .to_string()
57    }
58
59    fn generate_row_struct(
60        &self,
61        query_name: &str,
62        columns: &[ResolvedColumn],
63    ) -> Result<String, ScytheError> {
64        let struct_name = row_struct_name(query_name, &self.manifest.naming);
65        let mut out = String::new();
66        let _ = writeln!(out, "type {} struct {{", struct_name);
67        for col in columns {
68            let field = to_pascal_case(&col.field_name);
69            let json_tag = &col.field_name;
70            let _ = writeln!(out, "\t{} {} `json:\"{}\"`", field, col.full_type, json_tag);
71        }
72        let _ = write!(out, "}}");
73        Ok(out)
74    }
75
76    fn generate_model_struct(
77        &self,
78        table_name: &str,
79        columns: &[ResolvedColumn],
80    ) -> Result<String, ScytheError> {
81        let name = to_pascal_case(table_name);
82        self.generate_row_struct(&name, columns)
83    }
84
85    fn generate_query_fn(
86        &self,
87        analyzed: &AnalyzedQuery,
88        struct_name: &str,
89        columns: &[ResolvedColumn],
90        params: &[ResolvedParam],
91    ) -> Result<String, ScytheError> {
92        let func_name = fn_name(&analyzed.name, &self.manifest.naming);
93        let sql = super::clean_sql_oneline(&analyzed.sql);
94
95        let param_list = params
96            .iter()
97            .map(|p| {
98                let field = to_pascal_case(&p.field_name);
99                format!("{} {}", field, p.full_type)
100            })
101            .collect::<Vec<_>>()
102            .join(", ");
103        let sep = if param_list.is_empty() { "" } else { ", " };
104
105        let args = params
106            .iter()
107            .map(|p| to_pascal_case(&p.field_name).into_owned())
108            .collect::<Vec<_>>();
109
110        let mut out = String::new();
111
112        match &analyzed.command {
113            QueryCommand::Exec => {
114                // :exec - returns error only
115                let _ = writeln!(
116                    out,
117                    "func {}(ctx context.Context, db *pgxpool.Pool{}{}) error {{",
118                    func_name, sep, param_list
119                );
120                let args_str = if args.is_empty() {
121                    String::new()
122                } else {
123                    format!(", {}", args.join(", "))
124                };
125                let _ = writeln!(out, "\t_, err := db.Exec(ctx, \"{}\"{})", sql, args_str);
126                let _ = writeln!(out, "\treturn err");
127                let _ = write!(out, "}}");
128            }
129            QueryCommand::ExecResult | QueryCommand::ExecRows => {
130                // :exec_rows - returns affected row count
131                let _ = writeln!(
132                    out,
133                    "func {}(ctx context.Context, db *pgxpool.Pool{}{}) (int64, error) {{",
134                    func_name, sep, param_list
135                );
136                let args_str = if args.is_empty() {
137                    String::new()
138                } else {
139                    format!(", {}", args.join(", "))
140                };
141                let _ = writeln!(
142                    out,
143                    "\tresult, err := db.Exec(ctx, \"{}\"{})",
144                    sql, args_str
145                );
146                let _ = writeln!(out, "\tif err != nil {{");
147                let _ = writeln!(out, "\t\treturn 0, err");
148                let _ = writeln!(out, "\t}}");
149                let _ = writeln!(out, "\treturn result.RowsAffected(), nil");
150                let _ = write!(out, "}}");
151            }
152            QueryCommand::One => {
153                // :one - returns single struct
154                let _ = writeln!(
155                    out,
156                    "func {}(ctx context.Context, db *pgxpool.Pool{}{}) ({}, error) {{",
157                    func_name, sep, param_list, struct_name
158                );
159                let args_str = if args.is_empty() {
160                    String::new()
161                } else {
162                    format!(", {}", args.join(", "))
163                };
164                let _ = writeln!(out, "\trow := db.QueryRow(ctx, \"{}\"{})", sql, args_str);
165                let _ = writeln!(out, "\tvar r {}", struct_name);
166                let scan_fields: Vec<String> = columns
167                    .iter()
168                    .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
169                    .collect();
170                let _ = writeln!(out, "\terr := row.Scan({})", scan_fields.join(", "));
171                let _ = writeln!(out, "\treturn r, err");
172                let _ = write!(out, "}}");
173            }
174            QueryCommand::Many | QueryCommand::Batch => {
175                // :many - returns slice
176                let _ = writeln!(
177                    out,
178                    "func {}(ctx context.Context, db *pgxpool.Pool{}{}) ([]{}, error) {{",
179                    func_name, sep, param_list, struct_name
180                );
181                let args_str = if args.is_empty() {
182                    String::new()
183                } else {
184                    format!(", {}", args.join(", "))
185                };
186                let _ = writeln!(out, "\trows, err := db.Query(ctx, \"{}\"{})", sql, args_str);
187                let _ = writeln!(out, "\tif err != nil {{");
188                let _ = writeln!(out, "\t\treturn nil, err");
189                let _ = writeln!(out, "\t}}");
190                let _ = writeln!(out, "\tdefer rows.Close()");
191                let _ = writeln!(out, "\tvar result []{}", struct_name);
192                let _ = writeln!(out, "\tfor rows.Next() {{");
193                let _ = writeln!(out, "\t\tvar r {}", struct_name);
194                let scan_fields: Vec<String> = columns
195                    .iter()
196                    .map(|c| format!("&r.{}", to_pascal_case(&c.field_name)))
197                    .collect();
198                let _ = writeln!(
199                    out,
200                    "\t\tif err := rows.Scan({}); err != nil {{",
201                    scan_fields.join(", ")
202                );
203                let _ = writeln!(out, "\t\t\treturn nil, err");
204                let _ = writeln!(out, "\t\t}}");
205                let _ = writeln!(out, "\t\tresult = append(result, r)");
206                let _ = writeln!(out, "\t}}");
207                let _ = writeln!(out, "\treturn result, rows.Err()");
208                let _ = write!(out, "}}");
209            }
210        }
211
212        Ok(out)
213    }
214
215    fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
216        let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
217        let mut out = String::new();
218        let _ = writeln!(out, "type {} string", type_name);
219        let _ = writeln!(out);
220        let _ = writeln!(out, "const (");
221        for value in &enum_info.values {
222            let variant = enum_variant_name(value, &self.manifest.naming);
223            let _ = writeln!(
224                out,
225                "\t{}{} {} = \"{}\"",
226                type_name, variant, type_name, value
227            );
228        }
229        let _ = write!(out, ")");
230        Ok(out)
231    }
232
233    fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
234        let name = to_pascal_case(&composite.sql_name);
235        let mut out = String::new();
236        let _ = writeln!(out, "type {} struct {{", name);
237        if composite.fields.is_empty() {
238            // empty struct
239        } else {
240            for field in &composite.fields {
241                let field_name = to_pascal_case(&field.name);
242                let go_type = resolve_type(&field.neutral_type, &self.manifest, false)
243                    .map(|t| t.into_owned())
244                    .unwrap_or_else(|_| "interface{}".to_string());
245                let json_tag = &field.name;
246                let _ = writeln!(out, "\t{} {} `json:\"{}\"`", field_name, go_type, json_tag);
247            }
248        }
249        let _ = write!(out, "}}");
250        Ok(out)
251    }
252}
scythe_codegen/backends/go_pgx.rs

scythe_codegen/backends/
go_pgx.rs
