use crate::Result;
use crate::error::DbError;
use crate::udbc::connection::Connection;
use crate::udbc::driver::Driver;
use crate::udbc::mysql::connection::MysqlConnection;
use crate::udbc::{DEFAULT_DB_NAME, PoolOptions};
use async_trait::async_trait;
use mysql_async::{Opts, OptsBuilder, Pool, PoolConstraints, PoolOpts};
use std::time::Duration;
use tokio::time::timeout;
const MYSQL_TYPE: &str = "mysql";
pub struct MysqlDriver {
url: String,
name: String,
options: Option<PoolOptions>,
pool: Option<Pool>,
}
impl MysqlDriver {
pub fn new(url: impl Into<String>) -> Self {
Self {
name: DEFAULT_DB_NAME.to_string(),
url: url.into(),
options: None,
pool: None,
}
}
pub fn name(mut self, name: String) -> Self {
self.name = name;
self
}
pub fn options(mut self, options: PoolOptions) -> Self {
self.options = Some(options);
self
}
pub fn build(mut self) -> Result<Self> {
let opts = Opts::from_url(&self.url).map_err(|e| {
DbError::DbUrlError(format!("[{}] Invalid connection URL: {}", self.name, e))
})?;
let mut builder = OptsBuilder::from_opts(opts);
builder = builder.tcp_keepalive(Some(60_000u32));
if let Some(options) = &self.options {
if options.max_open_conns == 0 {
return Err(self.err_context(
"Invalid pool constraints: max_open_conns must be greater than 0",
));
}
let constraints = PoolConstraints::new(
options.max_idle_conns as usize,
options.max_open_conns as usize,
)
.ok_or_else(|| {
self.err_context(format!(
"Invalid pool constraints: max_idle_conns ({}) > max_open_conns ({})",
options.max_idle_conns, options.max_open_conns
))
})?;
let mut pool_opts = PoolOpts::default().with_constraints(constraints);
if options.max_lifetime > 0 {
pool_opts = pool_opts
.with_inactive_connection_ttl(Duration::from_secs(options.max_lifetime));
}
builder = builder.pool_opts(pool_opts);
}
let pool = Pool::new(builder);
self.pool = Some(pool);
Ok(self)
}
fn err_context<T: std::fmt::Display>(&self, msg: T) -> DbError {
DbError::DbError(format!("[{}] {}", self.name, msg))
}
}
#[async_trait]
impl Driver for MysqlDriver {
fn name(&self) -> &str {
&self.name
}
fn r#type(&self) -> &str {
MYSQL_TYPE
}
fn placeholder(&self, _param_seq: usize, _param_name: &str) -> String {
"?".to_string()
}
async fn acquire(&self) -> Result<Box<dyn Connection>> {
let pool = self.pool.as_ref().ok_or_else(|| {
self.err_context("Connection pool not initialized (call build() first)")
})?;
let get_conn_fut = pool.get_conn();
let conn = if let Some(options) = &self.options {
if options.timeout > 0 {
match timeout(Duration::from_secs(options.timeout), get_conn_fut).await {
Ok(result) => result,
Err(_) => {
return Err(self.err_context(format!(
"Connection acquisition timed out (timeout: {}s)",
options.timeout
)));
}
}
} else {
get_conn_fut.await
}
} else {
get_conn_fut.await
}
.map_err(|e| self.err_context(e))?;
Ok(Box::new(MysqlConnection::new(conn)))
}
async fn close(&self) -> Result<()> {
if let Some(pool) = &self.pool {
pool.clone()
.disconnect()
.await
.map_err(|e| self.err_context(format!("Failed to close pool: {}", e)))?;
}
Ok(())
}
}