use std::sync::Arc;
use super::{
backend::DatabaseBackend,
error::Result,
query_builder::{DeleteBuilder, InsertBuilder, SelectBuilder, UpdateBuilder},
};
#[cfg(feature = "postgres")]
use super::dialect::PostgresBackend;
#[cfg(feature = "postgres")]
const SQLSTATE_INVALID_CATALOG_NAME: &str = "3D000";
#[cfg(feature = "sqlite")]
use super::dialect::SqliteBackend;
#[cfg(feature = "mysql")]
use super::dialect::MySqlBackend;
#[derive(Clone)]
pub struct DatabaseConnection {
backend: Arc<dyn DatabaseBackend>,
is_cockroachdb: bool,
}
#[cfg(feature = "di")]
#[async_trait::async_trait]
impl reinhardt_di::Injectable for DatabaseConnection {
async fn inject(ctx: &reinhardt_di::InjectionContext) -> reinhardt_di::DiResult<Self> {
if let Some(conn) = ctx.get_singleton::<Self>() {
return Ok(std::sync::Arc::try_unwrap(conn).unwrap_or_else(|arc| (*arc).clone()));
}
if let Some(conn) = ctx.get_request::<Self>() {
return Ok(std::sync::Arc::try_unwrap(conn).unwrap_or_else(|arc| (*arc).clone()));
}
Err(reinhardt_di::DiError::NotRegistered {
type_name: std::any::type_name::<Self>().to_string(),
hint: "Use InjectionContextBuilder::singleton(db_connection) to register a \
DatabaseConnection. Create it with DatabaseConnection::connect_postgres(), \
connect_sqlite(), or connect_mysql()."
.to_string(),
})
}
async fn inject_uncached(ctx: &reinhardt_di::InjectionContext) -> reinhardt_di::DiResult<Self> {
Self::inject(ctx).await
}
}
impl DatabaseConnection {
pub fn new(backend: Arc<dyn DatabaseBackend>) -> Self {
Self::new_with_flavor(backend, false)
}
pub fn new_with_flavor(backend: Arc<dyn DatabaseBackend>, is_cockroachdb: bool) -> Self {
Self {
backend,
is_cockroachdb,
}
}
#[cfg(feature = "postgres")]
pub async fn connect_postgres(url: &str) -> Result<Self> {
Self::connect_postgres_with_pool_size(url, None).await
}
#[cfg(feature = "postgres")]
pub async fn connect_postgres_with_pool_size(
url: &str,
pool_size: Option<u32>,
) -> Result<Self> {
let pool = Self::build_postgres_pool(url, pool_size).await?;
let is_cockroachdb = Self::probe_cockroachdb(&pool).await;
Ok(Self {
backend: Arc::new(PostgresBackend::new(pool)),
is_cockroachdb,
})
}
#[cfg(feature = "postgres")]
async fn probe_cockroachdb(pool: &sqlx::PgPool) -> bool {
sqlx::query_scalar::<_, bool>("SELECT version() LIKE 'CockroachDB%'")
.fetch_one(pool)
.await
.unwrap_or(false)
}
#[cfg(feature = "postgres")]
pub async fn connect_postgres_or_create(url: &str) -> Result<Self> {
Self::connect_postgres_or_create_with_pool_size(url, None).await
}
#[cfg(feature = "postgres")]
async fn build_postgres_pool(
url: &str,
pool_size: Option<u32>,
) -> std::result::Result<sqlx::PgPool, sqlx::Error> {
use sqlx::postgres::PgPoolOptions;
use std::time::Duration;
let max_connections = pool_size
.or_else(|| {
std::env::var("DATABASE_POOL_MAX_CONNECTIONS")
.ok()
.and_then(|v| v.parse::<u32>().ok())
})
.unwrap_or(20);
PgPoolOptions::new()
.max_connections(max_connections)
.min_connections(1) .acquire_timeout(Duration::from_secs(10)) .idle_timeout(Some(Duration::from_secs(10))) .max_lifetime(Some(Duration::from_secs(30 * 60))) .connect(url)
.await
}
#[cfg(feature = "postgres")]
pub async fn connect_postgres_or_create_with_pool_size(
url: &str,
pool_size: Option<u32>,
) -> Result<Self> {
match Self::build_postgres_pool(url, pool_size).await {
Ok(pool) => {
let is_cockroachdb = Self::probe_cockroachdb(&pool).await;
return Ok(Self {
backend: Arc::new(PostgresBackend::new(pool)),
is_cockroachdb,
});
}
Err(e) => {
let is_db_not_found = matches!(
&e,
sqlx::Error::Database(db_err) if db_err.code().as_deref() == Some(SQLSTATE_INVALID_CATALOG_NAME)
);
if !is_db_not_found {
return Err(e.into());
}
}
}
let (admin_url, db_name) = Self::parse_postgres_url_for_creation(url)?;
use sqlx::postgres::PgPoolOptions;
use std::time::Duration;
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.acquire_timeout(Duration::from_secs(10))
.connect(&admin_url)
.await
.map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to connect to postgres database for auto-creation: {}",
e
))
})?;
let create_sql = format!("CREATE DATABASE \"{}\"", db_name.replace('"', "\"\""));
sqlx::query(&create_sql)
.execute(&admin_pool)
.await
.map_err(|e| {
super::error::DatabaseError::QueryError(format!(
"Failed to create database '{}': {}",
db_name, e
))
})?;
admin_pool.close().await;
Self::connect_postgres_with_pool_size(url, pool_size).await
}
#[cfg(feature = "postgres")]
fn parse_postgres_url_for_creation(url: &str) -> Result<(String, String)> {
let url_without_scheme = url
.strip_prefix("postgres://")
.or_else(|| url.strip_prefix("postgresql://"))
.ok_or_else(|| {
super::error::DatabaseError::ConnectionError(
"Invalid PostgreSQL URL: must start with postgres:// or postgresql://"
.to_string(),
)
})?;
let (path_part, query_part) = match url_without_scheme.find('?') {
Some(pos) => (&url_without_scheme[..pos], Some(&url_without_scheme[pos..])),
None => (url_without_scheme, None),
};
let last_slash_pos = path_part.rfind('/').ok_or_else(|| {
super::error::DatabaseError::ConnectionError(
"Invalid PostgreSQL URL: no database name found".to_string(),
)
})?;
let host_part = &path_part[..last_slash_pos];
let db_name = &path_part[last_slash_pos + 1..];
if db_name.is_empty() {
return Err(super::error::DatabaseError::ConnectionError(
"Invalid PostgreSQL URL: database name is empty".to_string(),
));
}
let admin_url = match query_part {
Some(params) => format!("postgres://{}/postgres{}", host_part, params),
None => format!("postgres://{}/postgres", host_part),
};
Ok((admin_url, db_name.to_string()))
}
#[cfg(feature = "sqlite")]
pub async fn connect_sqlite(url: &str) -> Result<Self> {
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool};
use std::path::Path;
use std::str::FromStr;
if url == "sqlite::memory:" {
let pool = SqlitePool::connect(url).await?;
return Ok(Self {
backend: Arc::new(SqliteBackend::new(pool)),
is_cockroachdb: false,
});
}
let file_path = if url.starts_with("sqlite:///") {
url.trim_start_matches("sqlite:///").to_string()
} else if url.starts_with("sqlite://") {
let rel_path = url.trim_start_matches("sqlite://");
std::env::current_dir()
.map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to get current directory: {}",
e
))
})?
.join(rel_path)
.to_string_lossy()
.to_string()
} else if url.starts_with("sqlite:") {
let rel_path = url.trim_start_matches("sqlite:");
std::env::current_dir()
.map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to get current directory: {}",
e
))
})?
.join(rel_path)
.to_string_lossy()
.to_string()
} else {
url.to_string()
};
let db_path = Path::new(&file_path);
let normalized_path = if db_path.exists() {
db_path.canonicalize().map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to canonicalize path {}: {}",
db_path.display(),
e
))
})?
} else {
if db_path.is_absolute() {
db_path.to_path_buf()
} else {
std::env::current_dir()
.map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to get current directory: {}",
e
))
})?
.join(db_path)
}
};
if let Some(parent) = normalized_path.parent()
&& !parent.as_os_str().is_empty()
&& !parent.exists()
{
std::fs::create_dir_all(parent).map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Failed to create database directory {}: {}",
parent.display(),
e
))
})?;
}
let path_str = normalized_path.to_string_lossy().replace('\\', "/");
let absolute_url = format!("sqlite:///{}", path_str);
let options = SqliteConnectOptions::from_str(&absolute_url)
.map_err(|e| {
super::error::DatabaseError::ConnectionError(format!(
"Invalid SQLite URL '{}': {}",
absolute_url, e
))
})?
.create_if_missing(true);
let pool = SqlitePool::connect_with(options).await?;
Ok(Self {
backend: Arc::new(SqliteBackend::new(pool)),
is_cockroachdb: false,
})
}
#[cfg(feature = "sqlite")]
pub fn from_sqlite_pool(pool: sqlx::SqlitePool) -> Self {
Self {
backend: Arc::new(SqliteBackend::new(pool)),
is_cockroachdb: false,
}
}
#[cfg(feature = "mysql")]
pub async fn connect_mysql(url: &str) -> Result<Self> {
use sqlx::MySqlPool;
let pool = MySqlPool::connect(url).await?;
Ok(Self {
backend: Arc::new(MySqlBackend::new(pool)),
is_cockroachdb: false,
})
}
pub fn backend(&self) -> Arc<dyn DatabaseBackend> {
self.backend.clone()
}
pub fn database_type(&self) -> super::types::DatabaseType {
self.backend.database_type()
}
pub fn is_cockroachdb(&self) -> bool {
self.is_cockroachdb
}
pub fn insert(&self, table: impl Into<String>) -> InsertBuilder {
InsertBuilder::new(self.backend.clone(), table)
}
pub fn update(&self, table: impl Into<String>) -> UpdateBuilder {
UpdateBuilder::new(self.backend.clone(), table)
}
pub fn select(&self) -> SelectBuilder {
SelectBuilder::new(self.backend.clone())
}
pub fn delete(&self, table: impl Into<String>) -> DeleteBuilder {
DeleteBuilder::new(self.backend.clone(), table)
}
#[cfg(feature = "settings")]
pub fn database_url_from<S>(settings: &S, env_override: Option<&str>) -> Result<String>
where
S: reinhardt_conf::HasCoreSettings + ?Sized,
{
if let Some(url) = env_override {
return Ok(url.to_string());
}
let core = settings.core();
let db_config = core.databases.get("default").ok_or_else(|| {
super::error::DatabaseError::ConnectionError(
"Database configuration `core.databases.default` not found in settings."
.to_string(),
)
})?;
Ok(db_config.to_url())
}
pub async fn execute(
&self,
sql: &str,
params: Vec<super::types::QueryValue>,
) -> Result<super::types::QueryResult> {
self.backend.execute(sql, params).await
}
pub async fn fetch_one(
&self,
sql: &str,
params: Vec<super::types::QueryValue>,
) -> Result<super::types::Row> {
self.backend.fetch_one(sql, params).await
}
pub async fn fetch_all(
&self,
sql: &str,
params: Vec<super::types::QueryValue>,
) -> Result<Vec<super::types::Row>> {
self.backend.fetch_all(sql, params).await
}
pub async fn fetch_optional(
&self,
sql: &str,
params: Vec<super::types::QueryValue>,
) -> Result<Option<super::types::Row>> {
self.backend.fetch_optional(sql, params).await
}
pub async fn begin(&self) -> Result<Box<dyn super::types::TransactionExecutor>> {
self.backend.begin().await
}
pub async fn begin_with_isolation(
&self,
level: super::types::IsolationLevel,
) -> Result<Box<dyn super::types::TransactionExecutor>> {
self.backend.begin_with_isolation(level).await
}
#[cfg(feature = "postgres")]
pub fn into_postgres(&self) -> Option<sqlx::PgPool> {
self.backend
.as_any()
.downcast_ref::<super::dialect::PostgresBackend>()
.map(|backend| backend.pool().clone())
}
#[cfg(feature = "sqlite")]
pub fn into_sqlite(&self) -> Option<sqlx::SqlitePool> {
self.backend
.as_any()
.downcast_ref::<super::dialect::SqliteBackend>()
.map(|backend| backend.pool().clone())
}
#[cfg(feature = "mysql")]
pub fn into_mysql(&self) -> Option<sqlx::MySqlPool> {
self.backend
.as_any()
.downcast_ref::<super::dialect::MySqlBackend>()
.map(|backend| backend.pool().clone())
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
fn build_create_database_sql(db_name: &str) -> String {
format!("CREATE DATABASE \"{}\"", db_name.replace('"', "\"\""))
}
#[rstest]
fn test_create_database_sql_normal_name() {
let db_name = "my_database";
let sql = build_create_database_sql(db_name);
assert_eq!(sql, "CREATE DATABASE \"my_database\"");
}
#[rstest]
fn test_create_database_sql_injection_with_double_quotes() {
let db_name = "test\"; DROP TABLE users; --";
let sql = build_create_database_sql(db_name);
assert_eq!(sql, "CREATE DATABASE \"test\"\"; DROP TABLE users; --\"");
}
#[rstest]
fn test_create_database_sql_injection_with_multiple_quotes() {
let db_name = "db\"\"injection";
let sql = build_create_database_sql(db_name);
assert_eq!(sql, "CREATE DATABASE \"db\"\"\"\"injection\"");
}
#[cfg(feature = "postgres")]
#[rstest]
fn test_parse_postgres_url_extracts_db_name() {
let url = "postgres://user:pass@localhost:5432/testdb";
let (admin_url, db_name) =
super::DatabaseConnection::parse_postgres_url_for_creation(url).unwrap();
assert_eq!(db_name, "testdb");
assert_eq!(admin_url, "postgres://user:pass@localhost:5432/postgres");
}
}