wasi_pg_client/query/
params.rs

1//! Parameterized query execution via the Extended Query Protocol.
2//!
3//! This module provides [`Connection::query_params`] and
4//! [`Connection::query_prepared`] for executing queries with parameters,
5//! using either text or binary encoding.
6//!
7//! # Wire protocol flow
8//!
9//! The extended query protocol sends Parse → Bind → Execute → Sync as a
10//! **batch** (no flush between individual messages), then flushes once before
11//! reading responses.  This is critical: if each message is flushed
12//! individually, Nagle's algorithm on the TCP layer may delay small writes,
13//! causing the server to wait for more data while the client waits for a
14//! response — a protocol-level deadlock.
15
16use std::sync::Arc;
17
18use crate::protocol::{BackendMessage, FrontendMessage, TransactionStatus};
19use crate::types::{Format, ToSql};
20
21use crate::connection::{Connection, ConnectionState};
22use crate::error::{PgError, PgServerError, Result};
23use crate::query::prepared::PreparedStatement;
24use crate::query::result::{CommandTag, ExecuteResult, QueryResult};
25use crate::query::row::{FieldDescription, Row};
26use crate::query::{read_data_row, read_row_description};
27use crate::transport::AsyncTransport;
28
29// ---------------------------------------------------------------------------
30// Parameter encoding helpers
31// ---------------------------------------------------------------------------
32
33/// Encode parameters as **text** (format 0).
34///
35/// Used for one-shot `query_params` where we don't yet know the server's
36/// inferred parameter types.  Each parameter is serialised as its text
37/// representation; `None` maps to a NULL (represented as `None` in the
38/// returned vector, which the Bind message encodes as length -1).
39pub(crate) fn encode_params_text(params: &[&dyn ToSql]) -> Result<Vec<Option<Vec<u8>>>> {
40    let mut values = Vec::with_capacity(params.len());
41    for p in params {
42        let mut buf = Vec::new();
43        // Text encoding ignores the specific Type, so we pass UNKNOWN.
44        let is_null = p.to_sql(&crate::types::Type::UNKNOWN, &mut buf, Format::Text)?;
45        match is_null {
46            crate::types::IsNull::Yes => values.push(None),
47            crate::types::IsNull::No => values.push(Some(buf)),
48        }
49    }
50    Ok(values)
51}
52
53/// Encode parameters as **binary** (format 1).
54///
55/// Used for `query_prepared` where the [`PreparedStatement`] already stores
56/// the parameter types.  NULL values are represented as `None` in the
57/// returned vector.
58pub(crate) fn encode_params_binary(
59    params: &[&dyn ToSql],
60    param_types: &[crate::types::Type],
61) -> Result<Vec<Option<Vec<u8>>>> {
62    if params.len() != param_types.len() {
63        return Err(PgError::Config(format!(
64            "parameter count mismatch: expected {}, got {}",
65            param_types.len(),
66            params.len()
67        )));
68    }
69
70    let mut values = Vec::with_capacity(params.len());
71    for (p, ty) in params.iter().zip(param_types.iter()) {
72        let mut buf = Vec::new();
73        let is_null = p.to_sql(ty, &mut buf, Format::Binary)?;
74        match is_null {
75            crate::types::IsNull::Yes => values.push(None),
76            crate::types::IsNull::No => values.push(Some(buf)),
77        }
78    }
79    Ok(values)
80}
81
82// ---------------------------------------------------------------------------
83// Connection methods
84// ---------------------------------------------------------------------------
85
86impl Connection {
87    /// Execute a parameterized query (one-shot extended query).
88    ///
89    /// Uses the unnamed prepared statement and unnamed portal. Parameters are
90    /// sent in **text** format (the server infers types from the SQL). Result
91    /// columns are requested in **binary** format where supported.
92    ///
93    /// # Wire protocol
94    ///
95    /// The messages Parse, Bind, Execute, and Sync are written into the
96    /// transport's write buffer **without flushing** between them, then a
97    /// single `flush()` is issued.  This avoids Nagle's-algorithm-induced
98    /// deadlocks where small individual writes are buffered by the kernel
99    /// while the server is waiting for the complete pipeline.
100    ///
101    /// # Example
102    /// ```ignore
103    /// let result = conn
104    ///     .query_params("SELECT * FROM users WHERE id = $1", &[&42i32])
105    ///     .await?;
106    /// ```
107    #[must_use = "query results should be checked for errors"]
108    pub async fn query_params(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<QueryResult> {
109        self.transition(ConnectionState::ActiveExtendedQuery)?;
110
111        let param_values = encode_params_text(params)?;
112
113        // ── Batch: Parse + Bind + Describe + Execute + Sync (no flush between) ──
114
115        // Parse (unnamed statement)
116        self.codec
117            .encode_and_write(
118                &mut self.transport,
119                &FrontendMessage::Parse {
120                    name: String::new(),
121                    sql: sql.to_string(),
122                    param_types: vec![],
123                },
124            )
125            .await?;
126
127        // Bind (unnamed portal, unnamed statement)
128        self.codec
129            .encode_and_write(
130                &mut self.transport,
131                &FrontendMessage::Bind {
132                    portal: String::new(),
133                    statement: String::new(),
134                    param_formats: vec![crate::protocol::FormatCode::Text],
135                    params: param_values,
136                    result_formats: vec![crate::protocol::FormatCode::Binary],
137                },
138            )
139            .await?;
140
141        // Describe the unnamed portal so the server sends RowDescription
142        // before any DataRow.  This gives us proper column metadata.
143        self.codec
144            .encode_and_write(
145                &mut self.transport,
146                &FrontendMessage::Describe {
147                    variant: b'P',
148                    name: String::new(),
149                },
150            )
151            .await?;
152
153        // Execute
154        self.codec
155            .encode_and_write(
156                &mut self.transport,
157                &FrontendMessage::Execute {
158                    portal: String::new(),
159                    max_rows: 0,
160                },
161            )
162            .await?;
163
164        // Sync
165        self.codec
166            .encode_and_write(&mut self.transport, &FrontendMessage::Sync)
167            .await?;
168
169        // ── Flush the entire batch ──
170        self.transport.flush().await.map_err(PgError::Transport)?;
171
172        // ── Read responses ──
173        let result = self.read_extended_query_result().await;
174        self.state = ConnectionState::Idle;
175        result
176    }
177
178    /// Execute a parameterized statement that does not return rows.
179    #[must_use = "execute results should be checked for errors"]
180    pub async fn execute_params(
181        &mut self,
182        sql: &str,
183        params: &[&dyn ToSql],
184    ) -> Result<ExecuteResult> {
185        let result = self.query_params(sql, params).await?;
186        Ok(ExecuteResult::new(result.command_tag().clone()))
187    }
188
189    /// Execute a parameterized query and return a streaming result.
190    ///
191    /// Unlike [`query_params`](Self::query_params), this does **not** buffer
192    /// all rows. Rows are fetched one at a time as the consumer calls
193    /// [`RowStream::next`](crate::RowStream::next).
194    ///
195    /// Uses the extended query protocol (Parse + Bind + Describe + Execute
196    /// + Sync). The ParseComplete, BindComplete, and RowDescription/NoData
197    ///   preamble is consumed eagerly so that parameter-binding errors are
198    ///   surfaced during this call rather than lazily during stream iteration.
199    ///
200    /// # Example
201    ///
202    /// ```ignore
203    /// let mut stream = conn
204    ///     .query_params_stream("SELECT id, name FROM users WHERE age > $1", &[&18i32])
205    ///     .await?;
206    /// while let Some(row) = stream.next().await? {
207    ///     let id: i32 = row.get(0)?;
208    /// }
209    /// ```
210    #[must_use = "stream errors should be checked"]
211    pub async fn query_params_stream(
212        &mut self,
213        sql: &str,
214        params: &[&dyn ToSql],
215    ) -> Result<crate::query::stream::RowStream<'_>> {
216        use crate::query::stream::RowStream;
217
218        #[cfg(feature = "tracing")]
219        tracing::debug!(
220            target: crate::tracing_ext::TARGET_QUERY,
221            sql_len = sql.len(),
222            sql_truncated = %crate::tracing_ext::truncate_str(sql, 200),
223            param_count = params.len(),
224            protocol = "extended",
225            "Executing parameterized query"
226        );
227        #[cfg(feature = "tracing")]
228        tracing::trace!(target: crate::tracing_ext::TARGET_QUERY, sql = %sql, "Full SQL text");
229
230        if self.needs_recovery {
231            self.recover().await?;
232        }
233
234        self.transition(ConnectionState::Streaming)?;
235
236        let param_values = encode_params_text(params)?;
237
238        // ── Batch: Parse + Bind + Describe + Execute + Sync (no flush between) ──
239
240        // Parse (unnamed statement)
241        self.codec
242            .encode_and_write(
243                &mut self.transport,
244                &FrontendMessage::Parse {
245                    name: String::new(),
246                    sql: sql.to_string(),
247                    param_types: vec![],
248                },
249            )
250            .await?;
251
252        // Bind (unnamed portal, unnamed statement)
253        self.codec
254            .encode_and_write(
255                &mut self.transport,
256                &FrontendMessage::Bind {
257                    portal: String::new(),
258                    statement: String::new(),
259                    param_formats: vec![crate::protocol::FormatCode::Text],
260                    params: param_values,
261                    result_formats: vec![crate::protocol::FormatCode::Binary],
262                },
263            )
264            .await?;
265
266        // Describe the unnamed portal so the server sends RowDescription
267        // before any DataRow.
268        self.codec
269            .encode_and_write(
270                &mut self.transport,
271                &FrontendMessage::Describe {
272                    variant: b'P',
273                    name: String::new(),
274                },
275            )
276            .await?;
277
278        // Execute
279        self.codec
280            .encode_and_write(
281                &mut self.transport,
282                &FrontendMessage::Execute {
283                    portal: String::new(),
284                    max_rows: 0,
285                },
286            )
287            .await?;
288
289        // Sync
290        self.codec
291            .encode_and_write(&mut self.transport, &FrontendMessage::Sync)
292            .await?;
293
294        // ── Flush the entire batch ──
295        self.transport.flush().await.map_err(PgError::Transport)?;
296
297        // ── Read the preamble: ParseComplete, BindComplete, RowDescription/NoData ──
298        // Errors during Parse/Bind are surfaced eagerly here rather than lazily
299        // in the stream.
300        let mut columns: Option<Arc<Vec<FieldDescription>>> = None;
301        loop {
302            let msg = self.codec.read_message(&mut self.transport).await?;
303            if self.handle_async_message(&msg) {
304                continue;
305            }
306            match msg {
307                BackendMessage::ParseComplete => {}
308                BackendMessage::BindComplete => {}
309                BackendMessage::RowDescription(body) => {
310                    columns = Some(Arc::new(read_row_description(body)?));
311                    break;
312                }
313                BackendMessage::NoData => {
314                    break;
315                }
316                BackendMessage::CommandComplete(_body) => {
317                    // Edge case: the server sent CommandComplete immediately
318                    // (e.g., an INSERT/UPDATE that returned no rows and no
319                    // RowDescription). This means there are no data rows.
320                    // The stream will handle this in ReceivingRows state.
321                    break;
322                }
323                BackendMessage::ReadyForQuery(body) => {
324                    // The query returned nothing at all (no rows, no
325                    // CommandComplete yet). Let the stream handle the
326                    // ReadyForQuery in the Finishing state.
327                    self.transaction_status = TransactionStatus::from_u8(body.status())
328                        .unwrap_or(TransactionStatus::Idle);
329                    break;
330                }
331                BackendMessage::ErrorResponse(body) => {
332                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
333                    self.read_until_ready().await?;
334                    self.state = ConnectionState::Idle;
335                    return Err(PgError::Server(Box::new(server_err)));
336                }
337                BackendMessage::EmptyQueryResponse => {
338                    break;
339                }
340                _ => {}
341            }
342        }
343
344        Ok(RowStream::new_extended_receiving(self, columns))
345    }
346
347    /// Execute a previously prepared statement with parameters.
348    ///
349    /// Parameters are encoded in **binary** format using the types stored in
350    /// the [`PreparedStatement`]. Results are requested in **binary** format.
351    #[must_use = "query results should be checked for errors"]
352    pub async fn query_prepared(
353        &mut self,
354        stmt: &PreparedStatement,
355        params: &[&dyn ToSql],
356    ) -> Result<QueryResult> {
357        self.transition(ConnectionState::ActiveExtendedQuery)?;
358
359        let param_values = encode_params_binary(params, &stmt.param_types)?;
360
361        // ── Batch: Bind + Describe + Execute + Sync ──
362
363        // Bind (unnamed portal, named statement)
364        self.codec
365            .encode_and_write(
366                &mut self.transport,
367                &FrontendMessage::Bind {
368                    portal: String::new(),
369                    statement: stmt.name.clone(),
370                    param_formats: vec![crate::protocol::FormatCode::Binary],
371                    params: param_values,
372                    result_formats: vec![crate::protocol::FormatCode::Binary],
373                },
374            )
375            .await?;
376
377        // Describe the unnamed portal so the server sends RowDescription
378        // before any DataRow.
379        self.codec
380            .encode_and_write(
381                &mut self.transport,
382                &FrontendMessage::Describe {
383                    variant: b'P',
384                    name: String::new(),
385                },
386            )
387            .await?;
388
389        // Execute
390        self.codec
391            .encode_and_write(
392                &mut self.transport,
393                &FrontendMessage::Execute {
394                    portal: String::new(),
395                    max_rows: 0,
396                },
397            )
398            .await?;
399
400        // Sync
401        self.codec
402            .encode_and_write(&mut self.transport, &FrontendMessage::Sync)
403            .await?;
404
405        // ── Flush the entire batch ──
406        self.transport.flush().await.map_err(PgError::Transport)?;
407
408        // ── Read responses ──
409        let result = self.read_extended_query_result().await;
410        self.state = ConnectionState::Idle;
411        result
412    }
413
414    /// Execute a previously prepared statement that does not return rows.
415    #[must_use = "execute results should be checked for errors"]
416    pub async fn execute_prepared(
417        &mut self,
418        stmt: &PreparedStatement,
419        params: &[&dyn ToSql],
420    ) -> Result<ExecuteResult> {
421        let result = self.query_prepared(stmt, params).await?;
422        Ok(ExecuteResult::new(result.command_tag().clone()))
423    }
424}
425
426// ---------------------------------------------------------------------------
427// Internal helpers
428// ---------------------------------------------------------------------------
429
430impl Connection {
431    /// Read an extended-query result set from the wire.
432    ///
433    /// Similar to `read_query_result` but also handles `ParseComplete`,
434    /// `BindComplete`, and `NoData` which are specific to the extended
435    /// query protocol.
436    async fn read_extended_query_result(&mut self) -> Result<QueryResult> {
437        let mut columns: Option<Arc<Vec<FieldDescription>>> = None;
438        let mut rows: Vec<Row> = Vec::new();
439        let mut tag = None;
440
441        loop {
442            let msg = self.codec.read_message(&mut self.transport).await?;
443            if self.handle_async_message(&msg) {
444                continue;
445            }
446            match msg {
447                // Extended-query acknowledgements — consume and continue.
448                BackendMessage::ParseComplete => {}
449                BackendMessage::BindComplete => {}
450                BackendMessage::NoData => {}
451
452                BackendMessage::RowDescription(body) => {
453                    columns = Some(Arc::new(read_row_description(body)?));
454                }
455                BackendMessage::DataRow(body) => {
456                    let values = read_data_row(body)?;
457                    if columns.is_none() {
458                        // The server did not send RowDescription (can happen
459                        // for unnamed portals). Synthesise column metadata
460                        // using the requested binary format.
461                        let synthetic: Vec<FieldDescription> = (0..values.len())
462                            .map(|i| {
463                                FieldDescription::new(
464                                    format!("col{}", i),
465                                    0,
466                                    0,
467                                    0, // UNKNOWN OID
468                                    -1,
469                                    -1,
470                                    1, // binary format
471                                )
472                            })
473                            .collect();
474                        columns = Some(Arc::new(synthetic));
475                    }
476                    rows.push(Row::new(columns.clone().unwrap(), values));
477                }
478                BackendMessage::CommandComplete(body) => {
479                    tag = Some(CommandTag::new(body.tag().unwrap_or("").into()));
480                }
481                BackendMessage::EmptyQueryResponse => {
482                    tag = Some(CommandTag::new("".into()));
483                }
484                BackendMessage::ErrorResponse(body) => {
485                    let server_err = PgServerError::from_error_body(&body).map_err(PgError::Io)?;
486                    // The server will send ReadyForQuery after ErrorResponse
487                    // in extended query mode (because we sent Sync).
488                    self.read_until_ready().await?;
489                    return Err(PgError::Server(Box::new(server_err)));
490                }
491                BackendMessage::ReadyForQuery(body) => {
492                    self.transaction_status = TransactionStatus::from_u8(body.status())
493                        .unwrap_or(TransactionStatus::Idle);
494                    break;
495                }
496                _ => {}
497            }
498        }
499
500        Ok(QueryResult::new(
501            rows,
502            tag.unwrap_or_default(),
503            columns.unwrap_or_default(),
504        ))
505    }
506}
507
508// ---------------------------------------------------------------------------
509// Tests
510// ---------------------------------------------------------------------------
511
512#[cfg(test)]
513mod tests {
514    use super::*;
515    use crate::auth::{Codec, ServerParams};
516    use crate::config::Config;
517    use crate::connection::ConnectionState;
518    use crate::transport::{BufferedTransport, ClientTransport, MockTransport, PgTransport};
519    use std::collections::VecDeque;
520
521    fn make_connection(read_data: Vec<u8>) -> Connection {
522        let transport = PgTransport::Plain(BufferedTransport::new(ClientTransport::Mock(
523            MockTransport::new(read_data),
524        )));
525        Connection {
526            transport,
527            codec: Codec::new(),
528            server_params: ServerParams::default(),
529            state: ConnectionState::Idle,
530            config: Config::new(),
531            transaction_status: TransactionStatus::Idle,
532            notification_queue: VecDeque::new(),
533            notice_handler: None,
534            statement_counter: 0,
535            needs_recovery: false,
536            health: crate::reconnect::session::ConnectionHealth::new(),
537            session_state: crate::reconnect::session::SessionState::new(),
538        }
539    }
540
541    fn build_parse_complete() -> Vec<u8> {
542        vec![b'1', 0, 0, 0, 4]
543    }
544
545    fn build_bind_complete() -> Vec<u8> {
546        vec![b'2', 0, 0, 0, 4]
547    }
548
549    fn build_row_description_msg(fields: &[(&str, u32)]) -> Vec<u8> {
550        let mut buf = vec![b'T'];
551        let mut body = Vec::new();
552        body.extend_from_slice(&(fields.len() as i16).to_be_bytes());
553        for (name, type_oid) in fields {
554            body.extend_from_slice(name.as_bytes());
555            body.push(0);
556            body.extend_from_slice(&0u32.to_be_bytes()); // table_oid
557            body.extend_from_slice(&0i16.to_be_bytes()); // column_id
558            body.extend_from_slice(&type_oid.to_be_bytes()); // type_oid
559            body.extend_from_slice(&(-1i16).to_be_bytes()); // type_size
560            body.extend_from_slice(&(-1i32).to_be_bytes()); // type_modifier
561            body.extend_from_slice(&0i16.to_be_bytes()); // format
562        }
563        let len = (body.len() + 4) as i32;
564        buf.extend_from_slice(&len.to_be_bytes());
565        buf.extend_from_slice(&body);
566        buf
567    }
568
569    fn build_data_row_msg(values: &[Option<&str>]) -> Vec<u8> {
570        let mut buf = vec![b'D'];
571        let mut body = Vec::new();
572        body.extend_from_slice(&(values.len() as i16).to_be_bytes());
573        for val in values {
574            match val {
575                Some(v) => {
576                    let bytes = v.as_bytes();
577                    body.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
578                    body.extend_from_slice(bytes);
579                }
580                None => {
581                    body.extend_from_slice(&(-1i32).to_be_bytes());
582                }
583            }
584        }
585        let len = (body.len() + 4) as i32;
586        buf.extend_from_slice(&len.to_be_bytes());
587        buf.extend_from_slice(&body);
588        buf
589    }
590
591    fn build_command_complete_msg(tag: &str) -> Vec<u8> {
592        let mut buf = vec![b'C'];
593        let mut body = Vec::new();
594        body.extend_from_slice(tag.as_bytes());
595        body.push(0);
596        let len = (body.len() + 4) as i32;
597        buf.extend_from_slice(&len.to_be_bytes());
598        buf.extend_from_slice(&body);
599        buf
600    }
601
602    fn build_ready_for_query(status: u8) -> Vec<u8> {
603        vec![b'Z', 0, 0, 0, 5, status]
604    }
605
606    #[tokio::test]
607    async fn test_query_params_select() {
608        let mut data = Vec::new();
609        data.extend_from_slice(&build_parse_complete());
610        data.extend_from_slice(&build_bind_complete());
611        data.extend_from_slice(&build_row_description_msg(&[(
612            "val",
613            crate::types::INT4_OID,
614        )]));
615        data.extend_from_slice(&build_data_row_msg(&[Some("42")]));
616        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
617        data.extend_from_slice(&build_ready_for_query(b'I'));
618
619        let mut conn = make_connection(data);
620        let result = conn.query_params("SELECT $1", &[&42i32]).await.unwrap();
621        assert_eq!(result.len(), 1);
622        let v: i32 = result.rows()[0].get(0).unwrap();
623        assert_eq!(v, 42);
624    }
625
626    #[tokio::test]
627    async fn test_query_params_insert() {
628        let mut data = Vec::new();
629        data.extend_from_slice(&build_parse_complete());
630        data.extend_from_slice(&build_bind_complete());
631        data.extend_from_slice(&build_command_complete_msg("INSERT 0 1"));
632        data.extend_from_slice(&build_ready_for_query(b'I'));
633
634        let mut conn = make_connection(data);
635        let result = conn
636            .execute_params("INSERT INTO t (id) VALUES ($1)", &[&1i32])
637            .await
638            .unwrap();
639        assert_eq!(result.rows_affected(), Some(1));
640    }
641
642    #[tokio::test]
643    async fn test_query_params_insert_two_params() {
644        let mut data = Vec::new();
645        data.extend_from_slice(&build_parse_complete());
646        data.extend_from_slice(&build_bind_complete());
647        data.extend_from_slice(&build_command_complete_msg("INSERT 0 1"));
648        data.extend_from_slice(&build_ready_for_query(b'I'));
649
650        let mut conn = make_connection(data);
651        let result = conn
652            .execute_params(
653                "INSERT INTO t (id, name) VALUES ($1, $2)",
654                &[&200i32, &"param_insert"],
655            )
656            .await
657            .unwrap();
658        assert_eq!(result.rows_affected(), Some(1));
659    }
660
661    #[tokio::test]
662    async fn test_query_prepared_select() {
663        let mut data = Vec::new();
664        data.extend_from_slice(&build_bind_complete());
665        data.extend_from_slice(&build_row_description_msg(&[(
666            "val",
667            crate::types::INT4_OID,
668        )]));
669        data.extend_from_slice(&build_data_row_msg(&[Some("99")]));
670        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
671        data.extend_from_slice(&build_ready_for_query(b'I'));
672
673        let mut conn = make_connection(data);
674        let stmt = PreparedStatement {
675            name: "__pg_stmt_1".into(),
676            sql: "SELECT $1".into(),
677            param_types: vec![crate::types::Type::INT4],
678            columns: Arc::new(vec![]),
679        };
680        let result = conn.query_prepared(&stmt, &[&99i32]).await.unwrap();
681        assert_eq!(result.len(), 1);
682        let v: i32 = result.rows()[0].get(0).unwrap();
683        assert_eq!(v, 99);
684    }
685
686    #[tokio::test]
687    async fn test_query_params_error() {
688        let mut data = Vec::new();
689        let mut err = vec![b'E', 0, 0, 0, 22];
690        err.extend_from_slice(b"S");
691        err.extend_from_slice(b"ERROR\0");
692        err.extend_from_slice(b"M");
693        err.extend_from_slice(b"syntax error\0");
694        err.push(0);
695        data.extend_from_slice(&err);
696        data.extend_from_slice(&build_ready_for_query(b'I'));
697
698        let mut conn = make_connection(data);
699        let result = conn.query_params("BAD $1", &[&1i32]).await;
700        assert!(result.is_err());
701        assert!(conn.is_idle());
702    }
703
704    fn build_error_response_msg(code: &str, message: &str) -> Vec<u8> {
705        let mut buf = vec![b'E'];
706        let mut body = Vec::new();
707        // Severity
708        body.push(b'S');
709        body.extend_from_slice(b"ERROR");
710        body.push(0);
711        // Code (SQLSTATE)
712        body.push(b'C');
713        body.extend_from_slice(code.as_bytes());
714        body.push(0);
715        // Message
716        body.push(b'M');
717        body.extend_from_slice(message.as_bytes());
718        body.push(0);
719        // Terminator
720        body.push(0);
721        let len = (body.len() + 4) as i32;
722        buf.extend_from_slice(&len.to_be_bytes());
723        buf.extend_from_slice(&body);
724        buf
725    }
726
727    #[tokio::test]
728    async fn test_query_params_stream() {
729        let mut data = Vec::new();
730        data.extend_from_slice(&build_parse_complete());
731        data.extend_from_slice(&build_bind_complete());
732        data.extend_from_slice(&build_row_description_msg(&[
733            ("id", crate::types::INT4_OID),
734            ("name", crate::types::TEXT_OID),
735        ]));
736        data.extend_from_slice(&build_data_row_msg(&[Some("1"), Some("Alice")]));
737        data.extend_from_slice(&build_data_row_msg(&[Some("2"), Some("Bob")]));
738        data.extend_from_slice(&build_command_complete_msg("SELECT 2"));
739        data.extend_from_slice(&build_ready_for_query(b'I'));
740
741        let mut conn = make_connection(data);
742        let mut stream = conn
743            .query_params_stream("SELECT * FROM users WHERE id > $1", &[&0i32])
744            .await
745            .unwrap();
746
747        let row1 = stream.next().await.unwrap().unwrap();
748        assert_eq!(row1.get::<i32>(0).unwrap(), 1);
749        assert_eq!(row1.get::<String>(1).unwrap(), "Alice");
750
751        let row2 = stream.next().await.unwrap().unwrap();
752        assert_eq!(row2.get::<i32>(0).unwrap(), 2);
753        assert_eq!(row2.get::<String>(1).unwrap(), "Bob");
754
755        assert!(stream.next().await.unwrap().is_none());
756        assert!(stream.is_done());
757    }
758
759    #[tokio::test]
760    async fn test_query_params_stream_error() {
761        let mut data = Vec::new();
762        data.extend_from_slice(&build_parse_complete());
763        data.extend_from_slice(&build_bind_complete());
764        // ErrorResponse instead of RowDescription
765        data.extend_from_slice(&build_error_response_msg("23505", "duplicate key value"));
766        data.extend_from_slice(&build_ready_for_query(b'I'));
767
768        let mut conn = make_connection(data);
769        let result = conn
770            .query_params_stream("INSERT INTO t VALUES ($1)", &[&1i32])
771            .await;
772        assert!(result.is_err());
773        match result {
774            Err(PgError::Server(e)) => {
775                assert_eq!(e.code(), "23505");
776            }
777            _ => panic!("expected server error"),
778        }
779    }
780
781    #[tokio::test]
782    async fn test_query_params_stream_no_data() {
783        // INSERT/UPDATE/DELETE that returns NoData (no RowDescription)
784        let mut data = Vec::new();
785        data.extend_from_slice(&build_parse_complete());
786        data.extend_from_slice(&build_bind_complete());
787        // NoData: the server signifies no rows will be returned
788        let nodata = vec![b'n', 0, 0, 0, 4];
789        data.extend_from_slice(&nodata);
790        data.extend_from_slice(&build_command_complete_msg("INSERT 0 1"));
791        data.extend_from_slice(&build_ready_for_query(b'I'));
792
793        let mut conn = make_connection(data);
794        let mut stream = conn
795            .query_params_stream("INSERT INTO t VALUES ($1)", &[&1i32])
796            .await
797            .unwrap();
798
799        assert!(stream.next().await.unwrap().is_none());
800        assert!(stream.is_done());
801    }
802}
wasi_pg_client/query/params.rs

wasi_pg_client/query/
params.rs
