Skip to main content

sqlx_gen/typemap/
mod.rs

1pub mod mysql;
2pub mod postgres;
3pub mod sqlite;
4
5use std::collections::HashMap;
6
7use crate::cli::DatabaseKind;
8use crate::introspect::{ColumnInfo, SchemaInfo};
9
10/// Resolved Rust type with its required imports.
11#[derive(Debug, Clone)]
12pub struct RustType {
13    pub path: String,
14    pub needs_import: Option<String>,
15}
16
17impl RustType {
18    pub fn simple(path: &str) -> Self {
19        Self {
20            path: path.to_string(),
21            needs_import: None,
22        }
23    }
24
25    pub fn with_import(path: &str, import: &str) -> Self {
26        Self {
27            path: path.to_string(),
28            needs_import: Some(import.to_string()),
29        }
30    }
31
32    pub fn wrap_option(self) -> Self {
33        Self {
34            path: format!("Option<{}>", self.path),
35            needs_import: self.needs_import,
36        }
37    }
38
39    pub fn wrap_vec(self) -> Self {
40        Self {
41            path: format!("Vec<{}>", self.path),
42            needs_import: self.needs_import,
43        }
44    }
45}
46
47pub fn map_column(
48    col: &ColumnInfo,
49    db_kind: DatabaseKind,
50    schema_info: &SchemaInfo,
51    overrides: &HashMap<String, String>,
52) -> RustType {
53    // Check type overrides first
54    if let Some(override_type) = overrides.get(&col.udt_name) {
55        let rt = RustType::simple(override_type);
56        return if col.is_nullable { rt.wrap_option() } else { rt };
57    }
58
59    let base = match db_kind {
60        DatabaseKind::Postgres => postgres::map_type(&col.udt_name, schema_info),
61        DatabaseKind::Mysql => mysql::map_type(&col.data_type, &col.udt_name),
62        DatabaseKind::Sqlite => sqlite::map_type(&col.udt_name),
63    };
64
65    if col.is_nullable {
66        base.wrap_option()
67    } else {
68        base
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use crate::introspect::SchemaInfo;
76    use std::collections::HashMap;
77
78    fn make_col(udt_name: &str, data_type: &str, nullable: bool) -> ColumnInfo {
79        ColumnInfo {
80            name: "test".to_string(),
81            data_type: data_type.to_string(),
82            udt_name: udt_name.to_string(),
83            is_nullable: nullable,
84            is_primary_key: false,
85            ordinal_position: 0,
86            schema_name: "public".to_string(),
87        }
88    }
89
90    // --- RustType::simple ---
91
92    #[test]
93    fn test_simple_creates_without_import() {
94        let rt = RustType::simple("i32");
95        assert_eq!(rt.path, "i32");
96        assert!(rt.needs_import.is_none());
97    }
98
99    #[test]
100    fn test_simple_path_correct() {
101        let rt = RustType::simple("String");
102        assert_eq!(rt.path, "String");
103    }
104
105    #[test]
106    fn test_simple_no_import() {
107        let rt = RustType::simple("bool");
108        assert_eq!(rt.needs_import, None);
109    }
110
111    // --- RustType::with_import ---
112
113    #[test]
114    fn test_with_import_creates_with_import() {
115        let rt = RustType::with_import("Uuid", "use uuid::Uuid;");
116        assert_eq!(rt.path, "Uuid");
117        assert_eq!(rt.needs_import, Some("use uuid::Uuid;".to_string()));
118    }
119
120    #[test]
121    fn test_with_import_path_correct() {
122        let rt = RustType::with_import("DateTime<Utc>", "use chrono::{DateTime, Utc};");
123        assert_eq!(rt.path, "DateTime<Utc>");
124    }
125
126    #[test]
127    fn test_with_import_import_present() {
128        let rt = RustType::with_import("Value", "use serde_json::Value;");
129        assert!(rt.needs_import.is_some());
130    }
131
132    // --- RustType::wrap_option ---
133
134    #[test]
135    fn test_wrap_option_wraps_path() {
136        let rt = RustType::simple("i32").wrap_option();
137        assert_eq!(rt.path, "Option<i32>");
138    }
139
140    #[test]
141    fn test_wrap_option_preserves_import() {
142        let rt = RustType::with_import("Uuid", "use uuid::Uuid;").wrap_option();
143        assert_eq!(rt.path, "Option<Uuid>");
144        assert_eq!(rt.needs_import, Some("use uuid::Uuid;".to_string()));
145    }
146
147    #[test]
148    fn test_wrap_option_double_wrap() {
149        let rt = RustType::simple("i32").wrap_option().wrap_option();
150        assert_eq!(rt.path, "Option<Option<i32>>");
151    }
152
153    // --- RustType::wrap_vec ---
154
155    #[test]
156    fn test_wrap_vec_wraps_path() {
157        let rt = RustType::simple("i32").wrap_vec();
158        assert_eq!(rt.path, "Vec<i32>");
159    }
160
161    #[test]
162    fn test_wrap_vec_preserves_import() {
163        let rt = RustType::with_import("Uuid", "use uuid::Uuid;").wrap_vec();
164        assert_eq!(rt.path, "Vec<Uuid>");
165        assert_eq!(rt.needs_import, Some("use uuid::Uuid;".to_string()));
166    }
167
168    // --- map_column ---
169
170    #[test]
171    fn test_override_takes_precedence() {
172        let col = make_col("uuid", "uuid", false);
173        let schema = SchemaInfo::default();
174        let mut overrides = HashMap::new();
175        overrides.insert("uuid".to_string(), "MyUuid".to_string());
176        let rt = map_column(&col, DatabaseKind::Postgres, &schema, &overrides);
177        assert_eq!(rt.path, "MyUuid");
178        assert!(rt.needs_import.is_none());
179    }
180
181    #[test]
182    fn test_override_with_nullable() {
183        let col = make_col("uuid", "uuid", true);
184        let schema = SchemaInfo::default();
185        let mut overrides = HashMap::new();
186        overrides.insert("uuid".to_string(), "MyUuid".to_string());
187        let rt = map_column(&col, DatabaseKind::Postgres, &schema, &overrides);
188        assert_eq!(rt.path, "Option<MyUuid>");
189    }
190
191    #[test]
192    fn test_no_override_dispatches_postgres() {
193        let col = make_col("int4", "integer", false);
194        let schema = SchemaInfo::default();
195        let overrides = HashMap::new();
196        let rt = map_column(&col, DatabaseKind::Postgres, &schema, &overrides);
197        assert_eq!(rt.path, "i32");
198    }
199
200    #[test]
201    fn test_nullable_without_override() {
202        let col = make_col("int4", "integer", true);
203        let schema = SchemaInfo::default();
204        let overrides = HashMap::new();
205        let rt = map_column(&col, DatabaseKind::Postgres, &schema, &overrides);
206        assert_eq!(rt.path, "Option<i32>");
207    }
208}
209