laminardb/

pgwire.rs

//! Postgres wire endpoint. Trust by default; MD5 with `pgwire_users`;
//! TLS with `pgwire_tls_cert` + `pgwire_tls_key`. Non-loopback binds
//! require `pgwire_allow_remote = true`.

use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use async_trait::async_trait;
use futures::{stream, Sink, StreamExt};
use laminar_sql::parser::{
    parse_streaming_sql, ShowCommand, StreamingStatement, SubscribeStatement,
};
use pgwire::api::auth::md5pass::{hash_md5_password, Md5PasswordAuthStartupHandler};
use pgwire::api::auth::noop::NoopStartupHandler;
use pgwire::api::auth::{
    AuthSource, DefaultServerParameterProvider, LoginInfo, Password, StartupHandler,
};
use pgwire::api::portal::{Format, Portal};
use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
use pgwire::api::stmt::QueryParser;
use pgwire::api::store::PortalStore;
use pgwire::api::{ClientInfo, ClientPortalStore, PgWireServerHandlers, Type};
use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage};
use pgwire::tokio::process_socket;
use sqlparser::ast::{
    CloseCursor, Expr, FetchDirection, FunctionArguments, SelectItem, Set, SetExpr, Statement,
    Value as AstValue,
};
use tokio::net::TcpListener;
use tokio::sync::Mutex as TokioMutex;
use tracing::{info, warn};

use laminar_db::subscription::{PortalFrame, SubscribeStart, SubscriptionPortal};
use laminar_db::LaminarDB;

use crate::config::Secret;
use crate::server::ServerError;

pub struct LaminarPgwireHandler {
    db: Arc<LaminarDB>,
    /// Per-peer SimpleQuery cursor map. Entries are evicted at the start of
    /// every `do_query` call once their cursors are dead and no transaction
    /// is open — pgwire 0.39 doesn't give us a connection-close hook, so this
    /// is the cheapest way to keep stale state from leaking on port reuse.
    connections: parking_lot::Mutex<HashMap<SocketAddr, Arc<ConnState>>>,
}

impl LaminarPgwireHandler {
    fn new(db: Arc<LaminarDB>) -> Self {
        Self {
            db,
            connections: parking_lot::Mutex::new(HashMap::new()),
        }
    }

    fn conn_state(&self, peer: SocketAddr) -> Arc<ConnState> {
        let mut guard = self.connections.lock();
        Arc::clone(
            guard
                .entry(peer)
                .or_insert_with(|| Arc::new(ConnState::default())),
        )
    }

    fn evict_idle_peer(&self, peer: SocketAddr) {
        let mut guard = self.connections.lock();
        if let Some(state) = guard.get(&peer) {
            if state.prune_dead_and_check_idle() {
                guard.remove(&peer);
            }
        }
    }
}

#[async_trait]
impl NoopStartupHandler for LaminarPgwireHandler {
    async fn post_startup<C>(
        &self,
        client: &mut C,
        _message: PgWireFrontendMessage,
    ) -> PgWireResult<()>
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        info!(peer = %client.socket_addr(), "pgwire client connected");
        Ok(())
    }
}

#[async_trait]
impl SimpleQueryHandler for LaminarPgwireHandler {
    async fn do_query<C>(&self, client: &mut C, query: &str) -> PgWireResult<Vec<Response>>
    where
        C: ClientInfo + ClientPortalStore + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::PortalStore: PortalStore,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        if query.trim().is_empty() {
            return Ok(vec![Response::EmptyQuery]);
        }
        let peer = client.socket_addr();
        self.evict_idle_peer(peer);
        let stmts = parse_streaming_sql(query)
            .map_err(|e| user_error("42601", format!("parse error: {e}")))?;

        // SUBSCRIBE owns the socket for its lifetime; it can't share a
        // simple-query batch with earlier or later statements. Reject
        // up front so trailing statements aren't silently dropped.
        if stmts.len() > 1
            && stmts
                .iter()
                .any(|s| matches!(s, StreamingStatement::Subscribe(_)))
        {
            return Err(user_error(
                "0A000",
                "SUBSCRIBE must be the only statement in a simple query",
            ));
        }

        let mut out = Vec::with_capacity(stmts.len());
        for stmt in stmts {
            out.push(match stmt {
                StreamingStatement::Subscribe(s) => {
                    let portal = open_portal_for_subscribe(&self.db, &s).await?;
                    // Simple query is always text (no Bind result format).
                    stream_subscribe_flushing(client, portal, true, None).await?;
                    return Ok(Vec::new());
                }
                StreamingStatement::Show(cmd) => {
                    engine_metadata_response(&self.db, &show_sql(&cmd)).await?
                }
                StreamingStatement::DeclareCursorForSubscribe {
                    name, subscribe, ..
                } => {
                    let state = self.conn_state(peer);
                    handle_declare_cursor(&self.db, &state, &name.value, *subscribe).await?
                }
                StreamingStatement::Standard(s) => {
                    let state = self.conn_state(peer);
                    standard_or_cursor_response(&self.db, &state, *s)?
                }
                other => {
                    return Err(user_error(
                        "0A000",
                        format!("not supported on pgwire (use HTTP /api/v1/sql): {other:?}"),
                    ));
                }
            });
        }
        Ok(out)
    }
}

async fn open_portal_for_subscribe(
    db: &LaminarDB,
    s: &SubscribeStatement,
) -> PgWireResult<SubscriptionPortal> {
    let name = s.name.to_string();
    let start = match s.as_of_epoch {
        Some(n) => SubscribeStart::AsOfEpoch(n),
        None => SubscribeStart::Tail,
    };
    db.open_subscription(&name, s.filter_sql.as_deref(), start)
        .await
        .map_err(|e| user_error("42P01", format!("SUBSCRIBE '{name}': {e}")))
}

/// Stream a SUBSCRIBE, flushing the `Sink` after every batch.
///
/// Workaround: pgwire `feed()`s `DataRow`s and only flushes at
/// end-of-response (never, for an unbounded SUBSCRIBE) or at its ~8 KB
/// buffer, so a sparse stream stalls. Per-batch flush is unconditional
/// (fine — batches amortise; not per-row). Retire when pgwire flushes
/// streaming responses upstream. Both paths need it (psql=simple,
/// psycopg/JDBC=extended).
///
/// `send_row_desc`: simple query carries `RowDescription`; extended
/// already sent it via `Describe` (caller returns `Response::Execution`
/// for `CommandComplete`). Returns `Ok(())` only on pump exit.
async fn stream_subscribe_flushing<C>(
    client: &mut C,
    mut portal: SubscriptionPortal,
    send_row_desc: bool,
    result_format: Option<&Format>,
) -> PgWireResult<()>
where
    C: Sink<PgWireBackendMessage> + Unpin + Send,
    C::Error: Debug,
    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
{
    use futures::SinkExt;

    let schema = portal.schema();
    // Honour the extended client's per-column binary/text choice; pgwire's
    // `Describe` advertised it, so the `DataRow` encoding must match.
    let fields = std::sync::Arc::new(field_infos(&schema, result_format));

    if send_row_desc {
        // Equivalent to pgwire's crate-private `into_row_description`.
        let row_desc =
            pgwire::messages::data::RowDescription::new(fields.iter().map(Into::into).collect());
        client
            .feed(PgWireBackendMessage::RowDescription(row_desc))
            .await?;
        client.flush().await?;
    }

    let mut rows: usize = 0;
    loop {
        match portal.next_frame().await {
            Some(PortalFrame::Batch(b)) if b.num_rows() > 0 => {
                for row in encode_batch(&b, &fields) {
                    client.feed(PgWireBackendMessage::DataRow(row?)).await?;
                    rows += 1;
                }
                client.flush().await?;
            }
            Some(PortalFrame::Batch(_)) => {}
            // Checkpoint barriers have no Postgres wire representation.
            Some(PortalFrame::Barrier { .. }) => {}
            Some(PortalFrame::Lagged(n)) => {
                return Err(user_error(
                    "54000",
                    format!("subscription lagged: skipped {n} messages, terminating"),
                ));
            }
            None => {
                // Simple query owns the whole response, so emit
                // CommandComplete here. The extended path lets pgwire
                // emit it from the returned `Response::Execution`.
                if send_row_desc {
                    let tag = Tag::new("SUBSCRIBE").with_rows(rows);
                    client
                        .feed(PgWireBackendMessage::CommandComplete(tag.into()))
                        .await?;
                    client.flush().await?;
                }
                return Ok(());
            }
        }
    }
}

/// Wrap a `SubscriptionPortal` in a pgwire `Response::Query` so the
/// framework can chunk via `Execute(max_rows)` and emit PortalSuspended
/// automatically. Used by the chunked extended-query path.
fn subscription_query_response(
    portal: SubscriptionPortal,
    result_format: Option<&Format>,
) -> Response {
    use futures::stream;
    let schema = portal.schema();
    let fields = Arc::new(field_infos(&schema, result_format));
    struct State {
        portal: SubscriptionPortal,
        fields: Arc<Vec<FieldInfo>>,
        pending: VecDeque<PgWireResult<pgwire::messages::data::DataRow>>,
    }
    let init = State {
        portal,
        fields: Arc::clone(&fields),
        pending: VecDeque::new(),
    };
    let row_stream = stream::unfold(init, |mut s| async move {
        loop {
            if let Some(row) = s.pending.pop_front() {
                return Some((row, s));
            }
            match s.portal.next_frame().await {
                None => return None,
                Some(PortalFrame::Batch(b)) if b.num_rows() > 0 => {
                    s.pending.extend(encode_batch(&b, &s.fields));
                }
                Some(PortalFrame::Batch(_)) | Some(PortalFrame::Barrier { .. }) => {}
                Some(PortalFrame::Lagged(n)) => {
                    let err = user_error(
                        "54000",
                        format!("subscription lagged: skipped {n} messages, terminating"),
                    );
                    return Some((Err(err), s));
                }
            }
        }
    });
    let mut resp = QueryResponse::new(fields, row_stream);
    resp.set_command_tag("SUBSCRIBE");
    Response::Query(resp)
}

/// State that shares the cursor's lifetime: the portal, the leftover-row
/// buffer, and the exhausted flag. Held by `Arc` so a row stream can keep
/// reading after `ConnState::get` returns.
struct CursorInner {
    /// Tokio mutex because a FETCH stream holds it across `await` while
    /// pulling frames.
    portal: TokioMutex<SubscriptionPortal>,
    /// Rows encoded from a prior frame that the previous FETCH didn't
    /// consume. Without this, a multi-row batch + `FETCH 1` would drop the
    /// leftover rows when the response stream ends.
    pending: parking_lot::Mutex<VecDeque<PgWireResult<pgwire::messages::data::DataRow>>>,
    /// Flipped when the pump emits `None` or `Lagged`. The next `evict_idle_peer`
    /// pass reaps the cursor.
    exhausted: AtomicBool,
}

#[derive(Clone)]
struct ActiveCursor {
    inner: Arc<CursorInner>,
    schema: arrow_schema::SchemaRef,
}

#[derive(Default)]
struct ConnState {
    cursors: parking_lot::Mutex<HashMap<String, ActiveCursor>>,
    in_tx: AtomicBool,
}

impl ConnState {
    /// Cursor names follow PG identifier folding: unquoted → lowercase. We
    /// don't track `quote_style`, so quoted-mixed-case cursors collapse too —
    /// good enough for the `\set FETCH_COUNT` case this targets.
    fn key(name: &str) -> String {
        name.to_ascii_lowercase()
    }

    fn insert(&self, name: &str, cursor: ActiveCursor) {
        self.cursors.lock().insert(Self::key(name), cursor);
    }

    fn contains(&self, name: &str) -> bool {
        self.cursors.lock().contains_key(&Self::key(name))
    }

    fn get(&self, name: &str) -> Option<ActiveCursor> {
        self.cursors.lock().get(&Self::key(name)).cloned()
    }

    fn remove(&self, name: &str) -> bool {
        self.cursors.lock().remove(&Self::key(name)).is_some()
    }

    fn drop_all(&self) {
        self.cursors.lock().clear();
    }

    /// Drop dead cursors and report whether the connection is now idle
    /// (no cursors, no transaction). Single inner-lock acquisition.
    fn prune_dead_and_check_idle(&self) -> bool {
        let mut cursors = self.cursors.lock();
        cursors.retain(|_, c| !c.inner.exhausted.load(Ordering::Acquire));
        cursors.is_empty() && !self.in_tx.load(Ordering::Acquire)
    }
}

/// Open a SUBSCRIBE behind a cursor name. Rejects with 42P03 if the name is
/// already in use on this connection (matches PG; user must `CLOSE` first).
async fn handle_declare_cursor(
    db: &LaminarDB,
    state: &ConnState,
    cursor_name: &str,
    subscribe: SubscribeStatement,
) -> PgWireResult<Response> {
    if state.contains(cursor_name) {
        return Err(user_error(
            "42P03",
            format!("cursor \"{cursor_name}\" already exists"),
        ));
    }
    let portal = open_portal_for_subscribe(db, &subscribe).await?;
    let schema = portal.schema();
    state.insert(
        cursor_name,
        ActiveCursor {
            inner: Arc::new(CursorInner {
                portal: TokioMutex::new(portal),
                pending: parking_lot::Mutex::new(VecDeque::new()),
                exhausted: AtomicBool::new(false),
            }),
            schema,
        },
    );
    Ok(Response::Execution(Tag::new("DECLARE CURSOR")))
}

/// `FETCH NEXT` and bare `FETCH FORWARD` map to a single row, matching PG.
fn fetch_direction_count(dir: &FetchDirection) -> PgWireResult<FetchTarget> {
    match dir {
        FetchDirection::Next | FetchDirection::Forward { limit: None } => Ok(FetchTarget::Count(1)),
        FetchDirection::Count { limit } | FetchDirection::Forward { limit: Some(limit) } => {
            value_to_u64(limit).map(FetchTarget::Count)
        }
        FetchDirection::All | FetchDirection::ForwardAll => Ok(FetchTarget::All),
        FetchDirection::Prior
        | FetchDirection::First
        | FetchDirection::Last
        | FetchDirection::Absolute { .. }
        | FetchDirection::Relative { .. }
        | FetchDirection::Backward { .. }
        | FetchDirection::BackwardAll => Err(user_error(
            "0A000",
            "FETCH direction not supported (SUBSCRIBE cursors are forward-only): use FORWARD or NEXT",
        )),
    }
}

#[derive(Copy, Clone)]
enum FetchTarget {
    Count(u64),
    All,
}

fn value_to_u64(v: &AstValue) -> PgWireResult<u64> {
    match v {
        AstValue::Number(n, _) => n
            .parse::<u64>()
            .map_err(|_| user_error("22023", format!("invalid FETCH count: {n}"))),
        other => Err(user_error(
            "22023",
            format!("FETCH count must be an integer, got {other}"),
        )),
    }
}

fn handle_fetch(
    state: &ConnState,
    cursor_name: &str,
    target: FetchTarget,
) -> PgWireResult<Response> {
    let cursor = state
        .get(cursor_name)
        .ok_or_else(|| user_error("34000", format!("cursor \"{cursor_name}\" does not exist")))?;
    Ok(fetch_response(cursor, target))
}

fn handle_close(state: &ConnState, cursor: &CloseCursor) -> PgWireResult<Response> {
    match cursor {
        CloseCursor::All => {
            state.drop_all();
            Ok(Response::Execution(Tag::new("CLOSE CURSOR ALL")))
        }
        CloseCursor::Specific { name } => {
            if state.remove(&name.value) {
                Ok(Response::Execution(Tag::new("CLOSE CURSOR")))
            } else {
                Err(user_error(
                    "34000",
                    format!("cursor \"{}\" does not exist", name.value),
                ))
            }
        }
    }
}

/// Wraps the original `standard_response` and intercepts cursor / transaction
/// statements that need ConnState. Anything else falls through to the
/// existing handler unchanged.
fn standard_or_cursor_response(
    db: &LaminarDB,
    state: &ConnState,
    stmt: Statement,
) -> PgWireResult<Response> {
    match stmt {
        Statement::StartTransaction { .. } => {
            state.in_tx.store(true, Ordering::Release);
            Ok(Response::Execution(Tag::new("BEGIN")))
        }
        Statement::Commit { .. } => {
            state.drop_all();
            state.in_tx.store(false, Ordering::Release);
            Ok(Response::Execution(Tag::new("COMMIT")))
        }
        Statement::Rollback { .. } => {
            state.drop_all();
            state.in_tx.store(false, Ordering::Release);
            Ok(Response::Execution(Tag::new("ROLLBACK")))
        }
        Statement::Fetch {
            ref name,
            ref direction,
            ..
        } => {
            let target = fetch_direction_count(direction)?;
            handle_fetch(state, &name.value, target)
        }
        Statement::Close { ref cursor } => handle_close(state, cursor),
        Statement::Declare { .. } => Err(user_error(
            "0A000",
            "DECLARE on pgwire only supports CURSOR FOR SUBSCRIBE …",
        )),
        other => standard_response(db, other),
    }
}

/// Connection-setup statements: transaction control, `SET`, and a tiny set
/// of catalog probes drivers send during handshake. Anything DDL/DML hits
/// the "use HTTP" error.
fn standard_response(db: &LaminarDB, stmt: Statement) -> PgWireResult<Response> {
    match stmt {
        Statement::StartTransaction { .. } => Ok(Response::Execution(Tag::new("BEGIN"))),
        Statement::Commit { .. } => Ok(Response::Execution(Tag::new("COMMIT"))),
        Statement::Rollback { .. } => Ok(Response::Execution(Tag::new("ROLLBACK"))),
        Statement::Set(s) => apply_set(db, s),
        Statement::Query(q) => driver_select_response(*q),
        Statement::Insert { .. }
        | Statement::Update { .. }
        | Statement::Delete { .. }
        | Statement::CreateTable { .. }
        | Statement::CreateView { .. }
        | Statement::Drop { .. } => Err(user_error(
            "0A000",
            "DDL/DML is not supported on pgwire; use HTTP /api/v1/sql",
        )),
        other => Err(user_error(
            "0A000",
            format!("not supported on pgwire: {other}"),
        )),
    }
}

/// Handle the `SELECT`s drivers issue at connect time. Single literal,
/// `SELECT version()`, and `SELECT current_schema()` are answered inline.
/// Anything else is rejected — real queries belong on `/api/v1/sql`.
fn driver_select_response(query: sqlparser::ast::Query) -> PgWireResult<Response> {
    let SetExpr::Select(select) = *query.body else {
        return Err(unsupported_select());
    };
    if select.projection.len() != 1 || !select.from.is_empty() || select.selection.is_some() {
        return Err(unsupported_select());
    }
    let SelectItem::UnnamedExpr(expr) = &select.projection[0] else {
        return Err(unsupported_select());
    };

    match expr {
        Expr::Value(v) => match &v.value {
            sqlparser::ast::Value::Number(n, _) => {
                let parsed: i32 = n.parse().map_err(|_| unsupported_select())?;
                Ok(text_response("?column?", Type::INT4, parsed.to_string()))
            }
            sqlparser::ast::Value::SingleQuotedString(s) => {
                Ok(text_response("?column?", Type::VARCHAR, s.clone()))
            }
            _ => Err(unsupported_select()),
        },
        Expr::Function(func) => {
            // Only no-arg builtins. `func.args` is `FunctionArguments::None`
            // for `func()`, the only shape we accept.
            if !matches!(func.args, FunctionArguments::List(ref a) if a.args.is_empty())
                && !matches!(func.args, FunctionArguments::None)
            {
                return Err(unsupported_select());
            }
            let name = func.name.to_string().to_ascii_lowercase();
            let (col, ty, value) = match name.as_str() {
                "version" | "pg_catalog.version" => (
                    "version",
                    Type::VARCHAR,
                    format!("LaminarDB {} on pgwire", env!("CARGO_PKG_VERSION")),
                ),
                "current_schema" | "pg_catalog.current_schema" => {
                    ("current_schema", Type::VARCHAR, "public".to_string())
                }
                "current_database" | "pg_catalog.current_database" => {
                    ("current_database", Type::VARCHAR, "laminar".to_string())
                }
                "current_user" | "session_user" | "user" => {
                    ("current_user", Type::VARCHAR, "laminar".to_string())
                }
                _ => return Err(unsupported_select()),
            };
            Ok(text_response(col, ty, value))
        }
        _ => Err(unsupported_select()),
    }
}

fn unsupported_select() -> PgWireError {
    user_error(
        "0A000",
        "pgwire SELECT is limited to literals and connect-time builtins; use HTTP /api/v1/sql",
    )
}

/// `SET` handling. We thread plain `SET name = value` to the engine's
/// session-property store, and refuse `SET TRANSACTION`-class statements
/// since we don't honor isolation levels.
fn apply_set(db: &LaminarDB, set: Set) -> PgWireResult<Response> {
    match set {
        Set::SingleAssignment {
            variable, values, ..
        } => {
            let key = variable.to_string();
            let value = values
                .first()
                .map(ToString::to_string)
                .unwrap_or_default()
                .trim_matches('\'')
                .to_string();
            db.set_session_property(&key, &value);
            Ok(Response::Execution(Tag::new("SET")))
        }
        // Refuse anything that implies semantics we do not provide.
        Set::SetTransaction { .. } => Err(user_error(
            "0A000",
            "SET TRANSACTION is not supported (no transactional semantics)",
        )),
        // Lenient pass-through for the harmless catalog-style SETs drivers
        // issue (NAMES, TIME ZONE, ROLE...). We don't honor them, but failing
        // the connection is worse than silently accepting.
        _ => Ok(Response::Execution(Tag::new("SET"))),
    }
}

fn user_error(code: &str, msg: impl Into<String>) -> PgWireError {
    PgWireError::UserError(Box::new(ErrorInfo::new(
        "ERROR".into(),
        code.into(),
        msg.into(),
    )))
}

/// Reconstruct a single SHOW statement from the parsed variant. Used by the
/// pgwire dispatcher so a multi-statement query (`SHOW SOURCES; SHOW SINKS`)
/// re-executes only the matching statement, not the whole input string.
fn show_sql(cmd: &ShowCommand) -> String {
    match cmd {
        ShowCommand::Sources => "SHOW SOURCES".into(),
        ShowCommand::Sinks => "SHOW SINKS".into(),
        ShowCommand::Queries => "SHOW QUERIES".into(),
        ShowCommand::MaterializedViews => "SHOW MATERIALIZED VIEWS".into(),
        ShowCommand::Streams => "SHOW STREAMS".into(),
        ShowCommand::Tables => "SHOW TABLES".into(),
        ShowCommand::CheckpointStatus => "SHOW CHECKPOINT STATUS".into(),
        ShowCommand::CreateSource { name } => format!("SHOW CREATE SOURCE {name}"),
        ShowCommand::CreateSink { name } => format!("SHOW CREATE SINK {name}"),
    }
}

/// Run a SHOW through the engine and stream its `RecordBatch` to the wire.
async fn engine_metadata_response(db: &LaminarDB, sql: &str) -> PgWireResult<Response> {
    use laminar_db::ExecuteResult;
    let result = db
        .execute(sql)
        .await
        .map_err(|e| user_error("XX000", e.to_string()))?;
    let ExecuteResult::Metadata(batch) = result else {
        return Err(user_error("XX000", "SHOW did not return metadata"));
    };
    Ok(record_batch_response(batch))
}

/// Single-row `text` response with one column.
fn text_response(col: &str, ty: Type, value: String) -> Response {
    let schema = Arc::new(vec![FieldInfo::new(
        col.into(),
        None,
        None,
        ty,
        FieldFormat::Text,
    )]);
    let schema_for_row = Arc::clone(&schema);
    let row_stream = stream::iter(std::iter::once(Ok::<_, PgWireError>(()))).map(move |_| {
        let mut enc = DataRowEncoder::new(Arc::clone(&schema_for_row));
        enc.encode_field(&Some(value.as_str()))?;
        Ok(enc.take_row())
    });
    Response::Query(QueryResponse::new(schema, row_stream))
}

fn record_batch_response(batch: arrow_array::RecordBatch) -> Response {
    let fields = Arc::new(field_infos(&batch.schema(), None));
    let nrows = batch.num_rows();

    // Encode rows eagerly: SHOW outputs are tiny and this avoids the
    // !Send formatter dance.
    let mut rows = Vec::with_capacity(nrows);
    {
        let opts = arrow_cast::display::FormatOptions::default();
        let formatters: Vec<_> = batch
            .columns()
            .iter()
            .map(|c| arrow_cast::display::ArrayFormatter::try_new(c.as_ref(), &opts))
            .collect::<Result<_, _>>()
            .unwrap_or_default();
        for row in 0..nrows {
            rows.push(encode_row(&batch, row, &fields, &formatters));
        }
    }

    let row_stream = stream::iter(rows);
    Response::Query(QueryResponse::new(fields, row_stream))
}

/// Strict-PG FETCH: blocks until `target` rows are produced, the pump exits,
/// or the broadcast lags. Lag/exit flips `cursor.inner.exhausted` so the next
/// `evict_idle_peer` reaps the cursor. Text format only; SimpleQuery has no
/// binary. Leftover rows from a multi-row frame stay in `cursor.inner.pending`
/// so successive FETCHes consume the frame in order.
fn fetch_response(cursor: ActiveCursor, target: FetchTarget) -> Response {
    let fields = Arc::new(field_infos(&cursor.schema, None));
    let remaining = match target {
        FetchTarget::Count(n) => Some(n),
        FetchTarget::All => None,
    };

    struct State {
        cursor: ActiveCursor,
        fields: Arc<Vec<FieldInfo>>,
        remaining: Option<u64>,
    }

    let init = State {
        cursor,
        fields: Arc::clone(&fields),
        remaining,
    };

    let row_stream = stream::unfold(init, |mut s| async move {
        loop {
            if matches!(s.remaining, Some(0)) {
                return None;
            }
            // Pending rows from a prior FETCH come out first. Anything left
            // here when remaining hits 0 stays for the next call.
            let popped = s.cursor.inner.pending.lock().pop_front();
            if let Some(row) = popped {
                if let Some(n) = s.remaining.as_mut() {
                    *n = n.saturating_sub(1);
                }
                return Some((row, s));
            }
            if s.cursor.inner.exhausted.load(Ordering::Acquire) {
                return None;
            }

            let next = s.cursor.inner.portal.lock().await.next_frame().await;
            match next {
                None => {
                    s.cursor.inner.exhausted.store(true, Ordering::Release);
                    return None;
                }
                Some(PortalFrame::Batch(b)) if b.num_rows() > 0 => {
                    let encoded = encode_batch(&b, &s.fields);
                    s.cursor.inner.pending.lock().extend(encoded);
                }
                Some(PortalFrame::Batch(_)) => {}
                Some(PortalFrame::Barrier { .. }) => {
                    // Same as portal_to_response: PG has no out-of-band marker.
                }
                Some(PortalFrame::Lagged(n)) => {
                    s.cursor.inner.exhausted.store(true, Ordering::Release);
                    let err = user_error(
                        "54000",
                        format!("subscription lagged: skipped {n} messages, terminating cursor"),
                    );
                    return Some((Err(err), s));
                }
            }
        }
    });
    Response::Query(QueryResponse::new(fields, row_stream))
}

fn encode_batch(
    batch: &arrow_array::RecordBatch,
    fields: &Arc<Vec<FieldInfo>>,
) -> Vec<PgWireResult<pgwire::messages::data::DataRow>> {
    let opts = arrow_cast::display::FormatOptions::default();
    let formatters: Vec<_> = match batch
        .columns()
        .iter()
        .map(|c| arrow_cast::display::ArrayFormatter::try_new(c.as_ref(), &opts))
        .collect::<Result<_, _>>()
    {
        Ok(f) => f,
        Err(e) => {
            return vec![Err(user_error("XX000", format!("format column: {e}")))];
        }
    };
    (0..batch.num_rows())
        .map(|row| encode_row(batch, row, fields, &formatters))
        .collect()
}

/// Build pgwire `FieldInfo`s from an Arrow schema. `result_format` (from a
/// `Bind`) sets per-column text/binary; `None` defaults all-text.
fn field_infos(schema: &arrow_schema::Schema, result_format: Option<&Format>) -> Vec<FieldInfo> {
    schema
        .fields()
        .iter()
        .enumerate()
        .map(|(i, f)| {
            let format = result_format.map_or(FieldFormat::Text, |rf| rf.format_for(i));
            FieldInfo::new(
                f.name().clone(),
                None,
                None,
                arrow_to_pg_type(f.data_type()),
                format,
            )
        })
        .collect()
}

fn encode_row(
    batch: &arrow_array::RecordBatch,
    row: usize,
    fields: &Arc<Vec<FieldInfo>>,
    formatters: &[arrow_cast::display::ArrayFormatter<'_>],
) -> PgWireResult<pgwire::messages::data::DataRow> {
    let mut enc = DataRowEncoder::new(Arc::clone(fields));
    for (i, col) in batch.columns().iter().enumerate() {
        let info = &fields[i];
        match info.format() {
            FieldFormat::Text => encode_field_text(&mut enc, col.as_ref(), row, &formatters[i])?,
            FieldFormat::Binary => encode_field_binary(&mut enc, col.as_ref(), row, info.name())?,
        }
    }
    Ok(enc.take_row())
}

fn encode_field_text(
    enc: &mut DataRowEncoder,
    col: &dyn arrow_array::Array,
    row: usize,
    formatter: &arrow_cast::display::ArrayFormatter<'_>,
) -> PgWireResult<()> {
    use arrow_schema::DataType;
    if col.is_null(row) {
        return enc.encode_field(&None::<&str>);
    }
    // A TEXT[] column must serialize as a Postgres array literal `{..}`, not
    // Arrow's `[..]` display, so text-mode clients parse it as an array.
    if matches!(col.data_type(), DataType::List(f) if matches!(f.data_type(), DataType::Utf8 | DataType::LargeUtf8))
    {
        return enc.encode_field(&Some(pg_text_array_literal(&list_text_elements(col, row))));
    }
    enc.encode_field(&Some(formatter.value(row).to_string()))
}

/// Owned elements of a `List<Utf8|LargeUtf8>` row, NULLs preserved.
fn list_text_elements(col: &dyn arrow_array::Array, row: usize) -> Vec<Option<String>> {
    use arrow_array::cast::AsArray;
    use arrow_array::Array;
    use arrow_schema::DataType;
    let values = col.as_list::<i32>().value(row);
    if matches!(values.data_type(), DataType::LargeUtf8) {
        let s = values.as_string::<i64>();
        (0..s.len())
            .map(|i| (!s.is_null(i)).then(|| s.value(i).to_owned()))
            .collect()
    } else {
        let s = values.as_string::<i32>();
        (0..s.len())
            .map(|i| (!s.is_null(i)).then(|| s.value(i).to_owned()))
            .collect()
    }
}

/// Postgres `text[]` literal, e.g. `{"en","ja",NULL}`. Every element is quoted
/// (NULL excepted) so commas/braces/quotes in values are unambiguous.
fn pg_text_array_literal(elements: &[Option<String>]) -> String {
    let mut out = String::from("{");
    for (i, elem) in elements.iter().enumerate() {
        if i > 0 {
            out.push(',');
        }
        match elem {
            None => out.push_str("NULL"),
            Some(v) => {
                out.push('"');
                for ch in v.chars() {
                    if ch == '"' || ch == '\\' {
                        out.push('\\');
                    }
                    out.push(ch);
                }
                out.push('"');
            }
        }
    }
    out.push('}');
    out
}

/// Binary-encode a single Arrow value via `postgres-types` `ToSql`.
///
/// Coverage: Int{8,16,32,64}, UInt{8,16,32,64}, Float{32,64}, Bool,
/// Utf8/LargeUtf8, Timestamp (any unit, naive), Date32, Date64, and
/// `List<Utf8>` (as `text[]`). UInt64 is widened to INT8 with saturation
/// since Postgres has no unsigned 64. Any other column type yields `0A000`.
fn encode_field_binary(
    enc: &mut DataRowEncoder,
    col: &dyn arrow_array::Array,
    row: usize,
    name: &str,
) -> PgWireResult<()> {
    use arrow_array::{cast::AsArray, types::*};
    use arrow_schema::DataType;

    if col.is_null(row) {
        return enc.encode_field(&None::<&str>);
    }

    // Pull the typed Arrow value and pass it to `DataRowEncoder`, which
    // calls `postgres-types::ToSql` for the wire format. The `as $cast`
    // arm widens a narrower Arrow int to the matching Postgres OID (see
    // `arrow_to_pg_type`); only lossless `From` casts go through here.
    macro_rules! prim {
        ($ty:ty as $cast:ty) => {
            enc.encode_field(&Some(<$cast>::from(col.as_primitive::<$ty>().value(row))))
        };
        ($ty:ty) => {
            enc.encode_field(&Some(col.as_primitive::<$ty>().value(row)))
        };
    }

    match col.data_type() {
        DataType::Int8 => prim!(Int8Type as i32),
        DataType::Int16 => prim!(Int16Type as i32),
        DataType::Int32 => prim!(Int32Type),
        DataType::Int64 => prim!(Int64Type),
        DataType::UInt8 => prim!(UInt8Type as i32),
        DataType::UInt16 => prim!(UInt16Type as i32),
        DataType::UInt32 => prim!(UInt32Type as i64),
        DataType::UInt64 => {
            // PG has no unsigned 64; saturate so we never wrap.
            let v = col.as_primitive::<UInt64Type>().value(row);
            enc.encode_field(&Some(i64::try_from(v).unwrap_or(i64::MAX)))
        }
        DataType::Float32 => prim!(Float32Type as f64),
        DataType::Float64 => prim!(Float64Type),
        DataType::Boolean => enc.encode_field(&Some(col.as_boolean().value(row))),
        DataType::Utf8 => enc.encode_field(&Some(col.as_string::<i32>().value(row))),
        DataType::LargeUtf8 => enc.encode_field(&Some(col.as_string::<i64>().value(row))),
        DataType::Timestamp(unit, _tz) => {
            // Each unit has its own Arrow type — `PrimitiveArray<TimestampMicrosecondType>`
            // is *not* `PrimitiveArray<Int64Type>`, so the downcast must match the unit.
            use arrow_array::temporal_conversions::{
                timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_s_to_datetime,
                timestamp_us_to_datetime,
            };
            use arrow_schema::TimeUnit;
            let (raw, dt) = match unit {
                TimeUnit::Second => {
                    let v = col.as_primitive::<TimestampSecondType>().value(row);
                    (v, timestamp_s_to_datetime(v))
                }
                TimeUnit::Millisecond => {
                    let v = col.as_primitive::<TimestampMillisecondType>().value(row);
                    (v, timestamp_ms_to_datetime(v))
                }
                TimeUnit::Microsecond => {
                    let v = col.as_primitive::<TimestampMicrosecondType>().value(row);
                    (v, timestamp_us_to_datetime(v))
                }
                TimeUnit::Nanosecond => {
                    let v = col.as_primitive::<TimestampNanosecondType>().value(row);
                    (v, timestamp_ns_to_datetime(v))
                }
            };
            let dt =
                dt.ok_or_else(|| user_error("22008", format!("timestamp out of range: {raw}")))?;
            enc.encode_field(&Some(dt))
        }
        DataType::Date32 => {
            let v = col.as_primitive::<Date32Type>().value(row);
            let dt = arrow_array::temporal_conversions::date32_to_datetime(v)
                .ok_or_else(|| user_error("22008", format!("DATE out of range: {v}")))?;
            enc.encode_field(&Some(dt.date()))
        }
        DataType::Date64 => {
            let v = col.as_primitive::<Date64Type>().value(row);
            let dt = arrow_array::temporal_conversions::date64_to_datetime(v)
                .ok_or_else(|| user_error("22008", format!("DATE out of range: {v}")))?;
            enc.encode_field(&Some(dt.date()))
        }
        DataType::List(field)
            if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) =>
        {
            // `postgres-types` encodes Vec<Option<String>> as the binary
            // text[] wire format (the column's OID is TEXT_ARRAY).
            enc.encode_field(&Some(list_text_elements(col, row)))
        }
        other => Err(user_error(
            "0A000",
            format!("binary format not supported for column '{name}' (type {other:?})"),
        )),
    }
}

fn arrow_to_pg_type(dt: &arrow_schema::DataType) -> Type {
    use arrow_schema::DataType;
    match dt {
        DataType::Int8 | DataType::Int16 | DataType::Int32 => Type::INT4,
        DataType::Int64 | DataType::UInt32 | DataType::UInt64 => Type::INT8,
        DataType::UInt8 | DataType::UInt16 => Type::INT4,
        DataType::Float32 | DataType::Float64 => Type::FLOAT8,
        DataType::Utf8 | DataType::LargeUtf8 => Type::VARCHAR,
        DataType::Boolean => Type::BOOL,
        DataType::Timestamp(_, _) => Type::TIMESTAMP,
        DataType::Date32 | DataType::Date64 => Type::DATE,
        DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => Type::NUMERIC,
        DataType::List(field)
            if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) =>
        {
            Type::TEXT_ARRAY
        }
        _ => Type::TEXT,
    }
}

/// Per-call salt + stored credential for the MD5 challenge flow. The
/// stored value is either plaintext (legacy) or `md5<32-hex>`, the same
/// format Postgres' `pg_authid` uses, where the hex is `md5(password ‖
/// user)`. The pre-hashed form lets operators avoid plaintext at rest.
#[derive(Debug)]
struct LaminarAuthSource {
    users: Arc<HashMap<String, Secret>>,
}

/// If `stored` is a `pg_authid`-style pre-hash, return the inner hex
/// (the bit after the `md5` tag). Lowercase hex only; uppercase or
/// other lengths fall back to plaintext handling.
pub(crate) fn parse_pre_hashed_md5(stored: &str) -> Option<&str> {
    let inner = stored.strip_prefix("md5")?;
    if inner.len() == 32 && inner.chars().all(|c| matches!(c, '0'..='9' | 'a'..='f')) {
        Some(inner)
    } else {
        None
    }
}

/// MD5 challenge response when only the inner hash is known: the client
/// sends `md5{hex(md5(inner_hex || salt))}` and the server precomputes
/// the same string for comparison.
fn outer_md5_challenge(inner_hex: &str, salt: &[u8]) -> String {
    use md5::{Digest, Md5};
    let mut hasher = Md5::new();
    hasher.update(inner_hex.as_bytes());
    hasher.update(salt);
    format!("md5{:x}", hasher.finalize())
}

#[async_trait]
impl AuthSource for LaminarAuthSource {
    async fn get_password(&self, login: &LoginInfo) -> PgWireResult<Password> {
        let user = login.user().unwrap_or("");
        // Indistinguishable from a wrong-password failure: both branches must
        // surface the same wire error so a client can't probe which usernames
        // are configured. pgwire emits exactly this variant on bad password.
        let stored = self
            .users
            .get(user)
            .ok_or_else(|| PgWireError::InvalidPassword(user.to_string()))?;
        let salt: [u8; 4] = rand::random();
        let expected = match parse_pre_hashed_md5(stored.expose()) {
            Some(inner_hex) => outer_md5_challenge(inner_hex, &salt),
            None => hash_md5_password(user, stored.expose(), &salt),
        };
        Ok(Password::new(Some(salt.to_vec()), expected.into_bytes()))
    }
}

type Md5Handler = Md5PasswordAuthStartupHandler<LaminarAuthSource, DefaultServerParameterProvider>;

/// Startup-phase dispatch. `Md5` requires password auth; `Trust` accepts any
/// connection. Selected once at listener startup based on whether
/// `pgwire_users` is non-empty.
enum StartupAuth {
    Trust(Arc<LaminarPgwireHandler>),
    Md5(Arc<Md5Handler>),
}

#[async_trait]
impl StartupHandler for StartupAuth {
    async fn on_startup<C>(
        &self,
        client: &mut C,
        message: PgWireFrontendMessage,
    ) -> PgWireResult<()>
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        match self {
            Self::Trust(h) => h.on_startup(client, message).await,
            Self::Md5(h) => h.on_startup(client, message).await,
        }
    }
}

pub struct LaminarHandlerFactory {
    handler: Arc<LaminarPgwireHandler>,
    startup: Arc<StartupAuth>,
}

impl LaminarHandlerFactory {
    fn new(db: Arc<LaminarDB>, users: HashMap<String, Secret>) -> Self {
        let handler = Arc::new(LaminarPgwireHandler::new(db));
        let startup = if users.is_empty() {
            Arc::new(StartupAuth::Trust(Arc::clone(&handler)))
        } else {
            let auth = LaminarAuthSource {
                users: Arc::new(users),
            };
            let md5 = Md5PasswordAuthStartupHandler::new(
                Arc::new(auth),
                Arc::new(DefaultServerParameterProvider::default()),
            );
            Arc::new(StartupAuth::Md5(Arc::new(md5)))
        };
        Self { handler, startup }
    }
}

impl PgWireServerHandlers for LaminarHandlerFactory {
    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
        Arc::clone(&self.handler)
    }

    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
        Arc::clone(&self.handler)
    }

    fn startup_handler(&self) -> Arc<impl StartupHandler> {
        Arc::clone(&self.startup)
    }
}

/// Parsed statement carried through `Parse` → `Bind` → `Execute`.
#[derive(Clone, Debug)]
pub enum LaminarStmt {
    /// `SUBSCRIBE` with its schema resolved at parse time so `Describe` can
    /// answer before the portal is bound.
    Subscribe {
        name: String,
        filter_sql: Option<String>,
        as_of_epoch: Option<u64>,
        schema: arrow_schema::SchemaRef,
    },
    Show(ShowCommand),
    Standard(Box<Statement>),
}

/// Resolves SQL to `LaminarStmt`, looking up stream schemas against the
/// live `LaminarDB` so the extended-query `Describe` returns columns
/// without running the query.
#[derive(Clone)]
pub struct LaminarQueryParser {
    db: Arc<LaminarDB>,
}

#[async_trait]
impl QueryParser for LaminarQueryParser {
    type Statement = LaminarStmt;

    async fn parse_sql<C>(
        &self,
        _client: &C,
        sql: &str,
        _types: &[Option<Type>],
    ) -> PgWireResult<Self::Statement>
    where
        C: ClientInfo + Unpin + Send + Sync,
    {
        let mut stmts = parse_streaming_sql(sql)
            .map_err(|e| user_error("42601", format!("parse error: {e}")))?;
        let stmt = stmts
            .pop()
            .ok_or_else(|| user_error("42601", "empty statement"))?;
        if !stmts.is_empty() {
            return Err(user_error(
                "42601",
                "extended query: multiple statements per Parse are not supported",
            ));
        }

        match stmt {
            StreamingStatement::Subscribe(s) => {
                let name = s.name.to_string();
                let (schema, _) = self.db.lookup_subscription_schema(&name).ok_or_else(|| {
                    user_error("42P01", format!("SUBSCRIBE '{name}': stream not found"))
                })?;
                Ok(LaminarStmt::Subscribe {
                    name,
                    filter_sql: s.filter_sql,
                    as_of_epoch: s.as_of_epoch,
                    schema,
                })
            }
            StreamingStatement::Show(cmd) => Ok(LaminarStmt::Show(cmd)),
            StreamingStatement::Standard(s) => Ok(LaminarStmt::Standard(s)),
            other => Err(user_error(
                "0A000",
                format!("not supported on pgwire (use HTTP /api/v1/sql): {other:?}"),
            )),
        }
    }

    fn get_parameter_types(&self, _stmt: &Self::Statement) -> PgWireResult<Vec<Type>> {
        // SUBSCRIBE has no `$N` placeholders.
        Ok(Vec::new())
    }

    fn get_result_schema(
        &self,
        stmt: &Self::Statement,
        column_format: Option<&Format>,
    ) -> PgWireResult<Vec<FieldInfo>> {
        // SHOW and Standard are tiny single-row outputs whose schema only
        // materialises after execution; clients see it on Execute's
        // RowDescription instead.
        match stmt {
            LaminarStmt::Subscribe { schema, .. } => Ok(field_infos(schema, column_format)),
            LaminarStmt::Show(_) | LaminarStmt::Standard(_) => Ok(Vec::new()),
        }
    }
}

#[async_trait]
impl ExtendedQueryHandler for LaminarPgwireHandler {
    type Statement = LaminarStmt;
    type QueryParser = LaminarQueryParser;

    fn query_parser(&self) -> Arc<Self::QueryParser> {
        Arc::new(LaminarQueryParser {
            db: Arc::clone(&self.db),
        })
    }

    async fn do_query<C>(
        &self,
        client: &mut C,
        portal: &Portal<Self::Statement>,
        max_rows: usize,
    ) -> PgWireResult<Response>
    where
        C: ClientInfo + ClientPortalStore + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::PortalStore: PortalStore<Statement = Self::Statement>,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        match &portal.statement.statement {
            LaminarStmt::Subscribe {
                name,
                filter_sql,
                as_of_epoch,
                ..
            } => {
                let start = match as_of_epoch {
                    Some(n) => SubscribeStart::AsOfEpoch(*n),
                    None => SubscribeStart::Tail,
                };
                let sub = self
                    .db
                    .open_subscription(name, filter_sql.as_deref(), start)
                    .await
                    .map_err(|e| user_error("42P01", format!("SUBSCRIBE '{name}': {e}")))?;
                if max_rows == 0 {
                    // Unbounded fetch — pgwire would buffer infinitely.
                    // Drive the stream ourselves with per-batch flushing.
                    stream_subscribe_flushing(
                        client,
                        sub,
                        false,
                        Some(&portal.result_column_format),
                    )
                    .await?;
                    Ok(Response::Execution(Tag::new("SUBSCRIBE")))
                } else {
                    // Chunked (JDBC setFetchSize / tokio-postgres query_portal).
                    // Hand pgwire a row stream so it honours max_rows and
                    // emits PortalSuspended automatically.
                    Ok(subscription_query_response(
                        sub,
                        Some(&portal.result_column_format),
                    ))
                }
            }
            LaminarStmt::Show(cmd) => engine_metadata_response(&self.db, &show_sql(cmd)).await,
            LaminarStmt::Standard(s) => standard_response(&self.db, *s.clone()),
        }
    }

    /// Per-Sync portal cleanup: only the unnamed portal is destroyed.
    ///
    /// The pgwire 0.39 default `on_sync` calls `clear_portals()`, which wipes
    /// every named portal on the connection. PostgreSQL keeps named portals
    /// alive until `Close` or end-of-transaction, so the default would break
    /// any client that does `Bind named_portal; Sync; Execute named_portal;`
    /// — the standard JDBC / asyncpg / tokio-postgres pattern for chunked
    /// fetches via `setFetchSize` / `query_portal`.
    async fn on_sync<C>(
        &self,
        client: &mut C,
        _message: pgwire::messages::extendedquery::Sync,
    ) -> PgWireResult<()>
    where
        C: ClientInfo + ClientPortalStore + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::PortalStore: PortalStore<Statement = Self::Statement>,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        use futures::SinkExt;
        use pgwire::messages::response::ReadyForQuery;

        // Drop only the unnamed portal; named portals survive Sync.
        client.portal_store().rm_portal("");

        client
            .send(PgWireBackendMessage::ReadyForQuery(ReadyForQuery::new(
                client.transaction_status(),
            )))
            .await?;
        client.flush().await?;
        Ok(())
    }
}

pub struct TlsPaths<'a> {
    pub cert: &'a std::path::Path,
    pub key: &'a std::path::Path,
    pub min_version: TlsMinVersion,
    /// PEM bundle of CA roots; presence enables mTLS — every client must
    /// present a cert that chains to one of these roots.
    pub client_ca: Option<&'a std::path::Path>,
}

/// Owned counterpart to `TlsPaths` that the listener keeps for the
/// lifetime of `serve()` so the file watcher can rebuild the acceptor
/// without the original config still being in scope.
#[derive(Debug, Clone)]
struct TlsConfigPaths {
    cert: std::path::PathBuf,
    key: std::path::PathBuf,
    min_version: TlsMinVersion,
    client_ca: Option<std::path::PathBuf>,
}

impl TlsConfigPaths {
    fn from_paths(paths: &TlsPaths<'_>) -> Self {
        Self {
            cert: paths.cert.to_path_buf(),
            key: paths.key.to_path_buf(),
            min_version: paths.min_version,
            client_ca: paths.client_ca.map(|p| p.to_path_buf()),
        }
    }

    fn borrow(&self) -> TlsPaths<'_> {
        TlsPaths {
            cert: &self.cert,
            key: &self.key,
            min_version: self.min_version,
            client_ca: self.client_ca.as_deref(),
        }
    }
}

/// Live TLS acceptor + paths needed to rebuild it on cert rotation.
/// Reads on the accept path are a single mutex acquire and a cheap
/// `TlsAcceptor` clone; reloads are triggered by the file watcher.
pub struct TlsReloadState {
    paths: TlsConfigPaths,
    acceptor: parking_lot::Mutex<Arc<tokio_rustls::TlsAcceptor>>,
}

impl TlsReloadState {
    fn snapshot(&self) -> Arc<tokio_rustls::TlsAcceptor> {
        Arc::clone(&self.acceptor.lock())
    }
}

/// Rebuild the TLS acceptor from `state.paths` and atomically swap it in.
/// On any error the previous acceptor is left in place, so a bad rotation
/// (truncated file, expired cert) doesn't take TLS down.
#[allow(clippy::result_large_err)]
pub(crate) fn try_reload_tls(state: &TlsReloadState) -> Result<(), ServerError> {
    let new_acceptor = load_tls_acceptor(state.paths.borrow())?;
    *state.acceptor.lock() = Arc::new(new_acceptor);
    Ok(())
}

/// Watch the cert / key / client-CA files and call `try_reload_tls` after
/// debounced changes. Mirrors the pattern in `watcher.rs` (parent-dir
/// watch, debounce, then act). Runs until the channel closes; the caller
/// drives shutdown by aborting the task that owns this future.
async fn watch_tls_files(state: Arc<TlsReloadState>, debounce: std::time::Duration) {
    use crossfire::{mpsc, MTx};
    use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher};

    // Track raw + canonical paths so symlink-swap rotations and edits to
    // the symlink target both produce visible events.
    let mut raw_targets: Vec<std::path::PathBuf> = Vec::new();
    let mut canon_targets: Vec<std::path::PathBuf> = Vec::new();
    for path in [
        Some(state.paths.cert.clone()),
        Some(state.paths.key.clone()),
        state.paths.client_ca.clone(),
    ]
    .into_iter()
    .flatten()
    {
        match path.canonicalize() {
            Ok(canonical) => {
                canon_targets.push(canonical);
                raw_targets.push(path);
            }
            Err(e) => {
                warn!(
                    path = %path.display(),
                    error = %e,
                    "pgwire TLS watcher: cannot canonicalize path; reload disabled",
                );
                return;
            }
        }
    }
    let mut dirs: Vec<std::path::PathBuf> = raw_targets
        .iter()
        .chain(canon_targets.iter())
        .filter_map(|p| p.parent().map(|d| d.to_path_buf()))
        .collect();
    dirs.sort();
    dirs.dedup();

    let (tx, rx) = mpsc::bounded_async::<()>(16);
    let blocking_tx: MTx<_> = tx.clone().into_blocking();
    let watch_raw = raw_targets.clone();
    let watch_canon = canon_targets.clone();

    let mut watcher: RecommendedWatcher = match notify::recommended_watcher(
        move |result: Result<Event, notify::Error>| match result {
            Ok(event) => {
                let touched = event.paths.iter().any(|p| {
                    watch_raw.iter().any(|t| t == p)
                        || p.canonicalize()
                            .ok()
                            .as_ref()
                            .is_some_and(|c| watch_canon.contains(c))
                });
                if touched {
                    let _ = blocking_tx.send(());
                }
            }
            Err(e) => warn!(error = %e, "pgwire TLS watcher: notify error"),
        },
    ) {
        Ok(w) => w,
        Err(e) => {
            warn!(error = %e, "pgwire TLS watcher: failed to create watcher; reload disabled");
            return;
        }
    };

    for dir in &dirs {
        if let Err(e) = watcher.watch(dir, RecursiveMode::NonRecursive) {
            warn!(
                dir = %dir.display(),
                error = %e,
                "pgwire TLS watcher: failed to watch directory; reload disabled",
            );
            return;
        }
    }
    info!(
        files = ?raw_targets.iter().map(|p| p.display().to_string()).collect::<Vec<_>>(),
        "pgwire TLS watcher started",
    );

    loop {
        if rx.recv().await.is_err() {
            return;
        }
        // Debounce: sleep then drain so a burst of inotify events
        // (cert + key written separately) coalesces into one reload.
        tokio::time::sleep(debounce).await;
        while rx.try_recv().is_ok() {}

        match try_reload_tls(&state) {
            Ok(()) => tracing::info!(
                target: "audit",
                event = "pgwire.tls_reload",
                outcome = "ok",
            ),
            Err(e) => tracing::warn!(
                target: "audit",
                event = "pgwire.tls_reload",
                outcome = "failed",
                error = %e,
                "pgwire TLS reload failed; previous certificate kept",
            ),
        }
    }
}

/// Minimum TLS protocol version accepted on the pgwire listener. rustls
/// already disables TLS 1.0/1.1; this narrows further when an operator
/// needs TLS 1.3 only.
#[derive(Clone, Copy, Debug)]
pub enum TlsMinVersion {
    V1_2,
    V1_3,
}

impl TlsMinVersion {
    pub(crate) fn from_config_str(s: &str) -> Option<Self> {
        match s {
            "1.2" => Some(Self::V1_2),
            "1.3" => Some(Self::V1_3),
            _ => None,
        }
    }

    fn versions(self) -> &'static [&'static tokio_rustls::rustls::SupportedProtocolVersion] {
        use tokio_rustls::rustls::version::{TLS12, TLS13};
        static BOTH: &[&tokio_rustls::rustls::SupportedProtocolVersion] = &[&TLS12, &TLS13];
        static ONLY_13: &[&tokio_rustls::rustls::SupportedProtocolVersion] = &[&TLS13];
        match self {
            Self::V1_2 => BOTH,
            Self::V1_3 => ONLY_13,
        }
    }

    fn label(self) -> &'static str {
        match self {
            Self::V1_2 => "1.2",
            Self::V1_3 => "1.3",
        }
    }
}

/// Warn if the key file is group/other-readable.
#[cfg(unix)]
fn warn_if_key_world_readable(file: &std::fs::File, path: &std::path::Path) {
    use std::os::unix::fs::MetadataExt;
    if let Ok(meta) = file.metadata() {
        let mode = meta.mode();
        if mode & 0o077 != 0 {
            warn!(
                path = %path.display(),
                mode = format!("{:o}", mode & 0o777),
                "pgwire_tls_key permissions are too broad; tighten to 0600",
            );
        }
    }
}

#[cfg(not(unix))]
fn warn_if_key_world_readable(_file: &std::fs::File, _path: &std::path::Path) {}

/// Rolling-window auth-failure count per peer IP.
#[derive(Debug, Default)]
struct FailureTracker {
    inner: parking_lot::Mutex<
        HashMap<std::net::IpAddr, std::collections::VecDeque<std::time::Instant>>,
    >,
}

impl FailureTracker {
    fn is_blocked(&self, ip: std::net::IpAddr, limit: u32, window: std::time::Duration) -> bool {
        if limit == 0 {
            return false;
        }
        let cutoff = std::time::Instant::now() - window;
        let mut inner = self.inner.lock();
        let Some(failures) = inner.get_mut(&ip) else {
            return false;
        };
        while failures.front().is_some_and(|t| *t < cutoff) {
            failures.pop_front();
        }
        let blocked = failures.len() >= limit as usize;
        if failures.is_empty() {
            inner.remove(&ip);
        }
        blocked
    }

    fn record_failure(&self, ip: std::net::IpAddr) {
        let mut inner = self.inner.lock();
        // When full, evict the entry whose newest failure is oldest.
        if !inner.contains_key(&ip) && inner.len() >= MAX_TRACKED_IPS {
            if let Some(oldest) = inner
                .iter()
                .min_by_key(|(_, q)| q.back().copied())
                .map(|(k, _)| *k)
            {
                inner.remove(&oldest);
            }
        }
        inner
            .entry(ip)
            .or_default()
            .push_back(std::time::Instant::now());
    }
}

const MAX_TRACKED_IPS: usize = 4096;

/// Stable audit code for a session's exit status.
fn classify_outcome(result: &Result<(), std::io::Error>) -> &'static str {
    match result {
        Ok(()) => "ok",
        Err(e) => {
            let msg = e.to_string();
            if msg.contains("28P01") {
                "auth_failed"
            } else if msg.contains("HandshakeFailure")
                || msg.contains("rustls")
                || msg.contains("tls")
            {
                "tls_failed"
            } else {
                "error"
            }
        }
    }
}

/// Reject certs past `notAfter`; warn within 30 days.
#[allow(clippy::result_large_err)]
fn check_cert_expiry(
    der: &tokio_rustls::rustls::pki_types::CertificateDer<'_>,
    path: &std::path::Path,
) -> Result<(), ServerError> {
    use x509_parser::prelude::FromDer;
    let (_, cert) = x509_parser::certificate::X509Certificate::from_der(der.as_ref())
        .map_err(|e| ServerError::Http(format!("parse pgwire_tls_cert {}: {e}", path.display())))?;
    let now = x509_parser::time::ASN1Time::now();
    let not_after = cert.validity().not_after;
    if not_after < now {
        return Err(ServerError::Http(format!(
            "pgwire_tls_cert {} expired at {not_after}",
            path.display()
        )));
    }
    let remaining = not_after.to_datetime() - now.to_datetime();
    if remaining <= time::Duration::days(30) {
        warn!(
            path = %path.display(),
            expires_at = %not_after,
            "pgwire_tls_cert expires within 30 days; rotate before it lapses",
        );
    }
    Ok(())
}

/// Idempotent install of aws-lc-rs as rustls' default provider.
fn ensure_tls_provider() {
    let _ = tokio_rustls::rustls::crypto::aws_lc_rs::default_provider().install_default();
}

#[allow(clippy::result_large_err)]
fn load_tls_acceptor(paths: TlsPaths<'_>) -> Result<tokio_rustls::TlsAcceptor, ServerError> {
    use std::fs::File;
    use std::io::BufReader;

    ensure_tls_provider();

    let cert_file = File::open(paths.cert)
        .map_err(|e| ServerError::Http(format!("open pgwire_tls_cert: {e}")))?;
    let certs = rustls_pemfile::certs(&mut BufReader::new(cert_file))
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| ServerError::Http(format!("parse pgwire_tls_cert: {e}")))?;
    if certs.is_empty() {
        return Err(ServerError::Http(format!(
            "pgwire_tls_cert {} contains no certificates",
            paths.cert.display()
        )));
    }
    for cert in &certs {
        check_cert_expiry(cert, paths.cert)?;
    }

    let key_file = File::open(paths.key)
        .map_err(|e| ServerError::Http(format!("open pgwire_tls_key: {e}")))?;
    warn_if_key_world_readable(&key_file, paths.key);
    let key = rustls_pemfile::private_key(&mut BufReader::new(key_file))
        .map_err(|e| ServerError::Http(format!("parse pgwire_tls_key: {e}")))?
        .ok_or_else(|| {
            ServerError::Http(format!(
                "pgwire_tls_key {} contains no private key",
                paths.key.display()
            ))
        })?;

    let builder = tokio_rustls::rustls::ServerConfig::builder_with_protocol_versions(
        paths.min_version.versions(),
    );
    let builder = match paths.client_ca {
        Some(ca_path) => {
            let verifier = build_client_cert_verifier(ca_path)?;
            builder.with_client_cert_verifier(verifier)
        }
        None => builder.with_no_client_auth(),
    };
    let server_config = builder
        .with_single_cert(certs, key)
        .map_err(|e| ServerError::Http(format!("rustls server config: {e}")))?;
    Ok(tokio_rustls::TlsAcceptor::from(Arc::new(server_config)))
}

#[allow(clippy::result_large_err)]
fn build_client_cert_verifier(
    ca_path: &std::path::Path,
) -> Result<Arc<dyn tokio_rustls::rustls::server::danger::ClientCertVerifier>, ServerError> {
    use std::fs::File;
    use std::io::BufReader;
    use tokio_rustls::rustls::server::WebPkiClientVerifier;
    use tokio_rustls::rustls::RootCertStore;

    let file = File::open(ca_path)
        .map_err(|e| ServerError::Http(format!("open pgwire_tls_client_ca: {e}")))?;
    let mut roots = RootCertStore::empty();
    let mut added = 0usize;
    for cert in rustls_pemfile::certs(&mut BufReader::new(file)) {
        let cert =
            cert.map_err(|e| ServerError::Http(format!("parse pgwire_tls_client_ca: {e}")))?;
        roots
            .add(cert)
            .map_err(|e| ServerError::Http(format!("invalid CA in pgwire_tls_client_ca: {e}")))?;
        added += 1;
    }
    if added == 0 {
        return Err(ServerError::Http(format!(
            "pgwire_tls_client_ca {} contains no certificates",
            ca_path.display()
        )));
    }
    WebPkiClientVerifier::builder(Arc::new(roots))
        .build()
        .map_err(|e| ServerError::Http(format!("build client-cert verifier: {e}")))
}

pub async fn serve(
    db: Arc<LaminarDB>,
    bind: &str,
    users: HashMap<String, Secret>,
    allow_remote: bool,
    tls: Option<TlsPaths<'_>>,
    max_connections: usize,
    max_auth_failures_per_min: u32,
) -> Result<(SocketAddr, tokio::task::JoinHandle<()>), ServerError> {
    let addr: SocketAddr = bind
        .parse()
        .map_err(|e| ServerError::Http(format!("invalid pgwire_bind '{bind}': {e}")))?;

    let auth_mode = if users.is_empty() { "trust" } else { "md5" };
    let is_remote_bind = !addr.ip().is_loopback();
    match (auth_mode, is_remote_bind, allow_remote) {
        ("trust", true, _) => {
            return Err(ServerError::Http(format!(
                "pgwire_bind '{addr}' is not loopback and pgwire_users is empty (trust auth); \
             configure pgwire_users + pgwire_allow_remote=true, or bind to 127.0.0.1"
            )))
        }
        ("md5", true, false) => {
            return Err(ServerError::Http(format!(
                "pgwire_bind '{addr}' is not loopback; set pgwire_allow_remote=true to opt in"
            )))
        }
        _ => {}
    }

    let tls_min_label = tls.as_ref().map(|p| p.min_version.label());
    let mtls_on = tls.as_ref().is_some_and(|p| p.client_ca.is_some());
    let tls_state: Option<Arc<TlsReloadState>> = match tls {
        Some(paths) => {
            let acceptor = load_tls_acceptor(TlsPaths {
                cert: paths.cert,
                key: paths.key,
                min_version: paths.min_version,
                client_ca: paths.client_ca,
            })?;
            Some(Arc::new(TlsReloadState {
                paths: TlsConfigPaths::from_paths(&paths),
                acceptor: parking_lot::Mutex::new(Arc::new(acceptor)),
            }))
        }
        None => None,
    };

    let listener = TcpListener::bind(addr)
        .await
        .map_err(|e| ServerError::Http(format!("pgwire bind {addr}: {e}")))?;
    let local_addr = listener
        .local_addr()
        .map_err(|e| ServerError::Http(format!("pgwire local_addr: {e}")))?;

    let factory = Arc::new(LaminarHandlerFactory::new(db, users));
    let tls_mode = if tls_state.is_some() { "on" } else { "off" };
    let tls_min = tls_min_label.unwrap_or("-");
    let mtls = if mtls_on { "on" } else { "off" };
    if auth_mode == "trust" {
        warn!(
            addr = %local_addr,
            tls = tls_mode,
            tls_min,
            mtls,
            "pgwire listening with TRUST auth — any client reaching this address is admin",
        );
    } else {
        info!(
            addr = %local_addr,
            auth = auth_mode,
            tls = tls_mode,
            tls_min,
            mtls,
            "pgwire listening",
        );
    }

    // Track per-connection tasks so abort on the outer JoinHandle stops
    // active sessions in addition to the accept loop.
    let failures = Arc::new(FailureTracker::default());
    let watcher_state = tls_state.as_ref().map(Arc::clone);
    let watcher_disabled =
        std::env::var("LAMINAR_DISABLE_FILE_WATCH").is_ok_and(|v| v == "1" || v == "true");
    let handle = tokio::spawn(async move {
        let mut sessions: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
        // Watcher in its own JoinSet so it doesn't count toward max_connections.
        let mut watcher_set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
        if let (Some(state), false) = (watcher_state, watcher_disabled) {
            watcher_set.spawn(async move {
                watch_tls_files(state, std::time::Duration::from_millis(500)).await;
            });
        }
        loop {
            tokio::select! {
                Some(_) = sessions.join_next(), if !sessions.is_empty() => {
                    // Reap completed sessions; nothing to do with the result.
                }
                Some(_) = watcher_set.join_next(), if !watcher_set.is_empty() => {}
                accepted = listener.accept() => {
                    match accepted {
                        Ok((sock, peer)) => {
                            if sessions.len() >= max_connections {
                                tracing::info!(
                                    target: "audit",
                                    event = "pgwire.connection_rejected",
                                    peer = %peer,
                                    reason = "max_connections",
                                    in_flight = sessions.len(),
                                );
                                drop(sock);
                                continue;
                            }
                            if failures.is_blocked(
                                peer.ip(),
                                max_auth_failures_per_min,
                                std::time::Duration::from_secs(60),
                            ) {
                                tracing::warn!(
                                    target: "audit",
                                    event = "pgwire.connection_rejected",
                                    peer = %peer,
                                    reason = "auth_failure_throttle",
                                );
                                drop(sock);
                                continue;
                            }
                            let factory_ref = Arc::clone(&factory);
                            // Snapshot the live acceptor so that an in-flight
                            // handshake completes against whatever cert was
                            // current when the socket was accepted, even if a
                            // hot-reload swaps it under us.
                            let tls_ref: Option<tokio_rustls::TlsAcceptor> =
                                tls_state.as_ref().map(|s| (*s.snapshot()).clone());
                            let failures_ref = Arc::clone(&failures);
                            let peer_str = peer.to_string();
                            tracing::info!(
                                target: "audit",
                                event = "pgwire.connection_accepted",
                                peer = %peer,
                                auth = auth_mode,
                                tls = tls_mode,
                            );
                            let peer_ip = peer.ip();
                            sessions.spawn(async move {
                                let result = process_socket(sock, tls_ref, factory_ref).await;
                                let outcome = classify_outcome(&result);
                                if outcome == "auth_failed" {
                                    failures_ref.record_failure(peer_ip);
                                }
                                tracing::info!(
                                    target: "audit",
                                    event = "pgwire.connection_closed",
                                    peer = %peer_str,
                                    outcome,
                                );
                                if let Err(e) = result {
                                    warn!(peer = %peer_str, error = %e, "pgwire connection error");
                                }
                            });
                        }
                        Err(e) => {
                            warn!(error = %e, "pgwire accept failed");
                            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                        }
                    }
                }
            }
        }
    });
    Ok((local_addr, handle))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parse_one(sql: &str) -> StreamingStatement {
        parse_streaming_sql(sql)
            .unwrap()
            .into_iter()
            .next()
            .unwrap()
    }

    fn standard(sql: &str) -> Statement {
        match parse_one(sql) {
            StreamingStatement::Standard(s) => *s,
            other => panic!("expected Standard, got {other:?}"),
        }
    }

    #[test]
    fn pg_text_array_literal_quotes_nulls_and_escapes() {
        assert_eq!(pg_text_array_literal(&[]), "{}");
        assert_eq!(
            pg_text_array_literal(&[Some("en".into()), Some("ja".into())]),
            r#"{"en","ja"}"#
        );
        assert_eq!(
            pg_text_array_literal(&[None, Some("x".into())]),
            r#"{NULL,"x"}"#
        );
        // Embedded quote and backslash are escaped, not left ambiguous.
        assert_eq!(
            pg_text_array_literal(&[Some("a\"b\\c".into())]),
            r#"{"a\"b\\c"}"#
        );
    }

    #[tokio::test]
    async fn select_one_dispatches() {
        let db = LaminarDB::open().unwrap();
        for sql in ["SELECT 1", "select 1", "/* hint */ SELECT 1"] {
            standard_response(&db, standard(sql)).unwrap();
        }
    }

    #[tokio::test]
    async fn driver_select_builtins_dispatch() {
        let db = LaminarDB::open().unwrap();
        for sql in [
            "SELECT version()",
            "SELECT current_schema()",
            "SELECT current_database()",
            "SELECT current_user",
        ] {
            // current_user parses as Expr::Function with no parens in some versions;
            // we accept whatever the parser gives us.
            let _ = standard_response(&db, standard(sql));
        }
    }

    #[tokio::test]
    async fn select_with_from_is_rejected() {
        let db = LaminarDB::open().unwrap();
        let err = standard_response(&db, standard("SELECT 1 FROM foo")).unwrap_err();
        assert!(err.to_string().contains("limited to literals"));
    }

    #[tokio::test]
    async fn ddl_routed_to_http() {
        let db = LaminarDB::open().unwrap();
        let err = standard_response(&db, standard("CREATE TABLE foo (id INT)")).unwrap_err();
        assert!(err.to_string().contains("HTTP /api/v1/sql"));
    }

    #[tokio::test]
    async fn transaction_control_dispatches() {
        let db = LaminarDB::open().unwrap();
        for sql in [
            "BEGIN",
            "BEGIN TRANSACTION",
            "START TRANSACTION",
            "COMMIT",
            "ROLLBACK",
        ] {
            standard_response(&db, standard(sql)).unwrap();
        }
    }

    #[tokio::test]
    async fn set_writes_to_session_properties() {
        let db = LaminarDB::open().unwrap();
        standard_response(&db, standard("SET extra_float_digits = 3")).unwrap();
        assert_eq!(
            db.get_session_property("extra_float_digits").as_deref(),
            Some("3"),
        );
    }

    #[tokio::test]
    async fn set_transaction_isolation_is_rejected() {
        let db = LaminarDB::open().unwrap();
        let err = standard_response(
            &db,
            standard("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE"),
        )
        .unwrap_err();
        assert!(err.to_string().contains("SET TRANSACTION"));
    }

    #[test]
    fn multi_statement_parses() {
        let stmts = parse_streaming_sql("BEGIN; SELECT 1; COMMIT").unwrap();
        assert_eq!(stmts.len(), 3);
    }

    #[test]
    fn classify_outcome_buckets_errors() {
        use std::io::{Error, ErrorKind};
        assert_eq!(super::classify_outcome(&Ok(())), "ok");
        assert_eq!(
            super::classify_outcome(&Err(Error::other("FATAL: 28P01 bad pass"))),
            "auth_failed"
        );
        assert_eq!(
            super::classify_outcome(&Err(Error::other("rustls HandshakeFailure"))),
            "tls_failed"
        );
        assert_eq!(
            super::classify_outcome(&Err(Error::new(ErrorKind::BrokenPipe, "broken"))),
            "error"
        );
    }

    #[test]
    fn failure_tracker_blocks_after_threshold() {
        use std::net::{IpAddr, Ipv4Addr};
        use std::time::Duration;
        let ip: IpAddr = Ipv4Addr::LOCALHOST.into();
        let tracker = super::FailureTracker::default();
        let limit = 3;
        let window = Duration::from_secs(60);

        for _ in 0..limit {
            assert!(!tracker.is_blocked(ip, limit, window));
            tracker.record_failure(ip);
        }
        assert!(tracker.is_blocked(ip, limit, window));
    }

    #[test]
    fn failure_tracker_disabled_when_limit_zero() {
        use std::net::{IpAddr, Ipv4Addr};
        use std::time::Duration;
        let ip: IpAddr = Ipv4Addr::LOCALHOST.into();
        let tracker = super::FailureTracker::default();
        for _ in 0..100 {
            tracker.record_failure(ip);
        }
        assert!(!tracker.is_blocked(ip, 0, Duration::from_secs(60)));
    }

    #[test]
    fn failure_tracker_expires_old_entries() {
        use std::net::{IpAddr, Ipv4Addr};
        use std::time::Duration;
        let ip: IpAddr = Ipv4Addr::LOCALHOST.into();
        let tracker = super::FailureTracker::default();
        for _ in 0..5 {
            tracker.record_failure(ip);
        }
        // Window of 0 means every recorded failure is already expired.
        assert!(!tracker.is_blocked(ip, 5, Duration::from_secs(0)));
    }

    #[test]
    fn failure_tracker_caps_distinct_ips() {
        use std::net::{IpAddr, Ipv4Addr};
        let tracker = super::FailureTracker::default();
        // Push past the cap; map size must stay bounded.
        for i in 0..(super::MAX_TRACKED_IPS + 100) {
            #[allow(clippy::cast_possible_truncation)]
            let ip: IpAddr = Ipv4Addr::new(10, 0, (i / 256) as u8, (i % 256) as u8).into();
            tracker.record_failure(ip);
        }
        let len = tracker.inner.lock().len();
        assert!(
            len <= super::MAX_TRACKED_IPS,
            "tracker exceeded cap: {len} > {}",
            super::MAX_TRACKED_IPS
        );
    }

    #[tokio::test]
    async fn serve_rejects_remote_bind_in_trust_mode() {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        let err = serve(db, "0.0.0.0:0", HashMap::new(), false, None, 256, 10)
            .await
            .expect_err("trust + 0.0.0.0 must fail");
        assert!(err.to_string().contains("trust auth"), "got: {err}");
    }

    #[tokio::test]
    async fn serve_rejects_remote_bind_without_explicit_optin() {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        let mut users = HashMap::new();
        users.insert("alice".into(), Secret::new("wonderland-key"));
        let err = serve(db, "0.0.0.0:0", users, false, None, 256, 10)
            .await
            .expect_err("md5 + 0.0.0.0 without allow_remote must fail");
        assert!(
            err.to_string().contains("pgwire_allow_remote"),
            "got: {err}"
        );
    }
}

#[cfg(test)]
mod integration_tests {
    //! End-to-end pgwire driven by `tokio_postgres` against an in-process
    //! `LaminarDB`. Verifies the wire-protocol surface — handshake, SimpleQuery
    //! dispatch, error reporting — that unit tests can't reach. Engine-level
    //! row flow is covered in `laminar-db`'s `db::tests`.

    use std::collections::HashMap;
    use std::sync::Arc;

    use laminar_db::LaminarDB;
    use tokio_postgres::{NoTls, SimpleQueryMessage};

    use super::Secret;

    async fn spawn_server_with(
        users: HashMap<String, Secret>,
    ) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE trades (symbol VARCHAR, price DOUBLE)")
            .await
            .expect("create source");
        db.execute(
            "CREATE MATERIALIZED VIEW prices AS \
             SELECT symbol, price FROM trades",
        )
        .await
        .expect("create mv");
        db.start().await.expect("db starts");

        let (addr, handle) =
            super::serve(Arc::clone(&db), "127.0.0.1:0", users, false, None, 256, 10)
                .await
                .expect("pgwire serve");
        (addr, handle)
    }

    async fn spawn_server() -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
        spawn_server_with(HashMap::new()).await
    }

    async fn connect(addr: std::net::SocketAddr) -> tokio_postgres::Client {
        let conn_str = format!(
            "host={} port={} user=any dbname=laminardb",
            addr.ip(),
            addr.port()
        );
        let (client, conn) = tokio_postgres::connect(&conn_str, NoTls)
            .await
            .expect("pgwire connect");
        tokio::spawn(async move {
            let _ = conn.await;
        });
        client
    }

    fn first_row_value(messages: &[SimpleQueryMessage], col: usize) -> Option<&str> {
        messages.iter().find_map(|m| match m {
            SimpleQueryMessage::Row(r) => r.get(col),
            _ => None,
        })
    }

    #[tokio::test]
    async fn handshake_and_builtins() {
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let messages = client
            .simple_query("SELECT version()")
            .await
            .expect("version");
        let v = first_row_value(&messages, 0).expect("row");
        assert!(v.contains("LaminarDB"), "version: {v}");

        let messages = client
            .simple_query("SELECT current_database()")
            .await
            .expect("current_database");
        assert_eq!(first_row_value(&messages, 0), Some("laminar"));

        handle.abort();
    }

    #[tokio::test]
    async fn show_streams_runs() {
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        // No assertion on contents — just that the dispatch path returns rows
        // without error. Engine-level SHOW behavior is covered in laminar-db.
        client
            .simple_query("SHOW STREAMS")
            .await
            .expect("SHOW STREAMS");

        handle.abort();
    }

    #[tokio::test]
    async fn subscribe_unknown_name_returns_pg_error() {
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("SUBSCRIBE no_such_view")
            .await
            .expect_err("must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(
            db_err.message().contains("no_such_view"),
            "message: {}",
            db_err.message()
        );

        handle.abort();
    }

    #[tokio::test]
    async fn subscribe_with_valid_where_is_accepted() {
        // SUBSCRIBE never returns CommandComplete, so a successful compile
        // shows up as a timeout. Anything else — Ok(Ok) or Ok(Err) — is a
        // regression.
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let r = tokio::time::timeout(
            std::time::Duration::from_millis(500),
            client.simple_query("SUBSCRIBE prices WHERE symbol = 'AAPL'"),
        )
        .await;

        assert!(
            r.is_err(),
            "subscribe must stay open until timeout, got: {r:?}"
        );

        handle.abort();
    }

    #[tokio::test]
    async fn subscribe_with_unknown_column_in_where_returns_pg_error() {
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("SUBSCRIBE prices WHERE no_such_col > 1")
            .await
            .expect_err("must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(
            db_err.message().contains("no_such_col"),
            "filter error must name the bad column, got: {}",
            db_err.message()
        );

        handle.abort();
    }

    #[tokio::test]
    async fn subscribe_as_of_unretained_returns_pg_error() {
        // No retention configured on the `prices` MV from the default setup,
        // so AS OF EPOCH 1 must come back as a typed PG error.
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("SUBSCRIBE prices AS OF EPOCH 1")
            .await
            .expect_err("must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(
            db_err.message().contains("no longer retained"),
            "message: {}",
            db_err.message()
        );

        handle.abort();
    }

    /// SUBSCRIBE must actually stream emitted MV rows over the socket: bind a
    /// portal, push rows into the source, and read them back via the
    /// extended-query portal (the chunked path JDBC/asyncpg use).
    #[tokio::test]
    async fn subscribe_streams_emitted_rows_over_the_wire() {
        use std::time::Duration;

        use arrow_array::{Float64Array, RecordBatch, StringArray};

        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE trades (symbol VARCHAR, price DOUBLE)")
            .await
            .expect("create source");
        db.execute("CREATE MATERIALIZED VIEW prices AS SELECT symbol, price FROM trades")
            .await
            .expect("create mv");
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            256,
            10,
        )
        .await
        .expect("serve");
        let mut client = connect(addr).await;
        let txn = client.transaction().await.expect("begin");

        // The subscription opens when the first Execute runs, so push once the
        // read is in flight (Tail would otherwise miss earlier rows).
        let stmt = txn.prepare("SUBSCRIBE prices").await.expect("prepare");
        let portal = txn.bind(&stmt, &[]).await.expect("bind");

        let pusher = tokio::spawn({
            let db = Arc::clone(&db);
            async move {
                tokio::time::sleep(Duration::from_millis(300)).await;
                let src = db.source_untyped("trades").expect("source handle");
                let batch = RecordBatch::try_new(
                    src.schema().clone(),
                    vec![
                        Arc::new(StringArray::from(vec!["AAPL", "MSFT"])),
                        Arc::new(Float64Array::from(vec![100.0, 200.0])),
                    ],
                )
                .expect("batch");
                src.push_arrow(batch).expect("push");
            }
        });

        let rows = tokio::time::timeout(Duration::from_secs(10), txn.query_portal(&portal, 2))
            .await
            .expect("read did not time out")
            .expect("query_portal");
        pusher.await.expect("pusher");

        let mut symbols: Vec<String> = rows
            .iter()
            .map(|r| r.get::<_, &str>(0).to_string())
            .collect();
        symbols.sort();
        assert_eq!(
            symbols,
            ["AAPL", "MSFT"],
            "both emitted rows arrive over pgwire"
        );

        handle.abort();
    }

    /// A TEXT[] column must round-trip over the binary wire (asyncpg/JDBC
    /// request binary): the column advertises the _text OID and encodes as a
    /// Postgres array, so tokio_postgres decodes it into a Vec<String>.
    #[tokio::test]
    async fn subscribe_decodes_text_array_in_binary_format() {
        use std::time::Duration;

        use arrow_array::{Int64Array, RecordBatch};

        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE feed (id BIGINT)")
            .await
            .expect("create source");
        db.execute(
            "CREATE MATERIALIZED VIEW tagged AS SELECT id, make_array('en','ja') AS tags FROM feed",
        )
        .await
        .expect("create mv");
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            256,
            10,
        )
        .await
        .expect("serve");
        let mut client = connect(addr).await;
        let txn = client.transaction().await.expect("begin");
        let stmt = txn.prepare("SUBSCRIBE tagged").await.expect("prepare");
        let portal = txn.bind(&stmt, &[]).await.expect("bind");

        let pusher = tokio::spawn({
            let db = Arc::clone(&db);
            async move {
                tokio::time::sleep(Duration::from_millis(300)).await;
                let src = db.source_untyped("feed").expect("source handle");
                let batch = RecordBatch::try_new(
                    src.schema().clone(),
                    vec![Arc::new(Int64Array::from(vec![1_i64]))],
                )
                .expect("batch");
                src.push_arrow(batch).expect("push");
            }
        });

        let rows = tokio::time::timeout(Duration::from_secs(10), txn.query_portal(&portal, 1))
            .await
            .expect("read did not time out")
            .expect("query_portal");
        pusher.await.expect("pusher");

        assert_eq!(rows.len(), 1);
        let tags: Vec<String> = rows[0].get(1);
        assert_eq!(
            tags,
            vec!["en".to_string(), "ja".to_string()],
            "TEXT[] decoded over the binary wire"
        );

        handle.abort();
    }

    #[tokio::test]
    async fn ddl_returns_pg_error_pointing_at_http() {
        let (addr, handle) = spawn_server().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("CREATE SOURCE more_trades (sym VARCHAR)")
            .await
            .expect_err("DDL must be rejected");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(
            db_err.message().contains("/api/v1/sql"),
            "message: {}",
            db_err.message()
        );

        handle.abort();
    }

    async fn md5_users() -> HashMap<String, Secret> {
        let mut u = HashMap::new();
        u.insert("alice".to_string(), Secret::new(TEST_PASSWORD));
        u
    }

    const TEST_PASSWORD: &str = "wonderland-key";

    async fn connect_with_password(
        addr: std::net::SocketAddr,
        user: &str,
        password: &str,
    ) -> Result<tokio_postgres::Client, tokio_postgres::Error> {
        let conn_str = format!(
            "host={} port={} user={user} password={password} dbname=laminardb",
            addr.ip(),
            addr.port()
        );
        let (client, conn) = tokio_postgres::connect(&conn_str, NoTls).await?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        Ok(client)
    }

    #[tokio::test]
    async fn md5_auth_accepts_correct_password() {
        let (addr, handle) = spawn_server_with(md5_users().await).await;

        let client = connect_with_password(addr, "alice", TEST_PASSWORD)
            .await
            .expect("auth must succeed");

        let messages = client
            .simple_query("SELECT version()")
            .await
            .expect("query after auth");
        let v = first_row_value(&messages, 0).expect("row");
        assert!(v.contains("LaminarDB"), "version: {v}");

        handle.abort();
    }

    #[tokio::test]
    async fn md5_auth_rejects_wrong_password() {
        let (addr, handle) = spawn_server_with(md5_users().await).await;

        let err = connect_with_password(addr, "alice", "not-the-password")
            .await
            .expect_err("auth must fail");

        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "28P01", "got: {db_err:?}");

        handle.abort();
    }

    /// Pre-hashed pgwire_users entry: stored value is `md5{hex(md5(pw||user))}`,
    /// matching pg_authid. Plaintext never touches disk yet auth still succeeds.
    fn md5_users_prehashed(user: &str, password: &str) -> HashMap<String, Secret> {
        use md5::{Digest, Md5};
        let mut h = Md5::new();
        h.update(password.as_bytes());
        h.update(user.as_bytes());
        let inner = format!("{:x}", h.finalize());
        let mut u = HashMap::new();
        u.insert(user.to_string(), Secret::new(format!("md5{inner}")));
        u
    }

    #[tokio::test]
    async fn md5_auth_accepts_correct_password_against_prehash() {
        let (addr, handle) = spawn_server_with(md5_users_prehashed("alice", TEST_PASSWORD)).await;
        let client = connect_with_password(addr, "alice", TEST_PASSWORD)
            .await
            .expect("auth must succeed against pre-hashed entry");
        let messages = client
            .simple_query("SELECT version()")
            .await
            .expect("query after auth");
        let v = first_row_value(&messages, 0).expect("row");
        assert!(v.contains("LaminarDB"), "version: {v}");
        handle.abort();
    }

    #[tokio::test]
    async fn md5_auth_rejects_wrong_password_against_prehash() {
        let (addr, handle) = spawn_server_with(md5_users_prehashed("alice", TEST_PASSWORD)).await;
        let err = connect_with_password(addr, "alice", "not-the-password")
            .await
            .expect_err("auth must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "28P01", "got: {db_err:?}");
        handle.abort();
    }

    #[test]
    fn parse_pre_hashed_md5_strict_format() {
        // 32 lowercase hex after the tag → accepted.
        let inner = "5d41402abc4b2a76b9719d911017c592";
        assert_eq!(
            super::parse_pre_hashed_md5(&format!("md5{inner}")),
            Some(inner),
        );
        // Wrong length, uppercase hex, missing prefix, or non-hex → rejected.
        assert_eq!(super::parse_pre_hashed_md5("md5short"), None);
        assert_eq!(
            super::parse_pre_hashed_md5("md55D41402ABC4B2A76B9719D911017C592"),
            None,
        );
        assert_eq!(super::parse_pre_hashed_md5(inner), None);
        assert_eq!(
            super::parse_pre_hashed_md5("md5zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"),
            None,
        );
    }

    #[tokio::test]
    async fn md5_auth_rejects_unknown_user() {
        let (addr, handle) = spawn_server_with(md5_users().await).await;

        let err = connect_with_password(addr, "mallory", "anything")
            .await
            .expect_err("auth must fail");

        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "28P01", "got: {db_err:?}");

        handle.abort();
    }

    #[tokio::test]
    async fn connection_cap_drops_excess_clients() {
        // Cap of 1; first client occupies the slot, second is dropped.
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE trades (symbol VARCHAR, price DOUBLE)")
            .await
            .expect("create source");
        db.execute("CREATE MATERIALIZED VIEW prices AS SELECT symbol, price FROM trades")
            .await
            .expect("create mv");
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            1,
            10,
        )
        .await
        .expect("pgwire serve");

        // First client occupies the slot via SUBSCRIBE (stays open until drop).
        let first = connect(addr).await;
        let _bg = tokio::spawn(async move {
            let _ = first.simple_query("SUBSCRIBE prices").await;
        });
        // Give the server a moment to register the session in the JoinSet.
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;

        // Second connect: server accepts the TCP, then closes it because the
        // cap is hit. tokio_postgres surfaces this as an IO error during
        // startup. Exact string varies; just assert it failed.
        let conn_str = format!(
            "host={} port={} user=any dbname=laminardb",
            addr.ip(),
            addr.port()
        );
        let result = tokio_postgres::connect(&conn_str, NoTls).await;
        assert!(result.is_err(), "second connect must be refused");

        handle.abort();
    }

    /// Self-signed cert+key written to a tempdir for the duration of the
    /// test. `rcgen` is the well-maintained option for ad-hoc certs.
    fn self_signed_pem() -> (tempfile::TempDir, std::path::PathBuf, std::path::PathBuf) {
        let cert =
            rcgen::generate_simple_self_signed(vec!["localhost".into()]).expect("rcgen issue cert");
        let dir = tempfile::tempdir().expect("tempdir");
        let cert_path = dir.path().join("cert.pem");
        let key_path = dir.path().join("key.pem");
        std::fs::write(&cert_path, cert.cert.pem()).unwrap();
        std::fs::write(&key_path, cert.key_pair.serialize_pem()).unwrap();
        (dir, cert_path, key_path)
    }

    /// CA + client-leaf bundle for mTLS tests. The CA PEM is written to a
    /// tempfile so the server can be pointed at it via `pgwire_tls_client_ca`;
    /// the leaf cert+key are returned in DER form for direct use by a rustls
    /// `ClientConfig`.
    struct MintedClientPki {
        _dir: tempfile::TempDir,
        ca_pem_path: std::path::PathBuf,
        leaf_chain: Vec<tokio_rustls::rustls::pki_types::CertificateDer<'static>>,
        leaf_key: tokio_rustls::rustls::pki_types::PrivateKeyDer<'static>,
    }

    fn mint_ca_and_client_leaf(common_name: &str) -> MintedClientPki {
        use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};

        let mut ca_params = rcgen::CertificateParams::new(vec!["mtls-test-ca".into()]).unwrap();
        ca_params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
        let ca_key = rcgen::KeyPair::generate().unwrap();
        let ca_cert = ca_params.self_signed(&ca_key).unwrap();

        let mut leaf_params = rcgen::CertificateParams::new(vec![common_name.into()]).unwrap();
        leaf_params.extended_key_usages = vec![rcgen::ExtendedKeyUsagePurpose::ClientAuth];
        let leaf_key = rcgen::KeyPair::generate().unwrap();
        let leaf_cert = leaf_params.signed_by(&leaf_key, &ca_cert, &ca_key).unwrap();

        let dir = tempfile::tempdir().unwrap();
        let ca_pem_path = dir.path().join("ca.pem");
        std::fs::write(&ca_pem_path, ca_cert.pem()).unwrap();

        let leaf_chain = vec![CertificateDer::from(leaf_cert.der().to_vec())];
        let leaf_key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(leaf_key.serialize_der()));

        MintedClientPki {
            _dir: dir,
            ca_pem_path,
            leaf_chain,
            leaf_key,
        }
    }

    /// Builds a tokio_postgres TLS connector that trusts `server_cert_path`
    /// for the server hello and (optionally) presents a client cert for mTLS.
    fn make_client_tls(
        server_cert_path: &std::path::Path,
        client_auth: Option<(
            Vec<tokio_rustls::rustls::pki_types::CertificateDer<'static>>,
            tokio_rustls::rustls::pki_types::PrivateKeyDer<'static>,
        )>,
    ) -> tokio_postgres_rustls::MakeRustlsConnect {
        super::ensure_tls_provider();
        let cert_bytes = std::fs::read(server_cert_path).unwrap();
        let mut roots = tokio_rustls::rustls::RootCertStore::empty();
        for c in rustls_pemfile::certs(&mut std::io::Cursor::new(cert_bytes))
            .collect::<Result<Vec<_>, _>>()
            .unwrap()
        {
            roots.add(c).unwrap();
        }
        let builder = tokio_rustls::rustls::ClientConfig::builder().with_root_certificates(roots);
        let client_cfg = match client_auth {
            Some((chain, key)) => builder.with_client_auth_cert(chain, key).unwrap(),
            None => builder.with_no_client_auth(),
        };
        tokio_postgres_rustls::MakeRustlsConnect::new(client_cfg)
    }

    /// Self-signed cert with notAfter in the past, for the expiry test.
    fn expired_self_signed_pem() -> (tempfile::TempDir, std::path::PathBuf, std::path::PathBuf) {
        let mut params = rcgen::CertificateParams::new(vec!["localhost".into()]).unwrap();
        let one_year_ago = time::OffsetDateTime::now_utc() - time::Duration::days(365);
        params.not_before = one_year_ago - time::Duration::days(2);
        params.not_after = one_year_ago;
        let key = rcgen::KeyPair::generate().unwrap();
        let cert = params.self_signed(&key).unwrap();
        let dir = tempfile::tempdir().unwrap();
        let cert_path = dir.path().join("cert.pem");
        let key_path = dir.path().join("key.pem");
        std::fs::write(&cert_path, cert.pem()).unwrap();
        std::fs::write(&key_path, key.serialize_pem()).unwrap();
        (dir, cert_path, key_path)
    }

    #[tokio::test]
    async fn tls_load_rejects_expired_cert() {
        let (_dir, cert_path, key_path) = expired_self_signed_pem();
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let err = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_2,
                client_ca: None,
            }),
            256,
            10,
        )
        .await
        .expect_err("expired cert must be rejected");
        assert!(err.to_string().contains("expired"), "got: {err}");
    }

    #[tokio::test]
    async fn tls_min_1_3_rejects_tls_1_2_client() {
        let (_dir, cert_path, key_path) = self_signed_pem();
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_3,
                client_ca: None,
            }),
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        let cert_bytes = std::fs::read(&cert_path).unwrap();
        let mut roots = tokio_rustls::rustls::RootCertStore::empty();
        for c in rustls_pemfile::certs(&mut std::io::Cursor::new(cert_bytes))
            .collect::<Result<Vec<_>, _>>()
            .unwrap()
        {
            roots.add(c).unwrap();
        }
        super::ensure_tls_provider();
        // Client pinned to TLS 1.2 only — must be refused by a 1.3-min server.
        let client_cfg = tokio_rustls::rustls::ClientConfig::builder_with_protocol_versions(&[
            &tokio_rustls::rustls::version::TLS12,
        ])
        .with_root_certificates(roots)
        .with_no_client_auth();

        let conn_str = format!(
            "host=localhost hostaddr={} port={} user=any dbname=laminardb sslmode=require",
            addr.ip(),
            addr.port(),
        );
        let tls = tokio_postgres_rustls::MakeRustlsConnect::new(client_cfg);
        let err = match tokio_postgres::connect(&conn_str, tls).await {
            Ok(_) => panic!("TLS 1.2 client must be refused by a 1.3-min server"),
            Err(e) => e,
        };
        // tokio_postgres wraps the rustls error; flatten the chain so we can
        // assert against the version-mismatch token rustls emits.
        let chain = std::iter::successors(Some(&err as &dyn std::error::Error), |e| e.source())
            .map(|e| e.to_string())
            .collect::<Vec<_>>()
            .join(" | ");
        assert!(
            chain.contains("ProtocolVersion") || chain.contains("incompatible"),
            "expected a TLS version-mismatch error, got: {chain}"
        );

        handle.abort();
    }

    #[tokio::test]
    async fn tls_handshake_succeeds() {
        let (_dir, cert_path, key_path) = self_signed_pem();
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_2,
                client_ca: None,
            }),
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        // Build a client TLS config that trusts the same self-signed cert.
        let cert_bytes = std::fs::read(&cert_path).unwrap();
        let mut roots = tokio_rustls::rustls::RootCertStore::empty();
        for c in rustls_pemfile::certs(&mut std::io::Cursor::new(cert_bytes))
            .collect::<Result<Vec<_>, _>>()
            .unwrap()
        {
            roots.add(c).unwrap();
        }
        super::ensure_tls_provider();
        let client_cfg = tokio_rustls::rustls::ClientConfig::builder()
            .with_root_certificates(roots)
            .with_no_client_auth();

        let conn_str = format!(
            "host=localhost hostaddr={} port={} user=any dbname=laminardb sslmode=require",
            addr.ip(),
            addr.port(),
        );
        let tls = tokio_postgres_rustls::MakeRustlsConnect::new(client_cfg);
        let (client, conn) = tokio_postgres::connect(&conn_str, tls)
            .await
            .expect("TLS handshake + connect");
        tokio::spawn(async move {
            let _ = conn.await;
        });

        let messages = client
            .simple_query("SELECT version()")
            .await
            .expect("query over TLS");
        let v = first_row_value(&messages, 0).expect("row");
        assert!(v.contains("LaminarDB"), "version: {v}");

        handle.abort();
    }

    /// mTLS: with a client_ca configured, a client that presents no cert
    /// must be refused at handshake time.
    #[tokio::test]
    async fn mtls_rejects_client_without_cert() {
        let (_dir, cert_path, key_path) = self_signed_pem();
        let pki = mint_ca_and_client_leaf("alice");
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_2,
                client_ca: Some(&pki.ca_pem_path),
            }),
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        let tls = make_client_tls(&cert_path, None);
        let conn_str = format!(
            "host=localhost hostaddr={} port={} user=any dbname=laminardb sslmode=require",
            addr.ip(),
            addr.port(),
        );
        let err = match tokio_postgres::connect(&conn_str, tls).await {
            Ok(_) => panic!("client without a cert must be refused under mTLS"),
            Err(e) => e,
        };
        assert!(
            err_chain(&err).contains("CertificateRequired")
                || err_chain(&err).contains("HandshakeFailure")
                || err_chain(&err).contains("certificate required"),
            "expected a missing-client-cert error, got: {}",
            err_chain(&err),
        );
        handle.abort();
    }

    /// mTLS: a client cert signed by an unknown CA must be refused.
    #[tokio::test]
    async fn mtls_rejects_untrusted_client_cert() {
        let (_dir, cert_path, key_path) = self_signed_pem();
        let trusted = mint_ca_and_client_leaf("trusted");
        let stranger = mint_ca_and_client_leaf("stranger");
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_2,
                client_ca: Some(&trusted.ca_pem_path),
            }),
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        // Client presents a leaf signed by a CA the server doesn't know.
        let tls = make_client_tls(
            &cert_path,
            Some((stranger.leaf_chain.clone(), stranger.leaf_key.clone_key())),
        );
        let conn_str = format!(
            "host=localhost hostaddr={} port={} user=any dbname=laminardb sslmode=require",
            addr.ip(),
            addr.port(),
        );
        let err = match tokio_postgres::connect(&conn_str, tls).await {
            Ok(_) => panic!("untrusted client cert must be refused"),
            Err(e) => e,
        };
        // rustls maps a verifier-rejected client cert to a fatal alert; the
        // exact variant depends on the protocol version and verifier path
        // (UnknownCA / BadCertificate on 1.2, DecryptError or
        // CertificateUnknown on 1.3). We assert it failed at the TLS layer.
        let chain = err_chain(&err);
        assert!(
            chain.contains("UnknownCA")
                || chain.contains("BadCertificate")
                || chain.contains("CertificateUnknown")
                || chain.contains("DecryptError")
                || chain.contains("HandshakeFailure"),
            "expected a cert-rejection alert, got: {chain}",
        );
        handle.abort();
    }

    /// mTLS: a client cert signed by the configured CA is accepted, and a
    /// SimpleQuery completes over the encrypted+authenticated session.
    #[tokio::test]
    async fn mtls_accepts_trusted_client_cert() {
        let (_dir, cert_path, key_path) = self_signed_pem();
        let pki = mint_ca_and_client_leaf("alice");
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.start().await.expect("db starts");
        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            Some(super::TlsPaths {
                cert: &cert_path,
                key: &key_path,
                min_version: super::TlsMinVersion::V1_2,
                client_ca: Some(&pki.ca_pem_path),
            }),
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        let tls = make_client_tls(
            &cert_path,
            Some((pki.leaf_chain.clone(), pki.leaf_key.clone_key())),
        );
        let conn_str = format!(
            "host=localhost hostaddr={} port={} user=any dbname=laminardb sslmode=require",
            addr.ip(),
            addr.port(),
        );
        let (client, conn) = tokio_postgres::connect(&conn_str, tls)
            .await
            .expect("mTLS handshake + connect");
        tokio::spawn(async move {
            let _ = conn.await;
        });

        let messages = client
            .simple_query("SELECT version()")
            .await
            .expect("query over mTLS");
        let v = first_row_value(&messages, 0).expect("row");
        assert!(v.contains("LaminarDB"), "version: {v}");
        handle.abort();
    }

    /// Build a `TlsReloadState` directly for unit-testing the reload path
    /// without standing up a listener.
    fn build_reload_state(cert: &std::path::Path, key: &std::path::Path) -> super::TlsReloadState {
        let paths = super::TlsPaths {
            cert,
            key,
            min_version: super::TlsMinVersion::V1_2,
            client_ca: None,
        };
        let acceptor = super::load_tls_acceptor(super::TlsPaths {
            cert: paths.cert,
            key: paths.key,
            min_version: paths.min_version,
            client_ca: paths.client_ca,
        })
        .expect("initial acceptor loads");
        super::TlsReloadState {
            paths: super::TlsConfigPaths::from_paths(&paths),
            acceptor: parking_lot::Mutex::new(Arc::new(acceptor)),
        }
    }

    /// Hot-reload: writing a fresh cert+key over the configured paths and
    /// calling `try_reload_tls` swaps the acceptor under the mutex.
    #[test]
    fn tls_reload_swaps_acceptor_on_valid_pair() {
        let dir = tempfile::tempdir().unwrap();
        let cert_path = dir.path().join("cert.pem");
        let key_path = dir.path().join("key.pem");
        // Initial cert.
        let first = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
        std::fs::write(&cert_path, first.cert.pem()).unwrap();
        std::fs::write(&key_path, first.key_pair.serialize_pem()).unwrap();

        let state = build_reload_state(&cert_path, &key_path);
        let before = state.snapshot();

        // Rotate to a brand-new pair.
        let second = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
        std::fs::write(&cert_path, second.cert.pem()).unwrap();
        std::fs::write(&key_path, second.key_pair.serialize_pem()).unwrap();

        super::try_reload_tls(&state).expect("reload succeeds");
        let after = state.snapshot();
        assert!(
            !Arc::ptr_eq(&before, &after),
            "acceptor pointer must change after a successful reload",
        );
    }

    /// Hot-reload: a corrupt cert file leaves the previous acceptor in
    /// place — TLS doesn't go down on a bad rotation.
    #[test]
    fn tls_reload_keeps_old_acceptor_on_garbage() {
        let dir = tempfile::tempdir().unwrap();
        let cert_path = dir.path().join("cert.pem");
        let key_path = dir.path().join("key.pem");
        let first = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
        std::fs::write(&cert_path, first.cert.pem()).unwrap();
        std::fs::write(&key_path, first.key_pair.serialize_pem()).unwrap();

        let state = build_reload_state(&cert_path, &key_path);
        let before = state.snapshot();

        // Truncate cert.pem to non-PEM garbage.
        std::fs::write(&cert_path, b"this is not a certificate").unwrap();
        let err = super::try_reload_tls(&state).expect_err("reload must fail");
        let after = state.snapshot();
        assert!(
            Arc::ptr_eq(&before, &after),
            "acceptor must be unchanged on reload failure",
        );
        assert!(
            err.to_string().contains("pgwire_tls_cert"),
            "error should mention pgwire_tls_cert, got: {err}",
        );
    }

    /// Flatten an error and its `source()` chain to a single string for
    /// substring assertions.
    fn err_chain(err: &(dyn std::error::Error + 'static)) -> String {
        std::iter::successors(Some(err), |e| e.source())
            .map(|e| e.to_string())
            .collect::<Vec<_>>()
            .join(" | ")
    }

    /// Push one row into the `trades` source so subsequent SUBSCRIBE
    /// reads have something to drain. Returns the schema for tests
    /// that want to build their own batches.
    async fn push_one_trade(
        db: &Arc<LaminarDB>,
        symbol: &str,
        price: f64,
    ) -> arrow_schema::SchemaRef {
        let handle = db.source_untyped("trades").expect("source handle");
        let schema = handle.schema().clone();
        let batch = arrow_array::RecordBatch::try_new(
            Arc::clone(&schema),
            vec![
                Arc::new(arrow_array::StringArray::from(vec![symbol])),
                Arc::new(arrow_array::Float64Array::from(vec![price])),
            ],
        )
        .expect("batch");
        handle.push_arrow(batch).expect("push");
        schema
    }

    /// Ingest a row and return both the running server and the underlying db
    /// so tests can keep pushing rows after the listener is up.
    async fn spawn_with_data() -> (
        Arc<LaminarDB>,
        std::net::SocketAddr,
        tokio::task::JoinHandle<()>,
    ) {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE trades (symbol VARCHAR, price DOUBLE)")
            .await
            .expect("create source");
        db.execute(
            "CREATE MATERIALIZED VIEW prices AS \
             SELECT symbol, price FROM trades",
        )
        .await
        .expect("create mv");
        db.start().await.expect("db starts");

        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            256,
            10,
        )
        .await
        .expect("pgwire serve");
        (db, addr, handle)
    }

    /// Same as `spawn_with_data`, but `prices` is a STREAM with retained
    /// history. Lets cursor tests push rows *before* SUBSCRIBE attaches
    /// without losing them — the receiver replays on attach.
    async fn spawn_with_retained_data() -> (
        Arc<LaminarDB>,
        std::net::SocketAddr,
        tokio::task::JoinHandle<()>,
    ) {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        db.execute("CREATE SOURCE trades (symbol VARCHAR, price DOUBLE)")
            .await
            .expect("create source");
        db.execute(
            "CREATE STREAM prices AS SELECT symbol, price FROM trades \
             WITH ('retain_history' = '4mb')",
        )
        .await
        .expect("create stream");
        db.start().await.expect("db starts");

        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            256,
            10,
        )
        .await
        .expect("pgwire serve");
        (db, addr, handle)
    }

    /// `prepare()` triggers `Parse` + `Describe(Statement)`. Verifies the
    /// extended-query parser resolves stream schemas at parse time and
    /// returns column metadata to the client.
    #[tokio::test]
    async fn extended_query_describe_subscribe_returns_columns() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        let stmt = client
            .prepare("SUBSCRIBE prices")
            .await
            .expect("prepare SUBSCRIBE prices");

        let cols = stmt.columns();
        assert_eq!(cols.len(), 2, "expected 2 columns, got {}", cols.len());
        assert_eq!(cols[0].name(), "symbol");
        assert_eq!(cols[1].name(), "price");
        assert_eq!(cols[0].type_(), &tokio_postgres::types::Type::VARCHAR);
        assert_eq!(cols[1].type_(), &tokio_postgres::types::Type::FLOAT8);

        handle.abort();
    }

    /// Unknown stream → typed PG error at `Parse` time, before any rows
    /// are pulled.
    #[tokio::test]
    async fn extended_query_prepare_unknown_stream_errors() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        let err = client
            .prepare("SUBSCRIBE no_such_view")
            .await
            .expect_err("must fail at Parse");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(db_err.message().contains("no_such_view"));

        handle.abort();
    }

    /// Bind + Execute with `max_rows=1` against a portal returns one row at a
    /// time and `PortalSuspended`. Drives the binary-format encoders for
    /// VARCHAR + FLOAT8.
    #[tokio::test]
    async fn extended_query_binary_chunked_subscribe() {
        let (db, addr, handle) = spawn_with_data().await;
        let mut client = connect(addr).await;

        // tokio_postgres' `bind` + `query_portal` uses the extended-query
        // protocol with binary format for known column types — the path
        // JDBC and asyncpg take with prepared statements.
        let tx = client.transaction().await.expect("BEGIN");
        let stmt = tx.prepare("SUBSCRIBE prices").await.expect("prepare");
        let portal = tx.bind(&stmt, &[]).await.expect("bind portal");

        // The MV broadcast has no receiver until `Execute` reaches the
        // server and runs `do_query` → `open_subscription`. We can't push
        // from this task before query_portal because query_portal blocks
        // waiting for a row, so spawn the pushes from a sibling task with
        // a short head start for the receiver to attach. With cap=0
        // retention, a push that lands before the receiver is dropped.
        let pusher = {
            let db = Arc::clone(&db);
            tokio::spawn(async move {
                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                push_one_trade(&db, "AAPL", 150.5).await;
                push_one_trade(&db, "GOOG", 2700.25).await;
            })
        };

        let first = tokio::time::timeout(
            std::time::Duration::from_secs(3),
            tx.query_portal(&portal, 1),
        )
        .await
        .expect("first chunk arrives within 3s")
        .expect("query_portal #1");
        assert_eq!(first.len(), 1);
        let symbol: &str = first[0].get(0);
        let price: f64 = first[0].get(1);
        assert_eq!(symbol, "AAPL");
        assert!((price - 150.5).abs() < 1e-9);

        let second = tokio::time::timeout(
            std::time::Duration::from_secs(3),
            tx.query_portal(&portal, 1),
        )
        .await
        .expect("second chunk arrives within 3s")
        .expect("query_portal #2");
        assert_eq!(second.len(), 1);
        let symbol: &str = second[0].get(0);
        let price: f64 = second[0].get(1);
        assert_eq!(symbol, "GOOG");
        assert!((price - 2700.25).abs() < 1e-9);

        pusher.await.expect("push task");
        handle.abort();
    }

    /// Regression: binary encoding of `TIMESTAMP` columns must downcast
    /// the Arrow array as its unit-specific primitive type
    /// (`PrimitiveArray<TimestampMicrosecondType>`, not
    /// `PrimitiveArray<Int64Type>`). A bug in this branch would panic on
    /// the first row.
    #[tokio::test]
    async fn extended_query_binary_timestamp() {
        let db = Arc::new(LaminarDB::open().expect("db opens"));
        // `WATERMARK FOR ts AS ts - INTERVAL '0' SECOND` declares event time
        // so the streaming pipeline drives progress on the timestamp
        // column — without it, the MV stays empty.
        db.execute(
            "CREATE SOURCE events (ts TIMESTAMP, sym VARCHAR, \
             WATERMARK FOR ts AS ts - INTERVAL '0' SECOND)",
        )
        .await
        .expect("create source");
        db.execute("CREATE MATERIALIZED VIEW ev AS SELECT ts, sym FROM events")
            .await
            .expect("create mv");
        db.start().await.expect("db starts");

        let (addr, handle) = super::serve(
            Arc::clone(&db),
            "127.0.0.1:0",
            HashMap::new(),
            false,
            None,
            256,
            10,
        )
        .await
        .expect("pgwire serve");

        let mut client = connect(addr).await;
        let tx = client.transaction().await.expect("BEGIN");
        let stmt = tx.prepare("SUBSCRIBE ev").await.expect("prepare");
        let portal = tx.bind(&stmt, &[]).await.expect("bind");

        let expected = chrono::NaiveDate::from_ymd_opt(2026, 5, 9)
            .unwrap()
            .and_hms_opt(0, 0, 0)
            .unwrap();
        let ts_us = expected.and_utc().timestamp_micros();

        // Push from a sibling task after a short delay so the MV
        // broadcast receiver (created inside `Execute`) is attached
        // before send_batch fires. See the matching note in
        // `extended_query_binary_chunked_subscribe`.
        let pusher = {
            let db = Arc::clone(&db);
            tokio::spawn(async move {
                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                let src = db.source_untyped("events").expect("source");
                let batch = arrow_array::RecordBatch::try_new(
                    src.schema().clone(),
                    vec![
                        Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![ts_us])),
                        Arc::new(arrow_array::StringArray::from(vec!["AAPL"])),
                    ],
                )
                .expect("batch");
                src.push_arrow(batch).expect("push");
            })
        };

        let rows = tokio::time::timeout(
            std::time::Duration::from_secs(3),
            tx.query_portal(&portal, 1),
        )
        .await
        .expect("row arrives within 3s")
        .expect("query_portal");
        assert_eq!(rows.len(), 1);

        let ts: chrono::NaiveDateTime = rows[0].get(0);
        let sym: &str = rows[0].get(1);
        assert_eq!(ts, expected);
        assert_eq!(sym, "AAPL");

        pusher.await.expect("push task");
        handle.abort();
    }

    /// DDL on the extended-query path is refused at `Parse` with a typed
    /// 0A000 error pointing at the HTTP endpoint — same surface as the
    /// SimpleQuery path.
    #[tokio::test]
    async fn extended_query_ddl_rejected() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        let err = client
            .prepare("CREATE SOURCE more_trades (sym VARCHAR)")
            .await
            .expect_err("DDL must be rejected at Parse");
        let db_err = err.as_db_error().expect("typed PG error");
        assert!(
            db_err.message().contains("/api/v1/sql"),
            "message: {}",
            db_err.message()
        );

        handle.abort();
    }

    /// `\set FETCH_COUNT N` flow: BEGIN; DECLARE …; FETCH N FROM …; CLOSE; COMMIT.
    /// All over SimpleQuery — the path psql uses when `FETCH_COUNT` is set.
    /// Uses the retained-history variant so we can push before SUBSCRIBE.
    #[tokio::test]
    async fn cursor_declare_fetch_close_happy_path() {
        let (db, addr, handle) = spawn_with_retained_data().await;
        let client = connect(addr).await;

        for i in 0..4 {
            push_one_trade(&db, &format!("S{i}"), i as f64).await;
        }

        client.simple_query("BEGIN").await.expect("BEGIN");
        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");

        let messages = client
            .simple_query("FETCH 2 FROM c")
            .await
            .expect("FETCH 2");
        let row_count = messages
            .iter()
            .filter(|m| matches!(m, SimpleQueryMessage::Row(_)))
            .count();
        assert_eq!(row_count, 2, "expected exactly 2 rows from FETCH 2");

        client.simple_query("CLOSE c").await.expect("CLOSE");
        client.simple_query("COMMIT").await.expect("COMMIT");

        handle.abort();
    }

    /// COMMIT must close any open cursors. After COMMIT, FETCH against the
    /// same name returns "cursor does not exist".
    #[tokio::test]
    async fn cursor_commit_closes_cursors() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        client.simple_query("BEGIN").await.expect("BEGIN");
        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");
        client.simple_query("COMMIT").await.expect("COMMIT");

        let err = client
            .simple_query("FETCH 1 FROM c")
            .await
            .expect_err("FETCH after COMMIT must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "34000", "got {db_err:?}");

        handle.abort();
    }

    /// ROLLBACK closes cursors too — same reaper as COMMIT.
    #[tokio::test]
    async fn cursor_rollback_closes_cursors() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        client.simple_query("BEGIN").await.expect("BEGIN");
        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");
        client.simple_query("ROLLBACK").await.expect("ROLLBACK");

        let err = client
            .simple_query("FETCH 1 FROM c")
            .await
            .expect_err("FETCH after ROLLBACK must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "34000", "got {db_err:?}");

        handle.abort();
    }

    /// Explicit CLOSE works outside a transaction. PG allows DECLARE without
    /// BEGIN; we follow that for parity with `\set FETCH_COUNT 0` clients.
    #[tokio::test]
    async fn cursor_close_explicit() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");
        client.simple_query("CLOSE c").await.expect("CLOSE");

        let err = client
            .simple_query("FETCH 1 FROM c")
            .await
            .expect_err("FETCH after CLOSE must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "34000", "got {db_err:?}");

        handle.abort();
    }

    /// `SCROLL`, `BINARY`, `WITH HOLD` all rejected at parse time.
    #[tokio::test]
    async fn cursor_unsupported_modifiers_rejected() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        for sql in [
            "DECLARE c SCROLL CURSOR FOR SUBSCRIBE prices",
            "DECLARE c BINARY CURSOR FOR SUBSCRIBE prices",
            "DECLARE c CURSOR WITH HOLD FOR SUBSCRIBE prices",
            "DECLARE c INSENSITIVE CURSOR FOR SUBSCRIBE prices",
        ] {
            let err = client
                .simple_query(sql)
                .await
                .expect_err(&format!("{sql} must fail"));
            let db_err = err.as_db_error().expect("typed PG error");
            assert_eq!(
                db_err.code().code(),
                "42601",
                "{sql}: expected parse error, got {db_err:?}"
            );
        }

        handle.abort();
    }

    /// `FETCH BACKWARD` and other reverse / absolute directions are rejected
    /// because SUBSCRIBE is forward-only.
    #[tokio::test]
    async fn cursor_backward_directions_rejected() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");

        for sql in [
            "FETCH PRIOR FROM c",
            "FETCH BACKWARD 1 FROM c",
            "FETCH FIRST FROM c",
            "FETCH LAST FROM c",
            "FETCH ABSOLUTE 1 FROM c",
            "FETCH RELATIVE 1 FROM c",
        ] {
            let err = client
                .simple_query(sql)
                .await
                .expect_err(&format!("{sql} must fail"));
            let db_err = err.as_db_error().expect("typed PG error");
            assert_eq!(db_err.code().code(), "0A000", "{sql}: got {db_err:?}");
        }

        client.simple_query("CLOSE c").await.expect("CLOSE");
        handle.abort();
    }

    /// `DECLARE … CURSOR FOR <SELECT …>` (regular query, not SUBSCRIBE) is
    /// not supported on pgwire.
    #[tokio::test]
    async fn cursor_for_non_subscribe_rejected() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("DECLARE c CURSOR FOR SELECT 1")
            .await
            .expect_err("DECLARE FOR SELECT must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "0A000", "got {db_err:?}");

        handle.abort();
    }

    /// FETCH against a name we never declared returns 34000 (invalid_cursor_name).
    #[tokio::test]
    async fn cursor_fetch_unknown_name_errors() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        let err = client
            .simple_query("FETCH 1 FROM nope")
            .await
            .expect_err("must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "34000", "got {db_err:?}");

        handle.abort();
    }

    /// A multi-row batch with `FETCH 1` repeated must return each row in
    /// order — leftover rows persist on the cursor instead of being dropped
    /// when the response stream ends. With the bug, `FETCH 1` would consume
    /// the batch internally, return row[0], and discard row[1].
    #[tokio::test]
    async fn cursor_fetch_preserves_leftover_rows_in_one_batch() {
        let (db, addr, handle) = spawn_with_retained_data().await;
        let client = connect(addr).await;

        let src = db.source_untyped("trades").expect("source");
        let batch = arrow_array::RecordBatch::try_new(
            src.schema().clone(),
            vec![
                Arc::new(arrow_array::StringArray::from(vec!["AAPL", "GOOG"])),
                Arc::new(arrow_array::Float64Array::from(vec![1.0, 2.0])),
            ],
        )
        .expect("batch");
        src.push_arrow(batch).expect("push");

        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");

        let first = client
            .simple_query("FETCH 1 FROM c")
            .await
            .expect("FETCH 1");
        let r1: Vec<&str> = first
            .iter()
            .filter_map(|m| match m {
                SimpleQueryMessage::Row(r) => r.get(0),
                _ => None,
            })
            .collect();
        assert_eq!(r1, vec!["AAPL"]);

        let second = client
            .simple_query("FETCH 1 FROM c")
            .await
            .expect("FETCH 1");
        let r2: Vec<&str> = second
            .iter()
            .filter_map(|m| match m {
                SimpleQueryMessage::Row(r) => r.get(0),
                _ => None,
            })
            .collect();
        assert_eq!(r2, vec!["GOOG"]);

        client.simple_query("CLOSE c").await.expect("CLOSE");
        handle.abort();
    }

    /// Re-DECLAREing an open cursor name returns 42P03; user must CLOSE first.
    #[tokio::test]
    async fn cursor_duplicate_declare_rejected() {
        let (_db, addr, handle) = spawn_with_data().await;
        let client = connect(addr).await;

        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("first DECLARE");

        let err = client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect_err("duplicate DECLARE must fail");
        let db_err = err.as_db_error().expect("typed PG error");
        assert_eq!(db_err.code().code(), "42P03", "got {db_err:?}");

        // After CLOSE the name is free again.
        client.simple_query("CLOSE c").await.expect("CLOSE");
        client
            .simple_query("DECLARE c CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("re-DECLARE after CLOSE");
        client.simple_query("CLOSE c").await.expect("CLOSE again");

        handle.abort();
    }

    /// Cursor name lookup is case-insensitive (PG identifier folding rules).
    #[tokio::test]
    async fn cursor_name_case_insensitive() {
        let (db, addr, handle) = spawn_with_retained_data().await;
        let client = connect(addr).await;
        push_one_trade(&db, "AAPL", 1.0).await;

        client
            .simple_query("DECLARE MyCursor CURSOR FOR SUBSCRIBE prices")
            .await
            .expect("DECLARE");

        let messages = client
            .simple_query("FETCH 1 FROM mycursor")
            .await
            .expect("FETCH from lowercased name");
        let row_count = messages
            .iter()
            .filter(|m| matches!(m, SimpleQueryMessage::Row(_)))
            .count();
        assert_eq!(row_count, 1);

        client.simple_query("CLOSE MYCURSOR").await.expect("CLOSE");
        handle.abort();
    }
}