Skip to content
Clone
d6dc02f
252 lines 8.3 KB Rust
use std::time::Duration;
use sea_orm::ConnectOptions;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct DatabaseConfig {
pub url: Option<String>,
pub host: Option<String>,
pub port: Option<u16>,
pub username: Option<String>,
pub password: Option<String>,
pub database: Option<String>,
pub schema_search_path: Option<String>,
pub application_name: Option<String>,
pub ssl_mode: Option<SslMode>,
pub max_connections: Option<u32>,
pub min_connections: Option<u32>,
pub connect_timeout_secs: Option<u64>,
pub idle_timeout_secs: Option<u64>,
pub acquire_timeout_secs: Option<u64>,
pub max_lifetime_secs: Option<u64>,
pub sqlx_logging: Option<bool>,
pub replicas: Vec<DatabaseEndpoint>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct DatabaseEndpoint {
pub url: Option<String>,
pub host: Option<String>,
pub port: Option<u16>,
pub username: Option<String>,
pub password: Option<String>,
pub database: Option<String>,
pub application_name: Option<String>,
pub ssl_mode: Option<SslMode>,
pub max_connections: Option<u32>,
pub min_connections: Option<u32>,
pub connect_timeout_secs: Option<u64>,
pub idle_timeout_secs: Option<u64>,
pub acquire_timeout_secs: Option<u64>,
pub max_lifetime_secs: Option<u64>,
}
impl DatabaseEndpoint {
pub fn url(&self) -> String {
if let Some(url) = self.url.as_deref().filter(|s| !s.is_empty()) {
return url.to_string();
}
let host = self.host.as_deref().unwrap_or("localhost");
let port = self.port.unwrap_or(5432);
let user = self.username.as_deref().unwrap_or("postgres");
let db = self.database.as_deref().unwrap_or("postgres");
let password = match &self.password {
Some(p) if !p.is_empty() => format!(":{p}"),
_ => String::new(),
};
let mut url = format!("postgres://{user}{password}@{host}:{port}/{db}");
let mut params: Vec<String> = Vec::new();
if let Some(app) = self.application_name.as_deref().filter(|s| !s.is_empty()) {
params.push(format!("application_name={}", {
let value = app;
let mut out = String::with_capacity(value.len());
for b in value.as_bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(*b as char);
}
_ => out.push_str(&format!("%{:02X}", b)),
}
}
out
}));
}
if let Some(mode) = self.ssl_mode {
params.push(format!("sslmode={}", mode.as_query_value()));
}
if !params.is_empty() {
url.push('?');
url.push_str(&params.join("&"));
}
url
}
pub fn validate(&self) -> Result<(), String> {
if self.url.as_deref().is_some_and(|s| !s.is_empty()) {
return Ok(());
}
if self.host.as_deref().filter(|s| !s.is_empty()).is_none() {
return Err("endpoint.host must be set when endpoint.url is not provided".to_string());
}
if matches!(self.port, Some(0)) {
return Err("endpoint.port must be greater than 0".to_string());
}
Ok(())
}
pub fn apply_to(&self, opts: &mut ConnectOptions) -> Result<(), String> {
self.validate()?;
if let Some(v) = self.max_connections {
opts.max_connections(v);
}
if let Some(v) = self.min_connections {
opts.min_connections(v);
}
if let Some(v) = self.connect_timeout_secs {
opts.connect_timeout(Duration::from_secs(v));
}
if let Some(v) = self.idle_timeout_secs {
opts.idle_timeout(Duration::from_secs(v));
}
if let Some(v) = self.acquire_timeout_secs {
opts.acquire_timeout(Duration::from_secs(v));
}
if let Some(v) = self.max_lifetime_secs {
opts.max_lifetime(Duration::from_secs(v));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SslMode {
#[default]
Disable,
Prefer,
Require,
VerifyCa,
VerifyFull,
}
impl SslMode {
pub fn as_query_value(&self) -> &'static str {
match self {
SslMode::Disable => "disable",
SslMode::Prefer => "prefer",
SslMode::Require => "require",
SslMode::VerifyCa => "verify-ca",
SslMode::VerifyFull => "verify-full",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseRole {
Primary,
Replica,
}
impl DatabaseConfig {
pub fn primary_url(&self) -> String {
let endpoint = DatabaseEndpoint {
url: self.url.clone(),
host: self.host.clone(),
port: self.port,
username: self.username.clone(),
password: self.password.clone(),
database: self.database.clone(),
application_name: self.application_name.clone(),
ssl_mode: self.ssl_mode,
max_connections: self.max_connections,
min_connections: self.min_connections,
connect_timeout_secs: self.connect_timeout_secs,
idle_timeout_secs: self.idle_timeout_secs,
acquire_timeout_secs: self.acquire_timeout_secs,
max_lifetime_secs: self.max_lifetime_secs,
};
endpoint.url()
}
pub fn replica_count(&self) -> usize {
self.replicas.len()
}
pub fn has_replicas(&self) -> bool {
!self.replicas.is_empty()
}
pub fn validate(&self) -> Result<(), String> {
if self.url.as_deref().is_some_and(|s| !s.is_empty()) {
return Ok(());
}
if self.host.as_deref().filter(|s| !s.is_empty()).is_none() {
return Err("database.host must be set when database.url is not provided".to_string());
}
if matches!(self.port, Some(0)) {
return Err("database.port must be greater than 0".to_string());
}
for (i, replica) in self.replicas.iter().enumerate() {
replica
.validate()
.map_err(|e| format!("replicas[{i}]: {e}"))?;
}
Ok(())
}
pub fn primary_connect_options(&self) -> Result<ConnectOptions, String> {
self.validate()?;
let url = self.primary_url();
let mut opts = ConnectOptions::new(url);
if let Some(v) = self.schema_search_path.as_deref().filter(|s| !s.is_empty()) {
opts.set_schema_search_path(v);
}
if let Some(v) = self.sqlx_logging {
opts.sqlx_logging(v);
}
let endpoint = DatabaseEndpoint {
url: self.url.clone(),
host: self.host.clone(),
port: self.port,
username: self.username.clone(),
password: self.password.clone(),
database: self.database.clone(),
application_name: self.application_name.clone(),
ssl_mode: self.ssl_mode,
max_connections: self.max_connections,
min_connections: self.min_connections,
connect_timeout_secs: self.connect_timeout_secs,
idle_timeout_secs: self.idle_timeout_secs,
acquire_timeout_secs: self.acquire_timeout_secs,
max_lifetime_secs: self.max_lifetime_secs,
};
endpoint.apply_to(&mut opts)?;
Ok(opts)
}
pub fn replica_connect_options(&self, index: usize) -> Result<ConnectOptions, String> {
let replica = self
.replicas
.get(index)
.ok_or_else(|| format!("replica index {index} out of range"))?;
replica.validate()?;
let url = replica.url();
let mut opts = ConnectOptions::new(url);
if let Some(v) = self.schema_search_path.as_deref().filter(|s| !s.is_empty()) {
opts.set_schema_search_path(v);
}
if let Some(v) = self.sqlx_logging {
opts.sqlx_logging(v);
}
replica.apply_to(&mut opts)?;
Ok(opts)
}
}