1use crate::{
3 error::Error,
4 ident::{field_ident, to_pascal_case},
5 plugin::GenerateRequestView,
6 types::TypeMap,
7};
8
9pub struct EnumInfo {
10 pub schema: String,
11 pub pg_name: String,
12 pub rust_name: String,
13 pub type_name: String,
16 pub vals: Vec<String>,
18}
19
20pub struct CompositeField {
21 pub pg_name: String,
22 pub rust_ident: proc_macro2::Ident,
23 pub rust_type: String,
25}
26
27pub struct CompositeInfo {
28 pub schema: String,
29 pub pg_name: String,
30 pub rust_name: String,
31 pub type_name: String,
32 pub fields: Vec<CompositeField>,
33}
34
35pub struct CatalogInfo {
36 pub enums: Vec<EnumInfo>,
37 pub composites: Vec<CompositeInfo>,
38}
39
40pub fn walk(
42 request: &GenerateRequestView<'_>,
43 type_map: &mut TypeMap,
44) -> Result<CatalogInfo, Error> {
45 let mut enums = Vec::new();
46 let mut composites = Vec::new();
47
48 let catalog = match request.catalog.as_option() {
49 Some(c) => c,
50 None => return Ok(CatalogInfo { enums, composites }),
51 };
52
53 let default_schema = catalog.default_schema;
54
55 for schema in catalog.schemas.iter() {
57 for e in schema.enums.iter() {
58 let rust_name = make_rust_name(schema.name, e.name, default_schema);
59 let type_name = sqlx_type_name(schema.name, e.name, default_schema);
60 let vals: Vec<String> = e.vals.iter().map(|v| v.to_string()).collect();
62 let pg_key = if schema.name == default_schema || schema.name.is_empty() {
63 e.name.to_string()
64 } else {
65 format!("{}.{}", schema.name, e.name)
66 };
67 type_map.register(&pg_key, &rust_name, false);
68 enums.push(EnumInfo {
69 schema: schema.name.to_string(),
70 pg_name: e.name.to_string(),
71 rust_name,
72 type_name,
73 vals,
74 });
75 }
76 }
77
78 for schema in catalog.schemas.iter() {
84 let composite_names: std::collections::HashSet<&str> =
85 schema.composite_types.iter().map(|c| c.name).collect();
86 if composite_names.is_empty() {
87 continue;
88 }
89
90 for table in schema.tables.iter() {
91 let rel = match table.rel.as_option() {
92 Some(r) => r,
93 None => continue,
94 };
95 if !composite_names.contains(rel.name) {
96 continue;
97 }
98
99 let rust_name = make_rust_name(schema.name, rel.name, default_schema);
100 let type_name = sqlx_type_name(schema.name, rel.name, default_schema);
101
102 let mut fields = Vec::new();
103 for col in table.columns.iter() {
104 let pg_type = col.r#type.as_option().map(|t| t.name).unwrap_or("");
105 let array_dims = if col.array_dims > 0 {
106 col.array_dims as usize
107 } else {
108 usize::from(col.is_array)
109 };
110 let rust_type = type_map
113 .resolve_pg_type_dims(pg_type, true, array_dims)
114 .ok_or_else(|| {
115 Error::Codegen(format!(
116 "composite '{}.{}' field '{}' has unknown type '{pg_type}'",
117 schema.name, rel.name, col.name
118 ))
119 })?;
120 fields.push(CompositeField {
121 pg_name: col.name.to_string(),
122 rust_ident: field_ident(col.name),
123 rust_type: rust_type.rust_type,
124 });
125 }
126
127 let pg_key = if schema.name == default_schema || schema.name.is_empty() {
128 rel.name.to_string()
129 } else {
130 format!("{}.{}", schema.name, rel.name)
131 };
132 type_map.register(&pg_key, &rust_name, false);
133 composites.push(CompositeInfo {
134 schema: schema.name.to_string(),
135 pg_name: rel.name.to_string(),
136 rust_name,
137 type_name,
138 fields,
139 });
140 }
141 }
142
143 Ok(CatalogInfo { enums, composites })
144}
145
146fn make_rust_name(schema: &str, name: &str, default_schema: &str) -> String {
147 if schema == default_schema || schema.is_empty() {
148 to_pascal_case(name)
149 } else {
150 format!("{}{}", to_pascal_case(schema), to_pascal_case(name))
151 }
152}
153
154fn sqlx_type_name(schema: &str, name: &str, default_schema: &str) -> String {
155 if schema == default_schema || schema.is_empty() {
156 name.to_string()
157 } else {
158 format!("{}.{}", schema, name)
159 }
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165
166 #[test]
167 fn make_rust_name_public_schema() {
168 assert_eq!(make_rust_name("public", "status", "public"), "Status");
169 }
170
171 #[test]
172 fn make_rust_name_non_default_schema() {
173 assert_eq!(
174 make_rust_name("myschema", "status", "public"),
175 "MyschemaStatus"
176 );
177 }
178
179 #[test]
180 fn sqlx_type_name_public_schema() {
181 assert_eq!(sqlx_type_name("public", "status", "public"), "status");
182 }
183
184 #[test]
185 fn sqlx_type_name_non_default_schema() {
186 assert_eq!(
187 sqlx_type_name("myschema", "status", "public"),
188 "myschema.status"
189 );
190 }
191
192 #[test]
193 fn make_rust_name_empty_schema() {
194 assert_eq!(make_rust_name("", "status", "public"), "Status");
195 }
196
197 #[test]
198 fn sqlx_type_name_empty_schema() {
199 assert_eq!(sqlx_type_name("", "status", "public"), "status");
200 }
201
202 #[test]
203 fn make_rust_name_underscored_schema() {
204 assert_eq!(
205 make_rust_name("my_app_schema", "order_status", "public"),
206 "MyAppSchemaOrderStatus"
207 );
208 }
209}