Skip to main content

erdify_rs/
config.rs

1use crate::errors::ErdifyError;
2use clap::Parser;
3use std::env;
4
5/// Mermaid ER diagram generator from PostgreSQL.
6#[derive(Parser, Debug)]
7#[command(name = "erdify", version, about)]
8pub struct Args {
9    /// PostgreSQL connection URL (ex: postgresql://user:pass@host:5432/dbname)
10    #[arg(short, long)]
11    pub url: Option<String>,
12
13    /// Schemas to include, comma-separated (ex: public,extended)
14    #[arg(long)]
15    pub schema: Option<String>,
16
17    /// Tables to include, comma-separated (ex: users,orders)
18    #[arg(long, conflicts_with = "ignore_tables")]
19    pub table: Option<String>,
20
21    /// Tables to exclude, comma-separated (ex: logs,audit_trail)
22    #[arg(long, conflicts_with = "table")]
23    pub ignore_tables: Option<String>,
24
25    /// Minimal mode: columns only, without PK/FK metadata
26    #[arg(long, conflicts_with = "full")]
27    pub minimal: bool,
28
29    /// Full mode: columns + PK/FK/NOT NULL + relationships + constraints + indexes
30    #[arg(long, conflicts_with = "minimal")]
31    pub full: bool,
32
33    /// Output file (default: stdout)
34    #[arg(short, long)]
35    pub output: Option<String>,
36
37    /// Diagram title (default: derived from the database name)
38    #[arg(long)]
39    pub title: Option<String>,
40}
41
42/// Connection information extracted from a PostgreSQL URL.
43#[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    /// Parses comma-separated values into a vector of &str.
54    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    /// Determines the output mode.
66    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    /// Builds ConnectionInfo from --url or DATABASE_URL.
77    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
99/// Default PostgreSQL port, used when the url doesn't specify one.
100const DEFAULT_PORT: u16 = 5432;
101
102/// Parses a PostgreSQL URL in the `postgresql://user:pass@host:port/dbname` format.
103fn 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    // `port_or_known_default` doesn't know the postgresql scheme.
121    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    // Credentials are percent-encoded within a url: `p%40ss` must be
131    // passed to PostgreSQL as `p@ss`.
132    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
144/// Decodes the `%XX` sequences of a url component, leaving it unchanged
145/// if the result isn't valid UTF-8.
146fn 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/// Diagram output mode.
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154pub enum OutputMode {
155    /// Columns only, without PK/FK.
156    Minimal,
157    /// Columns + PK/FK (default).
158    Default,
159    /// Everything: columns + PK/FK/NOT NULL + relationships + constraints + indexes.
160    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        // Save and restore DATABASE_URL to avoid side effects.
275        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        // Restore.
285        if let Some(val) = orig {
286            unsafe {
287                env::set_var("DATABASE_URL", val);
288            }
289        }
290    }
291}