1use crate::errors::ErdifyError;
2use clap::Parser;
3use std::env;
4
5#[derive(Parser, Debug)]
7#[command(name = "erdify", version, about)]
8pub struct Args {
9 #[arg(short, long)]
11 pub url: Option<String>,
12
13 #[arg(long)]
15 pub schema: Option<String>,
16
17 #[arg(long, conflicts_with = "ignore_tables")]
19 pub table: Option<String>,
20
21 #[arg(long, conflicts_with = "table")]
23 pub ignore_tables: Option<String>,
24
25 #[arg(long, conflicts_with = "full")]
27 pub minimal: bool,
28
29 #[arg(long, conflicts_with = "minimal")]
31 pub full: bool,
32
33 #[arg(short, long)]
35 pub output: Option<String>,
36
37 #[arg(long)]
39 pub title: Option<String>,
40}
41
42#[derive(Debug)]
44pub struct ConnectionInfo {
45 pub host: String,
46 pub port: u16,
47 pub database: String,
48 pub user: String,
49 pub password: String,
50}
51
52impl Args {
53 pub fn parse_csv<'a>(&self, value: Option<&'a str>) -> Vec<&'a str> {
55 match value {
56 Some(v) if !v.is_empty() => v
57 .split(',')
58 .map(|s| s.trim())
59 .filter(|s| !s.is_empty())
60 .collect(),
61 _ => Vec::new(),
62 }
63 }
64
65 pub fn output_mode(&self) -> OutputMode {
67 if self.minimal {
68 OutputMode::Minimal
69 } else if self.full {
70 OutputMode::Full
71 } else {
72 OutputMode::Default
73 }
74 }
75
76 pub fn parse_url(&self) -> Result<ConnectionInfo, ErdifyError> {
78 let url_str = match &self.url {
79 Some(u) if !u.is_empty() => Some(u.clone()),
80 _ => None,
81 };
82
83 let url_str = match url_str {
84 Some(url) => url,
85 None => match env::var("DATABASE_URL") {
86 Ok(v) if !v.is_empty() => v,
87 _ => {
88 return Err(ErdifyError::InvalidUrl(
89 "no url provided; use --url or the DATABASE_URL variable".into(),
90 ));
91 }
92 },
93 };
94
95 parse_postgres_url(&url_str)
96 }
97}
98
99const DEFAULT_PORT: u16 = 5432;
101
102fn parse_postgres_url(url: &str) -> Result<ConnectionInfo, ErdifyError> {
104 let url = url::Url::parse(url)
105 .map_err(|e| ErdifyError::InvalidUrl(format!("invalid url format: {e}")))?;
106
107 let scheme = url.scheme();
108 if scheme != "postgresql" && scheme != "postgres" && scheme != "pg" {
109 return Err(ErdifyError::InvalidUrl(format!(
110 "expected url scheme: postgresql/postgres/pg, got: {scheme}"
111 )));
112 }
113
114 let host = url
115 .host_str()
116 .filter(|h| !h.is_empty())
117 .ok_or_else(|| ErdifyError::InvalidUrl("no host in the url".into()))?
118 .to_string();
119
120 let port = url.port().unwrap_or(DEFAULT_PORT);
122
123 let database = url
124 .path_segments()
125 .and_then(|mut segs| segs.next())
126 .filter(|s| !s.is_empty())
127 .map(percent_decode)
128 .ok_or_else(|| ErdifyError::InvalidUrl("no database name in the url".into()))?;
129
130 let user = percent_decode(url.username());
133 let password = url.password().map(percent_decode).unwrap_or_default();
134
135 Ok(ConnectionInfo {
136 host,
137 port,
138 database,
139 user,
140 password,
141 })
142}
143
144fn percent_decode(raw: &str) -> String {
147 percent_encoding::percent_decode_str(raw)
148 .decode_utf8()
149 .map_or_else(|_| raw.to_string(), |s| s.into_owned())
150}
151
152#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154pub enum OutputMode {
155 Minimal,
157 Default,
159 Full,
161}
162
163#[cfg(test)]
164mod tests {
165 use super::*;
166
167 #[test]
168 fn test_parse_csv_empty() {
169 let args = Args::parse_from(["erdify"]);
170 assert!(args.parse_csv(None).is_empty());
171 assert!(args.parse_csv(Some("")).is_empty());
172 }
173
174 #[test]
175 fn test_parse_csv_single() {
176 let args = Args::parse_from(["erdify"]);
177 let result = args.parse_csv(Some("public"));
178 assert_eq!(result, vec!["public"]);
179 }
180
181 #[test]
182 fn test_parse_csv_multiple() {
183 let args = Args::parse_from(["erdify"]);
184 let result = args.parse_csv(Some("public,extended,custom"));
185 assert_eq!(result, vec!["public", "extended", "custom"]);
186 }
187
188 #[test]
189 fn test_parse_csv_with_spaces() {
190 let args = Args::parse_from(["erdify"]);
191 let result = args.parse_csv(Some(" public , extended "));
192 assert_eq!(result, vec!["public", "extended"]);
193 }
194
195 #[test]
196 fn test_table_and_ignore_tables_conflict() {
197 let result =
198 Args::try_parse_from(["erdify", "--table", "users", "--ignore-tables", "logs"]);
199 assert!(result.is_err());
200 }
201
202 #[test]
203 fn test_output_mode_default() {
204 let args = Args::parse_from(["erdify"]);
205 assert_eq!(args.output_mode(), OutputMode::Default);
206 }
207
208 #[test]
209 fn test_output_mode_minimal() {
210 let args = Args::parse_from(["erdify", "--minimal"]);
211 assert_eq!(args.output_mode(), OutputMode::Minimal);
212 }
213
214 #[test]
215 fn test_output_mode_full() {
216 let args = Args::parse_from(["erdify", "--full"]);
217 assert_eq!(args.output_mode(), OutputMode::Full);
218 }
219
220 #[test]
221 fn test_parse_url_valid() {
222 let args = Args::parse_from([
223 "erdify",
224 "--url",
225 "postgresql://admin:secret@localhost:5432/mydb",
226 ]);
227 let info = args.parse_url().unwrap();
228 assert_eq!(info.host, "localhost");
229 assert_eq!(info.port, 5432);
230 assert_eq!(info.database, "mydb");
231 assert_eq!(info.user, "admin");
232 assert_eq!(info.password, "secret");
233 }
234
235 #[test]
236 fn test_parse_url_default_port() {
237 let args = Args::parse_from([
238 "erdify",
239 "--url",
240 "postgresql://user@db.example.com/production",
241 ]);
242 let info = args.parse_url().unwrap();
243 assert_eq!(info.port, 5432);
244 assert_eq!(info.database, "production");
245 }
246
247 #[test]
248 fn test_parse_url_percent_encoded_credentials() {
249 let args = Args::parse_from([
250 "erdify",
251 "--url",
252 "postgresql://ad%40min:p%40ss%2Fword@localhost:5432/mydb",
253 ]);
254 let info = args.parse_url().unwrap();
255 assert_eq!(info.user, "ad@min");
256 assert_eq!(info.password, "p@ss/word");
257 }
258
259 #[test]
260 fn test_parse_url_missing_database() {
261 let args = Args::parse_from(["erdify", "--url", "postgresql://user@localhost:5432/"]);
262 assert!(args.parse_url().is_err());
263 }
264
265 #[test]
266 fn test_parse_url_missing_host() {
267 let args = Args::parse_from(["erdify", "--url", "postgresql:///dbname"]);
268 let result = args.parse_url();
269 assert!(result.is_err());
270 }
271
272 #[test]
273 fn test_parse_url_no_url_no_env() {
274 let orig = env::var("DATABASE_URL").ok();
276 unsafe {
277 env::remove_var("DATABASE_URL");
278 }
279
280 let args = Args::parse_from(["erdify"]);
281 let result = args.parse_url();
282 assert!(result.is_err());
283
284 if let Some(val) = orig {
286 unsafe {
287 env::set_var("DATABASE_URL", val);
288 }
289 }
290 }
291}