Dionisio
fc78dcd35d
- Replace sql.chars().collect::<Vec<char>>() with direct byte-level iteration over sql.as_bytes() - Saves ~140 KB heap allocation for 35 KB schema (4× input size) - Remove unused _start variable - SQL is ASCII-safe so byte comparison is sufficient for all delimiters
327 lines
10 KiB
Rust
Executable File
327 lines
10 KiB
Rust
Executable File
use crate::common::config::AppConfig;
|
||
use sqlx::{PgPool, postgres::PgPoolOptions};
|
||
use std::time::Duration;
|
||
|
||
/// Database initialization error.
|
||
#[derive(Debug, thiserror::Error)]
|
||
#[error("{0}")]
|
||
pub struct DbError(String);
|
||
|
||
type Result<T> = std::result::Result<T, DbError>;
|
||
|
||
/// Segmented database pools.
|
||
///
|
||
/// `primary` is used for all user-facing request paths (REST, WebDAV, CalDAV,
|
||
/// CardDAV). `maintenance` is a smaller, isolated pool reserved for
|
||
/// background / batch operations (verify_integrity, garbage_collect,
|
||
/// update_all_users_storage_usage, trash cleanup) so they can never starve
|
||
/// interactive requests.
|
||
pub struct DbPools {
|
||
/// Pool for user-facing request paths.
|
||
pub primary: PgPool,
|
||
/// Pool for background / batch maintenance tasks.
|
||
pub maintenance: PgPool,
|
||
}
|
||
|
||
/// Create both the primary and maintenance database pools.
|
||
///
|
||
/// The schema is applied once via the primary pool. The maintenance pool
|
||
/// shares the same connection string but has its own, smaller budget.
|
||
pub async fn create_database_pools(config: &AppConfig) -> Result<DbPools> {
|
||
tracing::info!(
|
||
"Initializing PostgreSQL connections with URL: {}",
|
||
config
|
||
.database
|
||
.connection_string
|
||
.replace("postgres://", "postgres://[user]:[pass]@")
|
||
);
|
||
|
||
// --- primary pool ---
|
||
let primary = create_pool_with_retries(
|
||
&config.database.connection_string,
|
||
config.database.max_connections,
|
||
config.database.min_connections,
|
||
config.database.connect_timeout_secs,
|
||
config.database.idle_timeout_secs,
|
||
config.database.max_lifetime_secs,
|
||
"primary",
|
||
)
|
||
.await?;
|
||
|
||
// Apply schema through the primary pool (idempotent)
|
||
tracing::info!("Applying database schema...");
|
||
if let Err(e) = apply_schema(&primary).await {
|
||
return Err(DbError(format!(
|
||
"Database schema could not be applied: {}. \
|
||
Run manually: psql -f db/schema.sql",
|
||
e
|
||
)));
|
||
}
|
||
tracing::info!("Database schema applied successfully");
|
||
|
||
// --- maintenance pool ---
|
||
let maintenance = create_pool_with_retries(
|
||
&config.database.connection_string,
|
||
config.database.maintenance_max_connections,
|
||
config.database.maintenance_min_connections,
|
||
config.database.connect_timeout_secs,
|
||
config.database.idle_timeout_secs,
|
||
config.database.max_lifetime_secs,
|
||
"maintenance",
|
||
)
|
||
.await?;
|
||
|
||
tracing::info!(
|
||
"Database pools ready — primary: {} max / {} min, maintenance: {} max / {} min",
|
||
config.database.max_connections,
|
||
config.database.min_connections,
|
||
config.database.maintenance_max_connections,
|
||
config.database.maintenance_min_connections,
|
||
);
|
||
|
||
Ok(DbPools {
|
||
primary,
|
||
maintenance,
|
||
})
|
||
}
|
||
|
||
/// Internal helper: create a single pool with retry logic.
|
||
async fn create_pool_with_retries(
|
||
connection_string: &str,
|
||
max_connections: u32,
|
||
min_connections: u32,
|
||
connect_timeout_secs: u64,
|
||
idle_timeout_secs: u64,
|
||
max_lifetime_secs: u64,
|
||
label: &str,
|
||
) -> Result<PgPool> {
|
||
let mut attempt = 0;
|
||
const MAX_ATTEMPTS: usize = 5;
|
||
|
||
while attempt < MAX_ATTEMPTS {
|
||
attempt += 1;
|
||
tracing::info!(
|
||
"PostgreSQL {} pool connection attempt #{}/{}",
|
||
label,
|
||
attempt,
|
||
MAX_ATTEMPTS
|
||
);
|
||
|
||
match PgPoolOptions::new()
|
||
.max_connections(max_connections)
|
||
.min_connections(min_connections)
|
||
.acquire_timeout(Duration::from_secs(connect_timeout_secs))
|
||
.idle_timeout(Duration::from_secs(idle_timeout_secs))
|
||
.max_lifetime(Duration::from_secs(max_lifetime_secs))
|
||
.connect(connection_string)
|
||
.await
|
||
{
|
||
Ok(pool) => match sqlx::query("SELECT 1").execute(&pool).await {
|
||
Ok(_) => {
|
||
tracing::info!("PostgreSQL {} pool established successfully", label);
|
||
return Ok(pool);
|
||
}
|
||
Err(e) => {
|
||
tracing::error!("Error verifying {} pool connection: {}", label, e);
|
||
if attempt >= MAX_ATTEMPTS {
|
||
return Err(DbError(format!(
|
||
"Error verifying PostgreSQL {} pool connection: {}",
|
||
label, e
|
||
)));
|
||
}
|
||
}
|
||
},
|
||
Err(e) => {
|
||
tracing::error!(
|
||
"Error connecting to PostgreSQL {} pool (attempt {}/{}): {}",
|
||
label,
|
||
attempt,
|
||
MAX_ATTEMPTS,
|
||
e
|
||
);
|
||
if attempt >= MAX_ATTEMPTS {
|
||
return Err(DbError(format!(
|
||
"Error in PostgreSQL {} pool connection: {}",
|
||
label, e
|
||
)));
|
||
}
|
||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(DbError(format!(
|
||
"Could not establish PostgreSQL {} pool connection after {} attempts",
|
||
label, MAX_ATTEMPTS
|
||
)))
|
||
}
|
||
|
||
/// Apply the embedded schema.sql to the database.
|
||
/// First tries `raw_sql` (simple query protocol). If that fails, falls back
|
||
/// to splitting the SQL into individual statements and executing them one by one.
|
||
async fn apply_schema(pool: &PgPool) -> Result<()> {
|
||
let schema_sql = include_str!("../../db/schema.sql");
|
||
|
||
// Attempt 1: raw_sql sends the entire script via the simple query protocol
|
||
match sqlx::raw_sql(schema_sql).execute(pool).await {
|
||
Ok(_) => return Ok(()),
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
"raw_sql failed ({}), falling back to statement-by-statement execution",
|
||
e
|
||
);
|
||
}
|
||
}
|
||
|
||
// Attempt 2: split into individual statements respecting dollar-quoting
|
||
let statements = split_sql_statements(schema_sql);
|
||
for (i, stmt) in statements.iter().enumerate() {
|
||
let trimmed = stmt.trim();
|
||
if trimmed.is_empty() || trimmed == ";" {
|
||
continue;
|
||
}
|
||
if let Err(e) = sqlx::raw_sql(trimmed).execute(pool).await {
|
||
let preview = if trimmed.len() > 200 {
|
||
&trimmed[..200]
|
||
} else {
|
||
trimmed
|
||
};
|
||
tracing::error!(
|
||
"Schema statement {} failed: {}\n--- SQL ---\n{}\n-----------",
|
||
i + 1,
|
||
e,
|
||
preview
|
||
);
|
||
return Err(DbError(format!("Schema statement {} failed: {}", i + 1, e)));
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Split a SQL script into individual statements, correctly handling:
|
||
/// - Dollar-quoted blocks (`$BODY$...$BODY$`, `$$...$$`)
|
||
/// - Single-quoted strings (`'...'`)
|
||
/// - Line comments (`-- ...`)
|
||
/// - Block comments (`/* ... */`)
|
||
///
|
||
/// Uses byte-level iteration over the `&str` directly — no intermediate
|
||
/// `Vec<char>` allocation (saves ~4× the input size in heap memory).
|
||
/// SQL is ASCII-safe, so byte comparison is sufficient for all delimiters.
|
||
fn split_sql_statements(sql: &str) -> Vec<String> {
|
||
let mut statements = Vec::new();
|
||
let mut current = String::new();
|
||
let bytes = sql.as_bytes();
|
||
let len = bytes.len();
|
||
let mut i = 0;
|
||
|
||
while i < len {
|
||
let b = bytes[i];
|
||
|
||
// Line comment
|
||
if b == b'-' && i + 1 < len && bytes[i + 1] == b'-' {
|
||
while i < len && bytes[i] != b'\n' {
|
||
current.push(bytes[i] as char);
|
||
i += 1;
|
||
}
|
||
continue;
|
||
}
|
||
|
||
// Block comment
|
||
if b == b'/' && i + 1 < len && bytes[i + 1] == b'*' {
|
||
current.push('/');
|
||
current.push('*');
|
||
i += 2;
|
||
while i + 1 < len && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
|
||
current.push(bytes[i] as char);
|
||
i += 1;
|
||
}
|
||
if i + 1 < len {
|
||
current.push('*');
|
||
current.push('/');
|
||
i += 2;
|
||
}
|
||
continue;
|
||
}
|
||
|
||
// Single-quoted string
|
||
if b == b'\'' {
|
||
current.push('\'');
|
||
i += 1;
|
||
while i < len {
|
||
current.push(bytes[i] as char);
|
||
if bytes[i] == b'\'' {
|
||
if i + 1 < len && bytes[i + 1] == b'\'' {
|
||
current.push('\'');
|
||
i += 2;
|
||
} else {
|
||
i += 1;
|
||
break;
|
||
}
|
||
} else {
|
||
i += 1;
|
||
}
|
||
}
|
||
continue;
|
||
}
|
||
|
||
// Dollar-quoted string ($tag$...$tag$ or $$...$$)
|
||
if b == b'$' {
|
||
i += 1;
|
||
let mut tag = String::from("$");
|
||
while i < len && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') {
|
||
tag.push(bytes[i] as char);
|
||
i += 1;
|
||
}
|
||
if i < len && bytes[i] == b'$' {
|
||
tag.push('$');
|
||
i += 1;
|
||
// We have a dollar-quote tag, find the closing tag
|
||
current.push_str(&tag);
|
||
let tag_bytes = tag.as_bytes();
|
||
loop {
|
||
if i >= len {
|
||
break;
|
||
}
|
||
if bytes[i] == b'$' && i + tag_bytes.len() <= len
|
||
&& &bytes[i..i + tag_bytes.len()] == tag_bytes
|
||
{
|
||
current.push_str(&tag);
|
||
i += tag_bytes.len();
|
||
break;
|
||
}
|
||
current.push(bytes[i] as char);
|
||
i += 1;
|
||
}
|
||
} else {
|
||
// Not a valid dollar-quote, push what we consumed
|
||
current.push_str(&tag);
|
||
}
|
||
continue;
|
||
}
|
||
|
||
// Statement separator
|
||
if b == b';' {
|
||
current.push(';');
|
||
let trimmed = current.trim().to_string();
|
||
if !trimmed.is_empty() && trimmed != ";" {
|
||
statements.push(trimmed);
|
||
}
|
||
current.clear();
|
||
i += 1;
|
||
continue;
|
||
}
|
||
|
||
current.push(b as char);
|
||
i += 1;
|
||
}
|
||
|
||
// Trailing statement without semicolon
|
||
let trimmed = current.trim().to_string();
|
||
if !trimmed.is_empty() && trimmed != ";" {
|
||
statements.push(trimmed);
|
||
}
|
||
|
||
statements
|
||
}
|