use async_trait::async_trait;
use ironflow_core::error::OperationError;
use ironflow_core::operation::{Operation, OperationContext, TypedOperation};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::{Executor, PgPool};
use crate::helpers::{bind_json_param, pg_error, to_value};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecuteOutput {
pub rows_affected: u64,
}
pub struct Execute {
pool: PgPool,
sql: String,
params: Vec<Value>,
}
impl Execute {
pub fn new(pool: PgPool, sql: impl Into<String>, params: Vec<Value>) -> Self {
Self {
pool,
sql: sql.into(),
params,
}
}
pub async fn run(&self, _ctx: &OperationContext) -> Result<ExecuteOutput, OperationError> {
let mut query = sqlx::query(&self.sql);
for p in &self.params {
query = bind_json_param(query, p);
}
let result = query.execute(&self.pool).await.map_err(pg_error)?;
Ok(ExecuteOutput {
rows_affected: result.rows_affected(),
})
}
}
#[async_trait]
impl Operation for Execute {
fn kind(&self) -> &str {
"postgres"
}
async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
to_value(&self.run(ctx).await?)
}
fn input(&self) -> Option<Value> {
Some(serde_json::json!({ "sql": self.sql, "params": self.params }))
}
}
impl TypedOperation for Execute {
type Output = ExecuteOutput;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecuteBatchOutput {
pub rows_affected: Vec<u64>,
}
pub struct ExecuteBatch {
pool: PgPool,
statements: Vec<String>,
}
impl ExecuteBatch {
pub fn new(pool: PgPool, statements: Vec<String>) -> Self {
Self { pool, statements }
}
pub async fn run(&self, _ctx: &OperationContext) -> Result<ExecuteBatchOutput, OperationError> {
let mut results = Vec::with_capacity(self.statements.len());
for stmt in &self.statements {
let result = self.pool.execute(stmt.as_str()).await.map_err(pg_error)?;
results.push(result.rows_affected());
}
Ok(ExecuteBatchOutput {
rows_affected: results,
})
}
}
#[async_trait]
impl Operation for ExecuteBatch {
fn kind(&self) -> &str {
"postgres"
}
async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
to_value(&self.run(ctx).await?)
}
fn input(&self) -> Option<Value> {
Some(serde_json::json!({ "statements": self.statements }))
}
}
impl TypedOperation for ExecuteBatch {
type Output = ExecuteBatchOutput;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionOutput {
pub rows_affected: Vec<u64>,
}
pub struct Transaction {
pool: PgPool,
statements: Vec<String>,
}
impl Transaction {
pub fn new(pool: PgPool, statements: Vec<String>) -> Self {
Self { pool, statements }
}
pub async fn run(&self, _ctx: &OperationContext) -> Result<TransactionOutput, OperationError> {
let mut tx = self.pool.begin().await.map_err(pg_error)?;
let mut results = Vec::with_capacity(self.statements.len());
for stmt in &self.statements {
let result = tx.execute(stmt.as_str()).await.map_err(pg_error)?;
results.push(result.rows_affected());
}
tx.commit().await.map_err(pg_error)?;
Ok(TransactionOutput {
rows_affected: results,
})
}
}
#[async_trait]
impl Operation for Transaction {
fn kind(&self) -> &str {
"postgres"
}
async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
to_value(&self.run(ctx).await?)
}
fn input(&self) -> Option<Value> {
Some(serde_json::json!({ "statements": self.statements }))
}
}
impl TypedOperation for Transaction {
type Output = TransactionOutput;
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn execute_kind() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = Execute::new(pool, "INSERT INTO t VALUES (1)", vec![]);
assert_eq!(op.kind(), "postgres");
}
#[tokio::test]
async fn execute_batch_kind() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = ExecuteBatch::new(pool, vec!["SELECT 1".to_string()]);
assert_eq!(op.kind(), "postgres");
}
#[tokio::test]
async fn transaction_kind() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = Transaction::new(pool, vec!["SELECT 1".to_string()]);
assert_eq!(op.kind(), "postgres");
}
#[tokio::test]
async fn execute_input_no_secrets() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = Execute::new(pool, "INSERT INTO t VALUES ($1)", vec![]);
let input = op.input().unwrap();
let text = input.to_string();
assert!(!text.contains("postgres://"), "leaked URL: {text}");
}
#[tokio::test]
async fn execute_batch_input_no_secrets() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = ExecuteBatch::new(pool, vec!["SELECT 1".to_string()]);
let input = op.input().unwrap();
let text = input.to_string();
assert!(!text.contains("postgres://"), "leaked URL: {text}");
}
#[tokio::test]
async fn transaction_input_no_secrets() {
let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
let op = Transaction::new(pool, vec!["SELECT 1".to_string()]);
let input = op.input().unwrap();
let text = input.to_string();
assert!(!text.contains("postgres://"), "leaked URL: {text}");
}
}