wasi_pg_client/query/
cursor.rs

1//! Cursor support for fetching large result sets in batches.
2//!
3//! A [`Cursor`] executes a portal with a limited `max_rows` count and
4//! provides `fetch_next()` to retrieve subsequent batches until the result
5//! set is exhausted.
6//!
7//! A [`CursorStream`] wraps a cursor and yields rows one at a time,
8//! automatically fetching the next batch when the current batch is exhausted.
9
10use std::sync::Arc;
11
12use crate::protocol::{BackendMessage, FrontendMessage, TransactionStatus};
13
14use crate::connection::{Connection, ConnectionState};
15use crate::error::{PgError, PgServerError, Result};
16use crate::query::params::encode_params_text;
17use crate::query::result::CommandTag;
18use crate::query::row::{FieldDescription, Row};
19use crate::query::{read_data_row, read_row_description};
20use crate::transport::AsyncTransport;
21
22// ---------------------------------------------------------------------------
23// Cursor
24// ---------------------------------------------------------------------------
25
26/// A cursor for fetching a large result set in batches.
27///
28/// Created via [`Connection::query_cursor`]. Each call to [`Cursor::fetch_next`]
29/// returns the next batch of rows. The cursor is automatically closed when
30/// dropped, but explicit [`Cursor::close`] is recommended for clean shutdown.
31#[non_exhaustive]
32pub struct Cursor<'a> {
33    conn: &'a mut Connection,
34    portal_name: String,
35    columns: Arc<Vec<FieldDescription>>,
36    fetch_size: i32,
37    done: bool,
38    /// Whether this cursor started the transaction and should commit on close.
39    owns_transaction: bool,
40}
41
42impl<'a> Cursor<'a> {
43    /// Fetch the next batch of rows.
44    ///
45    /// Returns an empty vector when all rows have been consumed.
46    #[must_use = "cursor errors should be checked"]
47    pub async fn fetch_next(&mut self) -> Result<Vec<Row>> {
48        if self.done {
49            return Ok(Vec::new());
50        }
51
52        self.conn.transition(ConnectionState::ActiveExtendedQuery)?;
53
54        // Execute portal with limited max_rows
55        self.conn
56            .codec
57            .encode_and_write(
58                &mut self.conn.transport,
59                &FrontendMessage::Execute {
60                    portal: self.portal_name.clone(),
61                    max_rows: self.fetch_size,
62                },
63            )
64            .await?;
65
66        self.conn
67            .codec
68            .encode_and_write(&mut self.conn.transport, &FrontendMessage::Sync)
69            .await?;
70
71        // Flush the batch
72        self.conn
73            .transport
74            .flush()
75            .await
76            .map_err(PgError::Transport)?;
77
78        let mut rows = Vec::new();
79
80        loop {
81            let msg = self
82                .conn
83                .codec
84                .read_message(&mut self.conn.transport)
85                .await?;
86            if self.conn.handle_async_message(&msg) {
87                continue;
88            }
89            match msg {
90                BackendMessage::RowDescription(body) => {
91                    self.columns = Arc::new(read_row_description(body)?);
92                }
93                BackendMessage::DataRow(body) => {
94                    let values = read_data_row(body)?;
95                    rows.push(Row::new(self.columns.clone(), values));
96                }
97                BackendMessage::CommandComplete(_body) => {
98                    self.done = true;
99                }
100                BackendMessage::PortalSuspended => {
101                    // More rows available; portal remains open
102                }
103                BackendMessage::ReadyForQuery(body) => {
104                    self.conn.transaction_status = TransactionStatus::from_u8(body.status())
105                        .unwrap_or(TransactionStatus::Idle);
106                    self.conn.state = ConnectionState::Idle;
107                    break;
108                }
109                BackendMessage::ErrorResponse(body) => {
110                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
111                    self.conn.read_until_ready().await?;
112                    self.conn.state = ConnectionState::Idle;
113                    return Err(PgError::Server(Box::new(server_err)));
114                }
115                _ => {}
116            }
117        }
118
119        Ok(rows)
120    }
121
122    /// Close the cursor, releasing the portal on the server.
123    ///
124    /// If the cursor automatically started a transaction (because no
125    /// transaction was active when the cursor was created), the transaction
126    /// is committed.
127    #[must_use = "cursor close errors should be checked"]
128    pub async fn close(mut self) -> Result<()> {
129        self.conn.transition(ConnectionState::ActiveExtendedQuery)?;
130
131        self.conn
132            .codec
133            .encode_and_write(
134                &mut self.conn.transport,
135                &FrontendMessage::Close {
136                    variant: b'P',
137                    name: self.portal_name.clone(),
138                },
139            )
140            .await?;
141
142        self.conn
143            .codec
144            .encode_and_write(&mut self.conn.transport, &FrontendMessage::Sync)
145            .await?;
146
147        // Flush the batch
148        self.conn
149            .transport
150            .flush()
151            .await
152            .map_err(PgError::Transport)?;
153
154        loop {
155            let msg = self
156                .conn
157                .codec
158                .read_message(&mut self.conn.transport)
159                .await?;
160            if self.conn.handle_async_message(&msg) {
161                continue;
162            }
163            match msg {
164                BackendMessage::CloseComplete => {}
165                BackendMessage::ReadyForQuery(body) => {
166                    self.conn.transaction_status = TransactionStatus::from_u8(body.status())
167                        .unwrap_or(TransactionStatus::Idle);
168                    self.conn.state = ConnectionState::Idle;
169                    break;
170                }
171                BackendMessage::ErrorResponse(body) => {
172                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
173                    self.conn.read_until_ready().await?;
174                    self.conn.state = ConnectionState::Idle;
175                    return Err(PgError::Server(Box::new(server_err)));
176                }
177                _ => {}
178            }
179        }
180
181        // Commit the transaction if we started it
182        if self.owns_transaction {
183            self.conn.execute("COMMIT").await?;
184        }
185
186        self.done = true;
187        Ok(())
188    }
189
190    /// Returns true if all rows have been fetched.
191    pub fn is_done(&self) -> bool {
192        self.done
193    }
194}
195
196// ---------------------------------------------------------------------------
197// CursorStream
198// ---------------------------------------------------------------------------
199
200/// Internal state of the cursor stream.
201#[derive(Debug)]
202enum CursorStreamState {
203    /// Rows are available in the buffer or more can be fetched.
204    Active,
205    /// All rows have been consumed and the cursor is closed.
206    Done { command_tag: CommandTag },
207    /// An error occurred.
208    Error,
209}
210
211/// A streaming cursor that yields rows one at a time from a portal.
212///
213/// Unlike [`Cursor`] which returns batches of rows, `CursorStream` provides
214/// a row-by-row iterator interface. When the current batch is exhausted,
215/// it automatically fetches the next batch from the server.
216///
217/// `CursorStream` borrows the connection mutably. You cannot use the
218/// connection while iterating. When the stream is dropped (or fully consumed),
219/// the connection is available again.
220///
221/// If the stream is dropped before being fully consumed, the connection is
222/// left in an inconsistent state. The [`Connection::needs_recovery`] flag is
223/// set, and you must call [`Connection::recover`] before using the connection
224/// again.
225#[non_exhaustive]
226pub struct CursorStream<'a> {
227    conn: &'a mut Connection,
228    portal_name: String,
229    columns: Arc<Vec<FieldDescription>>,
230    fetch_size: i32,
231    state: CursorStreamState,
232    /// Rows buffered from the current batch, waiting to be yielded.
233    buffered_rows: Vec<Row>,
234    /// Whether this cursor started the transaction and should commit on close.
235    owns_transaction: bool,
236}
237
238impl<'a> CursorStream<'a> {
239    /// Create a new `CursorStream` from an already-set-up portal.
240    pub(crate) fn new(
241        conn: &'a mut Connection,
242        portal_name: String,
243        columns: Arc<Vec<FieldDescription>>,
244        fetch_size: i32,
245        owns_transaction: bool,
246    ) -> Self {
247        CursorStream {
248            conn,
249            portal_name,
250            columns,
251            fetch_size,
252            state: CursorStreamState::Active,
253            buffered_rows: Vec::new(),
254            owns_transaction,
255        }
256    }
257
258    /// Fetch the next row from the stream.
259    ///
260    /// Returns `Ok(Some(row))` when a row is available, `Ok(None)` when the
261    /// stream is exhausted, and `Err(...)` on a server or protocol error.
262    /// After returning `None` or `Err`, subsequent calls return `None`.
263    #[must_use = "cursor stream errors should be checked"]
264    pub async fn next(&mut self) -> Result<Option<Row>> {
265        loop {
266            match self.state {
267                CursorStreamState::Done { .. } | CursorStreamState::Error => {
268                    // Yield any remaining buffered rows before reporting done
269                    if let Some(row) = self.buffered_rows.pop() {
270                        return Ok(Some(row));
271                    }
272                    return Ok(None);
273                }
274
275                CursorStreamState::Active => {
276                    // First, yield from the buffer if we have rows from a previous fetch
277                    if let Some(row) = self.buffered_rows.pop() {
278                        return Ok(Some(row));
279                    }
280
281                    // Buffer is empty — fetch the next batch from the server
282                    self.conn.transition(ConnectionState::ActiveExtendedQuery)?;
283
284                    self.conn
285                        .codec
286                        .encode_and_write(
287                            &mut self.conn.transport,
288                            &FrontendMessage::Execute {
289                                portal: self.portal_name.clone(),
290                                max_rows: self.fetch_size,
291                            },
292                        )
293                        .await?;
294
295                    self.conn
296                        .codec
297                        .encode_and_write(&mut self.conn.transport, &FrontendMessage::Sync)
298                        .await?;
299
300                    self.conn
301                        .transport
302                        .flush()
303                        .await
304                        .map_err(PgError::Transport)?;
305
306                    // Read the batch response
307                    let mut command_tag: Option<CommandTag> = None;
308                    let mut rows: Vec<Row> = Vec::new();
309
310                    loop {
311                        let msg = self
312                            .conn
313                            .codec
314                            .read_message(&mut self.conn.transport)
315                            .await?;
316                        if self.conn.handle_async_message(&msg) {
317                            continue;
318                        }
319                        match msg {
320                            BackendMessage::RowDescription(body) => {
321                                self.columns = Arc::new(read_row_description(body)?);
322                            }
323                            BackendMessage::DataRow(body) => {
324                                let values = read_data_row(body)?;
325                                rows.push(Row::new(self.columns.clone(), values));
326                            }
327                            BackendMessage::CommandComplete(body) => {
328                                command_tag =
329                                    Some(CommandTag::new(body.tag().unwrap_or("").into()));
330                            }
331                            BackendMessage::PortalSuspended => {
332                                // More rows available; portal remains open
333                            }
334                            BackendMessage::ReadyForQuery(body) => {
335                                self.conn.transaction_status =
336                                    TransactionStatus::from_u8(body.status())
337                                        .unwrap_or(TransactionStatus::Idle);
338                                self.conn.state = ConnectionState::Idle;
339                                break;
340                            }
341                            BackendMessage::ErrorResponse(body) => {
342                                let server_err =
343                                    PgServerError::from_error_body(&body).map_err(PgError::Io)?;
344                                self.conn.read_until_ready().await?;
345                                self.conn.state = ConnectionState::Idle;
346                                self.state = CursorStreamState::Error;
347                                return Err(PgError::Server(Box::new(server_err)));
348                            }
349                            _ => {}
350                        }
351                    }
352
353                    // If CommandComplete was received, all rows are done
354                    if let Some(tag) = command_tag {
355                        self.state = CursorStreamState::Done { command_tag: tag };
356                    }
357
358                    // If no rows were returned and we're not done, loop to fetch again
359                    // (this can happen with PortalSuspended when fetch_size rows were
360                    // already consumed in a previous batch)
361                    if rows.is_empty() && !self.is_done() {
362                        continue;
363                    }
364
365                    // Reverse the buffer so we can pop() from the front efficiently
366                    rows.reverse();
367                    self.buffered_rows = rows;
368
369                    // Yield the first row from the buffer
370                    if let Some(row) = self.buffered_rows.pop() {
371                        return Ok(Some(row));
372                    }
373
374                    // No rows and done
375                    return Ok(None);
376                }
377            }
378        }
379    }
380
381    /// Get the column metadata for the current result set.
382    pub fn columns(&self) -> &[FieldDescription] {
383        &self.columns
384    }
385
386    /// Returns true if the stream has been fully consumed or encountered an error.
387    pub fn is_done(&self) -> bool {
388        matches!(
389            self.state,
390            CursorStreamState::Done { .. } | CursorStreamState::Error
391        )
392    }
393
394    /// Get the command tag after the stream ends.
395    pub fn command_tag(&self) -> Option<&CommandTag> {
396        match &self.state {
397            CursorStreamState::Done { command_tag } => Some(command_tag),
398            _ => None,
399        }
400    }
401
402    /// Consume the remaining rows in the stream, discarding them, and close
403    /// the cursor portal.
404    #[must_use = "consume errors should be checked"]
405    pub async fn consume(mut self) -> Result<CommandTag> {
406        while self.next().await?.is_some() {}
407        self.close_portal().await?;
408        match &self.state {
409            CursorStreamState::Done { command_tag } => Ok(command_tag.clone()),
410            _ => Ok(CommandTag::default()),
411        }
412    }
413
414    /// Close the portal on the server and commit the transaction if we own it.
415    async fn close_portal(&mut self) -> Result<()> {
416        if matches!(self.state, CursorStreamState::Done { .. }) {
417            // Already done — just commit if needed
418            if self.owns_transaction {
419                self.conn.execute("COMMIT").await?;
420            }
421            return Ok(());
422        }
423
424        self.conn.transition(ConnectionState::ActiveExtendedQuery)?;
425
426        self.conn
427            .codec
428            .encode_and_write(
429                &mut self.conn.transport,
430                &FrontendMessage::Close {
431                    variant: b'P',
432                    name: self.portal_name.clone(),
433                },
434            )
435            .await?;
436
437        self.conn
438            .codec
439            .encode_and_write(&mut self.conn.transport, &FrontendMessage::Sync)
440            .await?;
441
442        self.conn
443            .transport
444            .flush()
445            .await
446            .map_err(PgError::Transport)?;
447
448        loop {
449            let msg = self
450                .conn
451                .codec
452                .read_message(&mut self.conn.transport)
453                .await?;
454            if self.conn.handle_async_message(&msg) {
455                continue;
456            }
457            match msg {
458                BackendMessage::CloseComplete => {}
459                BackendMessage::ReadyForQuery(body) => {
460                    self.conn.transaction_status = TransactionStatus::from_u8(body.status())
461                        .unwrap_or(TransactionStatus::Idle);
462                    self.conn.state = ConnectionState::Idle;
463                    break;
464                }
465                BackendMessage::ErrorResponse(body) => {
466                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
467                    self.conn.read_until_ready().await?;
468                    self.conn.state = ConnectionState::Idle;
469                    return Err(PgError::Server(Box::new(server_err)));
470                }
471                _ => {}
472            }
473        }
474
475        // Commit the transaction if we started it
476        if self.owns_transaction {
477            self.conn.execute("COMMIT").await?;
478        }
479
480        self.state = CursorStreamState::Done {
481            command_tag: CommandTag::default(),
482        };
483        Ok(())
484    }
485}
486
487impl<'a> Drop for CursorStream<'a> {
488    fn drop(&mut self) {
489        if !self.is_done() {
490            self.conn.needs_recovery = true;
491        }
492    }
493}
494
495// ---------------------------------------------------------------------------
496// Connection method
497// ---------------------------------------------------------------------------
498
499impl Connection {
500    /// Open a cursor for a parameterized query.
501    ///
502    /// The query is parsed and bound to a named portal. The first batch of
503    /// rows is fetched via [`Cursor::fetch_next`].
504    ///
505    /// **Important:** Named portals only survive within a transaction
506    /// block. If no transaction is active, this method automatically
507    /// begins one so the portal remains valid across `fetch_next` calls.
508    /// The transaction is committed when the cursor is closed.
509    #[must_use = "cursor errors should be checked"]
510    pub async fn query_cursor(
511        &mut self,
512        sql: &str,
513        params: &[&dyn crate::types::ToSql],
514        fetch_size: i32,
515    ) -> Result<Cursor<'_>> {
516        // Named portals only survive within a transaction. Start one if
517        // we're not already in a transaction.
518        let need_transaction = self.transaction_status == crate::protocol::TransactionStatus::Idle;
519        if need_transaction {
520            // Use simple query for BEGIN — it's a single statement with no params
521            self.query("BEGIN").await?;
522        }
523
524        self.transition(ConnectionState::ActiveExtendedQuery)?;
525
526        let param_values = encode_params_text(params)?;
527        let portal_name = format!("__pg_portal_{}", self.statement_counter);
528        self.statement_counter += 1;
529
530        // Parse (unnamed statement)
531        self.codec
532            .encode_and_write(
533                &mut self.transport,
534                &FrontendMessage::Parse {
535                    name: String::new(),
536                    sql: sql.to_string(),
537                    param_types: vec![],
538                },
539            )
540            .await?;
541
542        // Bind (named portal)
543        self.codec
544            .encode_and_write(
545                &mut self.transport,
546                &FrontendMessage::Bind {
547                    portal: portal_name.clone(),
548                    statement: String::new(),
549                    param_formats: vec![crate::protocol::FormatCode::Text],
550                    params: param_values,
551                    result_formats: vec![crate::protocol::FormatCode::Binary],
552                },
553            )
554            .await?;
555
556        // Describe the named portal so the server sends RowDescription
557        // (or NoData for non-SELECT statements).
558        self.codec
559            .encode_and_write(
560                &mut self.transport,
561                &FrontendMessage::Describe {
562                    variant: b'P',
563                    name: portal_name.clone(),
564                },
565            )
566            .await?;
567
568        // Sync to complete the sub-protocol
569        self.codec
570            .encode_and_write(&mut self.transport, &FrontendMessage::Sync)
571            .await?;
572
573        // Flush the batch
574        self.transport.flush().await.map_err(PgError::Transport)?;
575
576        let mut columns: Option<Arc<Vec<FieldDescription>>> = None;
577
578        loop {
579            let msg = self.codec.read_message(&mut self.transport).await?;
580            if self.handle_async_message(&msg) {
581                continue;
582            }
583            match msg {
584                BackendMessage::ParseComplete => {}
585                BackendMessage::BindComplete => {}
586                BackendMessage::NoData => {
587                    // Non-SELECT query opened as cursor
588                }
589                BackendMessage::RowDescription(body) => {
590                    columns = Some(Arc::new(read_row_description(body)?));
591                }
592                BackendMessage::ReadyForQuery(body) => {
593                    self.transaction_status = TransactionStatus::from_u8(body.status())
594                        .unwrap_or(TransactionStatus::Idle);
595                    self.state = ConnectionState::Idle;
596                    break;
597                }
598                BackendMessage::ErrorResponse(body) => {
599                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
600                    self.read_until_ready().await?;
601                    self.state = ConnectionState::Idle;
602                    return Err(PgError::Server(Box::new(server_err)));
603                }
604                _ => {}
605            }
606        }
607
608        Ok(Cursor {
609            conn: self,
610            portal_name,
611            columns: columns.unwrap_or_default(),
612            fetch_size,
613            done: false,
614            owns_transaction: need_transaction,
615        })
616    }
617
618    /// Open a streaming cursor for a parameterized query.
619    ///
620    /// Like [`Connection::query_cursor`], but returns a [`CursorStream`] that
621    /// yields rows one at a time instead of in batches. When the current batch
622    /// is exhausted, the next batch is automatically fetched from the server.
623    ///
624    /// **Important:** Named portals only survive within a transaction
625    /// block. If no transaction is active, this method automatically
626    /// begins one so the portal remains valid. The transaction is committed
627    /// when the stream is fully consumed or closed.
628    ///
629    /// # Example
630    ///
631    /// ```ignore
632    /// let mut stream = conn.query_cursor_stream(
633    ///     "SELECT id, name FROM users WHERE active = $1",
634    ///     &[&true],
635    ///     100, // fetch 100 rows at a time
636    /// ).await?;
637    /// while let Some(row) = stream.next().await? {
638    ///     let id: i32 = row.get(0)?;
639    /// }
640    /// ```
641    #[must_use = "cursor stream errors should be checked"]
642    pub async fn query_cursor_stream(
643        &mut self,
644        sql: &str,
645        params: &[&dyn crate::types::ToSql],
646        fetch_size: i32,
647    ) -> Result<CursorStream<'_>> {
648        // Named portals only survive within a transaction. Start one if
649        // we're not already in a transaction.
650        let need_transaction = self.transaction_status == crate::protocol::TransactionStatus::Idle;
651        if need_transaction {
652            self.query("BEGIN").await?;
653        }
654
655        self.transition(ConnectionState::ActiveExtendedQuery)?;
656
657        let param_values = encode_params_text(params)?;
658        let portal_name = format!("__pg_portal_{}", self.statement_counter);
659        self.statement_counter += 1;
660
661        // Parse (unnamed statement)
662        self.codec
663            .encode_and_write(
664                &mut self.transport,
665                &FrontendMessage::Parse {
666                    name: String::new(),
667                    sql: sql.to_string(),
668                    param_types: vec![],
669                },
670            )
671            .await?;
672
673        // Bind (named portal)
674        self.codec
675            .encode_and_write(
676                &mut self.transport,
677                &FrontendMessage::Bind {
678                    portal: portal_name.clone(),
679                    statement: String::new(),
680                    param_formats: vec![crate::protocol::FormatCode::Text],
681                    params: param_values,
682                    result_formats: vec![crate::protocol::FormatCode::Binary],
683                },
684            )
685            .await?;
686
687        // Describe the named portal so the server sends RowDescription
688        // (or NoData for non-SELECT statements).
689        self.codec
690            .encode_and_write(
691                &mut self.transport,
692                &FrontendMessage::Describe {
693                    variant: b'P',
694                    name: portal_name.clone(),
695                },
696            )
697            .await?;
698
699        // Sync to complete the sub-protocol
700        self.codec
701            .encode_and_write(&mut self.transport, &FrontendMessage::Sync)
702            .await?;
703
704        // Flush the batch
705        self.transport.flush().await.map_err(PgError::Transport)?;
706
707        let mut columns: Option<Arc<Vec<FieldDescription>>> = None;
708
709        loop {
710            let msg = self.codec.read_message(&mut self.transport).await?;
711            if self.handle_async_message(&msg) {
712                continue;
713            }
714            match msg {
715                BackendMessage::ParseComplete => {}
716                BackendMessage::BindComplete => {}
717                BackendMessage::NoData => {
718                    // Non-SELECT query opened as cursor
719                }
720                BackendMessage::RowDescription(body) => {
721                    columns = Some(Arc::new(read_row_description(body)?));
722                }
723                BackendMessage::ReadyForQuery(body) => {
724                    self.transaction_status = TransactionStatus::from_u8(body.status())
725                        .unwrap_or(TransactionStatus::Idle);
726                    self.state = ConnectionState::Idle;
727                    break;
728                }
729                BackendMessage::ErrorResponse(body) => {
730                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
731                    self.read_until_ready().await?;
732                    self.state = ConnectionState::Idle;
733                    return Err(PgError::Server(Box::new(server_err)));
734                }
735                _ => {}
736            }
737        }
738
739        Ok(CursorStream::new(
740            self,
741            portal_name,
742            columns.unwrap_or_default(),
743            fetch_size,
744            need_transaction,
745        ))
746    }
747}
wasi_pg_client/query/cursor.rs

wasi_pg_client/query/
cursor.rs
