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    pin::pin,
10    sync::Arc,
11};
12
13/// A query ready to be executed by an [`Executor`].
14///
15/// Represents either raw SQL (`Raw`) or a backend prepared statement
16/// (`Prepared`) carrying driver-specific caching / parsing state.
17#[derive(Clone)]
18pub enum Query<D: Driver> {
19    /// Unprepared SQL text.
20    Raw(String),
21    /// Driver prepared handle.
22    Prepared(D::Prepared),
23}
24
25impl<D: Driver> Query<D> {
26    /// Execute the query streaming heterogeneous [`QueryResult`] items.
27    pub fn run<'e, Exec: Executor<Driver = D>>(
28        self,
29        executor: &mut Exec,
30    ) -> impl Stream<Item = Result<QueryResult>> + Send {
31        executor.run(self)
32    }
33    /// Fetch at most one labeled row.
34    pub fn fetch_one<Exec: Executor<Driver = D>>(
35        self,
36        executor: &mut Exec,
37    ) -> impl Future<Output = Result<Option<RowLabeled>>> + Send {
38        let stream = executor.fetch(self);
39        async move { pin!(stream).into_future().map(|(v, _)| v).await.transpose() }
40    }
41    /// Stream all labeled rows.
42    pub fn fetch_many<Exec: Executor<Driver = D>>(
43        self,
44        executor: &mut Exec,
45    ) -> impl Stream<Item = Result<RowLabeled>> + Send {
46        executor.fetch(self)
47    }
48}
49
50impl<D: Driver> From<&str> for Query<D> {
51    fn from(value: &str) -> Self {
52        Query::Raw(value.into())
53    }
54}
55
56impl<D: Driver> From<String> for Query<D> {
57    fn from(value: String) -> Self {
58        Query::Raw(value.into())
59    }
60}
61
62impl<D, P> From<P> for Query<D>
63where
64    D: Driver<Prepared = P>,
65    P: Prepared,
66{
67    fn from(value: P) -> Self {
68        Query::Prepared(value)
69    }
70}
71
72impl<D: Driver> Display for Query<D> {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        match self {
75            Query::Raw(query) => write!(f, "{}", truncate_long!(query)),
76            Query::Prepared(query) => query.fmt(f),
77        }
78    }
79}
80
81/// Metadata about modify operations (INSERT/UPDATE/DELETE).
82#[derive(Default, Debug, Clone, Copy)]
83pub struct RowsAffected {
84    /// Total number of rows impacted.
85    pub rows_affected: u64,
86    /// Backend-specific last inserted / affected identifier when available.
87    pub last_affected_id: Option<i64>,
88}
89
90/// Shared reference-counted column name list.
91pub type RowNames = Arc<[String]>;
92/// Owned row value slice matching `RowNames` length.
93pub type Row = Box<[Value]>;
94
95/// A result row with its corresponding column labels.
96#[derive(Debug, Clone)]
97pub struct RowLabeled {
98    /// Column names.
99    pub labels: RowNames,
100    /// Data values (aligned by index with `labels`).
101    pub values: Row,
102}
103
104impl RowLabeled {
105    pub fn new(names: RowNames, values: Row) -> Self {
106        Self {
107            labels: names,
108            values,
109        }
110    }
111    pub fn names(&self) -> &[String] {
112        &self.labels
113    }
114    pub fn values(&self) -> &[Value] {
115        &self.values
116    }
117    pub fn get_column(&self, name: &str) -> Option<&Value> {
118        self.labels
119            .iter()
120            .position(|v| v == name)
121            .map(|i| &self.values()[i])
122    }
123}
124
125/// Heterogeneous items emitted by `Executor::run` combining rows and modify results.
126#[derive(Debug)]
127pub enum QueryResult {
128    /// A labeled row.
129    Row(RowLabeled),
130    /// A modify effect aggregation.
131    Affected(RowsAffected),
132}
133
134impl Extend<RowsAffected> for RowsAffected {
135    fn extend<T: IntoIterator<Item = RowsAffected>>(&mut self, iter: T) {
136        for elem in iter {
137            self.rows_affected += elem.rows_affected;
138            if elem.last_affected_id.is_some() {
139                self.last_affected_id = elem.last_affected_id;
140            }
141        }
142    }
143}
144
145impl From<RowLabeled> for Row {
146    fn from(value: RowLabeled) -> Self {
147        value.values
148    }
149}
150
151impl<'a> From<&'a RowLabeled> for &'a Row {
152    fn from(value: &'a RowLabeled) -> Self {
153        &value.values
154    }
155}
156
157impl From<RowLabeled> for QueryResult {
158    fn from(value: RowLabeled) -> Self {
159        QueryResult::Row(value)
160    }
161}
162
163impl From<RowsAffected> for QueryResult {
164    fn from(value: RowsAffected) -> Self {
165        QueryResult::Affected(value)
166    }
167}