qail_pg/driver/pool/
fetch.rs

1//! Fetch methods for PooledConnection: uncached, fast, cached, typed, and pipelined-RLS variants.
2
3use super::connection::PooledConnection;
4use super::lifecycle::MAX_HOT_STATEMENTS;
5use crate::driver::{
6    PgConnection, PgError, PgResult, ResultFormat,
7    extended_flow::{ExtendedFlowConfig, ExtendedFlowTracker},
8    is_ignorable_session_message, unexpected_backend_message,
9};
10use std::sync::Arc;
11
12#[inline]
13fn rollback_cache_miss_statement_registration(
14    conn: &mut PgConnection,
15    is_cache_miss: bool,
16    sql_hash: u64,
17    stmt_name: &str,
18) {
19    if is_cache_miss {
20        conn.stmt_cache.remove(&sql_hash);
21        conn.prepared_statements.remove(stmt_name);
22        conn.column_info_cache.remove(&sql_hash);
23    }
24}
25
26#[inline]
27fn register_hot_statement_after_parse_success(
28    pool: &super::lifecycle::PgPoolInner,
29    sql_hash: u64,
30    stmt_name: &str,
31    sql: &str,
32) {
33    if let Ok(mut hot) = pool.hot_statements.write()
34        && (hot.contains_key(&sql_hash) || hot.len() < MAX_HOT_STATEMENTS)
35    {
36        hot.insert(sql_hash, (stmt_name.to_string(), sql.to_string()));
37    }
38}
39
40#[inline]
41fn return_with_desync<T>(conn: &mut PgConnection, err: PgError) -> PgResult<T> {
42    if matches!(
43        err,
44        PgError::Protocol(_) | PgError::Connection(_) | PgError::Timeout(_)
45    ) {
46        conn.mark_io_desynced();
47    }
48    Err(err)
49}
50
51#[inline]
52fn encoded_sql_str(sql_buf: &[u8]) -> PgResult<&str> {
53    std::str::from_utf8(sql_buf)
54        .map_err(|e| PgError::Encode(format!("encoded SQL is not UTF-8: {}", e)))
55}
56
57async fn drain_extended_responses_after_rls_setup_error(conn: &mut PgConnection) -> PgResult<()> {
58    loop {
59        let msg = conn.recv().await?;
60        match msg {
61            crate::protocol::BackendMessage::ReadyForQuery(_) => return Ok(()),
62            crate::protocol::BackendMessage::ErrorResponse(_) => {}
63            msg if is_ignorable_session_message(&msg) => {}
64            // Best-effort drain: consume everything until Sync's ReadyForQuery.
65            _ => {}
66        }
67    }
68}
69
70fn copy_export_table_sql(table: &str, columns: &[String]) -> PgResult<String> {
71    let cols: Vec<String> = columns
72        .iter()
73        .map(|c| crate::driver::copy::quote_copy_column_ident(c))
74        .collect::<PgResult<_>>()?;
75
76    Ok(format!(
77        "COPY {} ({}) TO STDOUT",
78        crate::driver::copy::quote_copy_table_ref(table)?,
79        cols.join(", ")
80    ))
81}
82
83impl PooledConnection {
84    /// Execute a QAIL command and fetch all rows (UNCACHED).
85    /// Returns rows with column metadata for JSON serialization.
86    pub async fn fetch_all_uncached(
87        &mut self,
88        cmd: &qail_core::ast::Qail,
89    ) -> PgResult<Vec<crate::driver::PgRow>> {
90        self.fetch_all_uncached_with_format(cmd, ResultFormat::Text)
91            .await
92    }
93
94    /// Execute raw SQL with bind parameters and return raw row data.
95    ///
96    /// Uses the Extended Query Protocol so parameters are never interpolated
97    /// into the SQL string. Intended for EXPLAIN or other SQL that can't be
98    /// represented as a `Qail` AST but still needs parameterized execution.
99    ///
100    /// Returns raw column bytes; callers must decode as needed.
101    pub async fn query_raw_with_params(
102        &mut self,
103        sql: &str,
104        params: &[Option<Vec<u8>>],
105    ) -> PgResult<Vec<Vec<Option<Vec<u8>>>>> {
106        let conn = self.conn_mut()?;
107        conn.query(sql, params).await
108    }
109
110    /// Execute raw SQL with bind parameters and return rows with column metadata.
111    ///
112    /// Uses the Extended Query Protocol so parameters are never interpolated
113    /// into the SQL string. Intended for compatibility paths that need
114    /// `PgRow` decoding and stable column names, not just raw bytes.
115    pub async fn query_rows_with_params(
116        &mut self,
117        sql: &str,
118        params: &[Option<Vec<u8>>],
119    ) -> PgResult<Vec<crate::driver::PgRow>> {
120        self.query_rows_with_params_with_format(sql, params, ResultFormat::Text)
121            .await
122    }
123
124    /// Execute raw SQL with bind parameters and explicit result format,
125    /// returning rows with column metadata.
126    pub async fn query_rows_with_params_with_format(
127        &mut self,
128        sql: &str,
129        params: &[Option<Vec<u8>>],
130        result_format: ResultFormat,
131    ) -> PgResult<Vec<crate::driver::PgRow>> {
132        let conn = self.conn_mut()?;
133        conn.query_rows_with_result_format(sql, params, result_format.as_wire_code())
134            .await
135    }
136
137    /// Execute raw SQL with explicit PostgreSQL parameter type OIDs and return
138    /// rows with column metadata.
139    pub async fn query_rows_with_param_types_with_format(
140        &mut self,
141        sql: &str,
142        param_types: &[u32],
143        params: &[Option<Vec<u8>>],
144        result_format: ResultFormat,
145    ) -> PgResult<Vec<crate::driver::PgRow>> {
146        let conn = self.conn_mut()?;
147        conn.query_rows_with_param_types_and_result_format(
148            sql,
149            param_types,
150            params,
151            result_format.as_wire_code(),
152        )
153        .await
154    }
155
156    /// Validate raw SQL with explicit PostgreSQL parameter type OIDs without
157    /// executing it.
158    pub async fn probe_query_with_param_types(
159        &mut self,
160        sql: &str,
161        param_types: &[u32],
162        params: &[Option<Vec<u8>>],
163    ) -> PgResult<()> {
164        let conn = self.conn_mut()?;
165        conn.probe_query_with_param_types(sql, param_types, params)
166            .await
167    }
168
169    /// Export data using AST-native COPY TO STDOUT and collect parsed rows.
170    pub async fn copy_export(&mut self, cmd: &qail_core::ast::Qail) -> PgResult<Vec<Vec<String>>> {
171        self.conn_mut()?.copy_export(cmd).await
172    }
173
174    /// Stream AST-native COPY TO STDOUT chunks with bounded memory usage.
175    pub async fn copy_export_stream_raw<F, Fut>(
176        &mut self,
177        cmd: &qail_core::ast::Qail,
178        on_chunk: F,
179    ) -> PgResult<()>
180    where
181        F: FnMut(Vec<u8>) -> Fut,
182        Fut: std::future::Future<Output = PgResult<()>>,
183    {
184        self.conn_mut()?.copy_export_stream_raw(cmd, on_chunk).await
185    }
186
187    /// Stream AST-native COPY TO STDOUT rows with bounded memory usage.
188    pub async fn copy_export_stream_rows<F>(
189        &mut self,
190        cmd: &qail_core::ast::Qail,
191        on_row: F,
192    ) -> PgResult<()>
193    where
194        F: FnMut(Vec<String>) -> PgResult<()>,
195    {
196        self.conn_mut()?.copy_export_stream_rows(cmd, on_row).await
197    }
198
199    /// Export a table using COPY TO STDOUT and collect raw bytes.
200    pub async fn copy_export_table(
201        &mut self,
202        table: &str,
203        columns: &[String],
204    ) -> PgResult<Vec<u8>> {
205        let sql = copy_export_table_sql(table, columns)?;
206        self.conn_mut()?.copy_out_raw(&sql).await
207    }
208
209    /// Stream a table export using COPY TO STDOUT with bounded memory usage.
210    pub async fn copy_export_table_stream<F, Fut>(
211        &mut self,
212        table: &str,
213        columns: &[String],
214        on_chunk: F,
215    ) -> PgResult<()>
216    where
217        F: FnMut(Vec<u8>) -> Fut,
218        Fut: std::future::Future<Output = PgResult<()>>,
219    {
220        let sql = copy_export_table_sql(table, columns)?;
221        self.conn_mut()?.copy_out_raw_stream(&sql, on_chunk).await
222    }
223
224    /// Execute a QAIL mutation command and return the affected row count.
225    ///
226    /// This is the pooled equivalent of [`crate::driver::PgDriver::execute`].
227    /// It uses the extended AST wire path and never interpolates values into SQL.
228    pub async fn execute(&mut self, cmd: &qail_core::ast::Qail) -> PgResult<u64> {
229        use crate::protocol::AstEncoder;
230
231        let conn = self.conn_mut()?;
232
233        AstEncoder::encode_cmd_reuse_into(
234            cmd,
235            &mut conn.sql_buf,
236            &mut conn.params_buf,
237            &mut conn.write_buf,
238        )
239        .map_err(|e| PgError::Encode(e.to_string()))?;
240
241        conn.flush_write_buf().await?;
242
243        let mut affected = 0u64;
244        let mut error: Option<PgError> = None;
245        let mut flow =
246            ExtendedFlowTracker::new(ExtendedFlowConfig::parse_bind_describe_portal_execute());
247
248        loop {
249            let msg = conn.recv().await?;
250            if let Err(err) = flow.validate(&msg, "pool execute mutation", error.is_some()) {
251                return return_with_desync(conn, err);
252            }
253            match msg {
254                crate::protocol::BackendMessage::ParseComplete
255                | crate::protocol::BackendMessage::BindComplete => {}
256                crate::protocol::BackendMessage::RowDescription(_) => {}
257                crate::protocol::BackendMessage::DataRow(_) => {}
258                crate::protocol::BackendMessage::NoData => {}
259                crate::protocol::BackendMessage::CommandComplete(tag) => {
260                    if error.is_none() {
261                        match crate::driver::parse_affected_rows(&tag) {
262                            Ok(parsed) => affected = parsed,
263                            Err(err) => return return_with_desync(conn, err),
264                        }
265                    }
266                }
267                crate::protocol::BackendMessage::ReadyForQuery(_) => {
268                    if let Some(err) = error {
269                        return Err(err);
270                    }
271                    return Ok(affected);
272                }
273                crate::protocol::BackendMessage::ErrorResponse(err) => {
274                    if error.is_none() {
275                        error = Some(PgError::QueryServer(err.into()));
276                    }
277                }
278                msg if is_ignorable_session_message(&msg) => {}
279                other => {
280                    return return_with_desync(
281                        conn,
282                        unexpected_backend_message("pool execute mutation", &other),
283                    );
284                }
285            }
286        }
287    }
288
289    /// Execute a QAIL command and fetch all rows (UNCACHED) with explicit result format.
290    pub async fn fetch_all_uncached_with_format(
291        &mut self,
292        cmd: &qail_core::ast::Qail,
293        result_format: ResultFormat,
294    ) -> PgResult<Vec<crate::driver::PgRow>> {
295        use crate::driver::ColumnInfo;
296        use crate::protocol::AstEncoder;
297
298        let conn = self.conn_mut()?;
299
300        AstEncoder::encode_cmd_reuse_into_with_result_format(
301            cmd,
302            &mut conn.sql_buf,
303            &mut conn.params_buf,
304            &mut conn.write_buf,
305            result_format.as_wire_code(),
306        )
307        .map_err(|e| PgError::Encode(e.to_string()))?;
308
309        conn.flush_write_buf().await?;
310
311        let mut rows: Vec<crate::driver::PgRow> = Vec::new();
312        let mut column_info: Option<Arc<ColumnInfo>> = None;
313        let mut error: Option<PgError> = None;
314        let mut flow =
315            ExtendedFlowTracker::new(ExtendedFlowConfig::parse_bind_describe_portal_execute());
316
317        loop {
318            let msg = conn.recv().await?;
319            if let Err(err) = flow.validate(&msg, "pool fetch_all execute", error.is_some()) {
320                return return_with_desync(conn, err);
321            }
322            match msg {
323                crate::protocol::BackendMessage::ParseComplete
324                | crate::protocol::BackendMessage::BindComplete => {}
325                crate::protocol::BackendMessage::RowDescription(fields) => {
326                    column_info = Some(Arc::new(ColumnInfo::from_fields(&fields)));
327                }
328                crate::protocol::BackendMessage::DataRow(data) => {
329                    if error.is_none() {
330                        rows.push(crate::driver::PgRow {
331                            columns: data,
332                            column_info: column_info.clone(),
333                        });
334                    }
335                }
336                crate::protocol::BackendMessage::NoData => {}
337                crate::protocol::BackendMessage::CommandComplete(_) => {}
338                crate::protocol::BackendMessage::ReadyForQuery(_) => {
339                    if let Some(err) = error {
340                        return Err(err);
341                    }
342                    return Ok(rows);
343                }
344                crate::protocol::BackendMessage::ErrorResponse(err) => {
345                    if error.is_none() {
346                        error = Some(PgError::QueryServer(err.into()));
347                    }
348                }
349                msg if is_ignorable_session_message(&msg) => {}
350                other => {
351                    return return_with_desync(
352                        conn,
353                        unexpected_backend_message("pool fetch_all execute", &other),
354                    );
355                }
356            }
357        }
358    }
359
360    /// Execute a QAIL command and fetch all rows (FAST VERSION).
361    /// Uses native AST-to-wire encoding and optimized recv_with_data_fast.
362    /// Skips column metadata for maximum speed.
363    pub async fn fetch_all_fast(
364        &mut self,
365        cmd: &qail_core::ast::Qail,
366    ) -> PgResult<Vec<crate::driver::PgRow>> {
367        self.fetch_all_fast_with_format(cmd, ResultFormat::Text)
368            .await
369    }
370
371    /// Execute a QAIL command and fetch all rows (FAST VERSION) with explicit result format.
372    pub async fn fetch_all_fast_with_format(
373        &mut self,
374        cmd: &qail_core::ast::Qail,
375        result_format: ResultFormat,
376    ) -> PgResult<Vec<crate::driver::PgRow>> {
377        use crate::protocol::AstEncoder;
378
379        let conn = self.conn_mut()?;
380
381        AstEncoder::encode_cmd_reuse_into_with_result_format(
382            cmd,
383            &mut conn.sql_buf,
384            &mut conn.params_buf,
385            &mut conn.write_buf,
386            result_format.as_wire_code(),
387        )
388        .map_err(|e| PgError::Encode(e.to_string()))?;
389
390        conn.flush_write_buf().await?;
391
392        let mut rows: Vec<crate::driver::PgRow> = Vec::with_capacity(32);
393        let mut error: Option<PgError> = None;
394        let mut flow = ExtendedFlowTracker::new(ExtendedFlowConfig::parse_bind_execute(true));
395
396        loop {
397            let res = conn.recv_with_data_fast().await;
398            match res {
399                Ok((msg_type, data)) => {
400                    if let Err(err) = flow.validate_msg_type(
401                        msg_type,
402                        "pool fetch_all_fast execute",
403                        error.is_some(),
404                    ) {
405                        return return_with_desync(conn, err);
406                    }
407                    match msg_type {
408                        b'D' => {
409                            if error.is_none()
410                                && let Some(columns) = data
411                            {
412                                rows.push(crate::driver::PgRow {
413                                    columns,
414                                    column_info: None,
415                                });
416                            }
417                        }
418                        b'Z' => {
419                            if let Some(err) = error {
420                                return Err(err);
421                            }
422                            return Ok(rows);
423                        }
424                        _ => {}
425                    }
426                }
427                Err(e) => {
428                    if matches!(&e, PgError::QueryServer(_)) {
429                        if error.is_none() {
430                            error = Some(e);
431                        }
432                        continue;
433                    }
434                    return Err(e);
435                }
436            }
437        }
438    }
439
440    /// Execute a QAIL command and fetch all rows (CACHED).
441    /// Uses prepared statement caching: Parse+Describe on first call,
442    /// then Bind+Execute only on subsequent calls with the same SQL shape.
443    /// This matches PostgREST's behavior for fair benchmarks.
444    pub async fn fetch_all_cached(
445        &mut self,
446        cmd: &qail_core::ast::Qail,
447    ) -> PgResult<Vec<crate::driver::PgRow>> {
448        self.fetch_all_cached_with_format(cmd, ResultFormat::Text)
449            .await
450    }
451
452    /// Execute a QAIL command and fetch all rows (CACHED) with explicit result format.
453    pub async fn fetch_all_cached_with_format(
454        &mut self,
455        cmd: &qail_core::ast::Qail,
456        result_format: ResultFormat,
457    ) -> PgResult<Vec<crate::driver::PgRow>> {
458        let mut retried = false;
459        loop {
460            match self
461                .fetch_all_cached_with_format_once(cmd, result_format)
462                .await
463            {
464                Ok(rows) => return Ok(rows),
465                Err(err)
466                    if !retried
467                        && (err.is_prepared_statement_retryable()
468                            || err.is_prepared_statement_already_exists()) =>
469                {
470                    retried = true;
471                    if err.is_prepared_statement_retryable()
472                        && let Some(conn) = self.conn.as_mut()
473                    {
474                        conn.clear_prepared_statement_state();
475                    }
476                }
477                Err(err) => return Err(err),
478            }
479        }
480    }
481
482    /// Execute a QAIL command and decode rows into typed structs (CACHED, text format).
483    pub async fn fetch_typed<T: crate::driver::row::QailRow>(
484        &mut self,
485        cmd: &qail_core::ast::Qail,
486    ) -> PgResult<Vec<T>> {
487        self.fetch_typed_with_format(cmd, ResultFormat::Text).await
488    }
489
490    /// Execute a QAIL command and decode rows into typed structs with explicit result format.
491    ///
492    /// Use [`ResultFormat::Binary`] for binary wire values; row decoders should use
493    /// metadata-aware helpers like `PgRow::try_get()` / `try_get_by_name()`.
494    pub async fn fetch_typed_with_format<T: crate::driver::row::QailRow>(
495        &mut self,
496        cmd: &qail_core::ast::Qail,
497        result_format: ResultFormat,
498    ) -> PgResult<Vec<T>> {
499        let rows = self
500            .fetch_all_cached_with_format(cmd, result_format)
501            .await?;
502        Ok(rows.iter().map(T::from_row).collect())
503    }
504
505    /// Execute a QAIL command and decode one typed row (CACHED, text format).
506    pub async fn fetch_one_typed<T: crate::driver::row::QailRow>(
507        &mut self,
508        cmd: &qail_core::ast::Qail,
509    ) -> PgResult<Option<T>> {
510        self.fetch_one_typed_with_format(cmd, ResultFormat::Text)
511            .await
512    }
513
514    /// Execute a QAIL command and decode one typed row with explicit result format.
515    pub async fn fetch_one_typed_with_format<T: crate::driver::row::QailRow>(
516        &mut self,
517        cmd: &qail_core::ast::Qail,
518        result_format: ResultFormat,
519    ) -> PgResult<Option<T>> {
520        let rows = self
521            .fetch_all_cached_with_format(cmd, result_format)
522            .await?;
523        Ok(rows.first().map(T::from_row))
524    }
525
526    async fn fetch_all_cached_with_format_once(
527        &mut self,
528        cmd: &qail_core::ast::Qail,
529        result_format: ResultFormat,
530    ) -> PgResult<Vec<crate::driver::PgRow>> {
531        use crate::driver::ColumnInfo;
532        use std::collections::hash_map::DefaultHasher;
533        use std::hash::{Hash, Hasher};
534
535        let pool = std::sync::Arc::clone(&self.pool);
536        let conn = self.conn.as_mut().ok_or_else(|| {
537            PgError::Connection("Connection already released back to pool".into())
538        })?;
539
540        conn.sql_buf.clear();
541        conn.params_buf.clear();
542
543        // Encode SQL + params to reusable buffers
544        match cmd.action {
545            qail_core::ast::Action::Get | qail_core::ast::Action::With => {
546                crate::protocol::ast_encoder::dml::encode_select(
547                    cmd,
548                    &mut conn.sql_buf,
549                    &mut conn.params_buf,
550                )?;
551            }
552            qail_core::ast::Action::Add => {
553                crate::protocol::ast_encoder::dml::encode_insert(
554                    cmd,
555                    &mut conn.sql_buf,
556                    &mut conn.params_buf,
557                )?;
558            }
559            qail_core::ast::Action::Set => {
560                crate::protocol::ast_encoder::dml::encode_update(
561                    cmd,
562                    &mut conn.sql_buf,
563                    &mut conn.params_buf,
564                )?;
565            }
566            qail_core::ast::Action::Del => {
567                crate::protocol::ast_encoder::dml::encode_delete(
568                    cmd,
569                    &mut conn.sql_buf,
570                    &mut conn.params_buf,
571                )?;
572            }
573            _ => {
574                // Fallback: unsupported actions go through uncached path
575                return self
576                    .fetch_all_uncached_with_format(cmd, result_format)
577                    .await;
578            }
579        }
580
581        let mut hasher = DefaultHasher::new();
582        conn.sql_buf.hash(&mut hasher);
583        let sql_hash = hasher.finish();
584
585        let is_cache_miss = !conn.stmt_cache.contains(&sql_hash);
586
587        conn.write_buf.clear();
588
589        let stmt_name = if let Some(name) = conn.stmt_cache.get(&sql_hash) {
590            name
591        } else {
592            let name = format!("qail_{:x}", sql_hash);
593
594            conn.evict_prepared_if_full();
595
596            let sql_str = encoded_sql_str(&conn.sql_buf)?;
597
598            use crate::protocol::PgEncoder;
599            let parse_msg = PgEncoder::try_encode_parse(&name, sql_str, &[])?;
600            let describe_msg = PgEncoder::try_encode_describe(false, &name)?;
601            conn.write_buf.extend_from_slice(&parse_msg);
602            conn.write_buf.extend_from_slice(&describe_msg);
603
604            conn.stmt_cache.put(sql_hash, name.clone());
605            conn.prepared_statements
606                .insert(name.clone(), sql_str.to_string());
607
608            name
609        };
610
611        use crate::protocol::PgEncoder;
612        if let Err(e) = PgEncoder::encode_bind_to_with_result_format(
613            &mut conn.write_buf,
614            &stmt_name,
615            &conn.params_buf,
616            result_format.as_wire_code(),
617        ) {
618            if is_cache_miss {
619                conn.stmt_cache.remove(&sql_hash);
620                conn.prepared_statements.remove(&stmt_name);
621                conn.column_info_cache.remove(&sql_hash);
622            }
623            return Err(PgError::Encode(e.to_string()));
624        }
625        PgEncoder::encode_execute_to(&mut conn.write_buf);
626        PgEncoder::encode_sync_to(&mut conn.write_buf);
627
628        if let Err(err) = conn.flush_write_buf().await {
629            if is_cache_miss {
630                conn.stmt_cache.remove(&sql_hash);
631                conn.prepared_statements.remove(&stmt_name);
632                conn.column_info_cache.remove(&sql_hash);
633            }
634            return Err(err);
635        }
636
637        let cached_column_info = conn.column_info_cache.get(&sql_hash).cloned();
638
639        let mut rows: Vec<crate::driver::PgRow> = Vec::with_capacity(32);
640        let mut column_info: Option<Arc<ColumnInfo>> = cached_column_info;
641        let mut error: Option<PgError> = None;
642        let mut flow = ExtendedFlowTracker::new(
643            ExtendedFlowConfig::parse_describe_statement_bind_execute(is_cache_miss),
644        );
645
646        loop {
647            let msg = match conn.recv().await {
648                Ok(msg) => msg,
649                Err(err) => {
650                    if is_cache_miss && !flow.saw_parse_complete() {
651                        conn.stmt_cache.remove(&sql_hash);
652                        conn.prepared_statements.remove(&stmt_name);
653                        conn.column_info_cache.remove(&sql_hash);
654                    }
655                    return Err(err);
656                }
657            };
658            if let Err(err) = flow.validate(&msg, "pool fetch_all_cached execute", error.is_some())
659            {
660                if is_cache_miss && !flow.saw_parse_complete() {
661                    conn.stmt_cache.remove(&sql_hash);
662                    conn.prepared_statements.remove(&stmt_name);
663                    conn.column_info_cache.remove(&sql_hash);
664                }
665                return return_with_desync(conn, err);
666            }
667            match msg {
668                crate::protocol::BackendMessage::ParseComplete => {}
669                crate::protocol::BackendMessage::BindComplete => {}
670                crate::protocol::BackendMessage::ParameterDescription(_) => {}
671                crate::protocol::BackendMessage::RowDescription(fields) => {
672                    let info = Arc::new(ColumnInfo::from_fields(&fields));
673                    if is_cache_miss {
674                        conn.column_info_cache.insert(sql_hash, Arc::clone(&info));
675                    }
676                    column_info = Some(info);
677                }
678                crate::protocol::BackendMessage::DataRow(data) => {
679                    if error.is_none() {
680                        rows.push(crate::driver::PgRow {
681                            columns: data,
682                            column_info: column_info.clone(),
683                        });
684                    }
685                }
686                crate::protocol::BackendMessage::CommandComplete(_) => {}
687                crate::protocol::BackendMessage::ReadyForQuery(_) => {
688                    if let Some(err) = error {
689                        if is_cache_miss
690                            && !flow.saw_parse_complete()
691                            && !err.is_prepared_statement_already_exists()
692                        {
693                            conn.stmt_cache.remove(&sql_hash);
694                            conn.prepared_statements.remove(&stmt_name);
695                            conn.column_info_cache.remove(&sql_hash);
696                        }
697                        return Err(err);
698                    }
699                    if is_cache_miss && !flow.saw_parse_complete() {
700                        conn.stmt_cache.remove(&sql_hash);
701                        conn.prepared_statements.remove(&stmt_name);
702                        conn.column_info_cache.remove(&sql_hash);
703                        return return_with_desync(
704                            conn,
705                            PgError::Protocol(
706                                "Cache miss query reached ReadyForQuery without ParseComplete"
707                                    .to_string(),
708                            ),
709                        );
710                    }
711                    if is_cache_miss && let Some(sql) = conn.prepared_statements.get(&stmt_name) {
712                        register_hot_statement_after_parse_success(
713                            &pool, sql_hash, &stmt_name, sql,
714                        );
715                    }
716                    return Ok(rows);
717                }
718                crate::protocol::BackendMessage::ErrorResponse(err) => {
719                    if error.is_none() {
720                        error = Some(PgError::QueryServer(err.into()));
721                    }
722                }
723                msg if is_ignorable_session_message(&msg) => {}
724                other => {
725                    if is_cache_miss && !flow.saw_parse_complete() {
726                        conn.stmt_cache.remove(&sql_hash);
727                        conn.prepared_statements.remove(&stmt_name);
728                        conn.column_info_cache.remove(&sql_hash);
729                    }
730                    return return_with_desync(
731                        conn,
732                        unexpected_backend_message("pool fetch_all_cached execute", &other),
733                    );
734                }
735            }
736        }
737    }
738
739    /// Execute a QAIL command with RLS context in a SINGLE roundtrip.
740    ///
741    /// Pipelines the RLS setup (BEGIN + set_config) and the query
742    /// (Parse/Bind/Execute/Sync) into one `write_all` syscall.
743    /// PG processes messages in order, so the BEGIN + set_config
744    /// completes before the query executes — security is preserved.
745    ///
746    /// Wire layout:
747    /// ```text
748    /// [SimpleQuery: "BEGIN; SET LOCAL...; SELECT set_config(...)"]
749    /// [Parse (if cache miss)]
750    /// [Describe (if cache miss)]
751    /// [Bind]
752    /// [Execute]
753    /// [Sync]
754    /// ```
755    ///
756    /// Response processing: consume 2× ReadyForQuery (SimpleQuery + Sync).
757    pub async fn fetch_all_with_rls(
758        &mut self,
759        cmd: &qail_core::ast::Qail,
760        rls_sql: &str,
761    ) -> PgResult<Vec<crate::driver::PgRow>> {
762        self.fetch_all_with_rls_with_format(cmd, rls_sql, ResultFormat::Text)
763            .await
764    }
765
766    /// Execute a QAIL command with RLS context in a SINGLE roundtrip with explicit result format.
767    pub async fn fetch_all_with_rls_with_format(
768        &mut self,
769        cmd: &qail_core::ast::Qail,
770        rls_sql: &str,
771        result_format: ResultFormat,
772    ) -> PgResult<Vec<crate::driver::PgRow>> {
773        let mut retried = false;
774        loop {
775            match self
776                .fetch_all_with_rls_with_format_once(cmd, rls_sql, result_format)
777                .await
778            {
779                Ok(rows) => return Ok(rows),
780                Err(err)
781                    if !retried
782                        && (err.is_prepared_statement_retryable()
783                            || err.is_prepared_statement_already_exists()) =>
784                {
785                    retried = true;
786                    if let Some(conn) = self.conn.as_mut() {
787                        if err.is_prepared_statement_retryable() {
788                            conn.clear_prepared_statement_state();
789                        }
790                        // Always rollback transaction state before a retried RLS pipeline
791                        // attempt, including 42P05 "prepared statement already exists".
792                        let _ = conn.execute_simple("ROLLBACK").await;
793                    }
794                    self.rls_dirty = false;
795                }
796                Err(err) => return Err(err),
797            }
798        }
799    }
800
801    async fn fetch_all_with_rls_with_format_once(
802        &mut self,
803        cmd: &qail_core::ast::Qail,
804        rls_sql: &str,
805        result_format: ResultFormat,
806    ) -> PgResult<Vec<crate::driver::PgRow>> {
807        use crate::driver::ColumnInfo;
808        use std::collections::hash_map::DefaultHasher;
809        use std::hash::{Hash, Hasher};
810
811        let pool = std::sync::Arc::clone(&self.pool);
812        let conn = self.conn.as_mut().ok_or_else(|| {
813            PgError::Connection("Connection already released back to pool".into())
814        })?;
815
816        if !crate::protocol::AstEncoder::encode_cacheable_cmd_sql_to(
817            cmd,
818            &mut conn.sql_buf,
819            &mut conn.params_buf,
820        )? {
821            // Fallback: RLS setup must happen synchronously for unsupported actions
822            conn.execute_simple(rls_sql).await?;
823            self.rls_dirty = true;
824            return self
825                .fetch_all_uncached_with_format(cmd, result_format)
826                .await;
827        }
828
829        let mut hasher = DefaultHasher::new();
830        conn.sql_buf.hash(&mut hasher);
831        let sql_hash = hasher.finish();
832
833        let is_cache_miss = !conn.stmt_cache.contains(&sql_hash);
834
835        conn.write_buf.clear();
836
837        // ── Prepend RLS Simple Query message ─────────────────────────
838        // NOTE: this is PostgreSQL SimpleQuery text, so the backend still
839        // parses this segment on every request. The optimization here is
840        // batching RLS + query protocol messages into one network flush.
841        let rls_msg = crate::protocol::PgEncoder::try_encode_query_string(rls_sql)?;
842        conn.write_buf.extend_from_slice(&rls_msg);
843
844        // ── Then append the query messages (same as fetch_all_cached) ──
845        let stmt_name = if let Some(name) = conn.stmt_cache.get(&sql_hash) {
846            name
847        } else {
848            let name = format!("qail_{:x}", sql_hash);
849
850            conn.evict_prepared_if_full();
851
852            let sql_str = encoded_sql_str(&conn.sql_buf)?;
853
854            use crate::protocol::PgEncoder;
855            let parse_msg = PgEncoder::try_encode_parse(&name, sql_str, &[])?;
856            let describe_msg = PgEncoder::try_encode_describe(false, &name)?;
857            conn.write_buf.extend_from_slice(&parse_msg);
858            conn.write_buf.extend_from_slice(&describe_msg);
859
860            conn.stmt_cache.put(sql_hash, name.clone());
861            conn.prepared_statements
862                .insert(name.clone(), sql_str.to_string());
863
864            name
865        };
866
867        use crate::protocol::PgEncoder;
868        if let Err(e) = PgEncoder::encode_bind_to_with_result_format(
869            &mut conn.write_buf,
870            &stmt_name,
871            &conn.params_buf,
872            result_format.as_wire_code(),
873        ) {
874            rollback_cache_miss_statement_registration(conn, is_cache_miss, sql_hash, &stmt_name);
875            return Err(PgError::Encode(e.to_string()));
876        }
877        PgEncoder::encode_execute_to(&mut conn.write_buf);
878        PgEncoder::encode_sync_to(&mut conn.write_buf);
879
880        // ── Single write_all for RLS + Query ────────────────────────
881        if let Err(err) = conn.flush_write_buf().await {
882            rollback_cache_miss_statement_registration(conn, is_cache_miss, sql_hash, &stmt_name);
883            return Err(err);
884        }
885
886        // Mark connection as RLS-dirty (needs COMMIT on release)
887        self.rls_dirty = true;
888
889        // ── Phase 1: Consume Simple Query responses (RLS setup) ─────
890        // Simple Query produces: CommandComplete × N, then ReadyForQuery.
891        // set_config results and BEGIN/SET LOCAL responses are all here.
892        let mut rls_error: Option<PgError> = None;
893        loop {
894            let msg = match conn.recv().await {
895                Ok(msg) => msg,
896                Err(err) => {
897                    rollback_cache_miss_statement_registration(
898                        conn,
899                        is_cache_miss,
900                        sql_hash,
901                        &stmt_name,
902                    );
903                    return Err(err);
904                }
905            };
906            match msg {
907                crate::protocol::BackendMessage::ReadyForQuery(_) => {
908                    // RLS setup done — break to Extended Query phase
909                    if let Some(err) = rls_error {
910                        rollback_cache_miss_statement_registration(
911                            conn,
912                            is_cache_miss,
913                            sql_hash,
914                            &stmt_name,
915                        );
916                        if let Err(drain_err) =
917                            drain_extended_responses_after_rls_setup_error(conn).await
918                        {
919                            tracing::warn!(
920                                error = %drain_err,
921                                "failed to drain pipelined extended responses after RLS setup error"
922                            );
923                        }
924                        return Err(err);
925                    }
926                    break;
927                }
928                crate::protocol::BackendMessage::ErrorResponse(err) => {
929                    if rls_error.is_none() {
930                        rls_error = Some(PgError::QueryServer(err.into()));
931                    }
932                }
933                // CommandComplete, DataRow (from set_config), RowDescription — ignore
934                crate::protocol::BackendMessage::CommandComplete(_)
935                | crate::protocol::BackendMessage::DataRow(_)
936                | crate::protocol::BackendMessage::RowDescription(_)
937                | crate::protocol::BackendMessage::ParseComplete
938                | crate::protocol::BackendMessage::BindComplete => {}
939                msg if is_ignorable_session_message(&msg) => {}
940                other => {
941                    rollback_cache_miss_statement_registration(
942                        conn,
943                        is_cache_miss,
944                        sql_hash,
945                        &stmt_name,
946                    );
947                    return return_with_desync(
948                        conn,
949                        unexpected_backend_message("pool rls setup", &other),
950                    );
951                }
952            }
953        }
954
955        // ── Phase 2: Consume Extended Query responses (actual data) ──
956        let cached_column_info = conn.column_info_cache.get(&sql_hash).cloned();
957
958        let mut rows: Vec<crate::driver::PgRow> = Vec::with_capacity(32);
959        let mut column_info: Option<std::sync::Arc<ColumnInfo>> = cached_column_info;
960        let mut error: Option<PgError> = None;
961        let mut flow = ExtendedFlowTracker::new(
962            ExtendedFlowConfig::parse_describe_statement_bind_execute(is_cache_miss),
963        );
964
965        loop {
966            let msg = match conn.recv().await {
967                Ok(msg) => msg,
968                Err(err) => {
969                    if is_cache_miss && !flow.saw_parse_complete() {
970                        rollback_cache_miss_statement_registration(
971                            conn,
972                            is_cache_miss,
973                            sql_hash,
974                            &stmt_name,
975                        );
976                    }
977                    return Err(err);
978                }
979            };
980            if let Err(err) =
981                flow.validate(&msg, "pool fetch_all_with_rls execute", error.is_some())
982            {
983                if is_cache_miss && !flow.saw_parse_complete() {
984                    rollback_cache_miss_statement_registration(
985                        conn,
986                        is_cache_miss,
987                        sql_hash,
988                        &stmt_name,
989                    );
990                }
991                return return_with_desync(conn, err);
992            }
993            match msg {
994                crate::protocol::BackendMessage::ParseComplete => {}
995                crate::protocol::BackendMessage::BindComplete => {}
996                crate::protocol::BackendMessage::ParameterDescription(_) => {}
997                crate::protocol::BackendMessage::RowDescription(fields) => {
998                    let info = std::sync::Arc::new(ColumnInfo::from_fields(&fields));
999                    if is_cache_miss {
1000                        conn.column_info_cache
1001                            .insert(sql_hash, std::sync::Arc::clone(&info));
1002                    }
1003                    column_info = Some(info);
1004                }
1005                crate::protocol::BackendMessage::DataRow(data) => {
1006                    if error.is_none() {
1007                        rows.push(crate::driver::PgRow {
1008                            columns: data,
1009                            column_info: column_info.clone(),
1010                        });
1011                    }
1012                }
1013                crate::protocol::BackendMessage::CommandComplete(_) => {}
1014                crate::protocol::BackendMessage::ReadyForQuery(_) => {
1015                    if let Some(err) = error {
1016                        if is_cache_miss
1017                            && !flow.saw_parse_complete()
1018                            && !err.is_prepared_statement_already_exists()
1019                        {
1020                            rollback_cache_miss_statement_registration(
1021                                conn,
1022                                is_cache_miss,
1023                                sql_hash,
1024                                &stmt_name,
1025                            );
1026                        }
1027                        return Err(err);
1028                    }
1029                    if is_cache_miss && !flow.saw_parse_complete() {
1030                        rollback_cache_miss_statement_registration(
1031                            conn,
1032                            is_cache_miss,
1033                            sql_hash,
1034                            &stmt_name,
1035                        );
1036                        return return_with_desync(
1037                            conn,
1038                            PgError::Protocol(
1039                                "Cache miss query reached ReadyForQuery without ParseComplete"
1040                                    .to_string(),
1041                            ),
1042                        );
1043                    }
1044                    if is_cache_miss && let Some(sql) = conn.prepared_statements.get(&stmt_name) {
1045                        register_hot_statement_after_parse_success(
1046                            &pool, sql_hash, &stmt_name, sql,
1047                        );
1048                    }
1049                    return Ok(rows);
1050                }
1051                crate::protocol::BackendMessage::ErrorResponse(err) => {
1052                    if error.is_none() {
1053                        error = Some(PgError::QueryServer(err.into()));
1054                    }
1055                }
1056                msg if is_ignorable_session_message(&msg) => {}
1057                other => {
1058                    if is_cache_miss && !flow.saw_parse_complete() {
1059                        rollback_cache_miss_statement_registration(
1060                            conn,
1061                            is_cache_miss,
1062                            sql_hash,
1063                            &stmt_name,
1064                        );
1065                    }
1066                    return return_with_desync(
1067                        conn,
1068                        unexpected_backend_message("pool fetch_all_with_rls execute", &other),
1069                    );
1070                }
1071            }
1072        }
1073    }
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078    use super::{copy_export_table_sql, encoded_sql_str, return_with_desync};
1079
1080    #[cfg(unix)]
1081    fn test_conn() -> crate::driver::PgConnection {
1082        use crate::driver::connection::StatementCache;
1083        use crate::driver::stream::PgStream;
1084        use bytes::BytesMut;
1085        use std::collections::{HashMap, VecDeque};
1086        use std::num::NonZeroUsize;
1087        use tokio::net::UnixStream;
1088
1089        let (unix_stream, _peer) = UnixStream::pair().expect("unix stream pair");
1090        crate::driver::PgConnection {
1091            stream: PgStream::Unix(unix_stream),
1092            buffer: BytesMut::with_capacity(1024),
1093            write_buf: BytesMut::with_capacity(1024),
1094            sql_buf: BytesMut::with_capacity(256),
1095            params_buf: Vec::new(),
1096            prepared_statements: HashMap::new(),
1097            stmt_cache: StatementCache::new(NonZeroUsize::new(2).expect("non-zero")),
1098            column_info_cache: HashMap::new(),
1099            process_id: 0,
1100            cancel_key_bytes: Vec::new(),
1101            requested_protocol_minor: crate::driver::PgConnection::default_protocol_minor(),
1102            negotiated_protocol_minor: crate::driver::PgConnection::default_protocol_minor(),
1103            notifications: VecDeque::new(),
1104            replication_stream_active: false,
1105            replication_mode_enabled: false,
1106            last_replication_wal_end: None,
1107            io_desynced: false,
1108            pending_statement_closes: Vec::new(),
1109            draining_statement_closes: false,
1110        }
1111    }
1112
1113    #[test]
1114    fn pool_copy_export_table_sql_preserves_schema_qualified_table() {
1115        let sql = copy_export_table_sql(
1116            "tenant_a.users",
1117            &["id".to_string(), "display\"name".to_string()],
1118        )
1119        .unwrap();
1120
1121        assert_eq!(
1122            sql,
1123            "COPY \"tenant_a\".\"users\" (\"id\", \"display\"\"name\") TO STDOUT"
1124        );
1125    }
1126
1127    #[test]
1128    fn pool_copy_export_table_sql_rejects_nul_bytes() {
1129        assert!(copy_export_table_sql("tenant\0.users", &["id".to_string()]).is_err());
1130        assert!(copy_export_table_sql("users", &["id\0".to_string()]).is_err());
1131    }
1132
1133    #[test]
1134    fn pool_encoded_sql_str_rejects_invalid_utf8() {
1135        let err = encoded_sql_str(&[0xff]).expect_err("invalid SQL UTF-8 must fail");
1136        assert!(err.to_string().contains("encoded SQL is not UTF-8"));
1137    }
1138
1139    #[cfg(unix)]
1140    #[tokio::test]
1141    async fn pool_return_with_desync_marks_protocol_error() {
1142        let mut conn = test_conn();
1143
1144        let err = return_with_desync::<()>(
1145            &mut conn,
1146            crate::driver::PgError::Protocol("bad response ordering".to_string()),
1147        )
1148        .expect_err("protocol error must be returned");
1149
1150        assert!(err.to_string().contains("bad response ordering"));
1151        assert!(conn.is_io_desynced());
1152    }
1153}
qail_pg/driver/pool/fetch.rs

qail_pg/driver/pool/
fetch.rs
