Skip to main content

tank_core/
executor.rs

1use crate::{
2    AsEntity, AsQuery, Driver, DynQuery, Query, QueryResult, RawQuery, Result, Row, RowsAffected,
3    stream::{Stream, StreamExt, TryStreamExt},
4    writer::SqlWriter,
5};
6use anyhow::anyhow;
7use convert_case::{Case, Casing};
8use std::{
9    future::{self, Future},
10    mem,
11};
12
13/// Async query execution.
14///
15/// Implemented by connections.
16pub trait Executor: Send {
17    /// Associated driver.
18    type Driver: Driver;
19
20    /// Checks if the driver supports multiple SQL statements in a single request.
21    fn accepts_multiple_statements(&self) -> bool {
22        true
23    }
24
25    /// Returns the driver instance associated with this executor.
26    fn driver(&self) -> Self::Driver {
27        Default::default()
28    }
29
30    /// Prepares a query for execution, returning a handle to the prepared statement.
31    fn prepare<'s>(
32        &'s mut self,
33        query: impl AsQuery<Self::Driver> + 's,
34    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
35        let mut query = query.as_query();
36        let query = mem::take(query.as_mut());
37        async {
38            match query {
39                Query::Raw(RawQuery(sql)) => self.do_prepare(sql).await,
40                Query::Prepared(..) => Ok(query),
41            }
42        }
43    }
44
45    /// Internal hook for implementing prepared statement support.
46    fn do_prepare(
47        &mut self,
48        _sql: String,
49    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
50        future::ready(Err(anyhow!(
51            "{} does not support prepare",
52            self.driver().name().to_case(Case::Pascal)
53        )))
54    }
55
56    /// Executes a query and streams the results (rows or affected counts).
57    fn run<'s>(
58        &'s mut self,
59        query: impl AsQuery<Self::Driver> + 's,
60    ) -> impl Stream<Item = Result<QueryResult>> + Send;
61
62    /// Executes a query and streams the resulting rows, ignoring affected counts.
63    fn fetch<'s>(
64        &'s mut self,
65        query: impl AsQuery<Self::Driver> + 's,
66    ) -> impl Stream<Item = Result<Row>> + Send {
67        self.run(query).filter_map(|v| async move {
68            match v {
69                Ok(QueryResult::Row(v)) => Some(Ok(v)),
70                Err(e) => Some(Err(e)),
71                _ => None,
72            }
73        })
74    }
75
76    /// Executes a query and returns the total number of affected rows.
77    fn execute<'s>(
78        &'s mut self,
79        query: impl AsQuery<Self::Driver> + 's,
80    ) -> impl Future<Output = Result<RowsAffected>> + Send {
81        self.run(query)
82            .filter_map(|v| async move {
83                match v {
84                    Ok(QueryResult::Affected(v)) => Some(Ok(v)),
85                    Err(e) => Some(Err(e)),
86                    _ => None,
87                }
88            })
89            .try_collect()
90    }
91
92    /// Efficiently inserts a collection of entities bypassing regular SQL execution when supported by the driver.
93    fn append<It>(&mut self, entities: It) -> impl Future<Output = Result<RowsAffected>> + Send
94    where
95        It: IntoIterator + Send,
96        It::IntoIter: Send,
97        It::Item: AsEntity,
98    {
99        let mut query = DynQuery::default();
100        self.driver()
101            .sql_writer()
102            .write_insert(&mut query, entities, false);
103        self.execute(query)
104    }
105}
106
107impl<S: Executor + ?Sized> Executor for &mut S {
108    type Driver = S::Driver;
109
110    fn accepts_multiple_statements(&self) -> bool {
111        (**self).accepts_multiple_statements()
112    }
113
114    fn driver(&self) -> Self::Driver {
115        (**self).driver()
116    }
117
118    fn prepare<'s>(
119        &'s mut self,
120        query: impl AsQuery<Self::Driver> + 's,
121    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
122        (**self).prepare(query)
123    }
124
125    fn do_prepare(
126        &mut self,
127        sql: String,
128    ) -> impl Future<Output = Result<Query<Self::Driver>>> + Send {
129        (**self).do_prepare(sql)
130    }
131
132    fn run<'s>(
133        &'s mut self,
134        query: impl AsQuery<Self::Driver> + 's,
135    ) -> impl Stream<Item = Result<QueryResult>> + Send {
136        (**self).run(query)
137    }
138
139    fn fetch<'s>(
140        &'s mut self,
141        query: impl AsQuery<Self::Driver> + 's,
142    ) -> impl Stream<Item = Result<Row>> + Send {
143        (**self).fetch(query)
144    }
145
146    fn execute<'s>(
147        &'s mut self,
148        query: impl AsQuery<Self::Driver> + 's,
149    ) -> impl Future<Output = Result<RowsAffected>> + Send {
150        (**self).execute(query)
151    }
152
153    fn append<It>(&mut self, entities: It) -> impl Future<Output = Result<RowsAffected>> + Send
154    where
155        It: IntoIterator + Send,
156        It::IntoIter: Send,
157        It::Item: AsEntity,
158    {
159        (**self).append(entities)
160    }
161}