|
| 1 | +use std::collections::HashMap; |
| 2 | + |
| 3 | +use anyhow::Context as _; |
| 4 | +use sqlx::postgres::{PgConnectOptions, PgSslMode}; |
| 5 | + |
| 6 | +#[derive(Debug)] |
| 7 | +pub(crate) struct HydraConfig { |
| 8 | + options: HashMap<String, String>, |
| 9 | +} |
| 10 | + |
| 11 | +impl HydraConfig { |
| 12 | + pub(crate) fn load() -> Self { |
| 13 | + let mut options = HashMap::new(); |
| 14 | + |
| 15 | + let path = match std::env::var("HYDRA_CONFIG") { |
| 16 | + Ok(p) if !p.is_empty() => p, |
| 17 | + _ => return Self { options }, |
| 18 | + }; |
| 19 | + |
| 20 | + let contents = match fs_err::read_to_string(&path) { |
| 21 | + Ok(c) => c, |
| 22 | + Err(e) => { |
| 23 | + tracing::warn!("could not read HYDRA_CONFIG at {path}: {e}"); |
| 24 | + return Self { options }; |
| 25 | + } |
| 26 | + }; |
| 27 | + |
| 28 | + for line in contents.lines() { |
| 29 | + // Strip comments |
| 30 | + let line = match line.find('#') { |
| 31 | + Some(pos) => &line[..pos], |
| 32 | + None => line, |
| 33 | + }; |
| 34 | + let line = line.trim(); |
| 35 | + |
| 36 | + let Some(eq) = line.find('=') else { |
| 37 | + continue; |
| 38 | + }; |
| 39 | + |
| 40 | + let key = line[..eq].trim(); |
| 41 | + let value = line[eq + 1..].trim(); |
| 42 | + |
| 43 | + if key.is_empty() { |
| 44 | + continue; |
| 45 | + } |
| 46 | + |
| 47 | + options.insert(key.to_owned(), value.to_owned()); |
| 48 | + } |
| 49 | + |
| 50 | + Self { options } |
| 51 | + } |
| 52 | + |
| 53 | + pub(crate) fn get_int(&self, key: &str, default: u64) -> u64 { |
| 54 | + self.options |
| 55 | + .get(key) |
| 56 | + .and_then(|v| v.parse().ok()) |
| 57 | + .unwrap_or(default) |
| 58 | + } |
| 59 | +} |
| 60 | + |
| 61 | +/// Parse a `HYDRA_DBI` environment variable into `PgConnectOptions`. |
| 62 | +/// |
| 63 | +/// Accepts strings like `dbi:Pg:dbname=hydra;host=localhost;port=5432`. |
| 64 | +pub(crate) fn parse_hydra_dbi() -> anyhow::Result<PgConnectOptions> { |
| 65 | + let dbi = std::env::var("HYDRA_DBI").unwrap_or_else(|_| "dbi:Pg:dbname=hydra;".to_owned()); |
| 66 | + parse_dbi(&dbi) |
| 67 | +} |
| 68 | + |
| 69 | +fn parse_dbi(dbi: &str) -> anyhow::Result<PgConnectOptions> { |
| 70 | + let params = dbi |
| 71 | + .strip_prefix("dbi:Pg:") |
| 72 | + .or_else(|| dbi.strip_prefix("DBI:Pg:")) |
| 73 | + .context("$HYDRA_DBI does not denote a PostgreSQL database")?; |
| 74 | + |
| 75 | + let mut opts = PgConnectOptions::new(); |
| 76 | + |
| 77 | + for pair in params.split(';').filter(|s| !s.is_empty()) { |
| 78 | + let (key, value) = pair |
| 79 | + .split_once('=') |
| 80 | + .with_context(|| format!("invalid DBI parameter: {pair}"))?; |
| 81 | + match key.trim() { |
| 82 | + "dbname" => opts = opts.database(value.trim()), |
| 83 | + "host" => opts = opts.host(value.trim()), |
| 84 | + "port" => { |
| 85 | + opts = opts.port( |
| 86 | + value |
| 87 | + .trim() |
| 88 | + .parse() |
| 89 | + .with_context(|| format!("invalid port: {value}"))?, |
| 90 | + ); |
| 91 | + } |
| 92 | + "user" => opts = opts.username(value.trim()), |
| 93 | + "password" => opts = opts.password(value.trim()), |
| 94 | + "application_name" => opts = opts.application_name(value.trim()), |
| 95 | + // The C++ evaluator used libpq which understood these |
| 96 | + // natively; without explicit handling TLS-required |
| 97 | + // deployments would fail. |
| 98 | + "sslmode" => { |
| 99 | + let mode = match value.trim() { |
| 100 | + "disable" => PgSslMode::Disable, |
| 101 | + "allow" => PgSslMode::Allow, |
| 102 | + "prefer" => PgSslMode::Prefer, |
| 103 | + "require" => PgSslMode::Require, |
| 104 | + "verify-ca" => PgSslMode::VerifyCa, |
| 105 | + "verify-full" => PgSslMode::VerifyFull, |
| 106 | + v => anyhow::bail!("invalid sslmode: {v}"), |
| 107 | + }; |
| 108 | + opts = opts.ssl_mode(mode); |
| 109 | + } |
| 110 | + "sslrootcert" => opts = opts.ssl_root_cert(value.trim()), |
| 111 | + "sslcert" => opts = opts.ssl_client_cert(value.trim()), |
| 112 | + "sslkey" => opts = opts.ssl_client_key(value.trim()), |
| 113 | + // Warn rather than bail on unknown parameters, to avoid |
| 114 | + // breaking on libpq keywords we haven't mapped yet. |
| 115 | + other => { |
| 116 | + tracing::warn!("ignoring unsupported DBI parameter: {other}"); |
| 117 | + } |
| 118 | + } |
| 119 | + } |
| 120 | + |
| 121 | + Ok(opts) |
| 122 | +} |
| 123 | + |
| 124 | +#[cfg(test)] |
| 125 | +mod tests { |
| 126 | + use super::*; |
| 127 | + |
| 128 | + #[test] |
| 129 | + fn parse_simple_dbi() { |
| 130 | + parse_dbi("dbi:Pg:dbname=hydra;host=localhost;port=5432").unwrap(); |
| 131 | + } |
| 132 | + |
| 133 | + #[test] |
| 134 | + fn parse_dbi_uppercase() { |
| 135 | + parse_dbi("DBI:Pg:dbname=testdb;").unwrap(); |
| 136 | + } |
| 137 | +} |
0 commit comments