tank_core/

executor.rs

use crate::{
    AsEntity, AsQuery, Driver, DynQuery, Query, QueryResult, RawQuery, Result, Row, RowsAffected,
    stream::{Stream, StreamExt},
    writer::SqlWriter,
};
use anyhow::anyhow;
use convert_case::{Case, Casing};
use std::{
    future::{self, Future},
    mem,
    pin::{Pin, pin},
    task::{self, Poll},
};

/// Stream adapter returned by [`Executor::fetch`], filtering [`Executor::run`] to rows.
struct Fetched<S>(S);

impl<S, R> Stream for Fetched<S>
where
    S: Stream<Item = Result<QueryResult<R>>>,
{
    type Item = Result<R>;

    #[inline(always)]
    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
        let mut inner = unsafe { self.map_unchecked_mut(|s| &mut s.0) };
        loop {
            match inner.as_mut().poll_next(cx) {
                Poll::Ready(Some(Ok(QueryResult::Row(row)))) => return Poll::Ready(Some(Ok(row))),
                Poll::Ready(Some(Ok(QueryResult::Affected(_)))) => {}
                Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
                Poll::Ready(None) => return Poll::Ready(None),
                Poll::Pending => return Poll::Pending,
            }
        }
    }

    #[inline(always)]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

/// Async query execution.
///
/// Implemented by connections.
pub trait Executor: Send {
    /// Associated driver.
    type Driver: Driver;

    /// Checks if the driver supports multiple SQL statements in a single request.
    fn accepts_multiple_statements(&self) -> bool {
        true
    }

    /// Returns the driver instance associated with this executor.
    fn driver(&self) -> Self::Driver {
        Default::default()
    }

    /// Prepares a query for execution, returning a handle to the prepared statement.
    fn prepare<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
        let mut query = query.as_query();
        let query = mem::take(query.as_mut());
        async {
            match query {
                Query::Raw(RawQuery(sql)) => self.do_prepare(sql).await,
                Query::Prepared(..) => Ok(query),
            }
        }
    }

    /// Internal hook for implementing prepared statement support.
    fn do_prepare(
        &mut self,
        _sql: String,
    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
        future::ready(Err(anyhow!(
            "{} does not support prepare",
            self.driver().name().to_case(Case::Pascal)
        )))
    }

    /// Executes a query and streams the results (rows or affected counts).
    fn run<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Stream<Item = Result<QueryResult<impl Row + 's>>> + Send;

    /// Executes a query and streams the resulting rows, ignoring affected counts.
    fn fetch<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Stream<Item = Result<impl Row + 's>> + Send {
        Fetched(self.run(query))
    }

    /// Executes a query and returns the total number of affected rows.
    fn execute<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Future<Output = Result<RowsAffected>> + Send {
        let stream = self.run(query);
        async move {
            let mut stream = pin!(stream);
            let mut affected = RowsAffected::default();
            while let Some(item) = stream.next().await {
                if let QueryResult::Affected(a) = item? {
                    affected.extend([a]);
                }
            }
            Ok(affected)
        }
    }

    /// Efficiently inserts a collection of entities bypassing regular SQL execution when supported by the driver.
    fn append<It>(&mut self, entities: It) -> impl Future<Output = Result<RowsAffected>> + Send
    where
        It: IntoIterator + Send,
        It::IntoIter: Send,
        It::Item: AsEntity,
    {
        let mut query = DynQuery::default();
        self.driver()
            .sql_writer()
            .write_insert(&mut query, entities, false);
        self.execute(query)
    }
}

impl<S: Executor + ?Sized> Executor for &mut S {
    type Driver = S::Driver;

    fn accepts_multiple_statements(&self) -> bool {
        (**self).accepts_multiple_statements()
    }

    fn driver(&self) -> Self::Driver {
        (**self).driver()
    }

    fn prepare<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
        (**self).prepare(query)
    }

    fn do_prepare(
        &mut self,
        sql: String,
    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
        (**self).do_prepare(sql)
    }

    fn run<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Stream<Item = Result<QueryResult<impl Row + 's>>> + Send {
        (**self).run(query)
    }

    fn fetch<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Stream<Item = Result<impl Row + 's>> + Send {
        (**self).fetch(query)
    }

    fn execute<'s>(
        &'s mut self,
        query: impl AsQuery<Self::Driver> + 's,
    ) -> impl Future<Output = Result<RowsAffected>> + Send {
        (**self).execute(query)
    }

    fn append<It>(&mut self, entities: It) -> impl Future<Output = Result<RowsAffected>> + Send
    where
        It: IntoIterator + Send,
        It::IntoIter: Send,
        It::Item: AsEntity,
    {
        (**self).append(entities)
    }
}