qail_pg/driver/

copy.rs

//! COPY protocol methods for PostgreSQL bulk operations.
//!

use super::{PgConnection, PgError, PgResult, parse_affected_rows};
use crate::protocol::{AstEncoder, BackendMessage, PgEncoder};
use bytes::BytesMut;
use qail_core::ast::{Action, Qail};
use tokio::io::AsyncWriteExt;

/// Quote a SQL identifier to prevent injection.
/// Wraps in double-quotes, escapes embedded double-quotes, and strips NUL bytes.
fn quote_ident(ident: &str) -> String {
    format!("\"{}\"", ident.replace('\0', "").replace('"', "\"\""))
}

impl PgConnection {
    /// **Fast** bulk insert using COPY protocol with zero-allocation encoding.
    /// Encodes all rows into a single buffer and writes with one syscall.
    /// ~2x faster than `copy_in_internal` due to batched I/O.
    pub(crate) async fn copy_in_fast(
        &mut self,
        table: &str,
        columns: &[String],
        rows: &[Vec<qail_core::ast::Value>],
    ) -> PgResult<u64> {
        use crate::protocol::encode_copy_batch;

        let cols: Vec<String> = columns.iter().map(|c| quote_ident(c)).collect();
        let sql = format!(
            "COPY {} ({}) FROM STDIN",
            quote_ident(table),
            cols.join(", ")
        );

        // Send COPY command
        let bytes = PgEncoder::encode_query_string(&sql);
        self.stream.write_all(&bytes).await?;

        // Wait for CopyInResponse
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyInResponse { .. } => break,
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }

        // Encode ALL rows into a single buffer (zero-allocation per value)
        let batch_data = encode_copy_batch(rows);

        // Single write for entire batch!
        self.send_copy_data(&batch_data).await?;

        // Send CopyDone
        self.send_copy_done().await?;

        // Wait for CommandComplete
        let mut affected = 0u64;
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CommandComplete(tag) => {
                    affected = parse_affected_rows(&tag);
                }
                BackendMessage::ReadyForQuery(_) => {
                    return Ok(affected);
                }
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }
    }

    /// **Fastest** bulk insert using COPY protocol with pre-encoded data.
    /// Accepts raw COPY text format bytes, no encoding needed.
    /// Use when caller has already encoded rows to COPY format.
    /// # Format
    /// Data should be tab-separated rows with newlines:
    /// `1\thello\t3.14\n2\tworld\t2.71\n`
    pub async fn copy_in_raw(
        &mut self,
        table: &str,
        columns: &[String],
        data: &[u8],
    ) -> PgResult<u64> {
        let cols: Vec<String> = columns.iter().map(|c| quote_ident(c)).collect();
        let sql = format!(
            "COPY {} ({}) FROM STDIN",
            quote_ident(table),
            cols.join(", ")
        );

        // Send COPY command
        let bytes = PgEncoder::encode_query_string(&sql);
        self.stream.write_all(&bytes).await?;

        // Wait for CopyInResponse
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyInResponse { .. } => break,
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }

        // Single write - data is already encoded!
        self.send_copy_data(data).await?;

        // Send CopyDone
        self.send_copy_done().await?;

        // Wait for CommandComplete
        let mut affected = 0u64;
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CommandComplete(tag) => {
                    affected = parse_affected_rows(&tag);
                }
                BackendMessage::ReadyForQuery(_) => {
                    return Ok(affected);
                }
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }
    }

    /// Send CopyData message (raw bytes).
    async fn send_copy_data(&mut self, data: &[u8]) -> PgResult<()> {
        // CopyData: 'd' + length + data
        let len = (data.len() + 4) as i32;
        let mut buf = BytesMut::with_capacity(1 + 4 + data.len());
        buf.extend_from_slice(b"d");
        buf.extend_from_slice(&len.to_be_bytes());
        buf.extend_from_slice(data);
        self.stream.write_all(&buf).await?;
        Ok(())
    }

    async fn send_copy_done(&mut self) -> PgResult<()> {
        // CopyDone: 'c' + length (4)
        self.stream.write_all(&[b'c', 0, 0, 0, 4]).await?;
        Ok(())
    }

    /// Export data using COPY TO STDOUT (AST-native).
    /// Takes a `Qail::Export` and returns rows as `Vec<Vec<String>>`.
    /// # Example
    /// ```ignore
    /// let cmd = Qail::export("users")
    ///     .columns(["id", "name"])
    ///     .filter("active", true);
    /// let rows = conn.copy_export(&cmd).await?;
    /// ```
    pub async fn copy_export(&mut self, cmd: &Qail) -> PgResult<Vec<Vec<String>>> {
        if cmd.action != Action::Export {
            return Err(PgError::Query(
                "copy_export requires Qail::Export action".to_string(),
            ));
        }

        // Encode command to SQL using AST encoder
        let (sql, _params) =
            AstEncoder::encode_cmd_sql(cmd).map_err(|e| PgError::Encode(e.to_string()))?;

        // Send COPY command
        let bytes = PgEncoder::encode_query_string(&sql);
        self.stream.write_all(&bytes).await?;

        // Wait for CopyOutResponse
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyOutResponse { .. } => break,
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }

        // Receive CopyData messages until CopyDone
        let mut rows = Vec::new();
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyData(data) => {
                    let line = String::from_utf8_lossy(&data);
                    let line = line.trim_end_matches('\n');
                    let cols: Vec<String> = line.split('\t').map(|s| s.to_string()).collect();
                    rows.push(cols);
                }
                BackendMessage::CopyDone => {}
                BackendMessage::CommandComplete(_) => {}
                BackendMessage::ReadyForQuery(_) => {
                    return Ok(rows);
                }
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }
    }

    /// Export data using raw COPY TO STDOUT, returning raw bytes.
    /// Format: tab-separated values, newline-terminated rows.
    /// Suitable for direct re-import via copy_in_raw.
    ///
    /// # Safety
    /// `pub(crate)` — not exposed externally because callers pass raw SQL.
    /// External code should use `copy_export()` with the AST encoder instead.
    pub(crate) async fn copy_out_raw(&mut self, sql: &str) -> PgResult<Vec<u8>> {
        // Send COPY command
        let bytes = PgEncoder::encode_query_string(sql);
        self.stream.write_all(&bytes).await?;

        // Wait for CopyOutResponse
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyOutResponse { .. } => break,
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }

        // Receive CopyData messages until CopyDone
        let mut data = Vec::new();
        loop {
            let msg = self.recv().await?;
            match msg {
                BackendMessage::CopyData(chunk) => {
                    data.extend_from_slice(&chunk);
                }
                BackendMessage::CopyDone => {}
                BackendMessage::CommandComplete(_) => {}
                BackendMessage::ReadyForQuery(_) => {
                    return Ok(data);
                }
                BackendMessage::ErrorResponse(err) => {
                    return Err(PgError::QueryServer(err.into()));
                }
                _ => {}
            }
        }
    }
}