qail_pg/driver/
copy.rs

1//! COPY protocol methods for PostgreSQL bulk operations.
2//!
3
4use super::{
5    PgConnection, PgError, PgResult, is_ignorable_session_message, parse_affected_rows,
6    unexpected_backend_message,
7};
8use crate::protocol::{AstEncoder, BackendMessage, PgEncoder};
9use bytes::BytesMut;
10use qail_core::ast::{Action, Qail};
11use std::future::Future;
12
13/// Quote a single SQL identifier atom for COPY statements.
14pub(crate) fn quote_copy_column_ident(ident: &str) -> PgResult<String> {
15    if ident.is_empty() {
16        return Err(PgError::Query(
17            "COPY column identifier is empty".to_string(),
18        ));
19    }
20    if ident.contains('\0') {
21        return Err(PgError::Query(
22            "COPY column identifier contains NUL byte".to_string(),
23        ));
24    }
25    Ok(format!("\"{}\"", ident.replace('"', "\"\"")))
26}
27
28/// Quote a COPY table reference, preserving schema-qualified names.
29pub(crate) fn quote_copy_table_ref(table: &str) -> PgResult<String> {
30    if table.is_empty() {
31        return Err(PgError::Query("COPY table identifier is empty".to_string()));
32    }
33    if table.contains('\0') {
34        return Err(PgError::Query(
35            "COPY table identifier contains NUL byte".to_string(),
36        ));
37    }
38
39    table
40        .split('.')
41        .map(|part| {
42            let part = part.trim();
43            if part.is_empty() {
44                return Err(PgError::Query(
45                    "COPY table identifier contains an empty path segment".to_string(),
46                ));
47            }
48            quote_copy_column_ident(part)
49        })
50        .collect::<PgResult<Vec<_>>>()
51        .map(|parts| parts.join("."))
52}
53
54fn parse_copy_text_row(line: &[u8]) -> PgResult<Vec<String>> {
55    let line = if line.ends_with(b"\r") {
56        &line[..line.len().saturating_sub(1)]
57    } else {
58        line
59    };
60
61    let mut fields = Vec::new();
62    let mut start = 0;
63    for (idx, byte) in line.iter().enumerate() {
64        if *byte == b'\t' {
65            fields.push(decode_copy_text_field(&line[start..idx])?);
66            start = idx + 1;
67        }
68    }
69    fields.push(decode_copy_text_field(&line[start..])?);
70    Ok(fields)
71}
72
73fn decode_copy_text_field(field: &[u8]) -> PgResult<String> {
74    if field == b"\\N" {
75        return Err(PgError::Protocol(
76            "COPY text NULL cannot be represented by Vec<String>; use copy_export_stream_raw for nullable exports"
77                .to_string(),
78        ));
79    }
80
81    let mut out = Vec::with_capacity(field.len());
82    let mut idx = 0;
83    while idx < field.len() {
84        if field[idx] != b'\\' {
85            out.push(field[idx]);
86            idx += 1;
87            continue;
88        }
89
90        let Some(&escaped) = field.get(idx + 1) else {
91            return Err(PgError::Protocol(
92                "COPY text field ends with incomplete backslash escape".to_string(),
93            ));
94        };
95
96        match escaped {
97            b'b' => {
98                out.push(0x08);
99                idx += 2;
100            }
101            b'f' => {
102                out.push(0x0c);
103                idx += 2;
104            }
105            b'n' => {
106                out.push(b'\n');
107                idx += 2;
108            }
109            b'r' => {
110                out.push(b'\r');
111                idx += 2;
112            }
113            b't' => {
114                out.push(b'\t');
115                idx += 2;
116            }
117            b'v' => {
118                out.push(0x0b);
119                idx += 2;
120            }
121            b'\\' => {
122                out.push(b'\\');
123                idx += 2;
124            }
125            b'0'..=b'7' => {
126                let mut value = 0u16;
127                let mut next = idx + 1;
128                for _ in 0..3 {
129                    let Some(&digit) = field.get(next) else {
130                        break;
131                    };
132                    if !(b'0'..=b'7').contains(&digit) {
133                        break;
134                    }
135                    value = (value * 8) + u16::from(digit - b'0');
136                    next += 1;
137                }
138                if value > u16::from(u8::MAX) {
139                    return Err(PgError::Protocol(format!(
140                        "COPY text octal escape is out of byte range: \\{:o}",
141                        value
142                    )));
143                }
144                out.push(value as u8);
145                idx = next;
146            }
147            b'x' => {
148                let mut value = 0u8;
149                let mut next = idx + 2;
150                let mut digits = 0;
151                while digits < 2 {
152                    let Some(&digit) = field.get(next) else {
153                        break;
154                    };
155                    let Some(nibble) = hex_nibble(digit) else {
156                        break;
157                    };
158                    value = (value << 4) | nibble;
159                    next += 1;
160                    digits += 1;
161                }
162                if digits == 0 {
163                    return Err(PgError::Protocol(
164                        "COPY text hex escape requires at least one hex digit".to_string(),
165                    ));
166                } else {
167                    out.push(value);
168                    idx = next;
169                }
170            }
171            other => {
172                out.push(other);
173                idx += 2;
174            }
175        }
176    }
177
178    String::from_utf8(out)
179        .map_err(|e| PgError::Protocol(format!("COPY text field is not valid UTF-8: {}", e)))
180}
181
182fn hex_nibble(byte: u8) -> Option<u8> {
183    match byte {
184        b'0'..=b'9' => Some(byte - b'0'),
185        b'a'..=b'f' => Some(byte - b'a' + 10),
186        b'A'..=b'F' => Some(byte - b'A' + 10),
187        _ => None,
188    }
189}
190
191#[inline]
192fn return_with_desync<T>(conn: &mut PgConnection, err: PgError) -> PgResult<T> {
193    if matches!(
194        err,
195        PgError::Protocol(_) | PgError::Connection(_) | PgError::Timeout(_)
196    ) {
197        conn.mark_io_desynced();
198    }
199    Err(err)
200}
201
202fn encode_copy_export_sql(cmd: &Qail) -> PgResult<String> {
203    if cmd.action != Action::Export {
204        return Err(PgError::Query(
205            "copy_export requires Qail::Export action".to_string(),
206        ));
207    }
208
209    let (sql, params) =
210        AstEncoder::encode_cmd_sql(cmd).map_err(|e| PgError::Encode(e.to_string()))?;
211    if !params.is_empty() {
212        return Err(PgError::Encode(format!(
213            "copy_export cannot encode parameterized export with {} bind parameter(s); use an unfiltered export, a prefiltered database view, or a raw COPY statement with trusted SQL",
214            params.len()
215        )));
216    }
217
218    Ok(sql)
219}
220
221fn drain_copy_text_rows<F>(pending: &mut Vec<u8>, chunk: &[u8], on_row: &mut F) -> PgResult<()>
222where
223    F: FnMut(Vec<String>) -> PgResult<()>,
224{
225    pending.extend_from_slice(chunk);
226    while let Some(pos) = pending.iter().position(|&b| b == b'\n') {
227        let line = pending[..pos].to_vec();
228        pending.drain(..=pos);
229        let row = parse_copy_text_row(&line)?;
230        on_row(row)?;
231    }
232    Ok(())
233}
234
235fn flush_pending_copy_text_row(pending: &[u8]) -> PgResult<()> {
236    if pending.is_empty() {
237        return Ok(());
238    }
239    Err(PgError::Protocol(
240        "COPY text stream ended with a truncated row without final newline".to_string(),
241    ))
242}
243
244impl PgConnection {
245    /// **Fast** bulk insert using COPY protocol with zero-allocation encoding.
246    /// Encodes all rows into a single buffer and writes with one syscall.
247    /// ~2x faster than `copy_in_internal` due to batched I/O.
248    pub(crate) async fn copy_in_fast(
249        &mut self,
250        table: &str,
251        columns: &[String],
252        rows: &[Vec<qail_core::ast::Value>],
253    ) -> PgResult<u64> {
254        use crate::protocol::try_encode_copy_batch;
255
256        let cols: Vec<String> = columns
257            .iter()
258            .map(|c| quote_copy_column_ident(c))
259            .collect::<PgResult<_>>()?;
260        let sql = format!(
261            "COPY {} ({}) FROM STDIN",
262            quote_copy_table_ref(table)?,
263            cols.join(", ")
264        );
265
266        // Encode before opening COPY mode so invalid AST data cannot leave the
267        // connection waiting for CopyFail/CopyDone cleanup.
268        let batch_data = try_encode_copy_batch(rows)?;
269
270        // Send COPY command
271        let bytes = PgEncoder::try_encode_query_string(&sql)?;
272        self.send_bytes(&bytes).await?;
273
274        // Wait for CopyInResponse
275        let mut startup_error: Option<PgError> = None;
276        loop {
277            let msg = self.recv().await?;
278            match msg {
279                BackendMessage::CopyInResponse { .. } => {
280                    if let Some(err) = startup_error {
281                        return return_with_desync(self, err);
282                    }
283                    break;
284                }
285                BackendMessage::ReadyForQuery(_) => {
286                    return return_with_desync(
287                        self,
288                        startup_error.unwrap_or_else(|| {
289                            PgError::Protocol(
290                                "COPY IN failed before CopyInResponse (unexpected ReadyForQuery)"
291                                    .to_string(),
292                            )
293                        }),
294                    );
295                }
296                BackendMessage::ErrorResponse(err) => {
297                    if startup_error.is_none() {
298                        startup_error = Some(PgError::QueryServer(err.into()));
299                    }
300                }
301                msg if is_ignorable_session_message(&msg) => {}
302                other => {
303                    return return_with_desync(
304                        self,
305                        unexpected_backend_message("copy-in startup", &other),
306                    );
307                }
308            }
309        }
310
311        // Single write for entire batch!
312        self.send_copy_data(&batch_data).await?;
313
314        // Send CopyDone
315        self.send_copy_done().await?;
316
317        // Wait for CommandComplete
318        let mut affected = 0u64;
319        let mut final_error: Option<PgError> = None;
320        let mut saw_command_complete = false;
321        loop {
322            let msg = self.recv().await?;
323            match msg {
324                BackendMessage::CommandComplete(tag) => {
325                    if saw_command_complete {
326                        return return_with_desync(
327                            self,
328                            PgError::Protocol(
329                                "COPY IN received duplicate CommandComplete".to_string(),
330                            ),
331                        );
332                    }
333                    saw_command_complete = true;
334                    if final_error.is_none() {
335                        match parse_affected_rows(&tag) {
336                            Ok(parsed) => affected = parsed,
337                            Err(err) => return return_with_desync(self, err),
338                        }
339                    }
340                }
341                BackendMessage::ReadyForQuery(_) => {
342                    if let Some(err) = final_error {
343                        return Err(err);
344                    }
345                    if !saw_command_complete {
346                        return return_with_desync(
347                            self,
348                            PgError::Protocol(
349                                "COPY IN completion missing CommandComplete before ReadyForQuery"
350                                    .to_string(),
351                            ),
352                        );
353                    }
354                    return Ok(affected);
355                }
356                BackendMessage::ErrorResponse(err) => {
357                    if final_error.is_none() {
358                        final_error = Some(PgError::QueryServer(err.into()));
359                    }
360                }
361                msg if is_ignorable_session_message(&msg) => {}
362                other => {
363                    return return_with_desync(
364                        self,
365                        unexpected_backend_message("copy-in completion", &other),
366                    );
367                }
368            }
369        }
370    }
371
372    /// **Fastest** bulk insert using COPY protocol with pre-encoded data.
373    /// Accepts raw COPY text format bytes, no encoding needed.
374    /// Use when caller has already encoded rows to COPY format.
375    /// # Format
376    /// Data should be tab-separated rows with newlines:
377    /// `1\thello\t3.14\n2\tworld\t2.71\n`
378    pub async fn copy_in_raw(
379        &mut self,
380        table: &str,
381        columns: &[String],
382        data: &[u8],
383    ) -> PgResult<u64> {
384        let cols: Vec<String> = columns
385            .iter()
386            .map(|c| quote_copy_column_ident(c))
387            .collect::<PgResult<_>>()?;
388        let sql = format!(
389            "COPY {} ({}) FROM STDIN",
390            quote_copy_table_ref(table)?,
391            cols.join(", ")
392        );
393
394        // Send COPY command
395        let bytes = PgEncoder::try_encode_query_string(&sql)?;
396        self.send_bytes(&bytes).await?;
397
398        // Wait for CopyInResponse
399        let mut startup_error: Option<PgError> = None;
400        loop {
401            let msg = self.recv().await?;
402            match msg {
403                BackendMessage::CopyInResponse { .. } => {
404                    if let Some(err) = startup_error {
405                        return return_with_desync(self, err);
406                    }
407                    break;
408                }
409                BackendMessage::ReadyForQuery(_) => {
410                    return return_with_desync(
411                        self,
412                        startup_error.unwrap_or_else(|| {
413                            PgError::Protocol(
414                                "COPY IN failed before CopyInResponse (unexpected ReadyForQuery)"
415                                    .to_string(),
416                            )
417                        }),
418                    );
419                }
420                BackendMessage::ErrorResponse(err) => {
421                    if startup_error.is_none() {
422                        startup_error = Some(PgError::QueryServer(err.into()));
423                    }
424                }
425                msg if is_ignorable_session_message(&msg) => {}
426                other => {
427                    return return_with_desync(
428                        self,
429                        unexpected_backend_message("copy-in raw startup", &other),
430                    );
431                }
432            }
433        }
434
435        // Single write - data is already encoded!
436        self.send_copy_data(data).await?;
437
438        // Send CopyDone
439        self.send_copy_done().await?;
440
441        // Wait for CommandComplete
442        let mut affected = 0u64;
443        let mut final_error: Option<PgError> = None;
444        let mut saw_command_complete = false;
445        loop {
446            let msg = self.recv().await?;
447            match msg {
448                BackendMessage::CommandComplete(tag) => {
449                    if saw_command_complete {
450                        return return_with_desync(
451                            self,
452                            PgError::Protocol(
453                                "COPY IN raw received duplicate CommandComplete".to_string(),
454                            ),
455                        );
456                    }
457                    saw_command_complete = true;
458                    if final_error.is_none() {
459                        match parse_affected_rows(&tag) {
460                            Ok(parsed) => affected = parsed,
461                            Err(err) => return return_with_desync(self, err),
462                        }
463                    }
464                }
465                BackendMessage::ReadyForQuery(_) => {
466                    if let Some(err) = final_error {
467                        return Err(err);
468                    }
469                    if !saw_command_complete {
470                        return return_with_desync(
471                            self,
472                            PgError::Protocol(
473                                "COPY IN raw completion missing CommandComplete before ReadyForQuery"
474                                    .to_string(),
475                            ),
476                        );
477                    }
478                    return Ok(affected);
479                }
480                BackendMessage::ErrorResponse(err) => {
481                    if final_error.is_none() {
482                        final_error = Some(PgError::QueryServer(err.into()));
483                    }
484                }
485                msg if is_ignorable_session_message(&msg) => {}
486                other => {
487                    return return_with_desync(
488                        self,
489                        unexpected_backend_message("copy-in raw completion", &other),
490                    );
491                }
492            }
493        }
494    }
495
496    /// Send CopyData message (raw bytes).
497    pub(crate) async fn send_copy_data(&mut self, data: &[u8]) -> PgResult<()> {
498        let total_len = data
499            .len()
500            .checked_add(4)
501            .ok_or_else(|| PgError::Protocol("CopyData frame length overflow".to_string()))?;
502        let len = i32::try_from(total_len)
503            .map_err(|_| PgError::Protocol("CopyData frame exceeds i32::MAX".to_string()))?;
504
505        // CopyData: 'd' + length + data
506        let mut buf = BytesMut::with_capacity(1 + 4 + data.len());
507        buf.extend_from_slice(b"d");
508        buf.extend_from_slice(&len.to_be_bytes());
509        buf.extend_from_slice(data);
510        self.send_bytes(&buf).await?;
511        Ok(())
512    }
513
514    async fn send_copy_done(&mut self) -> PgResult<()> {
515        // CopyDone: 'c' + length (4)
516        self.send_bytes(&[b'c', 0, 0, 0, 4]).await?;
517        Ok(())
518    }
519
520    async fn start_copy_out(&mut self, sql: &str, context: &str) -> PgResult<()> {
521        let bytes = PgEncoder::try_encode_query_string(sql)?;
522        self.send_bytes(&bytes).await?;
523
524        let mut startup_error: Option<PgError> = None;
525        loop {
526            let msg = self.recv().await?;
527            match msg {
528                BackendMessage::CopyOutResponse { .. } => {
529                    if let Some(err) = startup_error {
530                        return return_with_desync(self, err);
531                    }
532                    return Ok(());
533                }
534                BackendMessage::ReadyForQuery(_) => {
535                    return return_with_desync(
536                        self,
537                        startup_error.unwrap_or_else(|| {
538                            PgError::Protocol(format!(
539                                "{} failed before CopyOutResponse (unexpected ReadyForQuery)",
540                                context
541                            ))
542                        }),
543                    );
544                }
545                BackendMessage::ErrorResponse(err) => {
546                    if startup_error.is_none() {
547                        startup_error = Some(PgError::QueryServer(err.into()));
548                    }
549                }
550                msg if is_ignorable_session_message(&msg) => {}
551                other => {
552                    return return_with_desync(self, unexpected_backend_message(context, &other));
553                }
554            }
555        }
556    }
557
558    async fn stream_copy_out_chunks<F, Fut>(
559        &mut self,
560        context: &str,
561        mut on_chunk: F,
562    ) -> PgResult<()>
563    where
564        F: FnMut(Vec<u8>) -> Fut,
565        Fut: Future<Output = PgResult<()>>,
566    {
567        let mut stream_error: Option<PgError> = None;
568        let mut callback_error: Option<PgError> = None;
569        let mut saw_copy_done = false;
570        let mut saw_command_complete = false;
571
572        loop {
573            let msg = self.recv().await?;
574            match msg {
575                BackendMessage::CopyData(chunk) => {
576                    if saw_copy_done {
577                        return return_with_desync(
578                            self,
579                            PgError::Protocol(format!(
580                                "{} received CopyData after CopyDone",
581                                context
582                            )),
583                        );
584                    }
585                    if stream_error.is_none()
586                        && callback_error.is_none()
587                        && let Err(e) = on_chunk(chunk).await
588                    {
589                        callback_error = Some(e);
590                    }
591                }
592                BackendMessage::CopyDone => {
593                    if saw_copy_done {
594                        return return_with_desync(
595                            self,
596                            PgError::Protocol(format!("{} received duplicate CopyDone", context)),
597                        );
598                    }
599                    saw_copy_done = true;
600                }
601                BackendMessage::CommandComplete(_) => {
602                    if !saw_copy_done {
603                        return return_with_desync(
604                            self,
605                            PgError::Protocol(format!(
606                                "{} received CommandComplete before CopyDone",
607                                context
608                            )),
609                        );
610                    }
611                    if saw_command_complete {
612                        return return_with_desync(
613                            self,
614                            PgError::Protocol(format!(
615                                "{} received duplicate CommandComplete",
616                                context
617                            )),
618                        );
619                    }
620                    saw_command_complete = true;
621                }
622                BackendMessage::ReadyForQuery(_) => {
623                    if let Some(err) = stream_error {
624                        return Err(err);
625                    }
626                    if let Some(err) = callback_error {
627                        return Err(err);
628                    }
629                    if !saw_copy_done {
630                        return return_with_desync(
631                            self,
632                            PgError::Protocol(format!(
633                                "{} missing CopyDone before ReadyForQuery",
634                                context
635                            )),
636                        );
637                    }
638                    if !saw_command_complete {
639                        return return_with_desync(
640                            self,
641                            PgError::Protocol(format!(
642                                "{} missing CommandComplete before ReadyForQuery",
643                                context
644                            )),
645                        );
646                    }
647                    return Ok(());
648                }
649                BackendMessage::ErrorResponse(err) => {
650                    if stream_error.is_none() {
651                        stream_error = Some(PgError::QueryServer(err.into()));
652                    }
653                }
654                msg if is_ignorable_session_message(&msg) => {}
655                other => {
656                    return return_with_desync(self, unexpected_backend_message(context, &other));
657                }
658            }
659        }
660    }
661
662    /// Export data using COPY TO STDOUT (AST-native).
663    /// Takes a `Qail::Export` and returns rows as `Vec<Vec<String>>`.
664    /// # Example
665    /// ```ignore
666    /// let cmd = Qail::export("users")
667    ///     .columns(["id", "name"]);
668    /// let rows = conn.copy_export(&cmd).await?;
669    /// ```
670    pub async fn copy_export(&mut self, cmd: &Qail) -> PgResult<Vec<Vec<String>>> {
671        let mut rows = Vec::new();
672        self.copy_export_stream_rows(cmd, |row| {
673            rows.push(row);
674            Ok(())
675        })
676        .await?;
677        Ok(rows)
678    }
679
680    /// Stream COPY TO STDOUT chunks using an AST-native `Qail::Export` command.
681    ///
682    /// Chunks are forwarded as they arrive from PostgreSQL, so memory usage
683    /// stays bounded by network frame size and callback processing.
684    pub async fn copy_export_stream_raw<F, Fut>(&mut self, cmd: &Qail, on_chunk: F) -> PgResult<()>
685    where
686        F: FnMut(Vec<u8>) -> Fut,
687        Fut: Future<Output = PgResult<()>>,
688    {
689        let sql = encode_copy_export_sql(cmd)?;
690
691        self.copy_out_raw_stream(&sql, on_chunk).await
692    }
693
694    /// Stream COPY TO STDOUT rows using an AST-native `Qail::Export` command.
695    ///
696    /// Parses PostgreSQL COPY text lines into `Vec<String>` rows and invokes
697    /// `on_row` for each row without buffering the full result.
698    pub async fn copy_export_stream_rows<F>(&mut self, cmd: &Qail, mut on_row: F) -> PgResult<()>
699    where
700        F: FnMut(Vec<String>) -> PgResult<()>,
701    {
702        let mut pending = Vec::new();
703        self.copy_export_stream_raw(cmd, |chunk| {
704            let res = drain_copy_text_rows(&mut pending, &chunk, &mut on_row);
705            std::future::ready(res)
706        })
707        .await?;
708        flush_pending_copy_text_row(&pending)
709    }
710
711    /// Export data using raw COPY TO STDOUT, returning raw bytes.
712    /// Format: tab-separated values, newline-terminated rows.
713    /// Suitable for direct re-import via copy_in_raw.
714    ///
715    /// # Safety
716    /// `pub(crate)` — not exposed externally because callers pass raw SQL.
717    /// External code should use `copy_export()` with the AST encoder instead.
718    pub(crate) async fn copy_out_raw(&mut self, sql: &str) -> PgResult<Vec<u8>> {
719        let mut data = Vec::new();
720        self.copy_out_raw_stream(sql, |chunk| {
721            data.extend_from_slice(&chunk);
722            std::future::ready(Ok(()))
723        })
724        .await?;
725        Ok(data)
726    }
727
728    /// Stream raw COPY TO STDOUT bytes with bounded memory usage.
729    ///
730    /// # Safety
731    /// `pub(crate)` — callers pass raw SQL.
732    pub(crate) async fn copy_out_raw_stream<F, Fut>(
733        &mut self,
734        sql: &str,
735        on_chunk: F,
736    ) -> PgResult<()>
737    where
738        F: FnMut(Vec<u8>) -> Fut,
739        Fut: Future<Output = PgResult<()>>,
740    {
741        self.start_copy_out(sql, "copy-out raw startup").await?;
742        self.stream_copy_out_chunks("copy-out raw stream", on_chunk)
743            .await
744    }
745}
746
747#[cfg(test)]
748mod tests {
749    use super::{
750        drain_copy_text_rows, encode_copy_export_sql, flush_pending_copy_text_row,
751        parse_copy_text_row, quote_copy_column_ident, quote_copy_table_ref, return_with_desync,
752    };
753    use crate::driver::{PgConnection, PgError, PgResult};
754    use qail_core::ast::{Operator, Qail};
755
756    #[cfg(unix)]
757    fn test_conn() -> PgConnection {
758        use crate::driver::connection::StatementCache;
759        use crate::driver::stream::PgStream;
760        use bytes::BytesMut;
761        use std::collections::{HashMap, VecDeque};
762        use std::num::NonZeroUsize;
763        use tokio::net::UnixStream;
764
765        let (unix_stream, _peer) = UnixStream::pair().expect("unix stream pair");
766        PgConnection {
767            stream: PgStream::Unix(unix_stream),
768            buffer: BytesMut::with_capacity(1024),
769            write_buf: BytesMut::with_capacity(1024),
770            sql_buf: BytesMut::with_capacity(256),
771            params_buf: Vec::new(),
772            prepared_statements: HashMap::new(),
773            stmt_cache: StatementCache::new(NonZeroUsize::new(2).expect("non-zero")),
774            column_info_cache: HashMap::new(),
775            process_id: 0,
776            cancel_key_bytes: Vec::new(),
777            requested_protocol_minor: PgConnection::default_protocol_minor(),
778            negotiated_protocol_minor: PgConnection::default_protocol_minor(),
779            notifications: VecDeque::new(),
780            replication_stream_active: false,
781            replication_mode_enabled: false,
782            last_replication_wal_end: None,
783            io_desynced: false,
784            pending_statement_closes: Vec::new(),
785            draining_statement_closes: false,
786        }
787    }
788
789    #[test]
790    fn parse_copy_text_row_splits_tabs() {
791        let row = parse_copy_text_row(b"a\tb\tc").unwrap();
792        assert_eq!(row, vec!["a", "b", "c"]);
793    }
794
795    #[test]
796    fn parse_copy_text_row_trims_cr() {
797        let row = parse_copy_text_row(b"a\tb\r").unwrap();
798        assert_eq!(row, vec!["a", "b"]);
799    }
800
801    #[test]
802    fn parse_copy_text_row_unescapes_copy_text_values() {
803        let row = parse_copy_text_row(b"a\\tb\tline\\nnext\tc\\\\d").unwrap();
804        assert_eq!(row, vec!["a\tb", "line\nnext", "c\\d"]);
805    }
806
807    #[test]
808    fn parse_copy_text_row_rejects_copy_null_marker() {
809        let err = parse_copy_text_row(b"a\t\\N\tb").expect_err("COPY NULL must not be lossy");
810        assert!(
811            err.to_string()
812                .contains("COPY text NULL cannot be represented"),
813            "{err}"
814        );
815    }
816
817    #[test]
818    fn parse_copy_text_row_rejects_invalid_utf8() {
819        let err = parse_copy_text_row(&[0xff]).expect_err("invalid UTF-8 must fail");
820        assert!(
821            err.to_string()
822                .contains("COPY text field is not valid UTF-8")
823        );
824    }
825
826    #[test]
827    fn parse_copy_text_row_rejects_incomplete_escape() {
828        let err = parse_copy_text_row(b"bad\\").expect_err("trailing backslash must fail");
829        assert!(err.to_string().contains("incomplete backslash escape"));
830    }
831
832    #[test]
833    fn parse_copy_text_row_rejects_out_of_range_octal_escape() {
834        let err = parse_copy_text_row(br"\400").expect_err("octal escape > 377 must fail");
835        assert!(err.to_string().contains("out of byte range"));
836    }
837
838    #[test]
839    fn parse_copy_text_row_rejects_hex_escape_without_digits() {
840        let err = parse_copy_text_row(br"\xG").expect_err("hex escape without digits must fail");
841        assert!(err.to_string().contains("hex escape requires"));
842    }
843
844    #[test]
845    fn copy_table_quoting_preserves_schema_qualification() {
846        assert_eq!(
847            quote_copy_table_ref("tenant_a.users").unwrap(),
848            "\"tenant_a\".\"users\""
849        );
850    }
851
852    #[test]
853    fn copy_identifier_quoting_rejects_nul_bytes() {
854        assert!(quote_copy_table_ref("tenant\0.users").is_err());
855        assert!(quote_copy_column_ident("name\0").is_err());
856    }
857
858    #[test]
859    fn copy_export_rejects_parameterized_ast_before_streaming() {
860        let cmd = Qail::export("users").filter("active", Operator::Eq, true);
861        let err = encode_copy_export_sql(&cmd).expect_err("bind params cannot be ignored");
862
863        assert!(matches!(err, PgError::Encode(msg) if msg.contains("parameterized export")));
864    }
865
866    #[cfg(unix)]
867    #[tokio::test]
868    async fn copy_return_with_desync_marks_protocol_error() {
869        let mut conn = test_conn();
870
871        let err = return_with_desync::<()>(
872            &mut conn,
873            PgError::Protocol("copy protocol ordering broke".to_string()),
874        )
875        .expect_err("protocol error must be returned");
876
877        assert!(err.to_string().contains("copy protocol ordering broke"));
878        assert!(conn.is_io_desynced());
879    }
880
881    #[test]
882    fn drain_copy_text_rows_handles_chunk_boundaries() {
883        let mut pending = Vec::new();
884        let mut rows: Vec<Vec<String>> = Vec::new();
885
886        drain_copy_text_rows(&mut pending, b"a\tb\nc", &mut |row: Vec<String>| {
887            rows.push(row);
888            Ok(())
889        })
890        .unwrap();
891        assert_eq!(rows, vec![vec!["a".to_string(), "b".to_string()]]);
892        assert_eq!(pending, b"c");
893
894        drain_copy_text_rows(&mut pending, b"\td\n", &mut |row: Vec<String>| {
895            rows.push(row);
896            Ok(())
897        })
898        .unwrap();
899        assert_eq!(
900            rows,
901            vec![
902                vec!["a".to_string(), "b".to_string()],
903                vec!["c".to_string(), "d".to_string()]
904            ]
905        );
906        assert!(pending.is_empty());
907    }
908
909    #[test]
910    fn flush_pending_copy_text_row_rejects_final_partial_line() {
911        let pending = b"x\ty".to_vec();
912        let err = flush_pending_copy_text_row(&pending)
913            .expect_err("partial final COPY row must fail closed");
914        assert!(matches!(err, PgError::Protocol(msg) if msg.contains("truncated row")));
915        assert_eq!(pending, b"x\ty");
916    }
917
918    #[test]
919    fn callback_error_bubbles_from_row_drainer() {
920        let mut pending = Vec::new();
921        let mut on_row =
922            |_row: Vec<String>| -> PgResult<()> { Err(PgError::Query("fail".to_string())) };
923
924        let err = drain_copy_text_rows(&mut pending, b"a\tb\n", &mut on_row).unwrap_err();
925        assert!(matches!(err, PgError::Query(msg) if msg == "fail"));
926    }
927}
qail_pg/driver/copy.rs

qail_pg/driver/
copy.rs
