use std::path::Path;
use pgevolve_core::plan::raw_step::{RawStep, TransactionConstraint};
use crate::cluster_config::ClusterConfig;
pub async fn apply_cluster_steps(
steps: &[RawStep],
cfg: &ClusterConfig,
) -> Result<(), ClusterApplyError> {
if steps.is_empty() {
return Ok(());
}
let (mut client, connection) =
tokio_postgres::connect(&cfg.connection.dsn, tokio_postgres::NoTls)
.await
.map_err(|e| ClusterApplyError::Connection(e.to_string()))?;
tokio::spawn(async move {
if let Err(err) = connection.await {
tracing::debug!(?err, "cluster apply connection task ended");
}
});
for step in steps {
execute_step(&mut client, step).await?;
}
Ok(())
}
pub async fn apply_cluster_plan_dir(
plan_dir: &Path,
cfg: &ClusterConfig,
) -> Result<(), ClusterApplyError> {
let sql_path = plan_dir.join("plan.sql");
let sql = std::fs::read_to_string(&sql_path).map_err(|e| ClusterApplyError::Io {
path: sql_path.clone(),
source: e,
})?;
let statements = split_sql_statements(&sql);
if statements.is_empty() {
return Ok(());
}
let (mut client, connection) =
tokio_postgres::connect(&cfg.connection.dsn, tokio_postgres::NoTls)
.await
.map_err(|e| ClusterApplyError::Connection(e.to_string()))?;
tokio::spawn(async move {
if let Err(err) = connection.await {
tracing::debug!(?err, "cluster plan-dir apply connection task ended");
}
});
for stmt in &statements {
run_in_transaction(&mut client, stmt).await?;
}
Ok(())
}
async fn execute_step(
client: &mut tokio_postgres::Client,
step: &RawStep,
) -> Result<(), ClusterApplyError> {
match step.transactional {
TransactionConstraint::InTransaction => {
run_in_transaction(client, &step.sql).await?;
}
TransactionConstraint::OutsideTransaction => {
client.execute(step.sql.as_str(), &[]).await.map_err(|e| {
ClusterApplyError::StepFailed {
sql: step.sql.clone(),
source: e,
}
})?;
}
}
Ok(())
}
async fn run_in_transaction(
client: &mut tokio_postgres::Client,
sql: &str,
) -> Result<(), ClusterApplyError> {
let tx = client
.transaction()
.await
.map_err(|e| ClusterApplyError::StepFailed {
sql: sql.to_string(),
source: e,
})?;
tx.execute(sql, &[])
.await
.map_err(|e| ClusterApplyError::StepFailed {
sql: sql.to_string(),
source: e,
})?;
tx.commit()
.await
.map_err(|e| ClusterApplyError::StepFailed {
sql: sql.to_string(),
source: e,
})?;
Ok(())
}
fn split_sql_statements(sql: &str) -> Vec<String> {
let without_comments: String = sql
.lines()
.filter(|l| !l.trim_start().starts_with("--"))
.collect::<Vec<_>>()
.join("\n");
without_comments
.split(';')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
}
#[derive(Debug, thiserror::Error)]
pub enum ClusterApplyError {
#[error("connection error: {0}")]
Connection(String),
#[error("step execution failed (sql: {sql:?}): {source}")]
StepFailed {
sql: String,
#[source]
source: tokio_postgres::Error,
},
#[error("i/o reading {path}: {source}")]
Io {
path: std::path::PathBuf,
#[source]
source: std::io::Error,
},
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn split_discards_comment_lines() {
let sql = "-- @pgevolve step kind=create_role\nCREATE ROLE a;\n-- comment\nCREATE ROLE b;";
let stmts = split_sql_statements(sql);
assert_eq!(stmts, vec!["CREATE ROLE a", "CREATE ROLE b"]);
}
#[test]
fn split_discards_empty_fragments() {
let sql = "CREATE ROLE a;\n\n;";
let stmts = split_sql_statements(sql);
assert_eq!(stmts, vec!["CREATE ROLE a"]);
}
#[test]
fn split_single_statement_no_trailing_semicolon() {
let stmts = split_sql_statements("CREATE ROLE x");
assert_eq!(stmts, vec!["CREATE ROLE x"]);
}
}