tank_postgres/
rows.rs

1use crate::util::extract_value;
2use postgres_types::{FromSql, Type};
3use std::{
4    borrow::Cow,
5    cmp::min,
6    error::Error,
7    fmt::{self, Debug, Formatter},
8    future::Future,
9    mem,
10    pin::Pin,
11    task::{Context, Poll},
12};
13use tank_core::{Error as TankError, Result, Row, RowLabels, Value, stream::Stream};
14
15struct RawCell<'a>(Option<&'a [u8]>);
16
17impl<'a> FromSql<'a> for RawCell<'a> {
18    #[inline(always)]
19    fn from_sql(
20        _t: &Type,
21        raw: &'a [u8],
22    ) -> std::result::Result<Self, Box<dyn Error + Sync + Send>> {
23        Ok(Self(Some(raw)))
24    }
25
26    #[inline(always)]
27    fn from_sql_null(_t: &Type) -> std::result::Result<Self, Box<dyn Error + Sync + Send>> {
28        Ok(Self(None))
29    }
30
31    #[inline(always)]
32    fn from_sql_nullable(
33        _t: &Type,
34        raw: Option<&'a [u8]>,
35    ) -> std::result::Result<Self, Box<dyn Error + Sync + Send>> {
36        Ok(Self(raw))
37    }
38
39    #[inline(always)]
40    fn accepts(_t: &Type) -> bool {
41        true
42    }
43}
44
45/// A native row view borrowing the connection's cached column labels.
46pub struct PostgresRow<'a> {
47    pub(crate) row: tokio_postgres::Row,
48    pub(crate) labels: &'a RowLabels,
49}
50
51impl Debug for PostgresRow<'_> {
52    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
53        f.debug_struct("PostgresRow")
54            .field("labels", &self.labels)
55            .field("len", &self.len())
56            .finish()
57    }
58}
59
60impl PostgresRow<'_> {
61    #[cold]
62    #[inline(never)]
63    fn decode_fallback<'a>(ty: &Type, raw: &'a [u8]) -> Option<Value<'a>> {
64        extract_value(ty, Some(raw)).ok()
65    }
66}
67
68impl Row for PostgresRow<'_> {
69    #[inline(always)]
70    fn labels(&self) -> &[String] {
71        self.labels
72    }
73
74    #[inline(always)]
75    fn len(&self) -> usize {
76        min(self.labels.len(), self.row.len())
77    }
78
79    #[inline(always)]
80    fn value(&self, index: usize) -> Option<Cow<'_, Value<'_>>> {
81        if index >= self.len() {
82            return None;
83        }
84        let raw = match self.row.try_get::<usize, RawCell<'_>>(index) {
85            Ok(RawCell(Some(raw))) => raw,
86            Ok(RawCell(None)) => return Some(Cow::Owned(Value::Null)),
87            Err(e) => {
88                mem::forget(e);
89                return None;
90            }
91        };
92        // `BOOL` also places a 1-byte `Option<bool>` store at offset 1..2 of `Value`,
93        // splitting the 3-byte 1..4 gap into `i8` + `i16` so LLVM SROA promotes both to SSA.
94        match *self.row.columns()[index].type_() {
95            Type::BOOL => match *raw {
96                [b] => Some(Cow::Owned(Value::Boolean(Some(b != 0)))),
97                _ => Some(Cow::Owned(Value::Null)),
98            },
99            Type::INT4 => match raw.try_into() {
100                Ok(b) => Some(Cow::Owned(Value::Int32(Some(i32::from_be_bytes(b))))),
101                Err(_) => Some(Cow::Owned(Value::Null)),
102            },
103            Type::VARCHAR | Type::TEXT => match std::str::from_utf8(raw) {
104                Ok(text) => Some(Cow::Owned(Value::Varchar(Some(Cow::Borrowed(text))))),
105                Err(_) => Some(Cow::Owned(Value::Null)),
106            },
107            ref ty => match Self::decode_fallback(ty, raw) {
108                Some(v) => Some(Cow::Owned(v)),
109                None => Some(Cow::Owned(Value::Null)),
110            },
111        }
112    }
113}
114
115/// A streaming view over a `tokio_postgres::RowStream` and shared column labels.
116pub struct PostgresRows<'s> {
117    first: Option<tokio_postgres::Row>,
118    stream: tokio_postgres::RowStream,
119    labels: &'s RowLabels,
120    /// Set when the underlying `RowStream` has already been fully drained
121    /// (zero-row raw query pre-polled in `fetch_inner`): polling a completed
122    /// `tokio_postgres::RowStream` yields `Error::closed()`.
123    done: bool,
124}
125
126impl<'s> PostgresRows<'s> {
127    #[inline(always)]
128    pub(crate) fn new(
129        first: Option<tokio_postgres::Row>,
130        stream: tokio_postgres::RowStream,
131        labels: &'s RowLabels,
132    ) -> Self {
133        Self {
134            first,
135            stream,
136            labels,
137            done: false,
138        }
139    }
140
141    #[inline(always)]
142    pub(crate) fn new_drained(stream: tokio_postgres::RowStream, labels: &'s RowLabels) -> Self {
143        Self {
144            first: None,
145            stream,
146            labels,
147            done: true,
148        }
149    }
150
151    #[cold]
152    #[inline(never)]
153    fn row_error(e: tokio_postgres::Error) -> TankError {
154        let error = TankError::new(e);
155        log::error!("{error:#}");
156        error
157    }
158
159    #[inline(always)]
160    fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<PostgresRow<'s>>>> {
161        if self.done {
162            return Poll::Ready(None);
163        }
164        if let Some(row) = self.first.take() {
165            return Poll::Ready(Some(Ok(PostgresRow {
166                row,
167                labels: self.labels,
168            })));
169        }
170        match unsafe { Pin::new_unchecked(&mut self.stream) }.poll_next(cx) {
171            Poll::Ready(Some(Ok(row))) => Poll::Ready(Some(Ok(PostgresRow {
172                row,
173                labels: self.labels,
174            }))),
175            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(Self::row_error(e)))),
176            Poll::Ready(None) => Poll::Ready(None),
177            Poll::Pending => Poll::Pending,
178        }
179    }
180}
181
182/// Stream adapter that awaits the query execution future once, then yields rows from `PostgresRows`.
183pub(crate) struct PostgresStream<'s, F> {
184    fut: Option<F>,
185    rows: Option<PostgresRows<'s>>,
186}
187
188impl<'s, F> PostgresStream<'s, F> {
189    #[inline(always)]
190    pub(crate) fn new(fut: F) -> Self {
191        Self {
192            fut: Some(fut),
193            rows: None,
194        }
195    }
196
197    #[cold]
198    #[inline(never)]
199    fn poll_init(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>>
200    where
201        F: Future<Output = Result<PostgresRows<'s>>>,
202    {
203        let Some(fut) = self.fut.as_mut() else {
204            return Poll::Ready(Ok(()));
205        };
206        let Poll::Ready(res) = unsafe { Pin::new_unchecked(fut) }.poll(cx) else {
207            return Poll::Pending;
208        };
209        self.fut = None;
210        self.rows = Some(res?);
211        Poll::Ready(Ok(()))
212    }
213}
214
215impl<'s, F> Stream for PostgresStream<'s, F>
216where
217    F: Future<Output = Result<PostgresRows<'s>>>,
218{
219    type Item = Result<PostgresRow<'s>>;
220
221    #[inline(always)]
222    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
223        let this = unsafe { self.get_unchecked_mut() };
224        if this.rows.is_none() {
225            match this.poll_init(cx) {
226                Poll::Ready(Ok(())) => {}
227                Poll::Ready(Err(e)) => return Poll::Ready(Some(Err(e))),
228                Poll::Pending => return Poll::Pending,
229            }
230        }
231        unsafe { this.rows.as_mut().unwrap_unchecked() }.poll_next(cx)
232    }
233}
tank_postgres/rows.rs

tank_postgres/
rows.rs
