wasi_pg_client/query/
stream.rs

1//! Streaming API for query results.
2//!
3//! This module provides [`RowStream`] — an async stream of rows from a query
4//! result that processes rows one at a time without buffering the entire
5//! result set in memory.
6
7use std::sync::Arc;
8
9use crate::protocol::{BackendMessage, TransactionStatus};
10
11use crate::connection::{Connection, ConnectionState};
12use crate::error::{PgError, PgServerError, Result};
13use crate::query::result::CommandTag;
14use crate::query::row::{FieldDescription, Row};
15use crate::query::{read_data_row, read_row_description};
16
17#[cfg(feature = "tracing")]
18use crate::tracing_ext::TARGET_QUERY;
19
20/// Internal state of the row stream.
21#[derive(Debug)]
22enum RowStreamState {
23    WaitingForDescription,
24    ReceivingRows,
25    Finishing { command_tag: CommandTag },
26    Done { command_tag: CommandTag },
27    Error,
28}
29
30/// An async stream of rows from a query result.
31///
32/// Rows are fetched from the server one at a time as the consumer calls `next()`.
33/// This provides natural backpressure and O(1) memory per row.
34///
35/// `RowStream` borrows the connection mutably. You cannot use the connection
36/// while iterating the stream. When the stream is dropped (or fully consumed),
37/// the connection is available again.
38///
39/// If the stream is dropped before being fully consumed, the connection is
40/// left in an inconsistent state. The [`Connection::needs_recovery`] flag is
41/// set, and you must call [`Connection::recover`] before using the connection
42/// again. To avoid this, either consume the stream fully or call
43/// [`RowStream::consume`] before dropping.
44#[non_exhaustive]
45pub struct RowStream<'a> {
46    conn: &'a mut Connection,
47    columns: Option<Arc<Vec<FieldDescription>>>,
48    state: RowStreamState,
49    /// Whether this stream was created via the extended query protocol.
50    /// Used to determine protocol-specific behavior (e.g., portal-based
51    /// fetch with `max_rows` for cursor streaming in a future iteration).
52    #[allow(dead_code)]
53    extended_protocol: bool,
54    /// Number of rows fetched so far (for tracing).
55    rows_fetched: u64,
56    /// Time when the stream was created (for tracing duration).
57    #[cfg(feature = "tracing")]
58    started_at: std::time::Instant,
59}
60
61impl<'a> RowStream<'a> {
62    /// Create a new `RowStream` for a simple query.
63    pub(crate) fn new_simple(conn: &'a mut Connection) -> Self {
64        RowStream {
65            conn,
66            columns: None,
67            state: RowStreamState::WaitingForDescription,
68            extended_protocol: false,
69            rows_fetched: 0,
70            #[cfg(feature = "tracing")]
71            started_at: std::time::Instant::now(),
72        }
73    }
74
75    /// Create a new `RowStream` for an extended query.
76    #[cfg_attr(not(test), allow(dead_code))]
77    pub(crate) fn new_extended(conn: &'a mut Connection) -> Self {
78        RowStream {
79            conn,
80            columns: None,
81            state: RowStreamState::WaitingForDescription,
82            extended_protocol: true,
83            rows_fetched: 0,
84            #[cfg(feature = "tracing")]
85            started_at: std::time::Instant::now(),
86        }
87    }
88
89    /// Create a new `RowStream` for an extended query with known columns.
90    pub(crate) fn new_extended_with_columns(
91        conn: &'a mut Connection,
92        columns: Arc<Vec<FieldDescription>>,
93    ) -> Self {
94        RowStream {
95            conn,
96            columns: Some(columns),
97            state: RowStreamState::WaitingForDescription,
98            extended_protocol: true,
99            rows_fetched: 0,
100            #[cfg(feature = "tracing")]
101            started_at: std::time::Instant::now(),
102        }
103    }
104
105    /// Create a new `RowStream` for an extended query where the preamble
106    /// (ParseComplete, BindComplete, RowDescription/NoData) has already
107    /// been consumed. The stream starts directly in `ReceivingRows` state.
108    pub(crate) fn new_extended_receiving(
109        conn: &'a mut Connection,
110        columns: Option<Arc<Vec<FieldDescription>>>,
111    ) -> Self {
112        RowStream {
113            conn,
114            columns,
115            state: RowStreamState::ReceivingRows,
116            extended_protocol: true,
117            rows_fetched: 0,
118            #[cfg(feature = "tracing")]
119            started_at: std::time::Instant::now(),
120        }
121    }
122
123    /// Fetch the next row from the stream.
124    ///
125    /// Returns `Ok(Some(row))` when a row is available, `Ok(None)` when the
126    /// stream is exhausted, and `Err(...)` on a server or protocol error.
127    /// After returning `None` or `Err`, subsequent calls return `None`.
128    #[must_use = "stream errors should be checked"]
129    pub async fn next(&mut self) -> Result<Option<Row>> {
130        loop {
131            match self.state {
132                RowStreamState::Done { .. } | RowStreamState::Error => {
133                    return Ok(None);
134                }
135
136                RowStreamState::WaitingForDescription => {
137                    let msg = self
138                        .conn
139                        .codec
140                        .read_message(&mut self.conn.transport)
141                        .await?;
142                    match msg {
143                        BackendMessage::RowDescription(body) => {
144                            self.columns = Some(Arc::new(read_row_description(body)?));
145                            self.state = RowStreamState::ReceivingRows;
146                        }
147                        BackendMessage::CommandComplete(body) => {
148                            let tag = CommandTag::new(body.tag().unwrap_or("").into());
149                            self.state = RowStreamState::Finishing { command_tag: tag };
150                        }
151                        BackendMessage::EmptyQueryResponse => {
152                            self.state = RowStreamState::Finishing {
153                                command_tag: CommandTag::default(),
154                            };
155                        }
156                        BackendMessage::ErrorResponse(body) => {
157                            let server_err =
158                                PgServerError::from_error_body(&body).map_err(PgError::Io)?;
159                            self.conn.read_until_ready().await?;
160                            self.conn.state = ConnectionState::Idle;
161                            self.state = RowStreamState::Error;
162                            return Err(PgError::Server(Box::new(server_err)));
163                        }
164                        BackendMessage::NotificationResponse(body) => {
165                            self.conn.notification_queue.push_back(
166                                crate::notification::Notification {
167                                    process_id: body.process_id(),
168                                    channel: body.channel().unwrap_or("").to_string(),
169                                    payload: body.message().unwrap_or("").to_string(),
170                                },
171                            );
172                            continue;
173                        }
174                        BackendMessage::NoticeResponse(body) => {
175                            if let Ok(notice) = crate::query::Notice::from_fields(&body) {
176                                self.conn.handle_notice(&notice);
177                            }
178                            continue;
179                        }
180                        BackendMessage::ParameterStatus(body) => {
181                            if let (Ok(name), Ok(value)) = (body.name(), body.value()) {
182                                self.conn
183                                    .server_params
184                                    .params
185                                    .insert(name.to_string(), value.to_string());
186                            }
187                            continue;
188                        }
189                        BackendMessage::ParseComplete
190                        | BackendMessage::BindComplete
191                        | BackendMessage::NoData => {
192                            continue;
193                        }
194                        _ => continue,
195                    }
196                }
197
198                RowStreamState::ReceivingRows => {
199                    let msg = self
200                        .conn
201                        .codec
202                        .read_message(&mut self.conn.transport)
203                        .await?;
204                    match msg {
205                        BackendMessage::DataRow(body) => {
206                            let values = read_data_row(body)?;
207                            let cols = self.columns.clone().unwrap_or_default();
208                            self.rows_fetched += 1;
209                            return Ok(Some(Row::new(cols, values)));
210                        }
211                        BackendMessage::CommandComplete(body) => {
212                            let tag = CommandTag::new(body.tag().unwrap_or("").into());
213                            self.state = RowStreamState::Finishing { command_tag: tag };
214                        }
215                        BackendMessage::ErrorResponse(body) => {
216                            let server_err =
217                                PgServerError::from_error_body(&body).map_err(PgError::Io)?;
218                            self.conn.read_until_ready().await?;
219                            self.conn.state = ConnectionState::Idle;
220                            self.state = RowStreamState::Error;
221                            return Err(PgError::Server(Box::new(server_err)));
222                        }
223                        BackendMessage::NotificationResponse(body) => {
224                            self.conn.notification_queue.push_back(
225                                crate::notification::Notification {
226                                    process_id: body.process_id(),
227                                    channel: body.channel().unwrap_or("").to_string(),
228                                    payload: body.message().unwrap_or("").to_string(),
229                                },
230                            );
231                            continue;
232                        }
233                        BackendMessage::NoticeResponse(body) => {
234                            if let Ok(notice) = crate::query::Notice::from_fields(&body) {
235                                self.conn.handle_notice(&notice);
236                            }
237                            continue;
238                        }
239                        BackendMessage::ParameterStatus(body) => {
240                            if let (Ok(name), Ok(value)) = (body.name(), body.value()) {
241                                self.conn
242                                    .server_params
243                                    .params
244                                    .insert(name.to_string(), value.to_string());
245                            }
246                            continue;
247                        }
248                        BackendMessage::PortalSuspended => {
249                            continue;
250                        }
251                        _ => continue,
252                    }
253                }
254
255                RowStreamState::Finishing { .. } => {
256                    let msg = self
257                        .conn
258                        .codec
259                        .read_message(&mut self.conn.transport)
260                        .await?;
261                    match msg {
262                        BackendMessage::ReadyForQuery(body) => {
263                            let tag =
264                                match std::mem::replace(&mut self.state, RowStreamState::Error) {
265                                    RowStreamState::Finishing { command_tag } => command_tag,
266                                    _ => CommandTag::default(),
267                                };
268                            self.conn.transaction_status =
269                                TransactionStatus::from_u8(body.status())
270                                    .unwrap_or(TransactionStatus::Idle);
271                            self.conn.state = ConnectionState::Idle;
272                            self.state = RowStreamState::Done { command_tag: tag };
273                            #[cfg(feature = "tracing")]
274                            tracing::info!(
275                                target: TARGET_QUERY,
276                                rows_fetched = self.rows_fetched,
277                                command_tag = self.command_tag().map(|t| t.as_str()).unwrap_or(""),
278                                elapsed_ms = self.started_at.elapsed().as_millis() as u64,
279                                "Query stream completed"
280                            );
281                            return Ok(None);
282                        }
283                        BackendMessage::ErrorResponse(body) => {
284                            // The server sent an error after CommandComplete but
285                            // before ReadyForQuery. This can happen with certain
286                            // PostgreSQL error conditions (e.g., cursor errors,
287                            // constraint violations in multi-statement queries).
288                            // We must surface this error rather than silently
289                            // discarding it.
290                            let server_err =
291                                PgServerError::from_error_body(&body).map_err(PgError::Io)?;
292                            self.conn.read_until_ready().await?;
293                            self.conn.state = ConnectionState::Idle;
294                            self.state = RowStreamState::Error;
295                            return Err(PgError::Server(Box::new(server_err)));
296                        }
297                        BackendMessage::NotificationResponse(body) => {
298                            self.conn.notification_queue.push_back(
299                                crate::notification::Notification {
300                                    process_id: body.process_id(),
301                                    channel: body.channel().unwrap_or("").to_string(),
302                                    payload: body.message().unwrap_or("").to_string(),
303                                },
304                            );
305                            continue;
306                        }
307                        BackendMessage::NoticeResponse(body) => {
308                            if let Ok(notice) = crate::query::Notice::from_fields(&body) {
309                                self.conn.handle_notice(&notice);
310                            }
311                            continue;
312                        }
313                        BackendMessage::ParameterStatus(body) => {
314                            if let (Ok(name), Ok(value)) = (body.name(), body.value()) {
315                                self.conn
316                                    .server_params
317                                    .params
318                                    .insert(name.to_string(), value.to_string());
319                            }
320                            continue;
321                        }
322                        _ => continue,
323                    }
324                }
325            }
326        }
327    }
328
329    /// Get the column metadata for the current result set.
330    pub fn columns(&self) -> Option<&[FieldDescription]> {
331        self.columns.as_ref().map(|c| c.as_slice())
332    }
333
334    /// Get the command tag after the stream ends.
335    pub fn command_tag(&self) -> Option<&CommandTag> {
336        match &self.state {
337            RowStreamState::Done { command_tag } | RowStreamState::Finishing { command_tag } => {
338                Some(command_tag)
339            }
340            _ => None,
341        }
342    }
343
344    /// Returns true if the stream has been fully consumed or encountered an error.
345    pub fn is_done(&self) -> bool {
346        matches!(
347            self.state,
348            RowStreamState::Done { .. } | RowStreamState::Error
349        )
350    }
351
352    /// Consume the remaining rows in the stream, discarding them.
353    #[must_use = "consume errors should be checked"]
354    pub async fn consume(mut self) -> Result<CommandTag> {
355        while self.next().await?.is_some() {}
356        match &self.state {
357            RowStreamState::Done { command_tag } => Ok(command_tag.clone()),
358            _ => Ok(CommandTag::default()),
359        }
360    }
361}
362
363impl<'a> Drop for RowStream<'a> {
364    fn drop(&mut self) {
365        if !self.is_done() {
366            #[cfg(feature = "tracing")]
367            if !std::thread::panicking() {
368                tracing::warn!(target: TARGET_QUERY, "RowStream dropped without full consumption; connection may need recovery");
369            }
370            self.conn.needs_recovery = true;
371        }
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378    use crate::auth::{Codec, ServerParams};
379    use crate::config::Config;
380    use crate::connection::ConnectionState;
381    use crate::transport::{BufferedTransport, ClientTransport, MockTransport, PgTransport};
382    use std::collections::VecDeque;
383
384    fn make_connection(read_data: Vec<u8>) -> Connection {
385        let transport = PgTransport::Plain(BufferedTransport::new(ClientTransport::Mock(
386            MockTransport::new(read_data),
387        )));
388        Connection {
389            transport,
390            codec: Codec::new(),
391            server_params: ServerParams::default(),
392            state: ConnectionState::Idle,
393            config: Config::new(),
394            transaction_status: TransactionStatus::Idle,
395            notification_queue: VecDeque::new(),
396            notice_handler: None,
397            statement_counter: 0,
398            needs_recovery: false,
399            health: crate::reconnect::session::ConnectionHealth::new(),
400            session_state: crate::reconnect::session::SessionState::new(),
401        }
402    }
403
404    fn build_row_description_msg(fields: &[(&str, u32)]) -> Vec<u8> {
405        let mut buf = vec![b'T'];
406        let mut body = Vec::new();
407        body.extend_from_slice(&(fields.len() as i16).to_be_bytes());
408        for (name, type_oid) in fields {
409            body.extend_from_slice(name.as_bytes());
410            body.push(0);
411            body.extend_from_slice(&0u32.to_be_bytes());
412            body.extend_from_slice(&0i16.to_be_bytes());
413            body.extend_from_slice(&type_oid.to_be_bytes());
414            body.extend_from_slice(&(-1i16).to_be_bytes());
415            body.extend_from_slice(&(-1i32).to_be_bytes());
416            body.extend_from_slice(&0i16.to_be_bytes());
417        }
418        let len = (body.len() + 4) as i32;
419        buf.extend_from_slice(&len.to_be_bytes());
420        buf.extend_from_slice(&body);
421        buf
422    }
423
424    fn build_data_row_msg(values: &[Option<&str>]) -> Vec<u8> {
425        let mut buf = vec![b'D'];
426        let mut body = Vec::new();
427        body.extend_from_slice(&(values.len() as i16).to_be_bytes());
428        for val in values {
429            match val {
430                Some(v) => {
431                    let bytes = v.as_bytes();
432                    body.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
433                    body.extend_from_slice(bytes);
434                }
435                None => {
436                    body.extend_from_slice(&(-1i32).to_be_bytes());
437                }
438            }
439        }
440        let len = (body.len() + 4) as i32;
441        buf.extend_from_slice(&len.to_be_bytes());
442        buf.extend_from_slice(&body);
443        buf
444    }
445
446    fn build_command_complete_msg(tag: &str) -> Vec<u8> {
447        let mut buf = vec![b'C'];
448        let mut body = Vec::new();
449        body.extend_from_slice(tag.as_bytes());
450        body.push(0);
451        let len = (body.len() + 4) as i32;
452        buf.extend_from_slice(&len.to_be_bytes());
453        buf.extend_from_slice(&body);
454        buf
455    }
456
457    fn build_ready_for_query(status: u8) -> Vec<u8> {
458        vec![b'Z', 0, 0, 0, 5, status]
459    }
460
461    fn build_error_response_msg(severity: &str, code: &str, message: &str) -> Vec<u8> {
462        let mut buf = vec![b'E'];
463        let mut body = Vec::new();
464        // Severity
465        body.push(b'S');
466        body.extend_from_slice(severity.as_bytes());
467        body.push(0);
468        // Code (SQLSTATE)
469        body.push(b'C');
470        body.extend_from_slice(code.as_bytes());
471        body.push(0);
472        // Message
473        body.push(b'M');
474        body.extend_from_slice(message.as_bytes());
475        body.push(0);
476        // Terminator
477        body.push(0);
478        let len = (body.len() + 4) as i32;
479        buf.extend_from_slice(&len.to_be_bytes());
480        buf.extend_from_slice(&body);
481        buf
482    }
483
484    #[tokio::test]
485    async fn test_row_stream_basic() {
486        let mut data = Vec::new();
487        data.extend_from_slice(&build_row_description_msg(&[
488            ("id", crate::types::INT4_OID),
489            ("name", crate::types::TEXT_OID),
490        ]));
491        data.extend_from_slice(&build_data_row_msg(&[Some("1"), Some("alice")]));
492        data.extend_from_slice(&build_data_row_msg(&[Some("2"), Some("bob")]));
493        data.extend_from_slice(&build_data_row_msg(&[Some("3"), Some("charlie")]));
494        data.extend_from_slice(&build_command_complete_msg("SELECT 3"));
495        data.extend_from_slice(&build_ready_for_query(b'I'));
496        let mut conn = make_connection(data);
497        let mut stream = RowStream::new_simple(&mut conn);
498        assert!(stream.columns().is_none());
499        let row1 = stream.next().await.unwrap().unwrap();
500        let id1: i32 = row1.get(0).unwrap();
501        assert_eq!(id1, 1);
502        assert!(stream.columns().is_some());
503        let row2 = stream.next().await.unwrap().unwrap();
504        let name2: String = row2.get(1).unwrap();
505        assert_eq!(name2, "bob");
506        let row3 = stream.next().await.unwrap().unwrap();
507        let name3: String = row3.get(1).unwrap();
508        assert_eq!(name3, "charlie");
509        let done = stream.next().await.unwrap();
510        assert!(done.is_none());
511        assert!(stream.is_done());
512        assert_eq!(stream.command_tag().unwrap().as_str(), "SELECT 3");
513    }
514
515    #[tokio::test]
516    async fn test_row_stream_empty() {
517        let mut data = Vec::new();
518        data.extend_from_slice(&build_row_description_msg(&[(
519            "id",
520            crate::types::INT4_OID,
521        )]));
522        data.extend_from_slice(&build_command_complete_msg("SELECT 0"));
523        data.extend_from_slice(&build_ready_for_query(b'I'));
524        let mut conn = make_connection(data);
525        let mut stream = RowStream::new_simple(&mut conn);
526        let done = stream.next().await.unwrap();
527        assert!(done.is_none());
528        assert!(stream.is_done());
529    }
530
531    #[tokio::test]
532    async fn test_row_stream_no_row_description() {
533        let mut data = Vec::new();
534        data.extend_from_slice(&build_command_complete_msg("INSERT 0 1"));
535        data.extend_from_slice(&build_ready_for_query(b'I'));
536        let mut conn = make_connection(data);
537        let mut stream = RowStream::new_simple(&mut conn);
538        let done = stream.next().await.unwrap();
539        assert!(done.is_none());
540        assert!(stream.is_done());
541        assert_eq!(stream.command_tag().unwrap().as_str(), "INSERT 0 1");
542    }
543
544    #[tokio::test]
545    async fn test_row_stream_consume() {
546        let mut data = Vec::new();
547        data.extend_from_slice(&build_row_description_msg(&[(
548            "val",
549            crate::types::INT4_OID,
550        )]));
551        data.extend_from_slice(&build_data_row_msg(&[Some("10")]));
552        data.extend_from_slice(&build_data_row_msg(&[Some("20")]));
553        data.extend_from_slice(&build_data_row_msg(&[Some("30")]));
554        data.extend_from_slice(&build_command_complete_msg("SELECT 3"));
555        data.extend_from_slice(&build_ready_for_query(b'I'));
556        let mut conn = make_connection(data);
557        let mut stream = RowStream::new_simple(&mut conn);
558        let row1 = stream.next().await.unwrap().unwrap();
559        let v: i32 = row1.get(0).unwrap();
560        assert_eq!(v, 10);
561        let tag = stream.consume().await.unwrap();
562        assert_eq!(tag.as_str(), "SELECT 3");
563        assert!(!conn.needs_recovery());
564    }
565
566    #[tokio::test]
567    async fn test_row_stream_drop_sets_needs_recovery() {
568        let mut data = Vec::new();
569        data.extend_from_slice(&build_row_description_msg(&[(
570            "val",
571            crate::types::INT4_OID,
572        )]));
573        data.extend_from_slice(&build_data_row_msg(&[Some("10")]));
574        data.extend_from_slice(&build_data_row_msg(&[Some("20")]));
575        data.extend_from_slice(&build_command_complete_msg("SELECT 2"));
576        data.extend_from_slice(&build_ready_for_query(b'I'));
577        let mut conn = make_connection(data);
578        {
579            let mut stream = RowStream::new_simple(&mut conn);
580            let _row1 = stream.next().await.unwrap().unwrap();
581        }
582        assert!(conn.needs_recovery());
583    }
584
585    #[tokio::test]
586    async fn test_row_stream_extended_protocol() {
587        let mut data = Vec::new();
588        data.extend_from_slice(&[b'1', 0, 0, 0, 4]);
589        data.extend_from_slice(&[b'2', 0, 0, 0, 4]);
590        data.extend_from_slice(&build_row_description_msg(&[(
591            "val",
592            crate::types::INT4_OID,
593        )]));
594        data.extend_from_slice(&build_data_row_msg(&[Some("42")]));
595        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
596        data.extend_from_slice(&build_ready_for_query(b'I'));
597        let mut conn = make_connection(data);
598        let mut stream = RowStream::new_extended(&mut conn);
599        let row = stream.next().await.unwrap().unwrap();
600        let v: i32 = row.get(0).unwrap();
601        assert_eq!(v, 42);
602        let done = stream.next().await.unwrap();
603        assert!(done.is_none());
604    }
605
606    #[tokio::test]
607    async fn test_connection_recover_after_dropped_stream() {
608        // Simulate: stream is dropped early, then recover() drains remaining messages
609        let mut data = Vec::new();
610        data.extend_from_slice(&build_row_description_msg(&[(
611            "val",
612            crate::types::INT4_OID,
613        )]));
614        data.extend_from_slice(&build_data_row_msg(&[Some("10")]));
615        data.extend_from_slice(&build_data_row_msg(&[Some("20")]));
616        data.extend_from_slice(&build_command_complete_msg("SELECT 2"));
617        data.extend_from_slice(&build_ready_for_query(b'I'));
618        let mut conn = make_connection(data);
619        {
620            let mut stream = RowStream::new_simple(&mut conn);
621            let _row1 = stream.next().await.unwrap().unwrap();
622            // Drop stream without consuming remaining rows
623        }
624        assert!(
625            conn.needs_recovery(),
626            "should need recovery after dropping stream"
627        );
628
629        // Now recover — this should read until ReadyForQuery
630        conn.recover().await.unwrap();
631        assert!(
632            !conn.needs_recovery(),
633            "should not need recovery after recover"
634        );
635        assert_eq!(
636            conn.state(),
637            ConnectionState::Idle,
638            "should be idle after recover"
639        );
640    }
641
642    #[tokio::test]
643    async fn test_row_stream_error_in_finishing_state() {
644        // Simulate: server sends ErrorResponse after CommandComplete but
645        // before ReadyForQuery. This used to be silently discarded.
646        let mut data = Vec::new();
647        data.extend_from_slice(&build_row_description_msg(&[(
648            "val",
649            crate::types::INT4_OID,
650        )]));
651        data.extend_from_slice(&build_data_row_msg(&[Some("10")]));
652        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
653        // Server sends an error during the finishing phase
654        data.extend_from_slice(&build_error_response_msg(
655            "ERROR",
656            "57014", // query_canceled
657            "canceling statement due to user request",
658        ));
659        data.extend_from_slice(&build_ready_for_query(b'I'));
660        let mut conn = make_connection(data);
661        let mut stream = RowStream::new_simple(&mut conn);
662        // First row is fine
663        let row = stream.next().await.unwrap().unwrap();
664        let v: i32 = row.get(0).unwrap();
665        assert_eq!(v, 10);
666        // After CommandComplete, the stream enters Finishing state.
667        // The next call should encounter the ErrorResponse and return an error.
668        let result = stream.next().await;
669        assert!(
670            result.is_err(),
671            "ErrorResponse in Finishing state should be surfaced as an error"
672        );
673        let err = result.unwrap_err();
674        match err {
675            PgError::Server(server_err) => {
676                assert_eq!(
677                    server_err.code, "57014",
678                    "error code should be query_canceled"
679                );
680            }
681            other => panic!("expected PgError::Server, got {:?}", other),
682        }
683        // Drop the stream so we can access conn
684        drop(stream);
685        // Connection should be idle after the error is handled
686        assert_eq!(
687            conn.state(),
688            ConnectionState::Idle,
689            "connection should be idle after error in finishing state"
690        );
691    }
692}
wasi_pg_client/query/stream.rs

wasi_pg_client/query/
stream.rs
