nusadb 0.1.0

//! A single connection to a NusaDB server: TCP + handshake (trust or SCRAM-SHA-256), simple and
//! extended (parameterized / prepared) queries, and transaction control.

use std::collections::VecDeque;
use std::io::BufReader;
use std::net::TcpStream;
use std::sync::Arc;
use std::time::Duration;

use crate::error::{Error, Result};
use crate::protocol::{
    self, auth, backend, frontend, Frame, Reader, TypeTag, Writer, PROTOCOL_MAGIC, PROTOCOL_MAJOR,
    PROTOCOL_MINOR, SCRAM_MECHANISM,
};
use crate::scram::Scram;
use crate::tls::TlsConfig;
use crate::transport::Transport;
use crate::value::{ToSql, Value};

/// How to reach and authenticate to a server.
#[derive(Debug, Clone)]
pub struct Config {
    /// Host (default `127.0.0.1`).
    pub host: String,
    /// TCP port (default `5678`).
    pub port: u16,
    /// Authenticating role.
    pub user: String,
    /// Password for SCRAM, or `None` for a trust-on-startup server.
    pub password: Option<String>,
    /// Target database (default `nusadb`).
    pub database: String,
    /// Optional connect timeout.
    pub connect_timeout: Option<Duration>,
    /// TLS settings, or `None` for a plaintext connection. Requires the `tls` feature to take effect
    /// (a `Some` value without the feature is a [`Error::Config`] at connect time).
    pub tls: Option<TlsConfig>,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            host: "127.0.0.1".to_owned(),
            port: 5678,
            user: "nusa-root".to_owned(),
            password: Some("nusa-root".to_owned()),
            database: "nusadb".to_owned(),
            connect_timeout: None,
            tls: None,
        }
    }
}

impl Config {
    /// Parse a `nusadb://[user[:password]@]host[:port]/database` URL into a [`Config`]. The
    /// `nusadbs://` scheme enables TLS (with public-root verification; refine via [`Config::tls`]).
    pub fn from_url(url: &str) -> Result<Self> {
        let mut cfg = Self::default();
        let rest = if let Some(rest) = url.strip_prefix("nusadbs://") {
            cfg.tls = Some(TlsConfig::with_public_roots());
            rest
        } else {
            url.strip_prefix("nusadb://").ok_or_else(|| {
                Error::Config("URL must start with nusadb:// or nusadbs://".to_owned())
            })?
        };
        // Split optional `user:pass@` authority prefix.
        let (authority, hostpart) = match rest.split_once('@') {
            Some((auth, host)) => (Some(auth), host),
            None => (None, rest),
        };
        if let Some(auth) = authority {
            let (user, pass) = match auth.split_once(':') {
                Some((u, p)) => (u, Some(p.to_owned())),
                None => (auth, None),
            };
            if !user.is_empty() {
                cfg.user = decode_pct(user);
            }
            cfg.password = pass.map(|p| decode_pct(&p));
        }
        // `host[:port]/database`
        let (host_port, db) = match hostpart.split_once('/') {
            Some((hp, db)) => (hp, Some(db)),
            None => (hostpart, None),
        };
        if let Some((h, p)) = host_port.rsplit_once(':') {
            cfg.host = h.to_owned();
            cfg.port = p
                .parse()
                .map_err(|_| Error::Config(format!("invalid port {p:?}")))?;
        } else if !host_port.is_empty() {
            cfg.host = host_port.to_owned();
        }
        if let Some(db) = db {
            if !db.is_empty() {
                cfg.database = decode_pct(db);
            }
        }
        Ok(cfg)
    }
}

/// Minimal percent-decoding for URL components (`%XX`).
fn decode_pct(s: &str) -> String {
    let bytes = s.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'%' && i + 2 < bytes.len() {
            if let Ok(b) = u8::from_str_radix(&s[i + 1..i + 3], 16) {
                out.push(b);
                i += 3;
                continue;
            }
        }
        out.push(bytes[i]);
        i += 1;
    }
    String::from_utf8_lossy(&out).into_owned()
}

/// One result row: its column values, plus a shared handle to the column names for by-name access.
#[derive(Debug, Clone)]
pub struct Row {
    columns: Arc<Vec<String>>,
    values: Vec<Value>,
}

impl Row {
    /// The number of columns.
    #[must_use]
    pub fn len(&self) -> usize {
        self.values.len()
    }

    /// Whether the row has no columns.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    /// Borrow the raw [`Value`] at `idx`.
    #[must_use]
    pub fn value(&self, idx: usize) -> Option<&Value> {
        self.values.get(idx)
    }

    /// All raw values.
    #[must_use]
    pub fn values(&self) -> &[Value] {
        &self.values
    }

    /// Read column `idx` into `T`.
    pub fn get<T: crate::value::FromSql>(&self, idx: usize) -> Result<T> {
        let v = self
            .values
            .get(idx)
            .ok_or_else(|| Error::Conversion(format!("column index {idx} out of range")))?;
        T::from_sql(v)
    }

    /// Read the column named `name` into `T`.
    pub fn get_by_name<T: crate::value::FromSql>(&self, name: &str) -> Result<T> {
        let idx = self
            .columns
            .iter()
            .position(|c| c == name)
            .ok_or_else(|| Error::Conversion(format!("no column named {name:?}")))?;
        self.get(idx)
    }
}

/// The outcome of a query: column metadata, the rows, and the server's completion tag.
#[derive(Debug, Clone)]
pub struct QueryResult {
    /// Output column names, in order.
    pub columns: Vec<String>,
    /// Per-column declared type (1.1). All [`TypeTag::Unknown`] against a 1.0 server.
    pub column_types: Vec<TypeTag>,
    /// The result rows.
    pub rows: Vec<Row>,
    /// The `CommandComplete` tag (e.g. `"SELECT 3"`, `"INSERT 1"`).
    pub tag: String,
}

impl QueryResult {
    /// The affected/returned row count parsed from the trailing number of the command tag, if any.
    #[must_use]
    pub fn affected(&self) -> u64 {
        self.tag
            .rsplit(' ')
            .next()
            .and_then(|n| n.parse().ok())
            .unwrap_or(0)
    }
}

/// A prepared statement: a name the server holds a parsed plan under, reusable across executions.
#[derive(Debug, Clone)]
pub struct Prepared {
    name: String,
}

#[cfg(feature = "async")]
impl Prepared {
    pub(crate) fn new(name: String) -> Self {
        Self { name }
    }

    pub(crate) fn name(&self) -> &str {
        &self.name
    }
}

/// A live connection to a NusaDB server.
/// An asynchronous `LISTEN`/`NOTIFY` notification delivered by the server (§ pub/sub). Collected by
/// the connection while it is idle and returned from [`Connection::notifications`] /
/// [`Connection::poll_notification`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Notification {
    /// Backend pid of the connection that issued the `NOTIFY`.
    pub pid: u32,
    /// The channel the notification was sent on.
    pub channel: String,
    /// The payload (empty string when `NOTIFY` carried none).
    pub payload: String,
}

/// A synchronous connection to a NusaDB server.
pub struct Connection {
    stream: BufReader<Transport>,
    addr: (String, u16),
    backend_pid: u32,
    backend_secret: u32,
    stmt_counter: u64,
    /// Notifications received (buffered) while reading other frames; drained by `notifications()`.
    notifications: VecDeque<Notification>,
}

impl Connection {
    /// Open a connection from a `nusadb://…` (or `nusadbs://…` for TLS) URL.
    pub fn connect(url: &str) -> Result<Self> {
        Self::connect_config(&Config::from_url(url)?)
    }

    /// Open a connection from an explicit [`Config`].
    pub fn connect_config(config: &Config) -> Result<Self> {
        let tcp = match config.connect_timeout {
            Some(timeout) => {
                let addrs: Vec<_> =
                    std::net::ToSocketAddrs::to_socket_addrs(&(config.host.as_str(), config.port))?
                        .collect();
                let addr = addrs
                    .first()
                    .ok_or_else(|| Error::Config("could not resolve host".to_owned()))?;
                TcpStream::connect_timeout(addr, timeout)?
            }
            None => TcpStream::connect((config.host.as_str(), config.port))?,
        };
        tcp.set_nodelay(true).ok();
        let transport = build_transport(tcp, config)?;
        let mut conn = Self {
            stream: BufReader::new(transport),
            addr: (config.host.clone(), config.port),
            backend_pid: 0,
            backend_secret: 0,
            stmt_counter: 0,
            notifications: VecDeque::new(),
        };
        conn.handshake(config)?;
        Ok(conn)
    }

    fn send(&mut self, frame: &[u8]) -> Result<()> {
        protocol::write_frame(self.stream.get_mut(), frame)
    }

    /// Read one frame off the wire, transparently buffering any asynchronous `NotificationResponse`
    /// (`LISTEN`/`NOTIFY`) frames. Because the server delivers notifications while the connection is
    /// idle, a pending one can lead the next query's response — every read path goes through here so
    /// they are collected for `notifications()` rather than mistaken for an unexpected message.
    fn recv(&mut self) -> Result<Frame> {
        loop {
            let frame = protocol::read_frame(&mut self.stream)?;
            if frame.kind == backend::NOTIFICATION_RESPONSE {
                let note = decode_notification(&frame)?;
                self.notifications.push_back(note);
                continue;
            }
            return Ok(frame);
        }
    }

    fn handshake(&mut self, config: &Config) -> Result<()> {
        // Startup.
        let startup = Writer::new()
            .u32(PROTOCOL_MAGIC)
            .u16(PROTOCOL_MAJOR)
            .u16(PROTOCOL_MINOR)
            .str(&config.user)
            .str(&config.database)
            .frame(frontend::STARTUP);
        self.send(&startup)?;

        // Authentication: a single top-level message — either an immediate AuthOk (trust) or an
        // AuthSASL that begins SCRAM. `scram` drives the rest of the SASL exchange and consumes its
        // trailing AuthOk, so exactly one auth frame is read here before the post-auth stream.
        let frame = self.recv()?;
        match frame.kind {
            backend::AUTH => {
                let mut r = Reader::new(&frame.payload);
                match r.u32()? {
                    auth::OK => {}
                    auth::SASL => self.scram(config)?,
                    other => {
                        return Err(Error::Protocol(format!("unexpected auth sub-code {other}")))
                    }
                }
            }
            backend::ERROR => return Err(decode_error(&frame.payload)),
            other => {
                return Err(Error::Protocol(format!(
                    "unexpected message {:?} during auth",
                    other as char
                )))
            }
        }

        // Post-auth: BackendKeyData (K) and any ParameterStatus (S) frames, then ReadyForQuery (Z).
        loop {
            let frame = self.recv()?;
            match frame.kind {
                backend::BACKEND_KEY => {
                    let mut r = Reader::new(&frame.payload);
                    self.backend_pid = r.u32()?;
                    self.backend_secret = r.u32()?;
                }
                // Runtime parameters announced at startup carry nothing the driver needs.
                backend::PARAMETER_STATUS => {}
                backend::READY => return Ok(()),
                backend::ERROR => return Err(decode_error(&frame.payload)),
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} after auth",
                        other as char
                    )))
                }
            }
        }
    }

    fn scram(&mut self, config: &Config) -> Result<()> {
        let password = config
            .password
            .as_deref()
            .ok_or_else(|| Error::Auth("server requires a password (SCRAM)".to_owned()))?;
        let (mut state, client_first) = Scram::start(&config.user, password)?;

        // SASLInitialResponse: [mechanism:Str][data_len:u32][data:bytes].
        let initial = Writer::new()
            .str(SCRAM_MECHANISM)
            .u32(client_first.len() as u32)
            .raw(client_first.as_bytes())
            .frame(frontend::SASL_INITIAL);
        self.send(&initial)?;

        // AuthSaslContinue (R, sub-code 11): [11][server-first bytes].
        let frame = self.recv()?;
        let server_first = self.expect_auth_sub(frame, auth::SASL_CONTINUE)?;
        let client_final = state.client_final(&server_first)?;

        // SASLResponse: [client-final bytes] (raw, frame-delimited).
        let response = Writer::new()
            .raw(client_final.as_bytes())
            .frame(frontend::SASL_RESPONSE);
        self.send(&response)?;

        // AuthSaslFinal (R, sub-code 12): [12][server-final bytes]; verify mutual auth.
        let frame = self.recv()?;
        let server_final = self.expect_auth_sub(frame, auth::SASL_FINAL)?;
        if !state.verify(&server_final) {
            return Err(Error::Auth(
                "server signature verification failed".to_owned(),
            ));
        }
        // AuthOk (R, sub-code 0).
        let frame = self.recv()?;
        let _ = self.expect_auth_sub(frame, auth::OK)?;
        Ok(())
    }

    fn expect_auth_sub(&self, frame: Frame, expect: u32) -> Result<String> {
        if frame.kind == backend::ERROR {
            return Err(decode_error(&frame.payload));
        }
        if frame.kind != backend::AUTH {
            return Err(Error::Protocol(format!(
                "expected auth message, got {:?}",
                frame.kind as char
            )));
        }
        let mut r = Reader::new(&frame.payload);
        let sub = r.u32()?;
        if sub != expect {
            return Err(Error::Protocol(format!(
                "expected auth sub-code {expect}, got {sub}"
            )));
        }
        let rest = r.rest()?;
        String::from_utf8(rest.to_vec())
            .map_err(|_| Error::Protocol("invalid UTF-8 in SASL message".to_owned()))
    }

    /// Run a simple (text) query with no parameters.
    pub fn query(&mut self, sql: &str) -> Result<QueryResult> {
        let frame = Writer::new().str(sql).frame(frontend::QUERY);
        self.send(&frame)?;
        self.collect(true)
    }

    /// Run a query with positional `$1..$n` parameters (extended protocol; one-shot unnamed statement).
    pub fn query_params(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<QueryResult> {
        let encoded: Vec<Option<Vec<u8>>> = params.iter().map(|p| p.to_sql_text()).collect();
        self.send(&Writer::new().str("").str(sql).u16(0).frame(frontend::PARSE))?;
        self.send(&bind_frame("", "", &encoded))?;
        self.send(&describe_portal(""))?;
        self.send(&execute_portal("", 0))?;
        self.send(&Writer::new().frame(frontend::SYNC))?;
        self.collect(false)
    }

    /// Prepare a statement once for repeated execution (extended protocol). Reuse it with
    /// [`Self::query_prepared`].
    pub fn prepare(&mut self, sql: &str) -> Result<Prepared> {
        self.stmt_counter += 1;
        let name = format!("nusadb_stmt_{}", self.stmt_counter);
        self.send(
            &Writer::new()
                .str(&name)
                .str(sql)
                .u16(0)
                .frame(frontend::PARSE),
        )?;
        self.send(&Writer::new().frame(frontend::SYNC))?;
        // Drain to ReadyForQuery, surfacing any parse error.
        self.collect(false)?;
        Ok(Prepared { name })
    }

    /// Execute a previously [`Self::prepare`]d statement with parameters.
    pub fn query_prepared(
        &mut self,
        stmt: &Prepared,
        params: &[&dyn ToSql],
    ) -> Result<QueryResult> {
        let encoded: Vec<Option<Vec<u8>>> = params.iter().map(|p| p.to_sql_text()).collect();
        let portal = "";
        self.send(&bind_frame(portal, &stmt.name, &encoded))?;
        self.send(&describe_portal(portal))?;
        self.send(&execute_portal(portal, 0))?;
        self.send(&Writer::new().frame(frontend::SYNC))?;
        self.collect(false)
    }

    /// Run a statement and return the affected/returned row count (parsed from its command tag).
    pub fn execute(&mut self, sql: &str) -> Result<u64> {
        Ok(self.query(sql)?.affected())
    }

    /// Bulk-load a table via `COPY <table> FROM STDIN` (§12.1). Sends the COPY statement, then
    /// streams `source`'s bytes — already in the server's text format (tab-delimited, `\N` for NULL,
    /// matching the statement's `WITH` options) — as `CopyData` frames, finishing with `CopyDone`.
    /// Returns the number of rows loaded.
    ///
    /// On a read error from `source` the driver sends `CopyFail`, drains to `ReadyForQuery`, and
    /// returns the I/O error, leaving the connection usable. A COPY the server refuses (bad SQL, an
    /// RLS-protected table) surfaces as [`Error::Server`].
    pub fn copy_in<R: std::io::Read>(&mut self, sql: &str, source: &mut R) -> Result<u64> {
        self.send(&Writer::new().str(sql).frame(frontend::QUERY))?;
        self.await_copy_start(backend::COPY_IN)?;
        let mut buf = vec![0u8; 64 * 1024];
        loop {
            match source.read(&mut buf) {
                Ok(0) => break,
                Ok(n) => self.send(&Writer::new().raw(&buf[..n]).frame(frontend::COPY_DATA))?,
                Err(e) => {
                    // Abort the load but keep the connection in protocol sync.
                    let _ = self.send(
                        &Writer::new()
                            .str("nusadb: client read error during COPY")
                            .frame(frontend::COPY_FAIL),
                    );
                    self.drain_to_ready();
                    return Err(Error::Io(e));
                }
            }
        }
        self.send(&Writer::new().frame(frontend::COPY_DONE))?;
        self.finish_copy()
    }

    /// Bulk-export a table via `COPY <table> TO STDOUT` (§12.2). Sends the COPY statement, then
    /// writes every `CopyData` chunk the server streams into `sink` (raw text-format bytes —
    /// tab-delimited, `\N` for NULL). Returns the number of rows exported.
    pub fn copy_out<W: std::io::Write>(&mut self, sql: &str, sink: &mut W) -> Result<u64> {
        self.send(&Writer::new().str(sql).frame(frontend::QUERY))?;
        self.await_copy_start(backend::COPY_OUT)?;
        loop {
            let frame = self.recv()?;
            match frame.kind {
                backend::COPY_DATA => sink.write_all(&frame.payload)?,
                backend::COPY_DONE => break,
                backend::ERROR => {
                    let err = decode_error(&frame.payload);
                    self.drain_to_ready();
                    return Err(err);
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} during COPY TO STDOUT",
                        other as char
                    )))
                }
            }
        }
        self.finish_copy()
    }

    /// Read frames until the expected `CopyInResponse`/`CopyOutResponse` arrives. A COPY the server
    /// refuses comes back as `Error` then `ReadyForQuery`; surface that error, connection left ready.
    fn await_copy_start(&mut self, expected: u8) -> Result<()> {
        let mut err: Option<Error> = None;
        loop {
            let frame = self.recv()?;
            match frame.kind {
                k if k == expected => return Ok(()),
                backend::ERROR => err = Some(decode_error(&frame.payload)),
                backend::READY => {
                    return Err(err.unwrap_or_else(|| {
                        Error::Protocol("COPY: server did not start the copy".to_owned())
                    }))
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} starting COPY",
                        other as char
                    )))
                }
            }
        }
    }

    /// Drain the tail of a COPY exchange to `ReadyForQuery`, returning the row count parsed from the
    /// `COPY <n>` command tag (or a server error).
    fn finish_copy(&mut self) -> Result<u64> {
        let mut tag = String::new();
        let mut err: Option<Error> = None;
        loop {
            let frame = self.recv()?;
            match frame.kind {
                backend::COMMAND_COMPLETE => tag = Reader::new(&frame.payload).str()?,
                backend::COPY_DONE => {}
                backend::ERROR => err = Some(decode_error(&frame.payload)),
                backend::READY => {
                    if let Some(e) = err {
                        return Err(e);
                    }
                    return Ok(copy_count(&tag));
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} finishing COPY",
                        other as char
                    )))
                }
            }
        }
    }

    /// Best-effort resync: read and discard frames until `ReadyForQuery`.
    fn drain_to_ready(&mut self) {
        while let Ok(frame) = self.recv() {
            if frame.kind == backend::READY {
                break;
            }
        }
    }

    /// Open a transaction (`BEGIN`).
    pub fn begin(&mut self) -> Result<()> {
        self.query("BEGIN").map(|_| ())
    }

    /// Commit the open transaction (`COMMIT`).
    pub fn commit(&mut self) -> Result<()> {
        self.query("COMMIT").map(|_| ())
    }

    /// Roll back the open transaction (`ROLLBACK`).
    pub fn rollback(&mut self) -> Result<()> {
        self.query("ROLLBACK").map(|_| ())
    }

    /// Establish a savepoint in the open transaction (`SAVEPOINT name`) — a named marker you can
    /// later roll back to without ending the transaction.
    pub fn savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("SAVEPOINT {}", quote_identifier(name)))
            .map(|_| ())
    }

    /// Roll back to a savepoint (`ROLLBACK TO SAVEPOINT name`); the transaction stays open.
    pub fn rollback_to_savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("ROLLBACK TO SAVEPOINT {}", quote_identifier(name)))
            .map(|_| ())
    }

    /// Release (forget) a savepoint (`RELEASE SAVEPOINT name`), keeping its work.
    pub fn release_savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("RELEASE SAVEPOINT {}", quote_identifier(name)))
            .map(|_| ())
    }

    /// Subscribe to asynchronous notifications on `channel` (`LISTEN channel`). After this, any
    /// `NOTIFY channel` from another connection on the same database is delivered here; collect them
    /// with [`Connection::notifications`] / [`Connection::poll_notification`].
    pub fn listen(&mut self, channel: &str) -> Result<()> {
        self.query(&format!("LISTEN {}", quote_identifier(channel)))
            .map(|_| ())
    }

    /// Stop listening on `channel` (`UNLISTEN channel`).
    pub fn unlisten(&mut self, channel: &str) -> Result<()> {
        self.query(&format!("UNLISTEN {}", quote_identifier(channel)))
            .map(|_| ())
    }

    /// Stop listening on every channel (`UNLISTEN *`).
    pub fn unlisten_all(&mut self) -> Result<()> {
        self.query("UNLISTEN *").map(|_| ())
    }

    /// Send a notification on `channel` with an optional `payload` (`NOTIFY channel[, 'payload']`).
    /// Listeners on the same database — including this connection if it listens — receive it.
    pub fn notify(&mut self, channel: &str, payload: Option<&str>) -> Result<()> {
        let sql = match payload {
            Some(p) => format!("NOTIFY {}, {}", quote_identifier(channel), quote_literal(p)),
            None => format!("NOTIFY {}", quote_identifier(channel)),
        };
        self.query(&sql).map(|_| ())
    }

    /// Drain the notifications already received (buffered while reading other responses). Does not
    /// touch the socket — call [`Connection::poll_notification`] to wait for new ones.
    pub fn notifications(&mut self) -> Vec<Notification> {
        self.notifications.drain(..).collect()
    }

    /// Wait up to `timeout` for the next notification (a buffered one is returned immediately;
    /// `None` on timeout). `timeout` of `None` blocks until one arrives. Only meaningful after
    /// [`Connection::listen`]; the server pushes notifications while the connection is idle.
    pub fn poll_notification(&mut self, timeout: Option<Duration>) -> Result<Option<Notification>> {
        if let Some(note) = self.notifications.pop_front() {
            return Ok(Some(note));
        }
        self.stream.get_ref().tcp().set_read_timeout(timeout).ok();
        let result = protocol::read_frame(&mut self.stream);
        // Restore blocking reads so ordinary queries are unaffected.
        self.stream.get_ref().tcp().set_read_timeout(None).ok();
        match result {
            Ok(frame) if frame.kind == backend::NOTIFICATION_RESPONSE => {
                Ok(Some(decode_notification(&frame)?))
            }
            Ok(frame) => Err(Error::Protocol(format!(
                "unexpected message {:?} while polling for notifications",
                frame.kind as char
            ))),
            Err(Error::Io(e))
                if matches!(
                    e.kind(),
                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                ) =>
            {
                Ok(None)
            }
            Err(e) => Err(e),
        }
    }

    /// Run `work` inside a transaction: `BEGIN`, then `COMMIT` on `Ok` or `ROLLBACK` on `Err`
    /// (the error is returned). A rollback failure after a work error surfaces the original error.
    pub fn transaction<T, F>(&mut self, work: F) -> Result<T>
    where
        F: FnOnce(&mut Self) -> Result<T>,
    {
        self.begin()?;
        match work(self) {
            Ok(value) => {
                self.commit()?;
                Ok(value)
            }
            Err(e) => {
                let _ = self.rollback();
                Err(e)
            }
        }
    }

    /// This connection's `(pid, secret)` for out-of-band cancellation (§13).
    #[must_use]
    pub fn cancel_handle(&self) -> CancelHandle {
        CancelHandle {
            host: self.addr.0.clone(),
            port: self.addr.1,
            pid: self.backend_pid,
            secret: self.backend_secret,
        }
    }

    /// Politely close the connection (`Terminate`). Errors are ignored — the socket closes regardless.
    pub fn close(mut self) {
        let _ = self.send(&Writer::new().frame(frontend::TERMINATE));
    }

    /// Drain the backend response stream until `ReadyForQuery`, assembling a [`QueryResult`].
    fn collect(&mut self, _simple: bool) -> Result<QueryResult> {
        let mut acc = ResultAccumulator::new();
        loop {
            let frame = self.recv()?;
            if let Some(result) = acc.push(&frame)? {
                return Ok(result);
            }
        }
    }
}

/// Accumulates the backend response stream into a [`QueryResult`], one frame at a time. Shared by the
/// sync and async result loops so the wire-decoding logic lives in exactly one place.
pub(crate) struct ResultAccumulator {
    columns: Vec<String>,
    types: Vec<TypeTag>,
    col_arc: Arc<Vec<String>>,
    rows: Vec<Row>,
    tag: String,
    error: Option<Error>,
}

impl ResultAccumulator {
    pub(crate) fn new() -> Self {
        Self {
            columns: Vec::new(),
            types: Vec::new(),
            col_arc: Arc::new(Vec::new()),
            rows: Vec::new(),
            tag: String::new(),
            error: None,
        }
    }

    /// Fold one frame in. Returns `Ok(Some(result))` at `ReadyForQuery` (or the server error), or
    /// `Ok(None)` to keep reading.
    pub(crate) fn push(&mut self, frame: &Frame) -> Result<Option<QueryResult>> {
        match frame.kind {
            backend::ROW_DESCRIPTION => {
                let mut r = Reader::new(&frame.payload);
                let n = r.u16()? as usize;
                self.columns = Vec::with_capacity(n);
                self.types = Vec::with_capacity(n);
                for _ in 0..n {
                    self.columns.push(r.str()?);
                    self.types.push(TypeTag::Unknown);
                }
                self.col_arc = Arc::new(self.columns.clone());
            }
            backend::ROW_DESCRIPTION_TYPED => {
                let mut r = Reader::new(&frame.payload);
                let n = r.u16()? as usize;
                self.columns = Vec::with_capacity(n);
                self.types = Vec::with_capacity(n);
                for _ in 0..n {
                    self.columns.push(r.str()?);
                    self.types.push(TypeTag::from_byte(r.u8()?));
                }
                self.col_arc = Arc::new(self.columns.clone());
            }
            backend::DATA_ROW => {
                let mut r = Reader::new(&frame.payload);
                let fields = r.fields()?;
                let mut values = Vec::with_capacity(fields.len());
                for (i, field) in fields.into_iter().enumerate() {
                    let tag = self.types.get(i).copied().unwrap_or(TypeTag::Unknown);
                    values.push(Value::decode(field, tag)?);
                }
                self.rows.push(Row {
                    columns: Arc::clone(&self.col_arc),
                    values,
                });
            }
            backend::COMMAND_COMPLETE => {
                let mut r = Reader::new(&frame.payload);
                self.tag = r.str()?;
            }
            backend::ERROR => {
                self.error = Some(decode_error(&frame.payload));
            }
            backend::READY => {
                if let Some(e) = self.error.take() {
                    return Err(e);
                }
                return Ok(Some(QueryResult {
                    columns: std::mem::take(&mut self.columns),
                    column_types: std::mem::take(&mut self.types),
                    rows: std::mem::take(&mut self.rows),
                    tag: std::mem::take(&mut self.tag),
                }));
            }
            // Extended-query control replies: acknowledged, no state to add.
            backend::PARSE_COMPLETE
            | backend::BIND_COMPLETE
            | backend::CLOSE_COMPLETE
            | backend::PARAMETER_DESCRIPTION
            | backend::NO_DATA
            | backend::PORTAL_SUSPENDED => {}
            other => {
                return Err(Error::Protocol(format!(
                    "unexpected message {:?} in result stream",
                    other as char
                )))
            }
        }
        Ok(None)
    }
}

/// A connection's cancellation key, usable from another connection to abort an in-flight statement.
#[derive(Debug, Clone)]
pub struct CancelHandle {
    host: String,
    port: u16,
    pid: u32,
    secret: u32,
}

impl CancelHandle {
    #[cfg(feature = "async")]
    pub(crate) fn new(host: String, port: u16, pid: u32, secret: u32) -> Self {
        Self {
            host,
            port,
            pid,
            secret,
        }
    }

    /// Open a fresh connection and request cancellation of the target's in-flight statement (§13).
    /// Best-effort: a wrong/stale key cancels nothing.
    pub fn cancel(&self) -> Result<()> {
        let mut stream = TcpStream::connect((self.host.as_str(), self.port))?;
        let frame = Writer::new()
            .u32(self.pid)
            .u32(self.secret)
            .frame(frontend::CANCEL);
        std::io::Write::write_all(&mut stream, &frame)?;
        std::io::Write::flush(&mut stream)?;
        Ok(())
    }
}

/// Wrap a connected TCP socket in the configured [`Transport`] — plaintext, or (with the `tls`
/// feature) an implicit TLS 1.3 client stream handshaked before any frame is sent (§5.1).
fn build_transport(tcp: TcpStream, config: &Config) -> Result<Transport> {
    let Some(tls) = &config.tls else {
        return Ok(Transport::Plain(tcp));
    };
    #[cfg(feature = "tls")]
    {
        let client_config = crate::tls::client_config(tls)?;
        let server_name = crate::tls::server_name(&config.host)?;
        let conn = rustls::ClientConnection::new(client_config, server_name)
            .map_err(|e| Error::Protocol(format!("TLS handshake setup: {e}")))?;
        Ok(Transport::Tls(Box::new(rustls::StreamOwned::new(
            conn, tcp,
        ))))
    }
    #[cfg(not(feature = "tls"))]
    {
        let _ = tls;
        Err(Error::Config(
            "TLS requested but the `tls` feature is not enabled".to_owned(),
        ))
    }
}

/// Quotes a savepoint identifier so an arbitrary name is emitted safely. The server folds a quoted
/// identifier to its literal value, so the same quoting matches across `SAVEPOINT`/`ROLLBACK
/// TO`/`RELEASE`.
pub(crate) fn quote_identifier(name: &str) -> String {
    format!("\"{}\"", name.replace('"', "\"\""))
}

/// Quotes a string as a SQL literal (`'…'`, doubling embedded `'`) — used for a `NOTIFY` payload so
/// arbitrary text is sent safely.
pub(crate) fn quote_literal(text: &str) -> String {
    format!("'{}'", text.replace('\'', "''"))
}

/// Decode a `NotificationResponse` frame (`[pid:u32][channel:str][payload:str]`).
pub(crate) fn decode_notification(frame: &Frame) -> Result<Notification> {
    let mut r = Reader::new(&frame.payload);
    Ok(Notification {
        pid: r.u32()?,
        channel: r.str()?,
        payload: r.str()?,
    })
}

pub(crate) fn bind_frame(portal: &str, statement: &str, params: &[Option<Vec<u8>>]) -> Vec<u8> {
    // [portal][statement][params:Fields][fmt_count:u16=0] — empty format list ⇒ all columns text.
    Writer::new()
        .str(portal)
        .str(statement)
        .fields(params)
        .u16(0)
        .frame(frontend::BIND)
}

pub(crate) fn describe_portal(portal: &str) -> Vec<u8> {
    Writer::new().u8(b'P').str(portal).frame(frontend::DESCRIBE)
}

pub(crate) fn execute_portal(portal: &str, max_rows: u32) -> Vec<u8> {
    Writer::new()
        .str(portal)
        .u32(max_rows)
        .frame(frontend::EXECUTE)
}

/// Parse the row count from a `COPY <n>` command tag (`0` if absent/unparseable).
pub(crate) fn copy_count(tag: &str) -> u64 {
    tag.rsplit(' ')
        .next()
        .and_then(|n| n.parse().ok())
        .unwrap_or(0)
}

pub(crate) fn decode_error(payload: &[u8]) -> Error {
    let mut r = Reader::new(payload);
    let code = r.str().unwrap_or_else(|_| "XX000".to_owned());
    let message = r
        .str()
        .unwrap_or_else(|_| "unknown server error".to_owned());
    Error::Server { code, message }
}