use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use anyhow::Context;
use base64::{engine, Engine as _};
use chrono::Utc;
use futures::future::BoxFuture;
use futures::{FutureExt, StreamExt, TryStreamExt};
use itertools::Itertools;
use rust_decimal::{prelude::FromPrimitive, Decimal};
use serde_json::value::RawValue;
use serde_json::Map;
use serde_json::Value;
use tokio_postgres::Client;
use tokio_postgres::{types::ToSql, Row};
use tokio_postgres::{
    types::{FromSql, IsNull, Kind, Type},
    Column,
};
use uuid::Uuid;
use windmill_common::azure_workload_identity::WORKLOAD_IDENTITY_PASSWORD;
use windmill_common::error::to_anyhow;
use windmill_common::error::{self, Error};
use windmill_common::worker::{
    to_raw_value, Connection, SqlAnnotations, SqlResultCollectionStrategy, CLOUD_HOSTED,
};
use windmill_common::workspaces::{
    get_datatable_resource_from_db, parse_datatable_ref, parse_datatable_ref_for, DatatableAccess,
};
use windmill_common::{PgDatabase, PrepareQueryColumnInfo, PrepareQueryResult, DB};
use windmill_parser::{Arg, Typ};
use windmill_parser_sql::{
    parse_db_resource, parse_pg_statement_arg_positions, parse_pgsql_sig_with_typed_schema,
    parse_s3_mode, parse_sql_blocks,
};
use windmill_queue::{CanceledBy, MiniPulledJob};

use crate::agent_workers::get_datatable_resource_from_agent_http;
use crate::common::{
    build_args_values, get_reserved_variables, s3_mode_args_to_worker_data,
    s3_stream_and_upload_with_logs, sizeof_val, OccupancyMetrics, S3ModeWorkerData,
};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use crate::sql_s3_input::fetch_s3object_as_json_text;
use crate::sql_utils::remove_comments;
use crate::{max_sql_result_size, sql_result_too_large_error, to_raw_value_within};
use bytes::Buf;
use lazy_static::lazy_static;
use windmill_common::client::AuthedClient;
use windmill_types::s3::S3Object;

lazy_static! {
    pub static ref CACHE_HITS: AtomicU64 = AtomicU64::new(0);
}

/// Idle connections kept per cache key, checked out by a job for the duration
/// of its query and checked back in afterwards.
///
/// Invariant: the process never keeps an idle connection in a pool scope
/// ([`pg_pool_scope`]) while one of its jobs runs on a connection of its own in
/// that scope. Behind a session-mode pooler (PgBouncer, RDS Proxy pinning, …)
/// an idle connection holds a server slot, and the job on its own connection
/// waits for that slot with no timeout, forever if other jobs keep reusing the
/// idle connection.
static PG_CONNECTION_CACHE: std::sync::Mutex<PgConnectionCache> =
    std::sync::Mutex::new(PgConnectionCache { idle: Vec::new(), own: Vec::new(), next_own_id: 0 });
static PG_CONNECTION_CACHE_SWEEPER: std::sync::Once = std::sync::Once::new();
const PG_CONNECTION_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
/// Across all keys. Native mode runs 8 jobs at once in one process.
const PG_MAX_IDLE_CONNECTIONS: usize = 8;
/// A pooler can also share slots across pool scopes: two host names for one
/// pooler, or a per-database or per-user limit. A job on its own connection for
/// longer than this may be waiting on such a slot, so no idle connection is
/// kept, in any scope, until it finishes.
const PG_OWN_CONNECTION_STALL: Duration = Duration::from_secs(5);
/// The session reset run before reusing a cached connection is a few trivial
/// statements. A server that vanished without closing the socket never answers
/// it, and waiting for TCP to notice takes the keepalive budget (~2 min).
const PG_RESET_PROBE_TIMEOUT: Duration = Duration::from_secs(5);
/// How long a closed connection gets to wind down. An idle one only has its Terminate to
/// send; one closed with results still arriving is cut off here.
const PG_CLOSE_GRACE: Duration = Duration::from_secs(2);
/// How long the connection of a failed job gets to end its query and reset its session.
/// Past this it is closed regardless.
const PG_CLOSE_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5);
/// How long a failed job's last query gets to be answered before it is taken to be still
/// sending results, and cancelled.
const PG_UNREAD_RESULT_GRACE: Duration = Duration::from_millis(250);
/// Clears the session state a job can leave behind. Each statement is explained where a
/// cached connection is probed with it.
const PG_SESSION_RESET: &str = "RESET ALL; RESET SESSION AUTHORIZATION; UNLISTEN *; CLOSE ALL; \
     SELECT pg_advisory_unlock_all();";

struct PgConnectionCache {
    /// Least recently used first.
    idle: Vec<PgConnection>,
    /// Jobs of this process running on a connection of their own.
    own: Vec<OwnConnection>,
    next_own_id: u64,
}

struct OwnConnection {
    id: u64,
    scope: String,
    since: Instant,
}

impl PgConnectionCache {
    /// Registers a job on its own connection. Returns the idle connections of
    /// its scope, to be closed before it connects.
    fn register_own(&mut self, scope: &str) -> (u64, Vec<PgConnection>) {
        let id = self.next_own_id;
        self.next_own_id += 1;
        self.own
            .push(OwnConnection { id, scope: scope.to_string(), since: Instant::now() });
        let (evicted, kept) = std::mem::take(&mut self.idle)
            .into_iter()
            .partition(|c| c.scope == scope);
        self.idle = kept;
        (id, evicted)
    }

    fn stalled(&self) -> bool {
        self.own
            .iter()
            .any(|o| o.since.elapsed() >= PG_OWN_CONNECTION_STALL)
    }

    fn may_keep(&self, scope: &str) -> bool {
        !self.stalled() && self.own.iter().all(|o| o.scope != scope)
    }
}

fn pg_connection_cache() -> std::sync::MutexGuard<'static, PgConnectionCache> {
    PG_CONNECTION_CACHE
        .lock()
        .unwrap_or_else(|e| e.into_inner())
}

/// What a session pooler shares server slots across: PgBouncer pools by
/// database and user. Keys differing only in sslmode, TLS inputs, auth mode or
/// options share a scope. Hosts are compared by name, not address.
fn pg_pool_scope(database: &PgDatabase) -> String {
    format!(
        "{:?}",
        (
            database.host.trim().to_ascii_lowercase(),
            database.port.unwrap_or(5432),
            &database.dbname,
            database.login_name().trim(),
        )
    )
}

struct PgConnection {
    key: String,
    scope: String,
    client: Client,
    connection_task: Option<tokio::task::JoinHandle<()>>,
    last_used: Instant,
}

impl Drop for PgConnection {
    /// The connection task sends the protocol Terminate and ends by itself once the client is
    /// gone, right after this returns. Aborting it straight away closes the socket without
    /// one, and a transaction pooler may then hand the server connection to its next client
    /// as the job left it, inside a failed transaction.
    fn drop(&mut self) {
        let Some(mut task) = self.connection_task.take() else {
            return;
        };
        match tokio::runtime::Handle::try_current() {
            Ok(runtime) => {
                runtime.spawn(async move {
                    if tokio::time::timeout(PG_CLOSE_GRACE, &mut task)
                        .await
                        .is_err()
                    {
                        task.abort();
                    }
                });
            }
            Err(_) => task.abort(),
        }
    }
}

/// A job's claim on a connection. Dropping it checks the connection back in,
/// or closes it when its query did not finish or keeping it would break the
/// cache invariant.
struct PgConnectionLease {
    scope: String,
    /// Set while the job is on a connection of its own.
    own: Option<u64>,
    conn: Option<PgConnection>,
    /// Check the connection back in on drop. Set once the job's query has read
    /// all its results, so every other way out closes the connection.
    keep: bool,
}

impl PgConnectionLease {
    /// Takes the most recently used idle connection for `key`. Otherwise
    /// registers the job as on its own connection and closes the idle
    /// connections of its scope, so they cannot hold the server slot its own
    /// connection is about to need.
    fn checkout(key: &str, scope: String) -> Self {
        let (lease, evicted) = {
            let mut cache = pg_connection_cache();
            match cache.idle.iter().rposition(|c| c.key == key) {
                Some(i) => {
                    let conn = cache.idle.remove(i);
                    (
                        Self { scope, own: None, conn: Some(conn), keep: false },
                        vec![],
                    )
                }
                None => {
                    let (id, evicted) = cache.register_own(&scope);
                    (
                        Self { scope, own: Some(id), conn: None, keep: false },
                        evicted,
                    )
                }
            }
        };
        drop(evicted);
        lease
    }

    fn uncached() -> Self {
        Self { scope: String::new(), own: None, conn: None, keep: false }
    }

    /// Gives up a cached connection that failed its probe; the job then
    /// connects on its own.
    fn discard(&mut self) {
        let conn = self.conn.take();
        let evicted = if self.own.is_none() {
            let (id, evicted) = pg_connection_cache().register_own(&self.scope);
            self.own = Some(id);
            evicted
        } else {
            vec![]
        };
        drop(conn);
        drop(evicted);
    }

    fn client(&self) -> &Client {
        &self.conn.as_ref().expect("lease holds a connection").client
    }
}

impl Drop for PgConnectionLease {
    fn drop(&mut self) {
        let conn = self.conn.take();
        let (kept, closed) = {
            let mut cache = pg_connection_cache();
            if let Some(id) = self.own {
                cache.own.retain(|o| o.id != id);
            }
            match conn {
                Some(mut conn) if self.keep && cache.may_keep(&self.scope) => {
                    conn.last_used = Instant::now();
                    cache.idle.push(conn);
                    let over = cache.idle.len().saturating_sub(PG_MAX_IDLE_CONNECTIONS);
                    (true, cache.idle.drain(..over).collect::<Vec<_>>())
                }
                conn => (false, conn.into_iter().collect()),
            }
        };
        drop(closed);
        if kept {
            PG_CONNECTION_CACHE_SWEEPER.call_once(|| {
                tokio::spawn(sweep_idle_pg_connections());
            });
        }
    }
}

/// Closes the connection of a job that failed or left a transaction open, in the background
/// so the job ends without waiting. A pooler in transaction mode hands the server connection
/// to its next client the way this one leaves it, and one closed inside a transaction it may
/// not get back at all, so the transaction is ended and the session reset first.
///
/// `abandoned` is set for a query the timeout or a cancellation cut short. Closing does not
/// stop it, the server works on until it next writes to the client, so it is cancelled. So
/// is one still sending results the job stopped reading, which would otherwise be drained.
fn close_unclean_connection(
    lease: PgConnectionLease,
    database: PgDatabase,
    token_auth: bool,
    abandoned: bool,
) {
    if lease.conn.is_none() {
        return;
    }
    tokio::spawn(async move {
        let client = lease.client();
        let cleanup = async {
            let cancel = || database.cancel_query(client.cancel_token(), token_auth);
            if abandoned {
                cancel().await?;
            }
            // Answered once the server is done with the job's last query. A cancel request
            // is only sent by then, not acted on, and a pooler drops one whose client left.
            let settled = client.batch_execute("");
            tokio::pin!(settled);
            let answer = if abandoned {
                settled.await
            } else {
                match tokio::time::timeout(PG_UNREAD_RESULT_GRACE, &mut settled).await {
                    Ok(answer) => answer,
                    Err(_) => {
                        cancel().await?;
                        settled.await
                    }
                }
            };
            // A cancel request that found the query already over lands on what comes next.
            let cancelled = |e: &tokio_postgres::Error| {
                e.code() == Some(&tokio_postgres::error::SqlState::QUERY_CANCELED)
            };
            answer
                .or_else(|e| if cancelled(&e) { Ok(()) } else { Err(e) })
                .map_err(to_anyhow)?;
            let reset = || async {
                let in_transaction =
                    client.transaction_status() != tokio_postgres::TransactionStatus::Idle;
                let rollback = if in_transaction { "ROLLBACK; " } else { "" };
                client
                    .batch_execute(&format!("{rollback}{PG_SESSION_RESET}"))
                    .await
            };
            let mut done = reset().await;
            if done.as_ref().is_err_and(cancelled) {
                done = reset().await;
            }
            done.map_err(to_anyhow)?;
            Ok::<_, Error>(())
        };
        match tokio::time::timeout(PG_CLOSE_CLEANUP_TIMEOUT, cleanup).await {
            Ok(Ok(())) => {}
            // A connection that failed the job has nothing left to clean up.
            Ok(Err(e)) if abandoned => {
                tracing::warn!("could not cancel an abandoned postgres query: {e}")
            }
            Ok(Err(_)) => {}
            Err(_) => tracing::warn!(
                "a postgres connection did not settle within {}s of its job ending, closing it",
                PG_CLOSE_CLEANUP_TIMEOUT.as_secs()
            ),
        }
        drop(lease);
    });
}

async fn sweep_idle_pg_connections() {
    loop {
        tokio::time::sleep(Duration::from_secs(1)).await;
        let (stalled, closed) = {
            let mut cache = pg_connection_cache();
            let stalled = cache.stalled();
            let (closed, kept): (Vec<_>, Vec<_>) = std::mem::take(&mut cache.idle)
                .into_iter()
                .partition(|c| stalled || c.last_used.elapsed() > PG_CONNECTION_IDLE_TIMEOUT);
            cache.idle = kept;
            (stalled, closed)
        };
        if !closed.is_empty() {
            if stalled {
                tracing::info!(
                    "Closing {} cached pg executor connection(s): a job has been on its own connection for over {}s",
                    closed.len(),
                    PG_OWN_CONNECTION_STALL.as_secs()
                );
            } else {
                tracing::info!(
                    "Closing {} cached pg executor connection(s) due to inactivity",
                    closed.len()
                );
            }
        }
    }
}

/// PostgreSQL's role, database and server caps, by SQLSTATE 53300 since the
/// server translates the message under a non-English `lc_messages`, and
/// PgBouncer's client and user caps, whose messages are English only and come
/// with a generic code.
fn is_connection_cap_refusal(e: &Error) -> bool {
    if let Error::Anyhow { error, .. } = e {
        let postgres_cap = error
            .chain()
            .filter_map(|cause| cause.downcast_ref::<tokio_postgres::Error>())
            .any(|pg| pg.code() == Some(&tokio_postgres::error::SqlState::TOO_MANY_CONNECTIONS));
        if postgres_cap {
            return true;
        }
    }
    e.to_string().contains("no more connections allowed")
}

fn close_idle_pg_connections() -> usize {
    let idle = std::mem::take(&mut pg_connection_cache().idle);
    idle.len()
}

/// How far a job's statements have got, read by its stall warning.
struct PgProgress {
    start: std::time::Instant,
    statement: AtomicUsize,
    rows: AtomicU64,
    last_progress_ms: AtomicU64,
}

impl PgProgress {
    fn new() -> Self {
        Self {
            start: std::time::Instant::now(),
            statement: AtomicUsize::new(0),
            rows: AtomicU64::new(0),
            last_progress_ms: AtomicU64::new(0),
        }
    }

    fn elapsed_ms(&self) -> u64 {
        self.start.elapsed().as_millis() as u64
    }

    fn start_statement(&self, index: usize) {
        self.statement.store(index, Ordering::Relaxed);
        self.last_progress_ms
            .store(self.elapsed_ms(), Ordering::Relaxed);
    }

    fn row(&self) {
        self.rows.fetch_add(1, Ordering::Relaxed);
        self.last_progress_ms
            .store(self.elapsed_ms(), Ordering::Relaxed);
    }

    fn rows(&self) -> u64 {
        self.rows.load(Ordering::Relaxed)
    }
}

const PG_STALL_WARNING_AFTER: Duration = Duration::from_secs(5 * 60);
/// Steps faster than this are not logged: a log line is a write to the main DB,
/// and a script split into many statements would pay one per statement.
const PG_SLOW_STEP: Duration = Duration::from_secs(1);

fn connection_kind(fresh_connection: bool) -> &'static str {
    if fresh_connection {
        "new connection"
    } else {
        "worker's cached connection"
    }
}

/// Warns once per stall in the job log; never completes. A statement still
/// computing its first row also trips it, which is why it only warns: the job
/// timeout stays the one thing that stops a job.
async fn warn_on_stalled_statement(
    progress: &PgProgress,
    statement_count: usize,
    fresh_connection: bool,
    job_id: Uuid,
    workspace_id: &str,
    conn: &Connection,
) -> std::convert::Infallible {
    let mut warned_for = None;
    loop {
        tokio::time::sleep(Duration::from_secs(30)).await;
        let last = progress.last_progress_ms.load(Ordering::Relaxed);
        let stalled = Duration::from_millis(progress.elapsed_ms().saturating_sub(last));
        if stalled < PG_STALL_WARNING_AFTER || warned_for == Some(last) {
            continue;
        }
        warned_for = Some(last);
        windmill_queue::append_logs(
            &job_id,
            workspace_id,
            format!(
                "No new row for {} min on statement {}/{statement_count} ({} rows so far, {}). \
                 A query still running on the database is normal; if the database shows no \
                 active query for this connection, the job is stuck in the worker.\n",
                stalled.as_secs() / 60,
                progress.statement.load(Ordering::Relaxed) + 1,
                progress.rows(),
                connection_kind(fresh_connection),
            ),
            conn,
        )
        .await;
    }
}

pub async fn clear_pg_cache() {
    close_idle_pg_connections();
}

/// How the connection authenticates, which also keys the connection cache: a
/// connection established under one mode must never be handed to a request asking
/// for another.
#[derive(Clone, Copy, Debug, PartialEq)]
enum PgAuthMode {
    Password,
    /// AWS RDS IAM.
    Iam,
    /// Azure Entra ID, via the worker's federated identity.
    WorkloadIdentity,
}

impl PgAuthMode {
    fn of(database: &PgDatabase) -> error::Result<Self> {
        let workload_identity =
            database.password.as_deref().map(str::trim) == Some(WORKLOAD_IDENTITY_PASSWORD);
        match (database.use_iam_auth == Some(true), workload_identity) {
            (true, true) => Err(Error::BadRequest(
                "IAM RDS authentication cannot use the Azure workload identity password"
                    .to_string(),
            )),
            (true, false) => Ok(PgAuthMode::Iam),
            (false, true) => Ok(PgAuthMode::WorkloadIdentity),
            (false, false) => Ok(PgAuthMode::Password),
        }
    }

    /// What to announce in the job log. Password auth is the default and stays silent.
    /// The token modes name the login they present, which is the one thing the token
    /// itself does not carry.
    fn log_name(&self, database: &PgDatabase) -> Option<String> {
        match self {
            PgAuthMode::Password => None,
            PgAuthMode::Iam => Some(format!(
                "IAM RDS authentication (login {})",
                database.login_name()
            )),
            PgAuthMode::WorkloadIdentity => Some(match database.entra_login() {
                Ok(login) => format!("Azure Workload Identity (login {login})"),
                // Connecting rejects a missing login; do not invent one here.
                Err(_) => "Azure Workload Identity".to_string(),
            }),
        }
    }

    fn cache_key_segment(&self) -> &'static str {
        match self {
            // Workload identity needs no segment of its own: to_uri() carries the raw
            // password and the mode is a pure function of it, so two modes can never
            // share a key even though the mode is selected on the trimmed value.
            PgAuthMode::Password | PgAuthMode::WorkloadIdentity => "",
            PgAuthMode::Iam => "&iam=true",
        }
    }
}

async fn new_pg_connection(
    database: &PgDatabase,
    auth_mode: PgAuthMode,
    main_db: Option<&DB>,
) -> error::Result<(tokio_postgres::Client, tokio::task::JoinHandle<()>)> {
    let (client, connection) = match auth_mode {
        PgAuthMode::Iam => {
            #[cfg(all(feature = "enterprise", feature = "private"))]
            {
                database.connect_with_iam().await?
            }
            #[cfg(not(all(feature = "enterprise", feature = "private")))]
            {
                return Err(Error::ExecutionErr(
                    "IAM RDS authentication requires Windmill Enterprise Edition".to_string(),
                ));
            }
        }
        PgAuthMode::WorkloadIdentity => {
            #[cfg(feature = "enterprise")]
            {
                database.connect_with_workload_identity().await?
            }
            #[cfg(not(feature = "enterprise"))]
            {
                return Err(Error::ExecutionErr(
                    "Azure workload identity authentication requires Windmill Enterprise Edition"
                        .to_string(),
                ));
            }
        }
        PgAuthMode::Password => database.connect(main_db).await?,
    };
    let handle = tokio::spawn(async move {
        if let Err(e) = connection.await {
            tracing::error!("connection error: {}", e);
        }
    });
    Ok((client, handle))
}

/// `ToSql` / `FromSql` wrapper for a value whose Postgres wire format is plain
/// UTF-8 text regardless of the column's *type kind*. Vanilla
/// `tokio_postgres`'s `ToSql for String` / `FromSql for String` only accepts a
/// fixed list of base text types (TEXT/VARCHAR/BPCHAR/NAME/UNKNOWN + citext) —
/// they reject user-defined `Kind::Enum` and `Kind::Domain` even though
/// enum/domain wire format is just the variant name / the underlying base
/// type's text. This wrapper plugs that gap on both directions:
///
/// - **bind side** (prepare-fallback path): `INSERT INTO t VALUES
///   ($1::my_enum)` works end-to-end without users needing the
///   `CAST($1::text AS my_enum)` workaround.
/// - **read side** (`pg_cell_to_json_value`'s fallback): `SELECT
///   $1::my_enum`, `SELECT enum_col FROM t`, etc. round-trip into a JSON
///   string instead of erroring with "cannot convert Option<String> and the
///   Postgres type `my_enum`".
#[derive(Debug)]
struct AnyTextValue(String);

fn any_text_accepts(ty: &Type) -> bool {
    // Base text-like types, plus the citext extension type matched by name
    // (it's not in `tokio_postgres::types::Type`'s constants), plus
    // enum/domain kinds. We accept `Kind::Domain` unconditionally — the
    // server is responsible for parsing the bytes and any domain whose
    // base type accepts text on the wire (which is most of them) round-trips
    // naturally.
    matches!(
        *ty,
        Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::NAME | Type::UNKNOWN
    ) || ty.name() == "citext"
        || matches!(ty.kind(), Kind::Enum(_) | Kind::Domain(_))
}

impl ToSql for AnyTextValue {
    fn to_sql(
        &self,
        _ty: &Type,
        out: &mut bytes::BytesMut,
    ) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
        use bytes::BufMut;
        out.put_slice(self.0.as_bytes());
        Ok(IsNull::No)
    }

    fn accepts(ty: &Type) -> bool {
        any_text_accepts(ty)
    }

    tokio_postgres::types::to_sql_checked!();
}

impl<'a> FromSql<'a> for AnyTextValue {
    fn from_sql(
        _ty: &Type,
        raw: &'a [u8],
    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
        // Postgres' text wire format for enums / domains-over-text / the
        // base text types is the same: UTF-8 bytes of the value.
        Ok(AnyTextValue(std::str::from_utf8(raw)?.to_owned()))
    }

    fn accepts(ty: &Type) -> bool {
        any_text_accepts(ty)
    }
}

impl ResultFormatState {
    /// Decide whether to actually run the precision-loss check for this cell.
    /// Returns `true` for the first `NUMERIC_PRECISION_CHECK_BUDGET` calls,
    /// then `false` thereafter — and always `false` once the warning has
    /// already been triggered. Cheap on the hot path: an atomic load + an
    /// atomic decrement (Relaxed ordering), no allocation.
    fn should_check_precision(&self) -> bool {
        use std::sync::atomic::Ordering;
        if self.numeric_precision_loss.load(Ordering::Relaxed) {
            return false;
        }
        // `fetch_sub` returns the value BEFORE the decrement. When that's
        // > 0 we had budget left for this cell. After the budget reaches 0
        // the next call would wrap to `u32::MAX-1`; pin it back to 0.
        let prev = self
            .numeric_precision_check_budget
            .fetch_sub(1, Ordering::Relaxed);
        if prev == 0 {
            self.numeric_precision_check_budget
                .store(0, Ordering::Relaxed);
            false
        } else {
            true
        }
    }
}

/// Emit a single job-log warning if `state.numeric_precision_loss` flipped
/// during the row iteration. The detection itself is bounded by
/// `NUMERIC_PRECISION_CHECK_BUDGET` cells, so this only adds a constant-cost
/// log call at end-of-query.
async fn warn_on_numeric_precision_loss(
    state: &ResultFormatState,
    job_id: Uuid,
    workspace_id: &str,
    log_conn: &Connection,
) {
    use std::sync::atomic::Ordering;
    if state.numeric_precision_loss.load(Ordering::Relaxed) {
        windmill_queue::append_logs(
            &job_id,
            workspace_id,
            "warning: at least one `numeric` value in the result lost precision \
             when serialised as a JSON number (the JSON Number format goes through \
             f64, which has ~15-17 significant digits). To preserve full precision, \
             cast the column to text in your SQL — e.g. `SELECT col::text` — and \
             parse the string client-side with a Decimal library.\n",
            log_conn,
        )
        .await;
    }
}

/// Emit a one-shot warning naming each declared arg the user didn't supply a
/// value for. PG executor binds these as NULL for back-compat — without a
/// warning, a misspelled arg key in the args object silently produces a row
/// of NULLs, which is a notoriously hard DX bug to track down.
async fn warn_on_missing_args(
    missing: &[String],
    job_id: Uuid,
    workspace_id: &str,
    log_conn: &Connection,
) {
    if missing.is_empty() {
        return;
    }
    let names = missing
        .iter()
        .map(|n| format!("`{n}`"))
        .collect::<Vec<_>>()
        .join(", ");
    windmill_queue::append_logs(
        &job_id,
        workspace_id,
        format!(
            "warning: argument(s) {names} declared in the query but not provided in the \
             args object — bound as NULL. Add the value(s) to the job args, declare a \
             default in the SQL (`-- $1 name (type) = default`), or remove the \
             declaration if the arg isn't used.\n"
        ),
        log_conn,
    )
    .await;
}

/// Short stable label for the JSON value's variant — used in error messages
/// so users can see *what kind of value* hit a binding error.
fn json_value_kind(v: &Value) -> &'static str {
    match v {
        Value::Null => "null",
        Value::Bool(_) => "bool",
        Value::Number(_) => "number",
        Value::String(_) => "string",
        Value::Array(_) => "array",
        Value::Object(_) => "object",
    }
}

/// rust-postgres reports parameter encoding failures as
/// `error serializing parameter N: <inner>` with N being a 0-based index.
/// Pull N out so we can attach our own metadata.
fn parse_param_index_from_err_msg(msg: &str) -> Option<usize> {
    msg.strip_prefix("error serializing parameter ")
        .and_then(|rest| rest.split(':').next())
        .and_then(|n| n.parse::<usize>().ok())
}

/// Replace a rust-postgres encoder error with one that names the offending
/// arg, its JSON value kind, and the Postgres type we asserted, plus a hint
/// about how to fix it. Other errors are passed through unchanged.
fn wrap_param_encoding_error(
    err: tokio_postgres::Error,
    param_meta: &[(String, &'static str)],
    param_types: &[Type],
) -> Error {
    let msg = err.to_string();
    if let Some(idx) = parse_param_index_from_err_msg(&msg) {
        if let (Some((name, kind)), Some(t)) = (param_meta.get(idx), param_types.get(idx)) {
            return Error::ExecutionErr(format!(
                "Cannot bind arg `{name}` (JSON {kind}) as Postgres type `{t}` ({err}). \
                 Try adding an explicit cast in the SQL — e.g. `${pos}::<column_type>` \
                 or `CAST(${pos} AS <column_type>)` — or declare the type via \
                 `-- ${pos} {name} (<column_type>)`.",
                pos = idx + 1,
            ));
        }
    }
    to_anyhow(err).into()
}

/// Whether a statement may return a row stream, and so needs its column types described
/// before it runs. Anything not recognised as returning at most a row counts as a stream:
/// a wrong guess only costs a round trip, while a missed stream can hang the job.
fn can_stream_rows(query: &str) -> bool {
    let Some(keyword) = leading_keyword(query) else {
        return true;
    };
    let rowless = matches!(
        keyword.as_str(),
        "insert"
            | "update"
            | "delete"
            | "merge"
            | "create"
            | "alter"
            | "drop"
            | "truncate"
            | "grant"
            | "revoke"
            | "comment"
            | "set"
            | "reset"
            | "call"
            | "do"
            | "begin"
            | "commit"
            | "rollback"
            | "lock"
            | "vacuum"
            | "analyze"
            | "refresh"
    );
    !rowless || query.to_ascii_lowercase().contains("returning")
}

/// The first keyword of a statement, past whitespace and comments, following Postgres's
/// scanner: a `--` comment ends at CR or LF, and block comments nest, so
/// `/* /* a */ INSERT */ SELECT` starts with SELECT. `None` when the comments never end.
fn leading_keyword(query: &str) -> Option<String> {
    let mut rest = query;
    loop {
        rest = rest.trim_start();
        if let Some(after) = rest.strip_prefix("--") {
            rest = &after[after.find(['\n', '\r'])?..];
        } else if rest.starts_with("/*") {
            let mut depth = 0usize;
            let mut i = 0;
            let bytes = rest.as_bytes();
            loop {
                match bytes.get(i..i + 2)? {
                    b"/*" => depth += 1,
                    b"*/" => depth -= 1,
                    _ => {
                        i += 1;
                        continue;
                    }
                }
                i += 2;
                if depth == 0 {
                    break;
                }
            }
            rest = &rest[i..];
        } else {
            break;
        }
    }
    let end = rest
        .find(|c: char| !(c.is_ascii_alphanumeric() || c == '_' || c == '$'))
        .unwrap_or(rest.len());
    Some(rest[..end].to_ascii_lowercase())
}

fn otyp_to_pg_type(otyp: &str) -> error::Result<Type> {
    let base = otyp.trim_end_matches("[]");
    let is_array = otyp.ends_with("[]");

    let (scalar, array) = match base {
        "bool" | "boolean" => (Type::BOOL, Type::BOOL_ARRAY),
        "char" | "character" => (Type::CHAR, Type::CHAR_ARRAY),
        "smallint" | "smallserial" | "int2" | "serial2" => (Type::INT2, Type::INT2_ARRAY),
        "int" | "integer" | "int4" | "serial" => (Type::INT4, Type::INT4_ARRAY),
        "bigint" | "bigserial" | "int8" | "serial8" => (Type::INT8, Type::INT8_ARRAY),
        "real" | "float4" => (Type::FLOAT4, Type::FLOAT4_ARRAY),
        "double" | "double precision" | "float8" => (Type::FLOAT8, Type::FLOAT8_ARRAY),
        "numeric" | "decimal" => (Type::NUMERIC, Type::NUMERIC_ARRAY),
        "text" => (Type::TEXT, Type::TEXT_ARRAY),
        "varchar" | "character varying" => (Type::VARCHAR, Type::VARCHAR_ARRAY),
        "uuid" => (Type::UUID, Type::UUID_ARRAY),
        "date" => (Type::DATE, Type::DATE_ARRAY),
        "time" => (Type::TIME, Type::TIME_ARRAY),
        "timetz" => (Type::TIMETZ, Type::TIMETZ_ARRAY),
        "timestamp" => (Type::TIMESTAMP, Type::TIMESTAMP_ARRAY),
        "timestamptz" => (Type::TIMESTAMPTZ, Type::TIMESTAMPTZ_ARRAY),
        "json" => (Type::JSON, Type::JSON_ARRAY),
        "jsonb" => (Type::JSONB, Type::JSONB_ARRAY),
        "bytea" => (Type::BYTEA, Type::BYTEA_ARRAY),
        "oid" => (Type::OID, Type::OID_ARRAY),
        _ => {
            return Err(Error::ExecutionErr(format!(
                "Unsupported PostgreSQL type for typed schema: {}",
                otyp
            )))
        }
    };

    Ok(if is_array { array } else { scalar })
}

fn do_postgresql_inner<'a>(
    mut query: String,
    param_idx_to_arg_and_value: &HashMap<i32, (&Arg, Option<&Value>)>,
    client: &'a Client,
    column_order: Option<&'a mut Option<Vec<String>>>,
    siz: &'a AtomicUsize,
    skip_collect: bool,
    first_row_only: bool,
    s3: Option<S3ModeWorkerData>,
    job_id: Uuid,
    workspace_id: &'a str,
    log_conn: &'a Connection,
    raw_output: bool,
    progress: &'a PgProgress,
) -> error::Result<BoxFuture<'a, error::Result<Vec<Box<RawValue>>>>> {
    let mut query_params = vec![];
    let mut param_types: Vec<Type> = vec![];
    // Per-param metadata used to wrap rust-postgres `error serializing
    // parameter N` errors with actionable context (arg name, JSON value kind,
    // asserted Postgres type) — see error wrapping at the dispatch site.
    let mut param_meta: Vec<(String, &'static str)> = vec![];
    // Track whether every arg has a resolvable Postgres type. We need *both*
    // the parser-supplied otyp to be in `otyp_to_pg_type`'s map (so the arg
    // isn't a custom enum / extension type) *and* convert_val to produce a
    // (binding, type) pair that the encoder can actually serialize. If both
    // hold for every arg, we send the query as an unnamed prepared statement
    // (query_typed_raw) — see the dispatch comment below. Otherwise we fall
    // back to prepare + query_raw and let the server resolve the parameter
    // types from the SQL context.
    let mut all_types_resolved = true;

    // Single tokenizer pass — derive both the index set (for the param
    // dispatch loop below) and the byte ranges (for sparse renumbering) from
    // one walk over the SQL. Positions skip occurrences inside string
    // literals, comments, and dollar-quoted blocks, so the rewrite below
    // doesn't mangle a query like `SELECT 'price: $5' AS lbl, $5 FROM t`.
    let positions = parse_pg_statement_arg_positions(&query);
    let arg_indices: HashSet<i32> = positions.iter().map(|(i, _)| *i).collect();

    // Renumber sparse positional placeholders (e.g. $5, $50 → $1, $2) by
    // byte position, walking back-to-front so earlier positions don't shift.
    let renumber_mapping: HashMap<i32, usize> = arg_indices
        .iter()
        .sorted()
        .enumerate()
        .map(|(i, oidx)| (*oidx, i + 1))
        .collect();
    if renumber_mapping
        .iter()
        .any(|(oidx, new_i)| *oidx as usize != *new_i)
    {
        let mut positions = positions.clone();
        positions.sort_by_key(|(_, range)| std::cmp::Reverse(range.start));
        for (oidx, range) in positions {
            if let Some(new_i) = renumber_mapping.get(&oidx) {
                if oidx as usize != *new_i {
                    query.replace_range(range, &new_i.to_string());
                }
            }
        }
    }

    // Args the user didn't supply a value for — if their declaration doesn't
    // carry a default, we still bind NULL (back-compat with how the PG
    // executor has worked for years), but we collect them here to emit a
    // single one-shot warning to the job logs after query execution so a typo
    // / missing key doesn't silently turn into a row of NULLs.
    let mut missing_args: Vec<String> = Vec::new();
    // Stash declaration-default values so we can borrow them by reference
    // alongside user-supplied values — both paths feed `convert_val(&Value)`.
    let mut default_values: HashMap<i32, serde_json::Value> = HashMap::new();

    for oidx in arg_indices.iter().sorted() {
        if let Some((arg, value)) = param_idx_to_arg_and_value.get(&oidx) {
            // Resolve the value: explicit user value > declaration default > NULL.
            let value: &serde_json::Value = match (value, arg.default.as_ref()) {
                (Some(v), _) => *v,
                (None, Some(d)) => default_values.entry(*oidx).or_insert_with(|| d.clone()),
                (None, None) => {
                    if !arg.has_default && !missing_args.contains(&arg.name) {
                        missing_args.push(arg.name.clone());
                    }
                    &serde_json::Value::Null
                }
            };
            let arg_t = arg
                .otyp
                .as_ref()
                .ok_or_else(|| anyhow::anyhow!("Missing otyp for pg arg"))?;
            let typ = &arg.typ;
            let (param, natural_type) = convert_val(value, arg_t, typ, arg.otyp_inferred)?;
            query_params.push(param);
            param_meta.push((arg.name.clone(), json_value_kind(value)));
            if all_types_resolved {
                if otyp_to_pg_type(arg_t).is_ok() {
                    // The Type comes from convert_val (paired with the binding's
                    // concrete Rust type) rather than from `otyp_to_pg_type(arg_t)`
                    // — this prevents the parser-default "text" otyp from
                    // forcing an assertion that the encoder can't satisfy
                    // (e.g. Value::Bool with parser-defaulted text → Type::TEXT
                    // on a Box<bool>).
                    param_types.push(natural_type);
                } else {
                    all_types_resolved = false;
                    param_types.clear();
                }
            }
        }
    }

    let result_f = async move {
        let mut res: Vec<Box<serde_json::value::RawValue>> = vec![];

        // Always prefer query_typed_raw (unnamed prepared statement). It and its
        // describe are each a self-contained round-trip ending in Sync, so they
        // survive transaction-mode connection poolers (PgBouncer/Supabase
        // pooler/RDS Proxy) where named
        // statements ("s0", "s1", ...) can be reported missing because the prepare
        // and the execute land on different backend connections. Fall back to
        // prepare + query_raw only when an arg has a type unsupported by
        // otyp_to_pg_type (e.g. custom enum, geometry, …) — in that case we lose
        // pooler safety, but the query at least runs against a direct connection.
        let rows = if all_types_resolved {
            // query_typed_raw looks up result column types this connection has not seen
            // (enums, domains, extension types) while the rows already stream, and on a
            // large result that lookup waits behind them forever. Describing first resolves
            // them up front, still without a named statement.
            if can_stream_rows(&query) {
                if let Err(e) = client.describe_typed(&query, &param_types).await {
                    return Err(wrap_param_encoding_error(e, &param_meta, &param_types));
                }
            }
            let typed_params = query_params
                .iter()
                .zip(param_types.iter())
                .map(|(p, t)| (&**p as &(dyn ToSql + Sync), t.clone()));
            match client.query_typed_raw(&query, typed_params).await {
                Ok(rows) => rows,
                Err(e) => {
                    return Err(wrap_param_encoding_error(e, &param_meta, &param_types));
                }
            }
        } else {
            let query_params = query_params
                .iter()
                .map(|p| &**p as &(dyn ToSql + Sync))
                .collect_vec();
            let statement = client.prepare(&query).await.map_err(to_anyhow)?;
            client
                .query_raw(&statement, query_params)
                .await
                .map_err(to_anyhow)?
        };

        // One state object per query — `pg_cell_to_json_value_with_state`
        // flips `numeric_precision_loss` once if any `numeric` cell can't
        // round-trip through f64. We emit a single warning to the job log
        // after the iteration finishes, instead of error-by-error or per
        // row, and the per-row check short-circuits on the flag so the cost
        // is one branch after the first lossy value.
        let format_state = ResultFormatState::default();

        if skip_collect {
            futures::pin_mut!(rows);
            while rows.try_next().await.map_err(to_anyhow)?.is_some() {
                progress.row();
            }
        } else if let Some(ref s3) = s3 {
            let format_state_ref = &format_state;
            let rows_stream = rows.map_err(to_anyhow).map(move |row_result| {
                progress.row();
                row_result.and_then(|row| {
                    postgres_row_to_json_value_with_state(row, format_state_ref).map_err(to_anyhow)
                })
            });

            s3_stream_and_upload_with_logs(
                "PostgreSQL",
                rows_stream.boxed(),
                s3,
                job_id,
                workspace_id,
                log_conn,
            )
            .await?;

            warn_on_numeric_precision_loss(&format_state, job_id, workspace_id, log_conn).await;
            warn_on_missing_args(&missing_args, job_id, workspace_id, log_conn).await;

            return Ok(vec![to_raw_value(&s3.to_return_s3_obj())]);
        } else {
            let rows = if first_row_only {
                rows.take(1).boxed()
            } else {
                rows.boxed()
            };

            // The stream is consumed one row at a time so the cap below can still
            // refuse. Collecting it into a `Vec<Row>` first holds every wire buffer
            // and every converted row at once, and the worker is already past the
            // budget by the time the first check gets to run.
            futures::pin_mut!(rows);
            let max_result_size = max_sql_result_size();
            let mut envelope = raw_output
                .then(|| crate::pg_raw_output::RawOutputEnvelopeBuilder::new(max_result_size));
            let mut column_names: Option<Vec<String>> = None;

            while let Some(row) = rows.try_next().await.map_err(to_anyhow)? {
                progress.row();
                if column_names.is_none() {
                    column_names = Some(
                        row.columns()
                            .iter()
                            .map(|x| x.name().to_string())
                            .collect::<Vec<String>>(),
                    );
                }

                if let Some(envelope) = envelope.as_mut() {
                    envelope.push(row, &format_state, siz)?;
                    continue;
                }

                let v = postgres_row_to_json_value_with_state(row, &format_state)?;
                // Serialized under what is left of the budget: escaping can expand
                // a row that fit in memory past what remains, and a check placed
                // after the write happens once the allocation already did.
                let raw = to_raw_value_within(
                    &v,
                    max_result_size.saturating_sub(siz.load(Ordering::Relaxed)),
                )
                .ok_or_else(|| sql_result_too_large_error(max_result_size))?;
                // Both are proxies for what the row costs the worker, and neither
                // dominates: the value tree is wider than its JSON for small
                // scalars, narrower once escaping expands the text.
                siz.fetch_add(sizeof_val(&v).max(raw.get().len()), Ordering::Relaxed);
                if siz.load(Ordering::Relaxed) > max_result_size {
                    return Err(sql_result_too_large_error(max_result_size));
                }
                res.push(raw);
            }

            if let Some(column_order) = column_order {
                // A statement that returned no rows reports no columns.
                *column_order = Some(column_names.unwrap_or_default());
            }

            if let Some(envelope) = envelope {
                // The envelope is budgeted against the whole cap: its rows were
                // already charged as text on the way in, and this is that same
                // text serialized, so charging it against the remainder would
                // reject a result that passed every row-level check.
                res.push(
                    to_raw_value_within(&envelope.finish(), max_result_size)
                        .ok_or_else(|| sql_result_too_large_error(max_result_size))?,
                );
            }
        }

        warn_on_numeric_precision_loss(&format_state, job_id, workspace_id, log_conn).await;
        warn_on_missing_args(&missing_args, job_id, workspace_id, log_conn).await;

        Ok(res)
    };

    Ok(result_f.boxed())
}

pub async fn do_postgresql(
    job: &MiniPulledJob,
    client: &AuthedClient,
    query: &str,
    conn: &Connection,
    mem_peak: &mut i32,
    canceled_by: &mut Option<CanceledBy>,
    worker_name: &str,
    column_order: &mut Option<Vec<String>>,
    occupancy_metrics: &mut OccupancyMetrics,
    parent_runnable_path: Option<String>,
    run_inline: bool,
) -> error::Result<Box<RawValue>> {
    let mut pg_args = build_args_values(job, client, conn).await?;

    let inline_db_res_path = parse_db_resource(&query);

    let s3 = parse_s3_mode(&query)?.map(|s3| s3_mode_args_to_worker_data(s3, client.clone(), job));

    let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
        Some(
            client
                .get_resource_value_interpolated::<serde_json::Value>(
                    &inline_db_res_path,
                    Some(job.id.to_string()),
                )
                .await?,
        )
    } else {
        match pg_args.get("database").cloned() {
            Some(Value::String(db_str)) if db_str.starts_with("datatable://") => {
                let reference = db_str.trim_start_matches("datatable://");
                // The annotation wins: a generated query can carry a `?role=` in the reference it
                // was handed, but only the script's author writes the leading comment block.
                let annotated = SqlAnnotations::datatable_role(&query)?;
                Some(match conn {
                    Connection::Http(client) => {
                        let (name, uri_role) = parse_datatable_ref(reference)?;
                        get_datatable_resource_from_agent_http(
                            client,
                            name,
                            &job.workspace_id,
                            annotated.as_deref().or(uri_role),
                            &job.id,
                        )
                        .await?
                    }
                    Connection::Sql(db) => {
                        let (name, uri_role) =
                            parse_datatable_ref_for(db, &job.workspace_id, reference).await?;
                        get_datatable_resource_from_db(
                            db,
                            &job.workspace_id,
                            &name,
                            annotated.as_deref().or(uri_role.as_deref()),
                            DatatableAccess::PermissionedAs {
                                permissioned_as: &job.permissioned_as,
                                email: &job.permissioned_as_email,
                            },
                        )
                        .await?
                    }
                })
            }
            database => database,
        }
    };

    let database = if let Some(db) = db_arg {
        serde_json::from_value::<PgDatabase>(db.clone())
            .map_err(|e| Error::ExecutionErr(e.to_string()))?
    } else {
        return Err(Error::BadRequest("Missing database argument".to_string()));
    };

    // Surface in the job logs (not just the worker logs) when a verify-ca/verify-full
    // resource is connecting without actually verifying the server certificate, so the
    // person running the query can see and fix the misconfiguration.
    if database.verify_mode_skips_verification() {
        windmill_queue::append_logs(
            &job.id,
            &job.workspace_id,
            format!(
                "warning: sslmode={} but the server's TLS certificate is not being verified \
                 (accept_invalid_certs is enabled, or no root certificate is configured). Set \
                 accept_invalid_certs to false or provide root_certificate_pem to verify the \
                 server identity.\n",
                database.sslmode.as_deref().unwrap_or("")
            ),
            conn,
        )
        .await;
    }

    let annotations = windmill_common::worker::SqlAnnotations::parse(query);
    let collection_strategy = if annotations.raw_output || annotations.return_last_result {
        // raw_output emits a single envelope from the last statement, so the
        // multi-statement collection modes don't apply.
        SqlResultCollectionStrategy::LastStatementAllRows
    } else {
        annotations.result_collection
    };

    let auth_mode = PgAuthMode::of(&database)?;

    if let Some(mode) = auth_mode.log_name(&database) {
        windmill_queue::append_logs(&job.id, &job.workspace_id, format!("Using {mode}\n"), conn)
            .await;
    }

    // Include the auth mode in the cache key to distinguish connections to the same host
    // authenticated differently. The cache key is static (doesn't include the token), which
    // is correct because PostgreSQL connections remain valid after initial auth — fresh
    // tokens are generated on cache miss.
    //
    // to_uri() collapses require/verify-ca/verify-full to the same string, so the TLS verification
    // inputs are folded into the key separately. Without this a connection established under a
    // weaker sslmode (or a different root cert) could be reused for a stricter request, undoing
    // the verification configured in PgDatabase::configure_pg_tls_verification.
    let tls_disc = {
        use std::hash::{Hash, Hasher};
        let mut h = std::collections::hash_map::DefaultHasher::new();
        database.sslmode.hash(&mut h);
        database.root_certificate_pem.hash(&mut h);
        database.accept_invalid_certs.hash(&mut h);
        h.finish()
    };
    // to_uri() already ends with `?sslmode=...`, so append further key segments
    // with `&` to keep database_string a well-formed URI (it is only ever a cache
    // key, but a malformed one would mislead anyone who later logs or parses it).
    let database_string = format!(
        "{}{}&tls={tls_disc:x}",
        database.to_uri(),
        auth_mode.cache_key_segment()
    );

    let connect_started = std::time::Instant::now();
    let mut cached_connection_failed_reset = false;
    let mut reused_cached_connection = false;
    let mut lease = if *CLOUD_HOSTED {
        PgConnectionLease::uncached()
    } else {
        PgConnectionLease::checkout(&database_string, pg_pool_scope(&database))
    };

    if lease.conn.is_some() {
        // Probe the cached connection with a curated session reset before
        // reusing it. Each statement targets a specific class of state:
        //
        //   RESET ALL                     — GUC parameters (search_path,
        //                                   application_name, statement_
        //                                   timeout, transaction_*…). Note
        //                                   that this does NOT reset SET
        //                                   ROLE or SET SESSION
        //                                   AUTHORIZATION (security!).
        //   RESET SESSION AUTHORIZATION   — undoes both `SET SESSION
        //                                   AUTHORIZATION` and `SET ROLE`,
        //                                   restoring the connecting user.
        //                                   Without this a previous job
        //                                   leaving an elevated role
        //                                   active would silently leak
        //                                   permissions into the next.
        //   UNLISTEN *                    — drops LISTEN registrations.
        //   CLOSE ALL                     — closes open cursors.
        //   pg_advisory_unlock_all()      — releases any session-scoped
        //                                   advisory locks. Without this
        //                                   a job that called
        //                                   pg_advisory_lock and exited
        //                                   without unlocking would block
        //                                   later jobs holding the same
        //                                   key (DISCARD ALL covered this
        //                                   too).
        //
        // We deliberately do NOT use `DISCARD ALL`. DISCARD includes
        // `DEALLOCATE ALL`, which deallocates *all* prepared statements
        // server-side — including the typeinfo statements that
        // tokio_postgres caches per-Client to resolve custom enum/domain
        // Oids. After DISCARD, tokio_postgres still holds Statement
        // objects whose names the server has forgotten, so the next
        // custom-type query fails with `prepared statement "sN" does not
        // exist`. The trade-off: temp tables and user-PREPARE statements
        // may persist across cached-connection reuse (rare in datatable /
        // script workloads).
        //
        // Doubles as a liveness probe — if the connection is broken any
        // statement in the chain fails, or it never answers, and we replace it.
        let reset = lease.client().batch_execute(PG_SESSION_RESET);
        if matches!(
            tokio::time::timeout(PG_RESET_PROBE_TIMEOUT, reset).await,
            Ok(Ok(()))
        ) {
            tracing::info!("Using cached connection");
            CACHE_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            reused_cached_connection = true;
        } else {
            tracing::info!("Cached connection is stale, creating new one");
            cached_connection_failed_reset = true;
            lease.discard();
        }
    }
    if lease.conn.is_none() {
        let (client, connection_task) =
            match new_pg_connection(&database, auth_mode, conn.as_sql()).await {
                Ok(connected) => connected,
                // A server or pooler can also refuse outright past a per-login or
                // total connection cap, which pool scopes do not capture. Free the
                // idle connections of every scope and try once more.
                Err(e) => {
                    if !is_connection_cap_refusal(&e) || close_idle_pg_connections() == 0 {
                        return Err(e);
                    }
                    tokio::time::sleep(Duration::from_millis(250)).await;
                    new_pg_connection(&database, auth_mode, conn.as_sql()).await?
                }
            };
        lease.conn = Some(PgConnection {
            key: database_string.clone(),
            scope: lease.scope.clone(),
            client,
            connection_task: Some(connection_task),
            last_used: Instant::now(),
        });
    }

    let fresh_connection = !reused_cached_connection;
    let log_progress = !run_inline && !annotations.prepare;
    let connect_time = connect_started.elapsed();
    if log_progress && connect_time >= PG_SLOW_STEP {
        windmill_queue::append_logs(
            &job.id,
            &job.workspace_id,
            format!(
                "Getting a database connection took {} ms ({})\n",
                connect_time.as_millis(),
                if cached_connection_failed_reset {
                    "the cached connection failed its reset, then a new connection"
                } else {
                    connection_kind(fresh_connection)
                }
            ),
            conn,
        )
        .await;
    }

    let (mut sig, _) = parse_pgsql_sig_with_typed_schema(&query)
        .map_err(|x| Error::ExecutionErr(x.to_string()))?;

    // Materialize any `(s3object)` args into JSON text and rebind them as `jsonb` so
    // `otyp_to_pg_type` picks the right binding. Must run before the param map is
    // built below.
    let had_s3object_input = materialize_s3object_args(
        &mut sig.args,
        &mut pg_args,
        client,
        conn,
        job.id,
        &job.workspace_id,
    )
    .await?;

    let reserved_variables =
        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;

    let (query, _) =
        &sanitize_and_interpolate_unsafe_sql_args(query, &sig.args, &pg_args, &reserved_variables)?;

    let queries = parse_sql_blocks(query, true);

    let client = lease.client();

    let param_idx_to_arg_and_value = sig
        .args
        .iter()
        .filter_map(|x| x.oidx.map(|oidx| (oidx, (x, pg_args.get(&x.name)))))
        .collect::<HashMap<_, _>>();

    let size = AtomicUsize::new(0);
    let size_ref = &size;
    let progress = PgProgress::new();
    let progress_ref = &progress;
    let statement_count = queries.len();
    let result_f = async move {
        let mut results = vec![];
        // Session reset (DISCARD ALL) is now handled eagerly when validating
        // the cached connection — no per-query reset needed here.

        for (i, query) in queries.iter().enumerate() {
            if annotations.prepare {
                let query = remove_comments(query);
                // Used by the data table typechecker to set default schemas
                if query.starts_with("SET search_path") || query.starts_with("RESET search_path") {
                    let _ = client.execute(&query.to_string(), &[]).await;
                    continue;
                }
                let prepared = client.prepare(&query).await;
                let prepared = match prepared {
                    Ok(prepared) => {
                        let columns: Option<Vec<PrepareQueryColumnInfo>> = Some(
                            prepared
                                .columns()
                                .iter()
                                .map(|col| PrepareQueryColumnInfo {
                                    name: col.name().to_string(),
                                    type_name: col.type_().name().to_string(),
                                })
                                .collect(),
                        );
                        PrepareQueryResult { columns, error: None }
                    }
                    Err(e) => PrepareQueryResult { columns: None, error: Some(e.to_string()) },
                };
                results.push(vec![to_raw_value(&prepared)]);
                continue;
            }

            let skip_collect = collection_strategy.collect_last_statement_only(queries.len())
                && i < queries.len() - 1;
            let is_last = i == queries.len() - 1;
            progress_ref.start_statement(i);
            let rows_before = progress_ref.rows();
            let statement_started = std::time::Instant::now();
            let result = do_postgresql_inner(
                query.to_string(),
                &param_idx_to_arg_and_value,
                client,
                if is_last
                    && s3.is_none()
                    && collection_strategy.collect_last_statement_only(queries.len())
                    && !collection_strategy.collect_scalar()
                {
                    Some(column_order)
                } else {
                    None
                },
                size_ref,
                skip_collect,
                collection_strategy.collect_first_row_only(),
                s3.clone(),
                job.id,
                &job.workspace_id,
                conn,
                annotations.raw_output && is_last && !skip_collect,
                progress_ref,
            )?
            .await?;
            let statement_time = statement_started.elapsed();
            if log_progress && statement_time >= PG_SLOW_STEP {
                windmill_queue::append_logs(
                    &job.id,
                    &job.workspace_id,
                    format!(
                        "Statement {}/{statement_count}: {} rows in {} ms\n",
                        i + 1,
                        progress_ref.rows() - rows_before,
                        statement_time.as_millis()
                    ),
                    conn,
                )
                .await;
            }
            results.push(result);
        }

        if annotations.raw_output {
            // The raw_output envelope already aggregates the last statement's
            // result; skip the collection_strategy reshape that wraps rows in
            // a JSON array.
            Ok(crate::pg_raw_output::extract_envelope_or_empty(results))
        } else {
            collection_strategy.collect(results)
        }
    };
    let result_f = async {
        if log_progress {
            tokio::select! {
                result = result_f => result,
                never = warn_on_stalled_statement(progress_ref, statement_count, fresh_connection, job.id, &job.workspace_id, conn) => match never {},
            }
        } else {
            result_f.await
        }
    };
    // Still unset once the job is over: the query was abandoned mid-flight, by the timeout or
    // a cancellation, rather than ended by an error of its own.
    let query_ended = std::sync::atomic::AtomicBool::new(false);
    let result_f = async {
        let result = result_f.await;
        query_ended.store(true, Ordering::Relaxed);
        result
    };

    let result = if run_inline {
        result_f.await
    } else {
        run_future_with_polling_update_job_poller(
            job.id,
            job.timeout,
            conn,
            mem_peak,
            canceled_by,
            result_f,
            worker_name,
            &job.workspace_id,
            &mut Some(occupancy_metrics),
            Box::pin(futures::stream::once(async { 0 })),
        )
        .await
    };
    let result = match result {
        Ok(result) => result,
        Err(e) => {
            close_unclean_connection(
                lease,
                database,
                auth_mode != PgAuthMode::Password,
                !query_ended.load(Ordering::Relaxed),
            );
            return Err(map_s3object_jsonb_overflow(e, had_s3object_input));
        }
    };

    *mem_peak = size.load(Ordering::Relaxed) as i32;

    // A transaction the script left open would otherwise carry over into the
    // next job that reuses the connection, so it is rolled back and the connection
    // closed. The status comes with the last reply, so a script that ends cleanly
    // pays nothing for this check.
    if lease.client().transaction_status() != tokio_postgres::TransactionStatus::Idle {
        if !run_inline {
            windmill_queue::append_logs(
                &job.id,
                &job.workspace_id,
                "The script ended inside an open transaction, which was rolled back. \
                 End it with COMMIT to keep its changes.\n",
                conn,
            )
            .await;
        }
        close_unclean_connection(lease, database, auth_mode != PgAuthMode::Password, false);
    } else {
        // Keep the connection only after a query that read its results to the end.
        // An error may have stopped reading early (the result cap, a decode failure),
        // as a timeout, a dropped future or a first-row collection do: the server
        // keeps sending, and the next job's reset probe would wait on it.
        lease.keep = !*CLOUD_HOSTED && !collection_strategy.collect_first_row_only();
        drop(lease);
    }

    *mem_peak = (result.get().len() / 1000) as i32;
    // And then check that we got back the same string we sent over.
    return Ok(result);
}

/// For each `(s3object)` arg in `sig_args`: download the referenced file, decode it
/// to JSON text, then rewrite the arg to bind as `jsonb`. Mutates `args_map` in place
/// so the existing bind path picks up the materialized payload. Returns whether any
/// `(s3object)` arg was materialized, so the jsonb-cap error can be rewritten.
async fn materialize_s3object_args(
    sig_args: &mut [Arg],
    args_map: &mut HashMap<String, Value>,
    client: &AuthedClient,
    conn: &Connection,
    job_id: Uuid,
    workspace_id: &str,
) -> error::Result<bool> {
    let mut materialized_any = false;
    for arg in sig_args.iter_mut() {
        if arg.otyp.as_deref() != Some("s3object") {
            continue;
        }
        let raw = args_map.remove(&arg.name).unwrap_or(Value::Null);
        if matches!(raw, Value::Null) {
            return Err(Error::BadRequest(format!(
                "Missing S3Object value for arg `{}`",
                arg.name
            )));
        }
        let s3_obj: S3Object = serde_json::from_value(raw).map_err(|e| {
            Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
        })?;
        let json_text = fetch_s3object_as_json_text(client, conn, job_id, workspace_id, &s3_obj)
            .await
            .map_err(|e| {
                Error::ExecutionErr(format!(
                    "Failed to fetch S3 object for arg `{}`: {e}",
                    arg.name
                ))
            })?;
        materialized_any = true;
        // Parse to a Value so `convert_val`'s Array/Object → JSONB branches bind it
        // correctly. A bare String would mismatch the JSONB param type.
        let parsed: Value = serde_json::from_str(&json_text).map_err(|e| {
            Error::ExecutionErr(format!(
                "S3 object for arg `{}` is not valid JSON after decoding: {e}",
                arg.name
            ))
        })?;
        args_map.insert(arg.name.clone(), parsed);
        arg.otyp = Some("jsonb".to_string());
        arg.typ = Typ::Object(windmill_parser::ObjectType::new(None, Some(vec![])));
    }
    Ok(materialized_any)
}

/// A `(s3object)` input materializes the whole file into a single jsonb parameter, which
/// PostgreSQL caps at ~256MB (`total size of jsonb {array,object} elements exceeds the
/// maximum of 268435455 bytes`). A large input trips this with an opaque server error;
/// rewrite it into guidance pointing at DuckDB, which reads S3 natively and streams.
///
/// The attribution is hedged: the same error can also come from SQL constructing an
/// oversized jsonb at execution time, and with several inputs we can't tell which one
/// overflowed, so we point at `(s3object)` inputs as the likely cause rather than naming
/// a specific file. The DuckDB remediation is the same either way.
fn map_s3object_jsonb_overflow(e: Error, had_s3object_input: bool) -> Error {
    if !had_s3object_input {
        return e;
    }
    let msg = e.to_string();
    // Match only the jsonb byte-size cap ("total size of jsonb ... elements exceeds the
    // maximum of 268435455 bytes"), so the ~256 MB wording stays accurate. Excludes the
    // element-count cap and unrelated caps like "array size exceeds the maximum allowed".
    if !msg.contains("total size of jsonb") {
        return e;
    }
    Error::ExecutionErr(format!(
        "This query hit PostgreSQL's ~256 MB size limit for a single jsonb value. This is a \
         server-side database limit, not a worker-memory limit, so a larger worker will not raise \
         it. If a large `(s3object)` input is the cause: native SQL `(s3object)` inputs load the \
         whole file into one jsonb parameter and do not stream, so they only fit small files. For \
         large Parquet/CSV files, use a DuckDB script instead: it reads the file directly from S3 \
         and streams (e.g. `read_parquet(...)` / `read_csv_auto(...)`) rather than materializing \
         it.\n\nUnderlying error: {msg}",
    ))
}

/// Parse a date string in formats produced by chrono's Display or JS frontends.
fn parse_naive_date(s: &str) -> Result<chrono::NaiveDate, chrono::ParseError> {
    chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d")
        .or_else(|_| chrono::NaiveDate::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.fZ"))
        .or_else(|_| chrono::NaiveDate::parse_from_str(s, "%Y-%m-%dT%H:%M:%SZ"))
        .or_else(|_| chrono::NaiveDate::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f"))
}

/// Parse a time string in formats produced by chrono's Display or JS frontends.
fn parse_naive_time(s: &str) -> Result<chrono::NaiveTime, chrono::ParseError> {
    chrono::NaiveTime::parse_from_str(s, "%H:%M:%S%.f")
        .or_else(|_| chrono::NaiveTime::parse_from_str(s, "%H:%M:%S"))
        .or_else(|_| chrono::NaiveTime::parse_from_str(s, "%H:%M"))
        .or_else(|_| {
            // Handle full datetime strings by extracting the time part
            chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.fZ")
                .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f"))
                .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%SZ"))
                .map(|dt| dt.time())
        })
}

/// Parse a naive datetime string in formats produced by chrono's Display or JS frontends.
fn parse_naive_datetime(s: &str) -> Result<chrono::NaiveDateTime, chrono::ParseError> {
    chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f")
        .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S"))
        .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.fZ"))
        .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f"))
        .or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%SZ"))
}

/// Parse a timestamptz string in formats produced by chrono's Display or JS frontends.
fn parse_datetime_utc(s: &str) -> Result<chrono::DateTime<Utc>, chrono::ParseError> {
    s.parse::<chrono::DateTime<Utc>>()
        .or_else(|_| {
            // Handle numeric timezone offsets: "2024-01-15 10:30:00+00", "+00:00", "+0000"
            chrono::DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%#z")
                .or_else(|_| chrono::DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%#z"))
                .map(|dt| dt.with_timezone(&Utc))
        })
        .or_else(|_| {
            // Handle chrono's Display format: "2024-01-15 10:30:00 UTC"
            let trimmed = s.trim_end_matches(" UTC");
            chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%d %H:%M:%S%.f")
                .or_else(|_| chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%d %H:%M:%S"))
                .or_else(|_| {
                    chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%dT%H:%M:%S%.fZ")
                })
                .or_else(|_| chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%dT%H:%M:%S%.f"))
                .or_else(|_| chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%dT%H:%M:%SZ"))
                .map(|ndt| ndt.and_utc())
        })
}

fn map_as_single_type<T>(
    vec: Option<&Vec<Value>>,
    f: impl Fn(&Value) -> Option<T>,
) -> anyhow::Result<Option<Vec<Option<T>>>> {
    if let Some(vec) = vec {
        Ok(Some(
            vec.into_iter()
                .map(|v| {
                    // first option is if the value is of the right type (if none, will stop the collection and throw error)
                    // second option is if the value is null
                    // allow nulls in arrays
                    if matches!(v, Value::Null) {
                        Some(None)
                    } else {
                        f(v).map(Some)
                    }
                })
                .collect::<Option<Vec<Option<T>>>>()
                .ok_or_else(|| anyhow::anyhow!("Mixed types in array"))?,
        ))
    } else {
        Ok(None)
    }
}

/// A boxed `ToSql` value paired with the Postgres `Type` that matches its
/// concrete Rust type. Returned by `convert_val` / `convert_vec_val` so the
/// dispatch in `do_postgresql_inner` always asserts the type that the encoder
/// can actually produce — never a parser-derived guess that drifts from the
/// runtime binding.
type ConvertedParam = (Box<dyn ToSql + Sync + Send>, Type);

fn convert_vec_val(
    vec: Option<&Vec<Value>>,
    arg_t: &String,
) -> windmill_common::error::Result<ConvertedParam> {
    match arg_t.as_str() {
        // Each integer / bool array arm accepts both JSON-native values AND
        // stringified counterparts ("1", "true", …) — same coercion the
        // scalar `Value::String → <type>` arms in `convert_val` apply, so an
        // array passed via `JSON.stringify(BigInt(...))` or hand-quoted
        // values doesn't trip a confusing "Mixed types in array" error.
        "bool" | "boolean" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_bool()
                    .or_else(|| match v.as_str()?.to_ascii_lowercase().as_str() {
                        "true" | "t" | "yes" | "y" | "1" | "on" => Some(true),
                        "false" | "f" | "no" | "n" | "0" | "off" => Some(false),
                        _ => None,
                    })
            })?),
            Type::BOOL_ARRAY,
        )),
        "char" | "character" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_i64()
                    .map(|x| x as i8)
                    .or_else(|| v.as_str().and_then(|s| s.parse::<i8>().ok()))
            })?),
            Type::CHAR_ARRAY,
        )),
        "smallint" | "smallserial" | "int2" | "serial2" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_i64()
                    .map(|x| x as i16)
                    .or_else(|| v.as_str().and_then(|s| s.parse::<i16>().ok()))
            })?),
            Type::INT2_ARRAY,
        )),
        "int" | "integer" | "int4" | "serial" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_i64()
                    .map(|x| x as i32)
                    .or_else(|| v.as_str().and_then(|s| s.parse::<i32>().ok()))
            })?),
            Type::INT4_ARRAY,
        )),
        // Mirror the scalar `Value::String → numeric` parsing arm so an array
        // like `["1.5", "2.5"]` works against `$1::numeric[]` — useful for
        // bulk-loading via `unnest`. Without this the user would see an
        // unhelpful "Mixed types in array" error.
        "numeric" | "decimal" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                if v.is_i64() {
                    Decimal::from_i64(v.as_i64().unwrap())
                } else if v.is_f64() {
                    Decimal::from_f64(v.as_f64().unwrap())
                } else {
                    v.as_str().and_then(|s| s.parse::<Decimal>().ok())
                }
            })?),
            Type::NUMERIC_ARRAY,
        )),
        "oid" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_u64()
                    .map(|x| x as u32)
                    .or_else(|| v.as_str().and_then(|s| s.parse::<u32>().ok()))
            })?),
            Type::OID_ARRAY,
        )),
        "bigint" | "bigserial" | "int8" | "serial8" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_i64()
                    .or_else(|| v.as_u64().map(|x| x as i64))
                    .or_else(|| v.as_str().and_then(|s| s.parse::<i64>().ok()))
            })?),
            Type::INT8_ARRAY,
        )),
        "real" | "float4" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_f64()
                    .map(|x| x as f32)
                    .or_else(|| v.as_str().and_then(|s| s.parse::<f32>().ok()))
            })?),
            Type::FLOAT4_ARRAY,
        )),
        "double" | "double precision" | "float8" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_f64()
                    .or_else(|| v.as_str().and_then(|s| s.parse::<f64>().ok()))
            })?),
            Type::FLOAT8_ARRAY,
        )),
        "uuid" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().map(|x| Uuid::parse_str(x).ok()).flatten()
            })?),
            Type::UUID_ARRAY,
        )),
        "date" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().and_then(|x| parse_naive_date(x).ok())
            })?),
            Type::DATE_ARRAY,
        )),
        "time" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().and_then(|x| parse_naive_time(x).ok())
            })?),
            Type::TIME_ARRAY,
        )),
        "timetz" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().and_then(|x| parse_naive_time(x).ok())
            })?),
            // chrono's `NaiveTime` only encodes for `TIME` — same caveat as
            // the scalar `timetz` arm. Asserting `TIMETZ_ARRAY` here would
            // fail at the encoder. Postgres has an implicit `time → timetz`
            // assignment cast at the column site.
            Type::TIME_ARRAY,
        )),
        "timestamp" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().and_then(|x| parse_naive_datetime(x).ok())
            })?),
            Type::TIMESTAMP_ARRAY,
        )),
        "timestamptz" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().and_then(|x| parse_datetime_utc(x).ok())
            })?),
            Type::TIMESTAMPTZ_ARRAY,
        )),
        "jsonb" => Ok((
            Box::new(vec.map(|v| v.clone().into_iter().map(Some).collect_vec())),
            Type::JSONB_ARRAY,
        )),
        "json" => Ok((
            Box::new(vec.map(|v| v.clone().into_iter().map(Some).collect_vec())),
            Type::JSON_ARRAY,
        )),
        "bytea" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().map(|x| {
                    engine::general_purpose::STANDARD
                        .decode(x)
                        .unwrap_or(vec![])
                })
            })?),
            Type::BYTEA_ARRAY,
        )),
        "varchar" | "character varying" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().map(|x| x.to_string())
            })?),
            Type::VARCHAR_ARRAY,
        )),
        "text" => Ok((
            Box::new(map_as_single_type(vec, |v| {
                v.as_str().map(|x| x.to_string())
            })?),
            Type::TEXT_ARRAY,
        )),
        _ => Err(anyhow::anyhow!("Unsupported JSON array type"))?,
    }
}

fn convert_val(
    value: &Value,
    arg_t: &String,
    typ: &Typ,
    otyp_inferred: bool,
) -> windmill_common::error::Result<ConvertedParam> {
    // Helper: was the user's intent explicitly "text" / "varchar" / "char"?
    // True when the parser saw an inline `$N::text` cast or a `-- $N (text)`
    // declaration. False when the parser fell back to "text" because nothing
    // else was found (in which case the caller has no real target type
    // committed and we should bind the value's natural type).
    let explicit_text_target = !otyp_inferred
        && (matches!(typ, Typ::Str(_))
            && (arg_t == "text"
                || arg_t == "varchar"
                || arg_t == "character varying"
                || arg_t == "char"
                || arg_t == "character"));
    match value {
        Value::Array(vec) if arg_t.ends_with("[]") => {
            let arg_t = arg_t.trim_end_matches("[]").to_string();
            convert_vec_val(Some(vec), &arg_t)
        }
        Value::Null if arg_t.ends_with("[]") => {
            let arg_t = arg_t.trim_end_matches("[]").to_string();
            convert_vec_val(None, &arg_t)
        }
        Value::Null => match arg_t.as_str() {
            "bool" | "boolean" => Ok((Box::new(None::<bool>), Type::BOOL)),
            "char" | "character" => Ok((Box::new(None::<i8>), Type::CHAR)),
            "smallint" | "smallserial" | "int2" | "serial2" => {
                Ok((Box::new(None::<i16>), Type::INT2))
            }
            "int" | "integer" | "int4" | "serial" => Ok((Box::new(None::<i32>), Type::INT4)),
            "numeric" | "decimal" => Ok((Box::new(None::<Decimal>), Type::NUMERIC)),
            "oid" => Ok((Box::new(None::<u32>), Type::OID)),
            "bigint" | "bigserial" | "int8" | "serial8" => Ok((Box::new(None::<i64>), Type::INT8)),
            "real" | "float4" => Ok((Box::new(None::<f32>), Type::FLOAT4)),
            "double" | "double precision" | "float8" => Ok((Box::new(None::<f64>), Type::FLOAT8)),
            "uuid" => Ok((Box::new(None::<Uuid>), Type::UUID)),
            "date" => Ok((Box::new(None::<chrono::NaiveDate>), Type::DATE)),
            "time" => Ok((Box::new(None::<chrono::NaiveTime>), Type::TIME)),
            // chrono's NaiveTime has no timezone, so its ToSql impl only
            // accepts TIME. We assert TIME and rely on Postgres' implicit
            // assignment cast time → timetz at the use site.
            "timetz" => Ok((Box::new(None::<chrono::NaiveTime>), Type::TIME)),
            "timestamp" => Ok((Box::new(None::<chrono::NaiveDateTime>), Type::TIMESTAMP)),
            "timestamptz" => Ok((Box::new(None::<chrono::DateTime<Utc>>), Type::TIMESTAMPTZ)),
            "jsonb" => Ok((Box::new(None::<Value>), Type::JSONB)),
            "json" => Ok((Box::new(None::<Value>), Type::JSON)),
            "bytea" => Ok((Box::new(None::<Vec<u8>>), Type::BYTEA)),
            "varchar" | "character varying" => Ok((Box::new(None::<String>), Type::VARCHAR)),
            "text" => Ok((Box::new(None::<String>), Type::TEXT)),
            // Unrecognised arg_t — bind as TEXT NULL. The dispatch will fall
            // back to prepare + query_raw, where the server resolves the
            // actual column type and `Option<String>`'s ToSql will accept the
            // resolved Type for any text-like base; for enum/domain kinds
            // None is encoded as the literal NULL message body, so the
            // accepts() check is the only place that matters and we just need
            // a binding whose accepts() is permissive enough.
            _ => Ok((Box::new(None::<AnyTextValue>), Type::TEXT)),
        },
        // Bool / Number with an *explicitly* text-typed arg: coerce to
        // String. Used when the user wrote `-- $N (text)` or `$N::text` —
        // they committed to text and may rely on equality comparisons like
        // `WHERE text_col = $1`, which need a `text = text` operator (PG has
        // no implicit `bool/int → text` cast in expression context).
        Value::Bool(b) if explicit_text_target => {
            // `char` (Type::CHAR, OID 18) is single-byte and `to_string()` of
            // a bool is multi-byte ("true"/"false") — we can't bind it as
            // CHAR. Fail explicitly with an actionable hint rather than
            // silently sending BOOL (which the server then can't compare
            // against a CHAR column — `operator does not exist: bool = char`).
            // `character` (= bpchar, fixed-length text) has the same issue.
            // For text/varchar/character varying we coerce to a string.
            match arg_t.as_str() {
                "char" | "character" => Err(Error::ExecutionErr(format!(
                    "Cannot bind a JSON bool to a `{arg_t}` arg. \
                     `char` and `character` are single-byte / fixed-width text — \
                     pass the value as a string (e.g. \"t\" / \"f\") or change \
                     the arg type to `bool`."
                ))),
                "varchar" | "character varying" => Ok((Box::new(b.to_string()), Type::VARCHAR)),
                _ => Ok((Box::new(b.to_string()), Type::TEXT)),
            }
        }
        // Bool: bind as BOOL when no explicit text target. Postgres has an
        // implicit assignment cast bool→text, so INSERTs into text columns
        // still work — this only differs from the explicit-text branch above
        // for expression-context uses (WHERE clauses, etc.).
        Value::Bool(_) if arg_t == "jsonb" => Ok((Box::new(value.clone()), Type::JSONB)),
        Value::Bool(_) if arg_t == "json" => Ok((Box::new(value.clone()), Type::JSON)),
        Value::Bool(b) => Ok((Box::new(b.clone()), Type::BOOL)),
        // Number with an explicitly text-typed arg: coerce to String. Same
        // reasoning as the Bool branch — preserves pre-#8988 behaviour for
        // hand-written PG scripts that use `WHERE text_col = $1` with a
        // numeric value and an explicit text declaration.
        // Skip `char`/`character`: those go to the existing single-byte arm
        // below or the generic INT8 fallthrough.
        Value::Number(n)
            if explicit_text_target
                && (arg_t == "text" || arg_t == "varchar" || arg_t == "character varying") =>
        {
            let t = if arg_t == "varchar" || arg_t == "character varying" {
                Type::VARCHAR
            } else {
                Type::TEXT
            };
            Ok((Box::new(n.to_string()), t))
        }
        Value::Number(n) if arg_t == "char" && n.is_i64() => {
            Ok((Box::new(n.as_i64().unwrap() as i8), Type::CHAR))
        }
        Value::Number(n)
            if (arg_t == "smallint"
                || arg_t == "smallserial"
                || arg_t == "int2"
                || arg_t == "serial2")
                && n.is_i64() =>
        {
            Ok((Box::new(n.as_i64().unwrap() as i16), Type::INT2))
        }
        Value::Number(n)
            if (arg_t == "int" || arg_t == "integer" || arg_t == "int4" || arg_t == "serial")
                && n.is_i64() =>
        {
            Ok((Box::new(n.as_i64().unwrap() as i32), Type::INT4))
        }
        Value::Number(n) if (arg_t == "real" || arg_t == "float4") && n.as_f64().is_some() => {
            Ok((Box::new(n.as_f64().unwrap() as f32), Type::FLOAT4))
        }
        Value::Number(n)
            if (arg_t == "double" || arg_t == "double precision" || arg_t == "float8")
                && n.as_f64().is_some() =>
        {
            Ok((Box::new(n.as_f64().unwrap()), Type::FLOAT8))
        }
        Value::Number(n) if (arg_t == "numeric" || arg_t == "decimal") && n.is_i64() => Ok((
            Box::new(Decimal::from_i64(n.as_i64().unwrap()).unwrap_or_default()),
            Type::NUMERIC,
        )),
        Value::Number(n) if (arg_t == "numeric" || arg_t == "decimal") && n.is_f64() => Ok((
            Box::new(Decimal::from_f64(n.as_f64().unwrap()).unwrap_or_default()),
            Type::NUMERIC,
        )),
        Value::Number(n) if arg_t == "oid" && n.is_u64() => {
            Ok((Box::new(n.as_u64().unwrap() as u32), Type::OID))
        }
        Value::Number(n)
            if (arg_t == "bigint"
                || arg_t == "bigserial"
                || arg_t == "int8"
                || arg_t == "serial8")
                && n.is_u64() =>
        {
            Ok((Box::new(n.as_u64().unwrap() as i64), Type::INT8))
        }
        Value::Number(n) if n.is_i64() => Ok((Box::new(n.as_i64().unwrap()), Type::INT8)),
        Value::Number(n) => Ok((Box::new(n.as_f64().unwrap()), Type::FLOAT8)),
        Value::String(s) if arg_t == "uuid" => Ok((Box::new(Uuid::parse_str(s)?), Type::UUID)),
        Value::String(s)
            if arg_t == "smallint"
                || arg_t == "smallserial"
                || arg_t == "int2"
                || arg_t == "serial2" =>
        {
            s.parse::<i16>()
                .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::INT2))
                .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as smallint: {e}").into())
        }
        Value::String(s)
            if arg_t == "int" || arg_t == "integer" || arg_t == "int4" || arg_t == "serial" =>
        {
            s.parse::<i32>()
                .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::INT4))
                .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as integer: {e}").into())
        }
        Value::String(s)
            if arg_t == "bigint"
                || arg_t == "bigserial"
                || arg_t == "int8"
                || arg_t == "serial8" =>
        {
            s.parse::<i64>()
                .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::INT8))
                .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as bigint: {e}").into())
        }
        Value::String(s) if arg_t == "date" => {
            let date = parse_naive_date(s)
                .map_err(|e| Error::ExecutionErr(format!("Cannot parse '{s}' as date: {e}")))?;
            Ok((Box::new(date), Type::DATE))
        }
        Value::String(s) if arg_t == "time" => {
            let time = parse_naive_time(s)
                .map_err(|e| Error::ExecutionErr(format!("Cannot parse '{s}' as time: {e}")))?;
            Ok((Box::new(time), Type::TIME))
        }
        Value::String(s) if arg_t == "timetz" => {
            let time = parse_naive_time(s)
                .map_err(|e| Error::ExecutionErr(format!("Cannot parse '{s}' as time: {e}")))?;
            // See the timetz Null arm — assert TIME, server casts to TIMETZ.
            Ok((Box::new(time), Type::TIME))
        }
        Value::String(s) if arg_t == "timestamp" => {
            let datetime = parse_naive_datetime(s).map_err(|e| {
                Error::ExecutionErr(format!("Cannot parse '{s}' as timestamp: {e}"))
            })?;
            Ok((Box::new(datetime), Type::TIMESTAMP))
        }
        Value::String(s) if arg_t == "timestamptz" => {
            let datetime = parse_datetime_utc(s).map_err(|e| {
                Error::ExecutionErr(format!("Cannot parse '{s}' as timestamptz: {e}"))
            })?;
            Ok((Box::new(datetime), Type::TIMESTAMPTZ))
        }
        Value::String(s) if arg_t == "bytea" => {
            let bytes = engine::general_purpose::STANDARD
                .decode(s)
                .unwrap_or(vec![]);
            Ok((Box::new(bytes), Type::BYTEA))
        }
        // Parse Strings into the matching native Rust type for the remaining
        // recognised arg_ts that didn't have a dedicated arm. Without these,
        // a string value lands in the generic Value::String fallback below
        // (Box<String> + TEXT) and the server-side comparison
        // `<numeric|real|...> = text` fails since PG has no implicit cast.
        Value::String(s) if arg_t == "numeric" || arg_t == "decimal" => s
            .parse::<Decimal>()
            .map(|d| (Box::new(d) as Box<dyn ToSql + Sync + Send>, Type::NUMERIC))
            .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as numeric: {e}").into()),
        Value::String(s) if arg_t == "real" || arg_t == "float4" => s
            .parse::<f32>()
            .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::FLOAT4))
            .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as real: {e}").into()),
        Value::String(s)
            if arg_t == "double" || arg_t == "double precision" || arg_t == "float8" =>
        {
            s.parse::<f64>()
                .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::FLOAT8))
                .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as double: {e}").into())
        }
        Value::String(s) if arg_t == "oid" => s
            .parse::<u32>()
            .map(|n| (Box::new(n) as Box<dyn ToSql + Sync + Send>, Type::OID))
            .map_err(|e| anyhow::anyhow!("Cannot parse '{s}' as oid: {e}").into()),
        Value::String(s) if arg_t == "bool" || arg_t == "boolean" => {
            // Accept the same literals Postgres' boolin() does.
            let b = match s.to_ascii_lowercase().as_str() {
                "true" | "t" | "yes" | "y" | "1" | "on" => true,
                "false" | "f" | "no" | "n" | "0" | "off" => false,
                _ => {
                    return Err(
                        anyhow::anyhow!("Cannot parse '{s}' as bool: invalid literal").into(),
                    )
                }
            };
            Ok((Box::new(b), Type::BOOL))
        }
        Value::String(s) if arg_t == "varchar" || arg_t == "character varying" => {
            Ok((Box::new(s.clone()), Type::VARCHAR))
        }
        // For arg_t in (json, jsonb): bind a JSON-encodable Value with the
        // matching pg type. Falling through to TEXT here would assert TEXT
        // and break query_typed_raw's encoder check.
        // Object / Array (no `[]` suffix): bind as JSONB by default and
        // JSON-stringify when the target is text-like.
        //
        // Note the asymmetry vs the Bool/Number arms above: we coerce to
        // text on `matches!(typ, Typ::Str(_))` (which is true for both
        // explicit `(text)` decls AND parser-default text), not on
        // `explicit_text_target`. Reason: serialising a JSON object/array
        // as JSONB and binding against a parser-default-text arg would
        // assert `JSONB` for what could be a plain-text column. Postgres
        // has no implicit cast `jsonb → text` in expression context, so
        // `WHERE text_col = $1::JSONB` would fail. JSON-stringifying into
        // TEXT is what users almost always want for these JSON shapes
        // (and the result is itself valid JSON, so a `::jsonb` cast in
        // SQL still round-trips). Bool/Number don't need this safety
        // because `bool → text` and `int → text` have implicit assignment
        // casts; the asymmetry is therefore semantic, not a bug.
        Value::Array(_) if arg_t == "jsonb" => Ok((Box::new(value.clone()), Type::JSONB)),
        Value::Array(_) if arg_t == "json" => Ok((Box::new(value.clone()), Type::JSON)),
        Value::Array(_) if matches!(typ, Typ::Str(_)) => {
            let s = serde_json::to_string(value).map_err(|err| {
                Error::ExecutionErr(format!("Failed to convert JSON to text: {}", err))
            })?;
            let t = if arg_t == "varchar" {
                Type::VARCHAR
            } else {
                Type::TEXT
            };
            Ok((Box::new(s), t))
        }
        // Default for arrays without a [] suffix: bind as JSONB.
        Value::Array(_) => Ok((Box::new(value.clone()), Type::JSONB)),
        Value::Object(_) if arg_t == "json" => Ok((Box::new(value.clone()), Type::JSON)),
        Value::Object(_) if arg_t == "varchar" || arg_t == "character varying" => Ok((
            Box::new(serde_json::to_string(value).map_err(|err| {
                Error::ExecutionErr(format!("Failed to convert JSON to text: {}", err))
            })?),
            Type::VARCHAR,
        )),
        Value::Object(_) if arg_t == "text" || matches!(typ, Typ::Str(_)) => Ok((
            Box::new(serde_json::to_string(value).map_err(|err| {
                Error::ExecutionErr(format!("Failed to convert JSON to text: {}", err))
            })?),
            Type::TEXT,
        )),
        Value::Object(_) => Ok((Box::new(value.clone()), Type::JSONB)),
        // Generic String fallback. Use `AnyTextValue` (rather than plain
        // `String`) so the binding's `accepts()` covers `Kind::Enum` and
        // `Kind::Domain` in addition to the base text types — this is what
        // makes `INSERT INTO t VALUES ($1::my_enum)` work end-to-end without
        // users needing the `CAST($1::text AS my_enum)` workaround.
        //
        // We always assert `Type::TEXT` (not `Type::UNKNOWN`): tokio_postgres
        // sends parameter values in binary format, and Postgres rejects
        // binary-formatted bytes for `UNKNOWN` parameters in operator
        // contexts ("incorrect binary data format in bind parameter N"). The
        // trade-off is that bare `$1` against a non-text column still needs
        // an explicit cast (`$1::my_enum`), but the failure mode is a clear
        // server error rather than a cryptic protocol mismatch.
        Value::String(s) => Ok((Box::new(AnyTextValue(s.clone())), Type::TEXT)),
    }
}

/// Hard cap on how many `numeric` cells we test for f64-precision loss per
/// query. The check is `Decimal -> f64 -> Decimal` round-trip + `==` (~tens
/// of ns each); on a query returning millions of numeric cells, an
/// unbounded check would add measurable latency. After this many "fits
/// fine" observations we assume the rest do too — the pathological case
/// (rows 1..N fit, row N+1 loses precision) goes silently truncated, but
/// users who care about precision in such results can `::text`-cast their
/// SQL anyway. The first cell that does NOT fit short-circuits the budget
/// (the warning fires once and the per-row check stops immediately).
const NUMERIC_PRECISION_CHECK_BUDGET: u32 = 256;

/// Per-query state carried through result formatting. Currently used to
/// detect precision loss on the first `numeric` cell that doesn't round-trip
/// through f64, so the caller can emit a single warning per job rather than
/// silently truncating every row. Uses atomics (rather than `Cell`) so the
/// s3-streaming path — which moves the closure across futures and requires
/// `Send` — can borrow it.
pub struct ResultFormatState {
    /// `true` once we've observed a `numeric` value that loses precision when
    /// converted via f64. Once flipped, the per-row check short-circuits.
    pub numeric_precision_loss: std::sync::atomic::AtomicBool,
    /// Decremented for each `numeric` cell we actually check. When it hits 0
    /// the per-row check is skipped (along with the precision-loss flag) for
    /// the rest of the query — see the rationale on
    /// `NUMERIC_PRECISION_CHECK_BUDGET`.
    numeric_precision_check_budget: std::sync::atomic::AtomicU32,
}

impl Default for ResultFormatState {
    fn default() -> Self {
        Self {
            numeric_precision_loss: std::sync::atomic::AtomicBool::new(false),
            numeric_precision_check_budget: std::sync::atomic::AtomicU32::new(
                NUMERIC_PRECISION_CHECK_BUDGET,
            ),
        }
    }
}

pub fn pg_cell_to_json_value(
    row: &Row,
    column: &Column,
    column_i: usize,
) -> Result<JSONValue, Error> {
    pg_cell_to_json_value_with_state(row, column, column_i, &ResultFormatState::default())
}

pub fn pg_cell_to_json_value_with_state(
    row: &Row,
    column: &Column,
    column_i: usize,
    state: &ResultFormatState,
) -> Result<JSONValue, Error> {
    // JSON has no encoding for NaN / +Inf / -Inf, but Postgres `float4` /
    // `float8` (and `numeric`, via the special `'NaN'` value) do return them.
    // Pre-fix the worker errored with "invalid json-float", failing the
    // entire query. Round-trip these as JSON strings ("NaN", "Infinity",
    // "-Infinity") so the rest of the row still comes through; users who
    // need numeric semantics can filter them out client-side.
    let f64_to_json_number = |raw_val: f64| -> Result<JSONValue, Error> {
        if raw_val.is_nan() {
            return Ok(JSONValue::String("NaN".to_string()));
        }
        if raw_val.is_infinite() {
            return Ok(JSONValue::String(if raw_val > 0.0 {
                "Infinity".to_string()
            } else {
                "-Infinity".to_string()
            }));
        }
        let temp =
            serde_json::Number::from_f64(raw_val).ok_or(anyhow::anyhow!("invalid json-float"))?;
        Ok(JSONValue::Number(temp))
    };
    Ok(match *column.type_() {
        // for rust-postgres <> postgres type-mappings: https://docs.rs/postgres/latest/postgres/types/trait.FromSql.html#types
        // for postgres types: https://www.postgresql.org/docs/7.4/datatype.html#DATATYPE-TABLE

        // single types
        Type::BOOL => get_basic(row, column, column_i, |a: bool| Ok(JSONValue::Bool(a)))?,
        Type::BIT => get_basic(row, column, column_i, |a: bit_vec::BitVec| match a.len() {
            1 => Ok(JSONValue::Bool(a.get(0).unwrap())),
            _ => Ok(JSONValue::String(
                a.iter()
                    .map(|x| if x { "1" } else { "0" })
                    .collect::<String>(),
            )),
        })?,
        Type::INT2 => get_basic(row, column, column_i, |a: i16| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::INT4 => get_basic(row, column, column_i, |a: i32| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::INT8 => get_basic(row, column, column_i, |a: i64| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::TEXT | Type::VARCHAR => {
            get_basic(row, column, column_i, |a: String| Ok(JSONValue::String(a)))?
        }
        // ISO-8601 / RFC-3339 for temporal types so values round-trip through
        // JS / Python clients (`new Date(s)`, `datetime.fromisoformat(s)`)
        // without manual parsing. chrono's default `to_string()` returns
        // space-separated for naive datetimes and " UTC" suffix for tz-aware,
        // neither of which is parseable as ISO 8601.
        Type::TIMESTAMP => get_basic(row, column, column_i, |a: chrono::NaiveDateTime| {
            Ok(JSONValue::String(format_naive_datetime_iso(&a)))
        })?,
        Type::DATE => get_basic(row, column, column_i, |a: chrono::NaiveDate| {
            // chrono's `NaiveDate::to_string` is already ISO-8601 (`%Y-%m-%d`).
            Ok(JSONValue::String(a.to_string()))
        })?,
        Type::TIME => get_basic(row, column, column_i, |a: chrono::NaiveTime| {
            // `NaiveTime::to_string` is already ISO-8601 (`%H:%M:%S` with
            // optional `.f`).
            Ok(JSONValue::String(a.to_string()))
        })?,
        Type::TIMETZ => get_basic(row, column, column_i, |a: TimeTZStr| {
            // TimeTZStr's `from_sql` already formats as ISO-8601 (see impl
            // below).
            Ok(JSONValue::String(a.0))
        })?,
        Type::TIMESTAMPTZ => get_basic(row, column, column_i, |a: chrono::DateTime<Utc>| {
            Ok(JSONValue::String(a.to_rfc3339()))
        })?,
        Type::UUID => get_basic(row, column, column_i, |a: uuid::Uuid| {
            Ok(JSONValue::String(a.to_string()))
        })?,
        Type::INET | Type::CIDR => get_basic(row, column, column_i, |a: InetStr| {
            Ok(JSONValue::String(a.0))
        })?,
        Type::INTERVAL => get_basic(row, column, column_i, |a: IntervalStr| {
            Ok(JSONValue::String(a.0))
        })?,
        Type::JSON | Type::JSONB => get_basic(row, column, column_i, |a: JSONValue| Ok(a))?,
        Type::FLOAT4 => get_basic(row, column, column_i, |a: f32| {
            Ok(f64_to_json_number(a.into())?)
        })?,
        // Pre-existing behaviour: `numeric` is serialised as a JSON Number
        // via `Decimal::serialize`, which goes through f64 and silently
        // truncates past ~15-17 significant digits. Switching to JSON String
        // would preserve precision but break any user script doing arithmetic
        // / comparison on numeric column results (`row.amount + 1` becomes
        // string concat, `row.amount > 100` is lexicographic). Left as Number
        // for back-compat. Instead, on the FIRST cell whose decimal
        // representation can't round-trip through f64, we flip
        // `state.numeric_precision_loss` so the caller can emit a single
        // job-log warning recommending a `::text` cast. The check is bounded
        // by `NUMERIC_PRECISION_CHECK_BUDGET` cells (see comment there) and
        // short-circuits on the first lossy value, so the hot path on a
        // numeric-heavy result set is two atomic loads + an early return.
        Type::NUMERIC => get_basic(row, column, column_i, |a: Decimal| {
            if state.should_check_precision() && !decimal_fits_f64_losslessly(&a) {
                state
                    .numeric_precision_loss
                    .store(true, std::sync::atomic::Ordering::Relaxed);
            }
            Ok(serde_json::to_value(a)
                .map_err(|_| anyhow::anyhow!("Cannot convert decimal to json"))?)
        })?,
        Type::FLOAT8 => get_basic(row, column, column_i, |a: f64| f64_to_json_number(a))?,
        Type::BYTEA => get_basic(row, column, column_i, |a: Vec<u8>| {
            Ok(JSONValue::String(format!("\\x{}", hex::encode(a))))
        })?,
        // these types require a custom StringCollector struct as an intermediary (see struct at bottom)
        Type::TS_VECTOR => get_basic(row, column, column_i, |a: StringCollector| {
            Ok(JSONValue::String(a.0))
        })?,
        Type::OID => get_basic(row, column, column_i, |a: u32| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        // array types
        Type::BOOL_ARRAY => get_array(row, column, column_i, |a: bool| Ok(JSONValue::Bool(a)))?,
        Type::BIT_ARRAY => get_array(row, column, column_i, |a: bit_vec::BitVec| match a.len() {
            1 => Ok(JSONValue::Bool(a.get(0).unwrap())),
            _ => Ok(JSONValue::String(
                a.iter()
                    .map(|x| if x { "1" } else { "0" })
                    .collect::<String>(),
            )),
        })?,
        Type::INT2_ARRAY => get_array(row, column, column_i, |a: i16| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::INT4_ARRAY => get_array(row, column, column_i, |a: i32| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::INT8_ARRAY => get_array(row, column, column_i, |a: i64| {
            Ok(JSONValue::Number(serde_json::Number::from(a)))
        })?,
        Type::TEXT_ARRAY | Type::VARCHAR_ARRAY => {
            get_array(row, column, column_i, |a: String| Ok(JSONValue::String(a)))?
        }
        Type::JSON_ARRAY | Type::JSONB_ARRAY => {
            get_array(row, column, column_i, |a: JSONValue| Ok(a))?
        }
        Type::FLOAT4_ARRAY => get_array(row, column, column_i, |a: f32| {
            Ok(f64_to_json_number(a.into())?)
        })?,
        Type::FLOAT8_ARRAY => {
            get_array(row, column, column_i, |a: f64| Ok(f64_to_json_number(a)?))?
        }
        // See scalar NUMERIC arm — kept as JSON Number for back-compat,
        // with bounded precision-loss detection.
        Type::NUMERIC_ARRAY => get_array(row, column, column_i, |a: Decimal| {
            if state.should_check_precision() && !decimal_fits_f64_losslessly(&a) {
                state
                    .numeric_precision_loss
                    .store(true, std::sync::atomic::Ordering::Relaxed);
            }
            Ok(serde_json::to_value(a)
                .map_err(|_| anyhow::anyhow!("Cannot convert decimal to json"))?)
        })?,
        // these types require a custom StringCollector struct as an intermediary (see struct at bottom)
        Type::TS_VECTOR_ARRAY => get_array(row, column, column_i, |a: StringCollector| {
            Ok(JSONValue::String(a.0))
        })?,
        // Same ISO-8601 formatting as the scalar arms above.
        Type::TIMESTAMP_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveDateTime| {
            Ok(JSONValue::String(format_naive_datetime_iso(&a)))
        })?,
        Type::DATE_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveDate| {
            Ok(JSONValue::String(a.to_string()))
        })?,
        Type::TIME_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveTime| {
            Ok(JSONValue::String(a.to_string()))
        })?,
        Type::TIMETZ_ARRAY => get_array(row, column, column_i, |a: TimeTZStr| {
            Ok(JSONValue::String(a.0))
        })?,
        Type::TIMESTAMPTZ_ARRAY => get_array(row, column, column_i, |a: chrono::DateTime<Utc>| {
            Ok(JSONValue::String(a.to_rfc3339()))
        })?,
        Type::BYTEA_ARRAY => get_array(row, column, column_i, |a: Vec<u8>| {
            Ok(JSONValue::String(format!("\\x{}", hex::encode(a))))
        })?,
        Type::INET_ARRAY | Type::CIDR_ARRAY => get_array(row, column, column_i, |a: InetStr| {
            Ok(JSONValue::String(a.0))
        })?,
        Type::VOID => JSONValue::Null,
        // Default fallback for unhandled column types: read as text. We use
        // `AnyTextValue` instead of plain `String` so that `Kind::Enum`,
        // `Kind::Domain`, and citext columns round-trip into JSON strings
        // rather than erroring with `cannot convert between Option<String>
        // and the Postgres type \`<custom>\``.
        _ => get_basic(row, column, column_i, |a: AnyTextValue| {
            Ok(JSONValue::String(a.0))
        })?,
    })
}

pub fn postgres_row_to_json_value(row: Row) -> Result<JSONValue, Error> {
    postgres_row_to_json_value_with_state(row, &ResultFormatState::default())
}

pub fn postgres_row_to_json_value_with_state(
    row: Row,
    state: &ResultFormatState,
) -> Result<JSONValue, Error> {
    let row_data = postgres_row_to_row_data_with_state(row, state)?;
    Ok(JSONValue::Object(row_data))
}

// some type-aliases I use in my project
pub type JSONValue = serde_json::Value;
pub type RowData = Map<String, JSONValue>;

pub fn postgres_row_to_row_data(row: Row) -> Result<RowData, Error> {
    postgres_row_to_row_data_with_state(row, &ResultFormatState::default())
}

pub fn postgres_row_to_row_data_with_state(
    row: Row,
    state: &ResultFormatState,
) -> Result<RowData, Error> {
    let mut result: Map<String, JSONValue> = Map::new();
    for (i, column) in row.columns().iter().enumerate() {
        let name = column.name();
        let json_value = pg_cell_to_json_value_with_state(&row, column, i, state)?;
        result.insert(name.to_string(), json_value);
    }
    Ok(result)
}

/// Returns true if the `Decimal` value can round-trip through `f64` without
/// losing precision. Used to detect when the user's `numeric` results are
/// being silently truncated by the JSON Number serialisation path so the
/// worker can log a one-shot warning recommending a `::text` cast in SQL.
fn decimal_fits_f64_losslessly(d: &Decimal) -> bool {
    use rust_decimal::prelude::ToPrimitive;
    match d.to_f64() {
        Some(f) if f.is_finite() => Decimal::from_f64(f).is_some_and(|round| &round == d),
        _ => false,
    }
}

fn get_basic<'a, T: FromSql<'a>>(
    row: &'a Row,
    column: &Column,
    column_i: usize,
    val_to_json_val: impl Fn(T) -> Result<JSONValue, Error>,
) -> Result<JSONValue, Error> {
    let raw_val = row.try_get::<_, Option<T>>(column_i).with_context(|| {
        format!(
            "conversion issue for value at column_name `{}` with type {:?}",
            column.name(),
            column.type_()
        )
    })?;
    raw_val.map_or(Ok(JSONValue::Null), val_to_json_val)
}

struct IntervalStr(String);

impl<'a> FromSql<'a> for IntervalStr {
    fn from_sql(
        _: &Type,
        mut raw: &'a [u8],
    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
        let microseconds = raw.get_i64();
        let days = raw.get_i32();
        let months = raw.get_i32();
        Ok(IntervalStr(format!(
            "{:?} months {:?} days {:?} ms",
            months, days, microseconds
        )))
    }

    fn accepts(ty: &Type) -> bool {
        matches!(ty, &Type::INTERVAL)
    }
}

struct TimeTZStr(String);
impl<'a> FromSql<'a> for TimeTZStr {
    fn from_sql(
        _: &Type,
        mut raw: &'a [u8],
    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
        let microsecond = raw.get_i64();
        let offset = raw.get_i32();
        let utc_sec = (microsecond / 1_000_000) + offset as i64;
        let utc = chrono::NaiveTime::from_num_seconds_from_midnight_opt(
            ((utc_sec + 3600 * 24) % (3600 * 24)) as u32,
            ((microsecond % 1_000_000) * 1_000) as u32,
        )
        .ok_or_else(|| anyhow::anyhow!("Invalid time value"))?;
        // ISO-8601: append `+00:00` since TIMETZ is normalised to UTC here.
        Ok(TimeTZStr(format!("{}+00:00", utc)))
    }

    fn accepts(ty: &Type) -> bool {
        matches!(ty, &Type::TIMETZ)
    }
}

/// Format a `NaiveDateTime` as ISO-8601 (`YYYY-MM-DDTHH:MM:SS[.fff…]`).
/// chrono's default `to_string` uses a space separator, which is not parseable
/// by `new Date(s)` in older JS engines or Python's `datetime.fromisoformat`
/// before 3.11. Use the explicit format string so output is portable.
fn format_naive_datetime_iso(dt: &chrono::NaiveDateTime) -> String {
    if dt.and_utc().timestamp_subsec_nanos() == 0 {
        dt.format("%Y-%m-%dT%H:%M:%S").to_string()
    } else {
        dt.format("%Y-%m-%dT%H:%M:%S%.f").to_string()
    }
}

fn get_array<'a, T: FromSql<'a>>(
    row: &'a Row,
    column: &Column,
    column_i: usize,
    val_to_json_val: impl Fn(T) -> Result<JSONValue, Error>,
) -> Result<JSONValue, Error> {
    let raw_val_array = row
        .try_get::<_, Option<Vec<Option<T>>>>(column_i)
        .with_context(|| {
            format!(
                "conversion issue for array at column_name `{}`",
                column.name()
            )
        })?;
    Ok(match raw_val_array {
        Some(val_array) => {
            let mut result = vec![];
            for val in val_array {
                result.push(
                    val.map(|v| val_to_json_val(v))
                        .transpose()?
                        .unwrap_or(Value::Null),
                );
            }
            JSONValue::Array(result)
        }
        None => JSONValue::Null,
    })
}

/// `inet` / `cidr` in Postgres's own text form. `IpAddr`'s `FromSql` drops the
/// prefix length, so a network such as `10.0.0.0/8` would read back as a host.
struct InetStr(String);
impl FromSql<'_> for InetStr {
    fn from_sql(
        ty: &Type,
        raw: &[u8],
    ) -> Result<InetStr, Box<dyn std::error::Error + Sync + Send>> {
        let inet = postgres_protocol::types::inet_from_sql(raw)?;
        let host_prefix = if inet.addr().is_ipv4() { 32 } else { 128 };
        Ok(InetStr(
            if *ty == Type::CIDR || inet.netmask() != host_prefix {
                format!("{}/{}", inet.addr(), inet.netmask())
            } else {
                inet.addr().to_string()
            },
        ))
    }
    fn accepts(ty: &Type) -> bool {
        matches!(*ty, Type::INET | Type::CIDR)
    }
}

// you can remove this section if not using TS_VECTOR (or other types requiring an intermediary `FromSQL` struct)
struct StringCollector(String);
impl FromSql<'_> for StringCollector {
    fn from_sql(
        _: &Type,
        raw: &[u8],
    ) -> Result<StringCollector, Box<dyn std::error::Error + Sync + Send>> {
        let result = std::str::from_utf8(raw)?;
        Ok(StringCollector(result.to_owned()))
    }
    fn accepts(_ty: &Type) -> bool {
        true
    }
}

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

    /// The sentinel password is the whole opt-in: nothing else marks the resource, so a
    /// resource carrying it must not fall through to password auth.
    #[test]
    fn test_workload_identity_password_selects_the_auth_mode() {
        let db = |password: &str| {
            PgDatabase::parse_uri(&format!("postgres://someuser:{password}@host:5432/db")).unwrap()
        };
        assert_eq!(
            PgAuthMode::of(&db(WORKLOAD_IDENTITY_PASSWORD)).unwrap(),
            PgAuthMode::WorkloadIdentity
        );
        assert_eq!(
            PgAuthMode::of(&db("hunter2")).unwrap(),
            PgAuthMode::Password
        );
        // A pasted sentinel keeps its surrounding whitespace, and an unrecognized one is
        // forwarded to the server as a real password instead of selecting the mode.
        assert_eq!(
            PgAuthMode::of(&db("%20ms_entraid%0A")).unwrap(),
            PgAuthMode::WorkloadIdentity
        );
    }

    /// The job log is the only place the presented login is visible, and the two token
    /// modes differ on whether a missing one has a default at all.
    #[test]
    fn test_log_name_reports_the_presented_login() {
        let db = |user: &str| {
            PgDatabase::parse_uri(&format!("postgres://{user}:pw@host:5432/db")).unwrap()
        };

        assert_eq!(PgAuthMode::Password.log_name(&db("someuser")), None);
        assert_eq!(
            PgAuthMode::Iam.log_name(&db("someuser")).unwrap(),
            "IAM RDS authentication (login someuser)"
        );
        assert_eq!(
            PgAuthMode::Iam.log_name(&db("")).unwrap(),
            "IAM RDS authentication (login postgres)"
        );
        assert_eq!(
            PgAuthMode::WorkloadIdentity
                .log_name(&db("someuser"))
                .unwrap(),
            "Azure Workload Identity (login someuser)"
        );
        // Entra has no default login, so none is named rather than implying `postgres`.
        assert_eq!(
            PgAuthMode::WorkloadIdentity.log_name(&db("")).unwrap(),
            "Azure Workload Identity"
        );
    }

    #[test]
    fn test_map_s3object_jsonb_overflow() {
        let pg_err = Error::ExecutionErr(
            "db error: ERROR: total size of jsonb array elements exceeds the maximum of 268435455 bytes".to_string(),
        );
        let mapped = map_s3object_jsonb_overflow(pg_err, true).to_string();
        assert!(mapped.contains("256 MB"));
        assert!(mapped.contains("larger worker will not"));
        assert!(mapped.contains("read_csv_auto"));
        assert!(mapped.contains("Underlying error"));

        // The element-count cap is not the byte-size cap, so the "256 MB" message would
        // mislabel it — it must pass through unchanged.
        let count_err = Error::ExecutionErr(
            "number of jsonb array elements exceeds the maximum of 268435455".to_string(),
        );
        assert_eq!(
            map_s3object_jsonb_overflow(count_err, true).to_string(),
            "number of jsonb array elements exceeds the maximum of 268435455",
        );

        // A different "exceeds the maximum" error must NOT be reclassified as jsonb overflow.
        let array_err =
            Error::ExecutionErr("array size exceeds the maximum allowed (134217727)".to_string());
        assert_eq!(
            map_s3object_jsonb_overflow(array_err, true).to_string(),
            "array size exceeds the maximum allowed (134217727)",
        );

        let other = Error::ExecutionErr("syntax error at or near \"SELCT\"".to_string());
        assert_eq!(
            map_s3object_jsonb_overflow(other, true).to_string(),
            "syntax error at or near \"SELCT\"",
        );

        // No `(s3object)` input → even a matching error is left alone.
        let pg_err2 = Error::ExecutionErr("jsonb array elements exceeds the maximum".to_string());
        assert_eq!(
            map_s3object_jsonb_overflow(pg_err2, false).to_string(),
            "jsonb array elements exceeds the maximum",
        );
    }

    #[test]
    fn inet_keeps_the_prefix_postgres_prints() {
        let inet = |netmask: u8, addr: [u8; 4]| {
            let mut raw = vec![2, netmask, 0, 4];
            raw.extend(addr);
            raw
        };
        let read = |ty: &Type, raw: &[u8]| InetStr::from_sql(ty, raw).unwrap().0;
        assert_eq!(read(&Type::INET, &inet(32, [10, 1, 2, 3])), "10.1.2.3");
        assert_eq!(read(&Type::INET, &inet(8, [10, 0, 0, 0])), "10.0.0.0/8");
        assert_eq!(read(&Type::CIDR, &inet(32, [10, 1, 2, 3])), "10.1.2.3/32");
    }

    #[test]
    fn test_parse_naive_date() {
        // chrono's NaiveDate::to_string() format
        let d = parse_naive_date("2024-01-15").unwrap();
        assert_eq!(d.to_string(), "2024-01-15");

        // JS ISO format
        let d = parse_naive_date("2024-01-15T00:00:00.000Z").unwrap();
        assert_eq!(d.to_string(), "2024-01-15");

        // ISO without fractional seconds
        let d = parse_naive_date("2024-01-15T00:00:00Z").unwrap();
        assert_eq!(d.to_string(), "2024-01-15");
    }

    #[test]
    fn test_parse_naive_time() {
        // chrono's NaiveTime::to_string() format
        let t = parse_naive_time("10:30:00").unwrap();
        assert_eq!(t.to_string(), "10:30:00");

        // With fractional seconds
        let t = parse_naive_time("10:30:00.123456").unwrap();
        assert_eq!(t.to_string(), "10:30:00.123456");

        // Short format
        let t = parse_naive_time("10:30").unwrap();
        assert_eq!(t.to_string(), "10:30:00");

        // From full datetime string (JS frontend)
        let t = parse_naive_time("1970-01-01T10:30:00.000Z").unwrap();
        assert_eq!(t.to_string(), "10:30:00");
    }

    #[test]
    fn test_parse_naive_datetime() {
        // chrono's NaiveDateTime::to_string() format
        let dt = parse_naive_datetime("2024-01-15 10:30:00").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00");

        // With fractional seconds
        let dt = parse_naive_datetime("2024-01-15 10:30:00.123456").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00.123456");

        // ISO format with Z
        let dt = parse_naive_datetime("2024-01-15T10:30:00.000Z").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00");

        // ISO format without Z
        let dt = parse_naive_datetime("2024-01-15T10:30:00.000").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00");

        // ISO without fractional seconds
        let dt = parse_naive_datetime("2024-01-15T10:30:00Z").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00");
    }

    #[test]
    fn test_parse_datetime_utc() {
        // chrono's DateTime<Utc>::to_string() format
        let dt = parse_datetime_utc("2024-01-15 10:30:00 UTC").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00 UTC");

        // With fractional seconds
        let dt = parse_datetime_utc("2024-01-15 10:30:00.123456 UTC").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00.123456 UTC");

        // RFC 3339
        let dt = parse_datetime_utc("2024-01-15T10:30:00Z").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00 UTC");

        // RFC 3339 with fractional seconds
        let dt = parse_datetime_utc("2024-01-15T10:30:00.123Z").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00.123 UTC");

        // Numeric timezone offset (PostgreSQL text representation)
        let dt = parse_datetime_utc("2026-04-14 18:09:00+00").unwrap();
        assert_eq!(dt.to_string(), "2026-04-14 18:09:00 UTC");

        // Numeric timezone offset with fractional seconds
        let dt = parse_datetime_utc("2024-01-15 10:30:00.123+02").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 08:30:00.123 UTC");

        // Full offset format +00:00
        let dt = parse_datetime_utc("2024-01-15 10:30:00+00:00").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00 UTC");

        // ISO without timezone (treated as UTC)
        let dt = parse_datetime_utc("2024-01-15T10:30:00.000").unwrap();
        assert_eq!(dt.to_string(), "2024-01-15 10:30:00 UTC");
    }

    #[test]
    fn test_roundtrip_timestamp_formats() {
        // Verify that the format produced by pg_cell_to_json_value can be parsed back
        let original = chrono::NaiveDateTime::parse_from_str(
            "2024-06-15 14:30:45.123",
            "%Y-%m-%d %H:%M:%S%.f",
        )
        .unwrap();
        let serialized = original.to_string();
        let parsed = parse_naive_datetime(&serialized).unwrap();
        assert_eq!(original, parsed);
    }

    #[test]
    fn test_roundtrip_timestamptz_formats() {
        let original = "2024-06-15T14:30:45Z"
            .parse::<chrono::DateTime<Utc>>()
            .unwrap();
        let serialized = original.to_string();
        let parsed = parse_datetime_utc(&serialized).unwrap();
        assert_eq!(original, parsed);
    }

    #[test]
    fn test_roundtrip_time_formats() {
        let original = chrono::NaiveTime::parse_from_str("14:30:45.123", "%H:%M:%S%.f").unwrap();
        let serialized = original.to_string();
        let parsed = parse_naive_time(&serialized).unwrap();
        assert_eq!(original, parsed);
    }

    // ---------------------------------------------------------------------
    // convert_val: exhaustive (Value × otyp) → Type matrix.
    //
    // For every (JSON Value, parser otyp) combination that can occur from
    // either windmill-client SDK output (TS or Python) or a hand-written
    // Postgres script, verify that convert_val returns a `(Box, Type)` pair
    // where the Type matches the Box's concrete Rust type. This is the core
    // invariant that makes `query_typed_raw` safe — if it ever drifts again
    // (the bug introduced by #8988), users get
    // `cannot convert between the Rust type X and the Postgres type Y`.
    //
    // We can't introspect the Box's Rust type at runtime, but we *can* feed
    // each (Box, Type) through `to_sql_checked` against the asserted Type —
    // that's exactly the codepath `query_typed_raw` uses, so any mismatch
    // surfaces here as a `ToSql` error.
    // ---------------------------------------------------------------------

    use bytes::BytesMut;
    use serde_json::json;
    use tokio_postgres::types::IsNull;

    /// Verify that convert_val for `(value, arg_t)` returns a binding whose
    /// Rust type matches the asserted Postgres `Type` — exactly the check
    /// `query_typed_raw` performs when serialising parameters.
    ///
    /// Defaults to `otyp_inferred = false` (= "user explicitly typed this").
    /// Tests that need the parser-default flavour use
    /// `assert_convert_val_consistent_inferred`.
    fn assert_convert_val_consistent(
        label: &str,
        value: Value,
        arg_t: &str,
        typ: Typ,
        expected_type: Type,
    ) {
        assert_convert_val_consistent_full(label, value, arg_t, typ, expected_type, false)
    }

    fn assert_convert_val_consistent_inferred(
        label: &str,
        value: Value,
        arg_t: &str,
        typ: Typ,
        expected_type: Type,
    ) {
        assert_convert_val_consistent_full(label, value, arg_t, typ, expected_type, true)
    }

    fn assert_convert_val_consistent_full(
        label: &str,
        value: Value,
        arg_t: &str,
        typ: Typ,
        expected_type: Type,
        otyp_inferred: bool,
    ) {
        let (boxed, t) = convert_val(&value, &arg_t.to_string(), &typ, otyp_inferred)
            .unwrap_or_else(|e| panic!("{label}: convert_val errored: {e}"));
        assert_eq!(
            t, expected_type,
            "{label}: expected Type {expected_type}, got {t}"
        );
        // Run the encoder check — this is what query_typed_raw does internally
        // when binding the param. A mismatch between the boxed Rust type and
        // the asserted Type fails here as a `WrongType` error.
        let mut buf = BytesMut::new();
        match boxed.to_sql_checked(&t, &mut buf) {
            Ok(IsNull::Yes) | Ok(IsNull::No) => {}
            Err(e) => panic!(
                "{label}: ToSql failed for value={value:?} arg_t={arg_t} (asserted {t}): {e}"
            ),
        }
    }

    fn typ_for(arg_t: &str) -> Typ {
        windmill_parser_sql::parse_pg_typ(arg_t)
    }

    #[test]
    fn only_rowless_statements_skip_the_describe() {
        for q in [
            "SELECT * FROM t",
            "-- $1 n (int)\nselect $1",
            "WITH x AS (DELETE FROM t RETURNING *) SELECT * FROM x",
            "INSERT INTO t VALUES (1) RETURNING id",
            "/* c */ UPDATE t SET a = 1\nreturning *",
            "UPDATE t SET note = $$a;b$$ RETURNING *",
            "/* /* inner */ INSERT */ SELECT * FROM t",
            "/* never closed INSERT",
            "-- INSERT",
            "-- header\rSELECT m, pad AS\nupdate FROM t",
            "insert_rows()",
            "TABLE t",
            "VALUES (1)",
            "EXPLAIN SELECT 1",
        ] {
            assert!(can_stream_rows(q), "{q}");
        }
        for q in [
            "INSERT INTO t VALUES (1)",
            "-- $1 n (int)\nUPDATE t SET a = $1",
            "delete from t",
            "CREATE TABLE t (a int)",
            "SET search_path TO x",
            "CALL p()",
            "/* a /* b */ c */\n-- d\nINSERT INTO t VALUES (1)",
        ] {
            assert!(!can_stream_rows(q), "{q}");
        }
    }

    #[test]
    fn convert_val_null_for_every_known_arg_t() {
        // `Value::Null` for every type the parser may resolve, plus an unknown
        // arg_t (custom enum / extension). Each must produce a matching Type
        // and serialise without error.
        let cases: &[(&str, Type)] = &[
            ("bool", Type::BOOL),
            ("boolean", Type::BOOL),
            ("char", Type::CHAR),
            ("character", Type::CHAR),
            ("smallint", Type::INT2),
            ("int2", Type::INT2),
            ("smallserial", Type::INT2),
            ("serial2", Type::INT2),
            ("int", Type::INT4),
            ("integer", Type::INT4),
            ("int4", Type::INT4),
            ("serial", Type::INT4),
            ("bigint", Type::INT8),
            ("int8", Type::INT8),
            ("bigserial", Type::INT8),
            ("serial8", Type::INT8),
            ("real", Type::FLOAT4),
            ("float4", Type::FLOAT4),
            ("double", Type::FLOAT8),
            ("double precision", Type::FLOAT8),
            ("float8", Type::FLOAT8),
            ("numeric", Type::NUMERIC),
            ("decimal", Type::NUMERIC),
            ("oid", Type::OID),
            ("uuid", Type::UUID),
            ("date", Type::DATE),
            ("time", Type::TIME),
            // chrono::NaiveTime can only encode as TIME — see the Null arm.
            ("timetz", Type::TIME),
            ("timestamp", Type::TIMESTAMP),
            ("timestamptz", Type::TIMESTAMPTZ),
            ("json", Type::JSON),
            ("jsonb", Type::JSONB),
            ("bytea", Type::BYTEA),
            ("text", Type::TEXT),
            ("varchar", Type::VARCHAR),
            ("character varying", Type::VARCHAR),
            // Unknown / custom type: convert_val falls back to TEXT NULL — the
            // dispatch then takes the prepare + query_raw path so the server
            // resolves the actual column type.
            ("my_custom_enum", Type::TEXT),
        ];
        for (arg_t, expected) in cases {
            assert_convert_val_consistent(
                &format!("Null/{arg_t}"),
                Value::Null,
                arg_t,
                typ_for(arg_t),
                expected.clone(),
            );
        }
    }

    #[test]
    fn convert_val_bool_against_every_arg_t() {
        // Value::Bool. Pre-#8988 this always produced Box<bool>, which
        // mismatched the asserted Type for parser-default "text" (the
        // regression we fix). Post-fix:
        //  - explicit text-like target (`-- $1 (text)` / `$1::text`):
        //    coerce to Box<String>+TEXT so `WHERE text_col = $1` works.
        //  - parser-default text (bare `$N`, no annotation): bind as BOOL
        //    natively, server casts at the use site.
        //  - any other target: bind as BOOL.
        let bool_targets = ["bool", "boolean"];
        let json_targets = [("json", Type::JSON), ("jsonb", Type::JSONB)];
        let explicit_text_targets = [
            ("text", Type::TEXT),
            ("varchar", Type::VARCHAR),
            ("character varying", Type::VARCHAR),
        ];
        let bind_as_bool = [
            "smallint",
            "int",
            "integer",
            "bigint",
            "int4",
            "int8",
            "real",
            "double",
            "double precision",
            "numeric",
            "uuid",
            "date",
            "time",
            "timestamp",
            "timestamptz",
            "bytea",
            "oid",
            // unknown — server resolves via prepare path
            "my_custom_enum",
        ];

        for v in [true, false] {
            for arg_t in &bool_targets {
                assert_convert_val_consistent(
                    &format!("Bool({v})/{arg_t}"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    Type::BOOL,
                );
            }
            for (arg_t, expected) in &json_targets {
                assert_convert_val_consistent(
                    &format!("Bool({v})/{arg_t}"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    expected.clone(),
                );
            }
            for (arg_t, expected) in &explicit_text_targets {
                // Explicit (otyp_inferred=false): coerce to text.
                assert_convert_val_consistent(
                    &format!("Bool({v})/{arg_t} explicit"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    expected.clone(),
                );
                // Inferred (parser-default): keep BOOL.
                assert_convert_val_consistent_inferred(
                    &format!("Bool({v})/{arg_t} inferred"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    Type::BOOL,
                );
            }
            // `char` and `character` (= bpchar) are single-byte / fixed-width
            // text. Explicit decl with a JSON bool errors with an actionable
            // hint instead of silently binding BOOL (which a CHAR column
            // can't compare against). Inferred-default still binds BOOL.
            for arg_t in &["char", "character"] {
                let err = convert_val(
                    &Value::Bool(v),
                    &arg_t.to_string(),
                    &typ_for(arg_t),
                    /* otyp_inferred = */ false,
                )
                .err()
                .unwrap_or_else(|| panic!("Bool({v})/{arg_t} explicit should error"));
                let msg = err.to_string();
                assert!(
                    msg.contains("Cannot bind a JSON bool"),
                    "Bool({v})/{arg_t} explicit error didn't have expected message: {msg}"
                );
                assert_convert_val_consistent_inferred(
                    &format!("Bool({v})/{arg_t} inferred"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    Type::BOOL,
                );
            }
            for arg_t in &bind_as_bool {
                assert_convert_val_consistent(
                    &format!("Bool({v})/{arg_t}"),
                    Value::Bool(v),
                    arg_t,
                    typ_for(arg_t),
                    Type::BOOL,
                );
            }
        }
    }

    #[test]
    fn convert_val_integer_number_against_every_arg_t() {
        // JSON integers. Each arg_t selects its matching encoder; for arg_ts
        // that don't have a numeric encoder (uuid, date, …), the value falls
        // through to the generic Number arm — Box<i64> bound as INT8 — and
        // server-side casts handle the rest if the SQL wants it.
        let cases: &[(&str, Type)] = &[
            ("char", Type::CHAR),
            // "character" (= bpchar in PG) doesn't have a Number arm, so it
            // falls through to generic Number → Box<i64> + INT8. Server
            // casts at the SQL site if the column is bpchar.
            ("character", Type::INT8),
            ("smallint", Type::INT2),
            ("smallserial", Type::INT2),
            ("int2", Type::INT2),
            ("serial2", Type::INT2),
            ("int", Type::INT4),
            ("integer", Type::INT4),
            ("int4", Type::INT4),
            ("serial", Type::INT4),
            ("bigint", Type::INT8),
            ("bigserial", Type::INT8),
            ("int8", Type::INT8),
            ("serial8", Type::INT8),
            ("oid", Type::OID),
            ("numeric", Type::NUMERIC),
            ("decimal", Type::NUMERIC),
            // Unknown arg_t falls through to generic Number → INT8.
            ("my_custom_enum", Type::INT8),
        ];
        for (arg_t, expected) in cases {
            assert_convert_val_consistent(
                &format!("Number(42)/{arg_t}"),
                json!(42),
                arg_t,
                typ_for(arg_t),
                expected.clone(),
            );
        }
        // Text targets: split between explicit (coerce to TEXT) and inferred
        // (parser-default, bind as INT8 — server casts at the use site).
        for (arg_t, expected_text) in [("text", Type::TEXT), ("varchar", Type::VARCHAR)] {
            assert_convert_val_consistent(
                &format!("Number(42)/{arg_t} explicit"),
                json!(42),
                arg_t,
                typ_for(arg_t),
                expected_text,
            );
            assert_convert_val_consistent_inferred(
                &format!("Number(42)/{arg_t} inferred"),
                json!(42),
                arg_t,
                typ_for(arg_t),
                Type::INT8,
            );
        }
        // Negative integer (is_u64 false → falls to generic i64 arm for bigint).
        assert_convert_val_consistent(
            "Number(-7)/bigint",
            json!(-7),
            "bigint",
            typ_for("bigint"),
            Type::INT8,
        );
        assert_convert_val_consistent("Number(0)/oid", json!(0), "oid", typ_for("oid"), Type::OID);
    }

    #[test]
    fn convert_val_float_number_against_every_arg_t() {
        let cases: &[(&str, Type)] = &[
            ("real", Type::FLOAT4),
            ("float4", Type::FLOAT4),
            ("double", Type::FLOAT8),
            ("double precision", Type::FLOAT8),
            ("float8", Type::FLOAT8),
            ("numeric", Type::NUMERIC),
            ("decimal", Type::NUMERIC),
            // Unknown arg_t falls through to generic → Box<f64>+FLOAT8.
            ("my_custom_enum", Type::FLOAT8),
        ];
        for (arg_t, expected) in cases {
            assert_convert_val_consistent(
                &format!("Number(3.14)/{arg_t}"),
                json!(3.14),
                arg_t,
                typ_for(arg_t),
                expected.clone(),
            );
        }
        // Text targets: split between explicit (coerce to TEXT/VARCHAR) and
        // inferred (parser-default, bind as FLOAT8 — server casts at use site).
        for (arg_t, expected_text) in [("text", Type::TEXT), ("varchar", Type::VARCHAR)] {
            assert_convert_val_consistent(
                &format!("Number(3.14)/{arg_t} explicit"),
                json!(3.14),
                arg_t,
                typ_for(arg_t),
                expected_text,
            );
            assert_convert_val_consistent_inferred(
                &format!("Number(3.14)/{arg_t} inferred"),
                json!(3.14),
                arg_t,
                typ_for(arg_t),
                Type::FLOAT8,
            );
        }
    }

    #[test]
    fn convert_val_string_against_every_arg_t() {
        // Strings parse into the matching Rust type when arg_t resolves to a
        // numeric / temporal / uuid / bytea type; otherwise they bind as TEXT.
        assert_convert_val_consistent(
            "String('42')/smallint",
            json!("42"),
            "smallint",
            typ_for("smallint"),
            Type::INT2,
        );
        assert_convert_val_consistent(
            "String('42')/int",
            json!("42"),
            "int",
            typ_for("int"),
            Type::INT4,
        );
        assert_convert_val_consistent(
            "String('42')/bigint",
            json!("42"),
            "bigint",
            typ_for("bigint"),
            Type::INT8,
        );
        assert_convert_val_consistent(
            "String(uuid)/uuid",
            json!("550e8400-e29b-41d4-a716-446655440000"),
            "uuid",
            typ_for("uuid"),
            Type::UUID,
        );
        assert_convert_val_consistent(
            "String(date)/date",
            json!("2024-01-15"),
            "date",
            typ_for("date"),
            Type::DATE,
        );
        assert_convert_val_consistent(
            "String(time)/time",
            json!("10:30:00"),
            "time",
            typ_for("time"),
            Type::TIME,
        );
        assert_convert_val_consistent(
            "String(time)/timetz",
            json!("10:30:00"),
            "timetz",
            typ_for("timetz"),
            Type::TIME,
        );
        assert_convert_val_consistent(
            "String(ts)/timestamp",
            json!("2024-01-15T10:30:00"),
            "timestamp",
            typ_for("timestamp"),
            Type::TIMESTAMP,
        );
        assert_convert_val_consistent(
            "String(tstz)/timestamptz",
            json!("2024-01-15T10:30:00Z"),
            "timestamptz",
            typ_for("timestamptz"),
            Type::TIMESTAMPTZ,
        );
        assert_convert_val_consistent(
            "String(b64)/bytea",
            json!("aGVsbG8="),
            "bytea",
            typ_for("bytea"),
            Type::BYTEA,
        );
        // Generic text arms.
        for (arg_t, expected) in [
            ("text", Type::TEXT),
            ("varchar", Type::VARCHAR),
            ("character varying", Type::VARCHAR),
            // Unknown → TEXT (prepare fallback in dispatch).
            ("my_custom_enum", Type::TEXT),
        ] {
            assert_convert_val_consistent(
                &format!("String('hello')/{arg_t}"),
                json!("hello"),
                arg_t,
                typ_for(arg_t),
                expected,
            );
        }
    }

    #[test]
    fn convert_val_object_against_every_arg_t() {
        // Object values: bind as JSONB (default), JSON if explicitly typed,
        // or JSON-stringify into TEXT/VARCHAR when arg_t is text-like.
        assert_convert_val_consistent(
            "Object/jsonb",
            json!({"k": 1}),
            "jsonb",
            typ_for("jsonb"),
            Type::JSONB,
        );
        assert_convert_val_consistent(
            "Object/json",
            json!({"k": 1}),
            "json",
            typ_for("json"),
            Type::JSON,
        );
        assert_convert_val_consistent(
            "Object/text",
            json!({"k": 1}),
            "text",
            typ_for("text"),
            Type::TEXT,
        );
        assert_convert_val_consistent(
            "Object/varchar",
            json!({"k": 1}),
            "varchar",
            typ_for("varchar"),
            Type::VARCHAR,
        );
        // Parser-default text (Typ::Str) still routes to TEXT-string via the
        // `matches!(typ, Typ::Str(_))` arm.
        assert_convert_val_consistent(
            "Object/parser-default-text",
            json!({"k": 1}),
            "text",
            Typ::Str(None),
            Type::TEXT,
        );
        // Unknown arg_t parses to Typ::Str (parser's catch-all), so the
        // text-coercion arm picks it up — Box<String> + TEXT. The dispatch
        // then takes the prepare + query_raw path because otyp_to_pg_type
        // returns Err for the unknown name, letting the server resolve the
        // actual column type (e.g. a custom enum that accepts JSON via cast).
        assert_convert_val_consistent(
            "Object/my_custom_enum",
            json!({"k": 1}),
            "my_custom_enum",
            typ_for("my_custom_enum"),
            Type::TEXT,
        );
    }

    #[test]
    fn convert_val_array_against_every_arg_t() {
        // arg_t with [] suffix routes to convert_vec_val.
        let int_array_cases: &[(&str, Type)] = &[
            ("int[]", Type::INT4_ARRAY),
            ("integer[]", Type::INT4_ARRAY),
            ("int4[]", Type::INT4_ARRAY),
            ("smallint[]", Type::INT2_ARRAY),
            ("bigint[]", Type::INT8_ARRAY),
        ];
        for (arg_t, expected) in int_array_cases {
            assert_convert_val_consistent(
                &format!("Array([1,2])/{arg_t}"),
                json!([1, 2]),
                arg_t,
                typ_for(arg_t),
                expected.clone(),
            );
        }
        assert_convert_val_consistent(
            "Array(strs)/text[]",
            json!(["a", "b"]),
            "text[]",
            typ_for("text[]"),
            Type::TEXT_ARRAY,
        );
        assert_convert_val_consistent(
            "Array(strs)/varchar[]",
            json!(["a", "b"]),
            "varchar[]",
            typ_for("varchar[]"),
            Type::VARCHAR_ARRAY,
        );
        assert_convert_val_consistent(
            "Array(bools)/bool[]",
            json!([true, false]),
            "bool[]",
            typ_for("bool[]"),
            Type::BOOL_ARRAY,
        );
        assert_convert_val_consistent(
            "Array(floats)/double[]",
            json!([1.5, 2.5]),
            "double[]",
            typ_for("double[]"),
            Type::FLOAT8_ARRAY,
        );
        assert_convert_val_consistent(
            "Array(uuids)/uuid[]",
            json!(["550e8400-e29b-41d4-a716-446655440000"]),
            "uuid[]",
            typ_for("uuid[]"),
            Type::UUID_ARRAY,
        );
        // `timetz[]` falls back to TIME_ARRAY for the same reason the scalar
        // `timetz` falls back to TIME — chrono's `NaiveTime` only encodes for
        // TIME. The encoder check (to_sql_checked) catches a mistakenly
        // asserted TIMETZ_ARRAY here.
        assert_convert_val_consistent(
            "Array(times)/timetz[] → TIME_ARRAY",
            json!(["10:30:00", "11:00:00"]),
            "timetz[]",
            typ_for("timetz[]"),
            Type::TIME_ARRAY,
        );
        assert_convert_val_consistent(
            "Array(times)/time[]",
            json!(["10:30:00", "11:00:00"]),
            "time[]",
            typ_for("time[]"),
            Type::TIME_ARRAY,
        );
        // Array without [] suffix on arg_t: bind as JSONB (or JSON / TEXT).
        assert_convert_val_consistent(
            "Array/jsonb (no [])",
            json!([1, 2, 3]),
            "jsonb",
            typ_for("jsonb"),
            Type::JSONB,
        );
        assert_convert_val_consistent(
            "Array/json (no [])",
            json!([1, 2, 3]),
            "json",
            typ_for("json"),
            Type::JSON,
        );
        assert_convert_val_consistent(
            "Array/text (parser-default)",
            json!([1, 2, 3]),
            "text",
            typ_for("text"),
            Type::TEXT,
        );
        // See Object/my_custom_enum: Typ::Str catch-all → TEXT-stringify;
        // dispatch falls back to prepare for the unknown arg_t.
        assert_convert_val_consistent(
            "Array/my_custom_enum",
            json!([1, 2, 3]),
            "my_custom_enum",
            typ_for("my_custom_enum"),
            Type::TEXT,
        );
        // NULL-array shape: Value::Null with arg_t ending in []
        assert_convert_val_consistent(
            "Null/int[]",
            Value::Null,
            "int[]",
            typ_for("int[]"),
            Type::INT4_ARRAY,
        );
    }

    /// Edge cases mirroring what the SDKs (TS / Python) and hand-written PG
    /// scripts can actually emit. Each case is a real input → encode round
    /// trip, and would have failed under #8988 if the asserted Type drifted
    /// from the binding's Rust type.
    #[test]
    fn convert_val_sdk_edge_cases() {
        // TS SDK shapes — `${val}` is auto-tagged with ::TYPE.
        assert_convert_val_consistent(
            "TS SDK ${42}",
            json!(42),
            "bigint",
            typ_for("bigint"),
            Type::INT8,
        );
        assert_convert_val_consistent(
            "TS SDK ${3.14}",
            json!(3.14),
            "double",
            typ_for("double"),
            Type::FLOAT8,
        );
        assert_convert_val_consistent(
            "TS SDK ${true}",
            json!(true),
            "boolean",
            typ_for("boolean"),
            Type::BOOL,
        );
        assert_convert_val_consistent(
            "TS SDK ${\"hello\"}",
            json!("hello"),
            "text",
            typ_for("text"),
            Type::TEXT,
        );
        assert_convert_val_consistent(
            "TS SDK ${{x:1}}",
            json!({"x": 1}),
            "json",
            typ_for("json"),
            Type::JSON,
        );
        assert_convert_val_consistent(
            "TS SDK ${[1,2,3]}",
            json!([1, 2, 3]),
            "json",
            typ_for("json"),
            Type::JSON,
        );
        assert_convert_val_consistent(
            "TS SDK ${null}",
            Value::Null,
            "text",
            typ_for("text"),
            Type::TEXT,
        );

        // CAST(${val} AS T) shape — the SDK strips its own ::TYPE here, so
        // the parser sees a bare $N and otyp defaults to "text" *with
        // otyp_inferred = true*. This is the original regression #8988
        // introduced; the inferred-default flag is what lets convert_val
        // bind the value's natural type rather than coerce to TEXT.
        assert_convert_val_consistent_inferred(
            "CAST AS bool / Bool true",
            json!(true),
            "text",
            typ_for("text"),
            Type::BOOL,
        );
        assert_convert_val_consistent_inferred(
            "CAST AS bool / Bool false",
            json!(false),
            "text",
            typ_for("text"),
            Type::BOOL,
        );
        // Object falls into the text-coercion arm (Object branch checks
        // `matches!(typ, Typ::Str(_))` regardless of otyp_inferred — JSON
        // serialisation is always safer than asserting JSONB for an
        // unannotated arg).
        assert_convert_val_consistent_inferred(
            "CAST AS jsonb / Object",
            json!({"a": 1, "b": [2, 3]}),
            "text",
            typ_for("text"),
            Type::TEXT,
        );
        assert_convert_val_consistent_inferred(
            "CAST AS int / Number",
            json!(7),
            "text",
            typ_for("text"),
            Type::INT8,
        );

        // Python SDK datatable shape — type sits in the declaration comment
        // (`-- $1 arg1 (BIGINT)`). Parser resolves otyp before convert_val.
        assert_convert_val_consistent(
            "Python decl bigint / Number",
            json!(42),
            "bigint",
            typ_for("bigint"),
            Type::INT8,
        );
        assert_convert_val_consistent(
            "Python decl text / String",
            json!("hello"),
            "text",
            typ_for("text"),
            Type::TEXT,
        );
        assert_convert_val_consistent(
            "Python decl jsonb / Object",
            json!({"k": [1, 2]}),
            "jsonb",
            typ_for("jsonb"),
            Type::JSONB,
        );

        // Mismatched-but-coercible JSON shape: JSON int 0/1 into a bool col
        // still works because tokio_postgres encodes Number as INT8 and
        // postgres has int→bool cast at the SQL site. Uses the inferred
        // path (parser-default text otyp).
        assert_convert_val_consistent_inferred(
            "Number(0)/bool (parser default)",
            json!(0),
            "text",
            typ_for("text"),
            Type::INT8,
        );
    }

    /// `decimal_fits_f64_losslessly` returns true for values that round-trip
    /// through f64 and false for values that don't. This is the predicate
    /// behind the one-shot precision-loss warning.
    #[test]
    fn decimal_fits_f64_losslessly_predicate() {
        use std::str::FromStr;
        // Values that fit f64 cleanly:
        for s in &["0", "1", "-1", "3.14", "1234.5", "-0.5", "10000000000"] {
            let d = Decimal::from_str(s).unwrap();
            assert!(
                decimal_fits_f64_losslessly(&d),
                "expected `{s}` to fit f64 losslessly"
            );
        }
        // Values past f64's ~15 significant-digit window lose precision:
        for s in &[
            "12345678901234.56789", // 19 sig digits
            "0.123456789012345678", // 18 sig digits past the decimal
            "99999999999999999999", // 20-digit integer
        ] {
            let d = Decimal::from_str(s).unwrap();
            assert!(
                !decimal_fits_f64_losslessly(&d),
                "expected `{s}` to NOT fit f64 losslessly"
            );
        }
    }

    /// `should_check_precision` returns `true` exactly
    /// `NUMERIC_PRECISION_CHECK_BUDGET` times, then `false` forever — and
    /// `false` immediately once the precision-loss flag has been set, so the
    /// hot path on a numeric-heavy result set is one cheap atomic load after
    /// the first lossy value is observed.
    #[test]
    fn precision_check_budget_caps_per_query_overhead() {
        use std::sync::atomic::Ordering;
        let state = ResultFormatState::default();
        let mut allowed = 0u32;
        let mut denied = 0u32;
        for _ in 0..(NUMERIC_PRECISION_CHECK_BUDGET + 100) {
            if state.should_check_precision() {
                allowed += 1;
            } else {
                denied += 1;
            }
        }
        assert_eq!(allowed, NUMERIC_PRECISION_CHECK_BUDGET);
        assert_eq!(denied, 100);
        // The flag short-circuits the budget — once set, no more checks run
        // even if the budget hadn't been spent.
        let state = ResultFormatState::default();
        state.numeric_precision_loss.store(true, Ordering::Relaxed);
        for _ in 0..10 {
            assert!(!state.should_check_precision());
        }
        // Budget untouched.
        assert_eq!(
            state.numeric_precision_check_budget.load(Ordering::Relaxed),
            NUMERIC_PRECISION_CHECK_BUDGET
        );
    }

    /// Sparse positional placeholders renumber to a contiguous 1..=N without
    /// substring collisions OR mangling string-literal/comment occurrences.
    /// The pre-existing `String::replace` chain turned `$50` into `$10` when
    /// oidx=5 was processed first; even the regex-with-greedy-digits approach
    /// (a regression of its own) walked through string literals. The current
    /// position-aware rewrite uses the parser's tokenizer to skip those.
    #[test]
    fn renumber_sparse_placeholders_no_collision_no_string_mangling() {
        fn renumber(input: &str, mapping: &HashMap<i32, usize>) -> String {
            let mut out = input.to_owned();
            let mut positions = windmill_parser_sql::parse_pg_statement_arg_positions(input);
            positions.sort_by_key(|(_, range)| std::cmp::Reverse(range.start));
            for (oidx, range) in positions {
                if let Some(new_i) = mapping.get(&oidx) {
                    if oidx as usize != *new_i {
                        out.replace_range(range, &new_i.to_string());
                    }
                }
            }
            out
        }

        let mapping: HashMap<i32, usize> = [(5, 1), (50, 2)].into_iter().collect();
        let cases = &[
            // Two placeholders, full rewrite (greedy-digit collision check).
            ("SELECT $5, $50", "SELECT $1, $2"),
            // Same input flipped — order independence.
            ("SELECT $50, $5", "SELECT $2, $1"),
            // Repeat use of an index — every site gets rewritten.
            (
                "SELECT $5 FROM t WHERE id = $5 OR ref = $50",
                "SELECT $1 FROM t WHERE id = $1 OR ref = $2",
            ),
            // Index outside the mapping is left intact.
            ("SELECT $5, $99", "SELECT $1, $99"),
            // String literal containing the same `$N` syntax must not be
            // rewritten — the tokenizer marks it as inside a string.
            (
                "SELECT 'price: $5' AS lbl, $5 FROM t",
                "SELECT 'price: $5' AS lbl, $1 FROM t",
            ),
            // Single-line comment must not be rewritten either.
            ("-- mention $5\nSELECT $5", "-- mention $5\nSELECT $1"),
            // Dollar-quoted block ($$ … $$) must not be rewritten.
            ("SELECT $$body with $5$$, $5", "SELECT $$body with $5$$, $1"),
        ];
        for (input, expected) in cases {
            assert_eq!(
                renumber(input, &mapping).as_str(),
                *expected,
                "input={input}"
            );
        }
    }

    /// Drift-prevention: `otyp_to_pg_type` and `convert_val` must agree on the
    /// Type for every recognised arg_t when the JSON value matches the arg_t's
    /// "natural" Rust kind. Fails if someone adds a new arg_t to one but not
    /// the other, or changes the Type returned by either.
    #[test]
    fn otyp_to_pg_type_and_convert_val_agree_for_recognised_types() {
        // (arg_t, natural-value, expected scalar Type)
        let cases: &[(&str, Value, Type)] = &[
            ("bool", json!(true), Type::BOOL),
            ("boolean", json!(false), Type::BOOL),
            ("char", json!(65), Type::CHAR),
            ("smallint", json!(1), Type::INT2),
            ("smallserial", json!(1), Type::INT2),
            ("int2", json!(1), Type::INT2),
            ("serial2", json!(1), Type::INT2),
            ("int", json!(1), Type::INT4),
            ("integer", json!(1), Type::INT4),
            ("int4", json!(1), Type::INT4),
            ("serial", json!(1), Type::INT4),
            ("bigint", json!(1), Type::INT8),
            ("int8", json!(1), Type::INT8),
            ("bigserial", json!(1), Type::INT8),
            ("serial8", json!(1), Type::INT8),
            ("real", json!(1.5), Type::FLOAT4),
            ("float4", json!(1.5), Type::FLOAT4),
            ("double", json!(1.5), Type::FLOAT8),
            ("double precision", json!(1.5), Type::FLOAT8),
            ("float8", json!(1.5), Type::FLOAT8),
            ("numeric", json!(1), Type::NUMERIC),
            ("decimal", json!(1), Type::NUMERIC),
            ("oid", json!(1), Type::OID),
            (
                "uuid",
                json!("550e8400-e29b-41d4-a716-446655440000"),
                Type::UUID,
            ),
            ("date", json!("2024-01-15"), Type::DATE),
            ("time", json!("10:30:00"), Type::TIME),
            // chrono::NaiveTime can only encode TIME — see the timetz arm.
            ("timetz", json!("10:30:00"), Type::TIME),
            ("timestamp", json!("2024-01-15T10:30:00"), Type::TIMESTAMP),
            (
                "timestamptz",
                json!("2024-01-15T10:30:00Z"),
                Type::TIMESTAMPTZ,
            ),
            ("json", json!({"k": 1}), Type::JSON),
            ("jsonb", json!({"k": 1}), Type::JSONB),
            ("bytea", json!("aGVsbG8="), Type::BYTEA),
            ("text", json!("hello"), Type::TEXT),
            ("varchar", json!("hello"), Type::VARCHAR),
            ("character varying", json!("hello"), Type::VARCHAR),
        ];
        for (arg_t, value, expected) in cases {
            // 1. The dispatch's "is recognised" gate must accept this arg_t.
            //    `timetz` is the one exception where we deliberately return
            //    TIME from convert_val (chrono limitation), but otyp_to_pg_type
            //    returns TIMETZ.
            let from_otyp = otyp_to_pg_type(arg_t)
                .unwrap_or_else(|e| panic!("otyp_to_pg_type lost arg_t `{arg_t}`: {e}"));
            if *arg_t != "timetz" {
                assert_eq!(
                    from_otyp, *expected,
                    "otyp_to_pg_type({arg_t}) drift: expected {expected}, got {from_otyp}"
                );
            }
            // 2. convert_val must produce a binding whose Type matches
            //    `expected`, AND whose Rust type successfully encodes against
            //    that Type (the to_sql_checked round-trip).
            assert_convert_val_consistent(
                &format!("meta/{arg_t}"),
                value.clone(),
                arg_t,
                typ_for(arg_t),
                expected.clone(),
            );
        }
    }
}