use crate::errors::ErdifyError;
use clap::Parser;
use std::env;
#[derive(Parser, Debug)]
#[command(name = "erdify", version, about)]
pub struct Args {
#[arg(short, long)]
pub url: Option<String>,
#[arg(long)]
pub schema: Option<String>,
#[arg(long, conflicts_with = "ignore_tables")]
pub table: Option<String>,
#[arg(long, conflicts_with = "table")]
pub ignore_tables: Option<String>,
#[arg(long, conflicts_with = "full")]
pub minimal: bool,
#[arg(long, conflicts_with = "minimal")]
pub full: bool,
#[arg(short, long)]
pub output: Option<String>,
#[arg(long)]
pub title: Option<String>,
}
#[derive(Debug)]
pub struct ConnectionInfo {
pub host: String,
pub port: u16,
pub database: String,
pub user: String,
pub password: String,
}
impl Args {
pub fn parse_csv<'a>(&self, value: Option<&'a str>) -> Vec<&'a str> {
match value {
Some(v) if !v.is_empty() => v
.split(',')
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.collect(),
_ => Vec::new(),
}
}
pub fn output_mode(&self) -> OutputMode {
if self.minimal {
OutputMode::Minimal
} else if self.full {
OutputMode::Full
} else {
OutputMode::Default
}
}
pub fn parse_url(&self) -> Result<ConnectionInfo, ErdifyError> {
let url_str = match &self.url {
Some(u) if !u.is_empty() => Some(u.clone()),
_ => None,
};
let url_str = match url_str {
Some(url) => url,
None => match env::var("DATABASE_URL") {
Ok(v) if !v.is_empty() => v,
_ => {
return Err(ErdifyError::InvalidUrl(
"no url provided; use --url or the DATABASE_URL variable".into(),
));
}
},
};
parse_postgres_url(&url_str)
}
}
const DEFAULT_PORT: u16 = 5432;
fn parse_postgres_url(url: &str) -> Result<ConnectionInfo, ErdifyError> {
let url = url::Url::parse(url)
.map_err(|e| ErdifyError::InvalidUrl(format!("invalid url format: {e}")))?;
let scheme = url.scheme();
if scheme != "postgresql" && scheme != "postgres" && scheme != "pg" {
return Err(ErdifyError::InvalidUrl(format!(
"expected url scheme: postgresql/postgres/pg, got: {scheme}"
)));
}
let host = url
.host_str()
.filter(|h| !h.is_empty())
.ok_or_else(|| ErdifyError::InvalidUrl("no host in the url".into()))?
.to_string();
let port = url.port().unwrap_or(DEFAULT_PORT);
let database = url
.path_segments()
.and_then(|mut segs| segs.next())
.filter(|s| !s.is_empty())
.map(percent_decode)
.ok_or_else(|| ErdifyError::InvalidUrl("no database name in the url".into()))?;
let user = percent_decode(url.username());
let password = url.password().map(percent_decode).unwrap_or_default();
Ok(ConnectionInfo {
host,
port,
database,
user,
password,
})
}
fn percent_decode(raw: &str) -> String {
percent_encoding::percent_decode_str(raw)
.decode_utf8()
.map_or_else(|_| raw.to_string(), |s| s.into_owned())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputMode {
Minimal,
Default,
Full,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_csv_empty() {
let args = Args::parse_from(["erdify"]);
assert!(args.parse_csv(None).is_empty());
assert!(args.parse_csv(Some("")).is_empty());
}
#[test]
fn test_parse_csv_single() {
let args = Args::parse_from(["erdify"]);
let result = args.parse_csv(Some("public"));
assert_eq!(result, vec!["public"]);
}
#[test]
fn test_parse_csv_multiple() {
let args = Args::parse_from(["erdify"]);
let result = args.parse_csv(Some("public,extended,custom"));
assert_eq!(result, vec!["public", "extended", "custom"]);
}
#[test]
fn test_parse_csv_with_spaces() {
let args = Args::parse_from(["erdify"]);
let result = args.parse_csv(Some(" public , extended "));
assert_eq!(result, vec!["public", "extended"]);
}
#[test]
fn test_table_and_ignore_tables_conflict() {
let result =
Args::try_parse_from(["erdify", "--table", "users", "--ignore-tables", "logs"]);
assert!(result.is_err());
}
#[test]
fn test_output_mode_default() {
let args = Args::parse_from(["erdify"]);
assert_eq!(args.output_mode(), OutputMode::Default);
}
#[test]
fn test_output_mode_minimal() {
let args = Args::parse_from(["erdify", "--minimal"]);
assert_eq!(args.output_mode(), OutputMode::Minimal);
}
#[test]
fn test_output_mode_full() {
let args = Args::parse_from(["erdify", "--full"]);
assert_eq!(args.output_mode(), OutputMode::Full);
}
#[test]
fn test_parse_url_valid() {
let args = Args::parse_from([
"erdify",
"--url",
"postgresql://admin:secret@localhost:5432/mydb",
]);
let info = args.parse_url().unwrap();
assert_eq!(info.host, "localhost");
assert_eq!(info.port, 5432);
assert_eq!(info.database, "mydb");
assert_eq!(info.user, "admin");
assert_eq!(info.password, "secret");
}
#[test]
fn test_parse_url_default_port() {
let args = Args::parse_from([
"erdify",
"--url",
"postgresql://user@db.example.com/production",
]);
let info = args.parse_url().unwrap();
assert_eq!(info.port, 5432);
assert_eq!(info.database, "production");
}
#[test]
fn test_parse_url_percent_encoded_credentials() {
let args = Args::parse_from([
"erdify",
"--url",
"postgresql://ad%40min:p%40ss%2Fword@localhost:5432/mydb",
]);
let info = args.parse_url().unwrap();
assert_eq!(info.user, "ad@min");
assert_eq!(info.password, "p@ss/word");
}
#[test]
fn test_parse_url_missing_database() {
let args = Args::parse_from(["erdify", "--url", "postgresql://user@localhost:5432/"]);
assert!(args.parse_url().is_err());
}
#[test]
fn test_parse_url_missing_host() {
let args = Args::parse_from(["erdify", "--url", "postgresql:///dbname"]);
let result = args.parse_url();
assert!(result.is_err());
}
#[test]
fn test_parse_url_no_url_no_env() {
let orig = env::var("DATABASE_URL").ok();
unsafe {
env::remove_var("DATABASE_URL");
}
let args = Args::parse_from(["erdify"]);
let result = args.parse_url();
assert!(result.is_err());
if let Some(val) = orig {
unsafe {
env::set_var("DATABASE_URL", val);
}
}
}
}