use std::future::Future;
use keelson_core::{FromValue, Query};
use crate::error::ExecError;
use crate::executor::{ExecResult, Executor, Statement};
use crate::row::{FromRow, Row};
pub trait Execute: Query {
fn fetch_all<T: FromRow>(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<Vec<T>, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move {
let rows = run_fetch(db, stmt?).await?;
rows.into_iter().map(|mut r| T::from_row(&mut r)).collect()
}
}
fn fetch_rows(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<Vec<Row>, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move { run_fetch(db, stmt?).await }
}
fn fetch_one<T: FromRow>(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<T, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move {
let mut rows = run_fetch(db, stmt?).await?;
match rows.len() {
0 => Err(ExecError::RowNotFound),
1 => T::from_row(&mut rows[0]),
_ => Err(ExecError::TooManyRows),
}
}
}
fn fetch_optional<T: FromRow>(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<Option<T>, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move {
let mut rows = run_fetch(db, stmt?).await?;
match rows.len() {
0 => Ok(None),
1 => T::from_row(&mut rows[0]).map(Some),
_ => Err(ExecError::TooManyRows),
}
}
}
fn fetch_scalar<T: FromValue>(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<T, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move {
let mut rows = run_fetch(db, stmt?).await?;
match rows.len() {
0 => Err(ExecError::RowNotFound),
1 => rows[0].take_at(0),
_ => Err(ExecError::TooManyRows),
}
}
}
fn fetch_scalars<T: FromValue>(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<Vec<T>, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move {
let rows = run_fetch(db, stmt?).await?;
rows.into_iter().map(|mut r| r.take_at(0)).collect()
}
}
fn execute(
&self,
db: &(impl Executor + ?Sized),
) -> impl Future<Output = Result<ExecResult, ExecError>> + Send {
let stmt = Statement::from_query(self);
async move { run_execute(db, stmt?).await }
}
}
impl<Q: Query + ?Sized> Execute for Q {}
pub(crate) async fn run_fetch<E: Executor + ?Sized>(
db: &E,
stmt: Statement,
) -> Result<Vec<Row>, ExecError> {
#[cfg(feature = "tracing")]
{
use tracing::Instrument as _;
let span = query_span(db, &stmt);
let res = db.fetch(stmt).instrument(span.clone()).await;
match &res {
Ok(rows) => span.record("keelson.rows", rows.len() as u64),
Err(e) => span.record("error", tracing::field::display(e)),
};
res
}
#[cfg(not(feature = "tracing"))]
{
db.fetch(stmt).await
}
}
pub(crate) async fn run_execute<E: Executor + ?Sized>(
db: &E,
stmt: Statement,
) -> Result<ExecResult, ExecError> {
#[cfg(feature = "tracing")]
{
use tracing::Instrument as _;
let span = query_span(db, &stmt);
let res = db.execute(stmt).instrument(span.clone()).await;
match &res {
Ok(done) => span.record("keelson.rows_affected", done.rows_affected),
Err(e) => span.record("error", tracing::field::display(e)),
};
res
}
#[cfg(not(feature = "tracing"))]
{
db.execute(stmt).await
}
}
#[cfg(feature = "tracing")]
fn query_span<E: Executor + ?Sized>(db: &E, stmt: &Statement) -> tracing::Span {
tracing::info_span!(
"keelson.query",
db.system = db.family().as_str(),
db.query.text = %stmt.sql,
keelson.query_type = %stmt.query_type,
keelson.args.count = stmt.args.len() as u64,
keelson.rows = tracing::field::Empty,
keelson.rows_affected = tracing::field::Empty,
error = tracing::field::Empty,
)
}