use crate::DbkitError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Backend {
Postgres,
MySql,
Sqlite,
}
impl Backend {
pub fn from_url(url: &str) -> Result<Self, DbkitError> {
match url.split(':').next().unwrap_or("") {
"postgres" | "postgresql" => Ok(Backend::Postgres),
"mysql" => Ok(Backend::MySql),
"sqlite" => Ok(Backend::Sqlite),
other => Err(DbkitError::UnsupportedBackend(other.to_string())),
}
}
fn scheme(self) -> &'static str {
match self {
Backend::Postgres => "postgres",
Backend::MySql => "mysql",
Backend::Sqlite => "sqlite",
}
}
fn default_port(self) -> u16 {
match self {
Backend::Postgres => 5432,
Backend::MySql => 3306,
Backend::Sqlite => 0,
}
}
}
#[derive(Debug, Clone)]
pub struct DbkitConfig {
pub url: String,
pub pool_size: usize,
pub connect_timeout_secs: u64,
pub idle_timeout_secs: u64,
pub auto_create_db: bool,
}
impl DbkitConfig {
pub fn from_url(url: &str) -> Self {
Self {
url: url.to_string(),
pool_size: 16,
connect_timeout_secs: 30,
idle_timeout_secs: 300,
auto_create_db: true,
}
}
pub fn builder() -> ConfigBuilder {
ConfigBuilder::default()
}
}
pub struct ConfigBuilder {
backend: Backend,
host: String,
port: Option<u16>,
database: String,
user: Option<String>,
password: Option<String>,
pool_size: usize,
connect_timeout_secs: u64,
idle_timeout_secs: u64,
auto_create_db: bool,
ssl_mode: SslMode,
}
#[derive(Debug, Clone, Copy, Default)]
pub enum SslMode {
#[default]
Disable,
Prefer,
Require,
}
impl Default for ConfigBuilder {
fn default() -> Self {
Self {
backend: Backend::Postgres,
host: "localhost".into(),
port: None,
database: "postgres".into(),
user: None,
password: None,
pool_size: 16,
connect_timeout_secs: 30,
idle_timeout_secs: 300,
auto_create_db: true,
ssl_mode: SslMode::default(),
}
}
}
impl ConfigBuilder {
pub fn backend(mut self, backend: Backend) -> Self {
self.backend = backend;
self
}
pub fn host(mut self, host: &str) -> Self {
self.host = host.to_string();
self
}
pub fn port(mut self, port: u16) -> Self {
self.port = Some(port);
self
}
pub fn database(mut self, database: &str) -> Self {
self.database = database.to_string();
self
}
pub fn user(mut self, user: &str) -> Self {
self.user = Some(user.to_string());
self
}
pub fn password(mut self, password: &str) -> Self {
self.password = Some(password.to_string());
self
}
pub fn pool_size(mut self, size: usize) -> Self {
self.pool_size = size;
self
}
pub fn connect_timeout_secs(mut self, secs: u64) -> Self {
self.connect_timeout_secs = secs;
self
}
pub fn idle_timeout_secs(mut self, secs: u64) -> Self {
self.idle_timeout_secs = secs;
self
}
pub fn auto_create_db(mut self, enabled: bool) -> Self {
self.auto_create_db = enabled;
self
}
pub fn ssl_mode(mut self, mode: SslMode) -> Self {
self.ssl_mode = mode;
self
}
pub fn build(self) -> DbkitConfig {
let url = match self.backend {
Backend::Sqlite => format!("sqlite://{}", self.database),
backend => {
use percent_encoding::{NON_ALPHANUMERIC, utf8_percent_encode};
let enc = |s: &str| utf8_percent_encode(s, NON_ALPHANUMERIC).to_string();
let auth = match (&self.user, &self.password) {
(Some(u), Some(p)) => format!("{}:{}@", enc(u), enc(p)),
(Some(u), None) => format!("{}@", enc(u)),
_ => String::new(),
};
let port = self.port.unwrap_or_else(|| backend.default_port());
let ssl_param = match (backend, self.ssl_mode) {
(Backend::MySql, SslMode::Disable) => "?ssl-mode=DISABLED",
(Backend::MySql, SslMode::Prefer) => "?ssl-mode=PREFERRED",
(Backend::MySql, SslMode::Require) => "?ssl-mode=REQUIRED",
(_, SslMode::Disable) => "?sslmode=disable",
(_, SslMode::Prefer) => "?sslmode=prefer",
(_, SslMode::Require) => "?sslmode=require",
};
format!(
"{}://{}{}:{}/{}{}",
backend.scheme(),
auth,
self.host,
port,
self.database,
ssl_param
)
}
};
DbkitConfig {
url,
pool_size: self.pool_size,
connect_timeout_secs: self.connect_timeout_secs,
idle_timeout_secs: self.idle_timeout_secs,
auto_create_db: self.auto_create_db,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_from_url() {
let config = DbkitConfig::from_url("postgres://localhost/mydb");
assert_eq!(config.url, "postgres://localhost/mydb");
assert_eq!(config.pool_size, 16);
assert!(config.auto_create_db);
}
#[test]
fn test_builder_full() {
let config = DbkitConfig::builder()
.host("db.example.com")
.port(5433)
.database("myapp")
.user("admin")
.password("secret")
.pool_size(32)
.connect_timeout_secs(10)
.ssl_mode(SslMode::Require)
.build();
assert_eq!(
config.url,
"postgres://admin:secret@db.example.com:5433/myapp?sslmode=require"
);
assert_eq!(config.pool_size, 32);
assert_eq!(config.connect_timeout_secs, 10);
}
#[test]
fn test_builder_minimal() {
let config = DbkitConfig::builder().database("test").build();
assert_eq!(config.url, "postgres://localhost:5432/test?sslmode=disable");
}
#[test]
fn test_builder_percent_encodes_credentials() {
let config = DbkitConfig::builder()
.host("localhost")
.port(5432)
.user("postgres")
.password("LexLuthern246!!??")
.database("sports_ai_baseball")
.build();
assert_eq!(
config.url,
"postgres://postgres:LexLuthern246%21%21%3F%3F@localhost:5432/sports_ai_baseball?sslmode=disable"
);
}
#[test]
fn test_builder_user_no_password() {
let config = DbkitConfig::builder()
.user("readonly")
.database("prod")
.build();
assert_eq!(config.url, "postgres://readonly@localhost:5432/prod?sslmode=disable");
}
#[test]
fn test_builder_mysql_default_port() {
let config = DbkitConfig::builder()
.backend(Backend::MySql)
.user("root")
.database("app")
.build();
assert_eq!(config.url, "mysql://root@localhost:3306/app?ssl-mode=DISABLED");
}
#[test]
fn test_builder_mysql_ssl_mode_spelling() {
let config = DbkitConfig::builder()
.backend(Backend::MySql)
.database("app")
.ssl_mode(SslMode::Require)
.build();
assert_eq!(config.url, "mysql://localhost:3306/app?ssl-mode=REQUIRED");
}
#[test]
fn test_builder_sqlite() {
let config = DbkitConfig::builder()
.backend(Backend::Sqlite)
.database("data/app.db")
.build();
assert_eq!(config.url, "sqlite://data/app.db");
}
#[test]
fn test_backend_from_url() {
assert_eq!(Backend::from_url("postgres://x/y").unwrap(), Backend::Postgres);
assert_eq!(Backend::from_url("postgresql://x/y").unwrap(), Backend::Postgres);
assert_eq!(Backend::from_url("mysql://x/y").unwrap(), Backend::MySql);
assert_eq!(Backend::from_url("sqlite://f.db").unwrap(), Backend::Sqlite);
assert!(Backend::from_url("oracle://x/y").is_err());
}
}