tank_core/
query.rs

1use crate::{
2    Driver, Executor, Prepared, Result, Value,
3    future::FutureExt,
4    printable_query,
5    stream::{Stream, StreamExt},
6};
7use std::{
8    fmt::{self, Display},
9    pin::pin,
10    sync::Arc,
11};
12
13#[derive(Clone)]
14pub enum Query<D: Driver> {
15    Raw(String),
16    Prepared(D::Prepared),
17}
18
19impl<D: Driver> Query<D> {
20    pub fn run<'e, Exec: Executor<Driver = D>>(
21        self,
22        executor: &mut Exec,
23    ) -> impl Stream<Item = Result<QueryResult>> + Send {
24        executor.run(self)
25    }
26    pub fn fetch_one<Exec: Executor<Driver = D>>(
27        self,
28        executor: &mut Exec,
29    ) -> impl Future<Output = Result<Option<RowLabeled>>> + Send {
30        let stream = executor.fetch(self);
31        async move { pin!(stream).into_future().map(|(v, _)| v).await.transpose() }
32    }
33    pub fn fetch_many<Exec: Executor<Driver = D>>(
34        self,
35        executor: &mut Exec,
36    ) -> impl Stream<Item = Result<RowLabeled>> + Send {
37        executor.fetch(self)
38    }
39}
40
41impl<D: Driver> From<&str> for Query<D> {
42    fn from(value: &str) -> Self {
43        Query::Raw(value.into())
44    }
45}
46
47impl<D: Driver> From<String> for Query<D> {
48    fn from(value: String) -> Self {
49        Query::Raw(value.into())
50    }
51}
52
53impl<D, P> From<P> for Query<D>
54where
55    D: Driver<Prepared = P>,
56    P: Prepared,
57{
58    fn from(value: P) -> Self {
59        Query::Prepared(value)
60    }
61}
62
63impl<D: Driver> Display for Query<D> {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        match self {
66            Query::Raw(query) => write!(f, "{}", printable_query!(query)),
67            Query::Prepared(query) => query.fmt(f),
68        }
69    }
70}
71
72#[derive(Default, Debug, Clone, Copy)]
73pub struct RowsAffected {
74    pub rows_affected: u64,
75    pub last_affected_id: Option<i64>,
76}
77
78pub type RowNames = Arc<[String]>;
79pub type Row = Box<[Value]>;
80
81#[derive(Debug, Clone)]
82pub struct RowLabeled {
83    pub labels: RowNames,
84    pub values: Row,
85}
86
87impl RowLabeled {
88    pub fn new(names: RowNames, values: Row) -> Self {
89        Self {
90            labels: names,
91            values,
92        }
93    }
94    pub fn names(&self) -> &[String] {
95        &self.labels
96    }
97    pub fn values(&self) -> &[Value] {
98        &self.values
99    }
100    pub fn get_column(&self, name: &str) -> Option<&Value> {
101        self.labels
102            .iter()
103            .position(|v| v == name)
104            .map(|i| &self.values()[i])
105    }
106}
107
108#[derive(Debug)]
109pub enum QueryResult {
110    RowLabeled(RowLabeled),
111    Affected(RowsAffected),
112}
113
114impl Extend<RowsAffected> for RowsAffected {
115    fn extend<T: IntoIterator<Item = RowsAffected>>(&mut self, iter: T) {
116        for elem in iter {
117            self.rows_affected += elem.rows_affected;
118            if elem.last_affected_id.is_some() {
119                self.last_affected_id = elem.last_affected_id;
120            }
121        }
122    }
123}
124
125impl From<RowLabeled> for Row {
126    fn from(value: RowLabeled) -> Self {
127        value.values
128    }
129}
130
131impl<'a> From<&'a RowLabeled> for &'a Row {
132    fn from(value: &'a RowLabeled) -> Self {
133        &value.values
134    }
135}
136
137impl From<RowLabeled> for QueryResult {
138    fn from(value: RowLabeled) -> Self {
139        QueryResult::RowLabeled(value)
140    }
141}
142
143impl From<RowsAffected> for QueryResult {
144    fn from(value: RowsAffected) -> Self {
145        QueryResult::Affected(value)
146    }
147}