tank_core/
query.rs

1use crate::{
2    Driver, Executor, Prepared, Result, Value,
3    future::FutureExt,
4    stream::{Stream, StreamExt},
5    truncate_long,
6};
7use std::{
8    fmt::{self, Display},
9    sync::Arc,
10};
11
12/// A query ready to be executed by an [`Executor`].
13///
14/// Represents either raw SQL (`Raw`) or a backend prepared statement
15/// (`Prepared`) carrying driver-specific caching / parsing state.
16pub enum Query<D: Driver> {
17    /// Unprepared SQL text.
18    Raw(String),
19    /// Driver prepared handle.
20    Prepared(D::Prepared),
21}
22
23impl<'d, D: Driver> Query<D>
24where
25    D: 'd,
26{
27    /// Execute the query streaming heterogeneous [`QueryResult`] items.
28    pub fn run<Exec: Executor<Driver = D>>(
29        self,
30        executor: &'d mut Exec,
31    ) -> impl Stream<Item = Result<QueryResult>> + Send {
32        executor.run(self)
33    }
34    /// Fetch at most one labeled row.
35    pub fn fetch_one<Exec: Executor<Driver = D>>(
36        self,
37        executor: &'d mut Exec,
38    ) -> impl Future<Output = Result<Option<RowLabeled>>> + Send {
39        // TODO: replace boxed with pin! once https://github.com/rust-lang/rust/issues/100013 is fixed
40        async {
41            executor
42                .fetch(self)
43                .boxed()
44                .into_future()
45                .map(|(v, _)| v.transpose())
46                .await
47        }
48    }
49    /// Stream all labeled rows.
50    pub fn fetch_many<Exec: Executor<Driver = D>>(
51        self,
52        executor: &'d mut Exec,
53    ) -> impl Stream<Item = Result<RowLabeled>> + Send {
54        executor.fetch(self)
55    }
56}
57
58pub trait AsQuery<D: Driver> {
59    type Output: AsMut<Query<D>> + Send;
60    fn as_query(self) -> Self::Output;
61}
62
63impl<D: Driver> AsQuery<D> for Query<D> {
64    type Output = Query<D>;
65    fn as_query(self) -> Self::Output {
66        self
67    }
68}
69
70impl<'q, D: Driver + 'q> AsQuery<D> for &'q mut Query<D> {
71    type Output = &'q mut Query<D>;
72    fn as_query(self) -> Self::Output {
73        self
74    }
75}
76
77impl<D: Driver> AsQuery<D> for String {
78    type Output = Query<D>;
79    fn as_query(self) -> Self::Output {
80        Query::Raw(self)
81    }
82}
83
84impl<D: Driver> AsQuery<D> for &str {
85    type Output = Query<D>;
86    fn as_query(self) -> Self::Output {
87        Query::Raw(self.to_owned())
88    }
89}
90
91impl<D: Driver> Default for Query<D> {
92    fn default() -> Self {
93        Self::Raw(Default::default())
94    }
95}
96
97impl<D: Driver> From<&str> for Query<D> {
98    fn from(value: &str) -> Self {
99        Query::Raw(value.into())
100    }
101}
102
103impl<D: Driver> From<String> for Query<D> {
104    fn from(value: String) -> Self {
105        Query::Raw(value.into())
106    }
107}
108
109impl<D, P> From<P> for Query<D>
110where
111    D: Driver<Prepared = P>,
112    P: Prepared,
113{
114    fn from(value: P) -> Self {
115        Query::Prepared(value)
116    }
117}
118
119impl<D: Driver> Display for Query<D> {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        match self {
122            Query::Raw(query) => write!(f, "{}", truncate_long!(query)),
123            Query::Prepared(query) => query.fmt(f),
124        }
125    }
126}
127
128impl<D: Driver> AsMut<Query<D>> for Query<D> {
129    fn as_mut(&mut self) -> &mut Query<D> {
130        self
131    }
132}
133
134/// Metadata about modify operations (INSERT/UPDATE/DELETE).
135#[derive(Default, Debug, Clone, Copy)]
136pub struct RowsAffected {
137    /// Total number of rows impacted.
138    pub rows_affected: u64,
139    /// Backend-specific last inserted / affected identifier when available.
140    pub last_affected_id: Option<i64>,
141}
142
143/// Shared reference-counted column name list.
144pub type RowNames = Arc<[String]>;
145/// Owned row value slice matching `RowNames` length.
146pub type Row = Box<[Value]>;
147
148/// A result row with its corresponding column labels.
149#[derive(Debug, Clone)]
150pub struct RowLabeled {
151    /// Column names.
152    pub labels: RowNames,
153    /// Data values (aligned by index with `labels`).
154    pub values: Row,
155}
156
157impl RowLabeled {
158    pub fn new(names: RowNames, values: Row) -> Self {
159        Self {
160            labels: names,
161            values,
162        }
163    }
164    pub fn names(&self) -> &[String] {
165        &self.labels
166    }
167    pub fn values(&self) -> &[Value] {
168        &self.values
169    }
170    pub fn get_column(&self, name: &str) -> Option<&Value> {
171        self.labels
172            .iter()
173            .position(|v| v == name)
174            .map(|i| &self.values()[i])
175    }
176}
177
178/// Heterogeneous items emitted by `Executor::run` combining rows and modify results.
179#[derive(Debug)]
180pub enum QueryResult {
181    /// A labeled row.
182    Row(RowLabeled),
183    /// A modify effect aggregation.
184    Affected(RowsAffected),
185}
186
187impl Extend<RowsAffected> for RowsAffected {
188    fn extend<T: IntoIterator<Item = RowsAffected>>(&mut self, iter: T) {
189        for elem in iter {
190            self.rows_affected += elem.rows_affected;
191            if elem.last_affected_id.is_some() {
192                self.last_affected_id = elem.last_affected_id;
193            }
194        }
195    }
196}
197
198impl From<RowLabeled> for Row {
199    fn from(value: RowLabeled) -> Self {
200        value.values
201    }
202}
203
204impl<'a> From<&'a RowLabeled> for &'a Row {
205    fn from(value: &'a RowLabeled) -> Self {
206        &value.values
207    }
208}
209
210impl From<RowLabeled> for QueryResult {
211    fn from(value: RowLabeled) -> Self {
212        QueryResult::Row(value)
213    }
214}
215
216impl From<RowsAffected> for QueryResult {
217    fn from(value: RowsAffected) -> Self {
218        QueryResult::Affected(value)
219    }
220}