Files
MercuryToolbox/crates/sqlshape/src/connection.rs
T

139 lines
3.9 KiB
Rust

use common::CliError;
use url::Url;
use crate::model::{ConnectionSource, Engine};
pub(crate) fn source_from_parts(
engine: Option<Engine>,
raw_url: String,
) -> Result<ConnectionSource, CliError> {
let source_redacted = redact_source(&raw_url);
let engine = engine.or_else(|| infer_engine(&raw_url)).ok_or_else(|| {
CliError::usage(format!(
"--engine is required for ambiguous SQL source '{source_redacted}'"
))
})?;
Ok(ConnectionSource {
engine,
raw_url,
source_redacted,
})
}
pub(crate) fn infer_engine(raw_url: &str) -> Option<Engine> {
Url::parse(raw_url)
.ok()
.and_then(|url| Engine::parse(url.scheme()))
}
pub(crate) fn redact_source(raw: &str) -> String {
if let Ok(mut url) = Url::parse(raw) {
if url.password().is_some() {
let _ = url.set_password(Some("***"));
}
if url
.query_pairs()
.any(|(key, _)| is_sensitive_connection_key(&key))
{
let pairs = url
.query_pairs()
.map(|(key, value)| {
let rendered = if is_sensitive_connection_key(&key) {
"***".into()
} else {
value
};
(key.into_owned(), rendered.into_owned())
})
.collect::<Vec<_>>();
url.query_pairs_mut().clear().extend_pairs(pairs);
}
return url.to_string();
}
redact_key_value_passwords(raw)
}
fn redact_key_value_passwords(raw: &str) -> String {
let mut changed = false;
let parts = raw
.split(';')
.map(|part| {
let Some((key, value)) = part.split_once('=') else {
return part.to_string();
};
let key_trimmed = key.trim();
if is_sensitive_connection_key(key_trimmed) {
changed = true;
format!("{key}=***")
} else {
format!("{key}={value}")
}
})
.collect::<Vec<_>>();
if changed {
parts.join(";")
} else {
raw.to_string()
}
}
fn is_sensitive_connection_key(key: &str) -> bool {
let normalized = key
.chars()
.filter(|character| !matches!(character, '_' | '-' | ' '))
.flat_map(char::to_lowercase)
.collect::<String>();
matches!(
normalized.as_str(),
"password"
| "pwd"
| "token"
| "accesstoken"
| "apikey"
| "secret"
| "clientsecret"
| "credential"
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn redacts_url_and_ado_passwords() {
assert_eq!(
redact_source("postgres://user:secret@example.test/db"),
"postgres://user:***@example.test/db"
);
assert_eq!(
redact_source("server=tcp:host,1433;User ID=sa;Password=secret;Database=db"),
"server=tcp:host,1433;User ID=sa;Password=***;Database=db"
);
}
#[test]
fn redacts_url_query_and_key_value_secrets() {
assert_eq!(
redact_source(
"postgres://example.test/db?password=secret&access_token=abc&sslmode=require"
),
"postgres://example.test/db?password=***&access_token=***&sslmode=require"
);
assert_eq!(
redact_source("server=tcp:host,1433;User ID=sa;Access Token=abc;ApiKey=key"),
"server=tcp:host,1433;User ID=sa;Access Token=***;ApiKey=***"
);
}
#[test]
fn infers_engines_from_url_schemes_only() {
assert_eq!(
infer_engine("postgres://localhost/db"),
Some(Engine::Postgres)
);
assert_eq!(infer_engine("duckdb://local/path"), Some(Engine::DuckDb));
assert_eq!(infer_engine(r".\local.db"), None);
}
}