tdsql 0.6.0

//! The connection: connect once, run many statements.
//!
//! The connection lives in a background task that owns the driver client.
//! [`Client`] is a handle that sends it requests over a channel and awaits the
//! reply. That indirection buys two things: the task, not the caller's future,
//! reads every response -- to its end, or once the caller has gone, up to the
//! server's acknowledgement of a cancel -- so a dropped future cannot desync
//! the protocol stream; and work can be *enqueued* without `await`, which is
//! what lets [`Transaction`](crate::Transaction) roll back from `Drop`.

use std::borrow::Cow;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use futures::StreamExt;
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot};
use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};

use crate::command::{Command, Prepared};
use crate::config::Config;
use crate::dataset::{DataSet, DataTable};
use crate::error::{Error, Result, StatementKind};
use crate::row::{Column, Row};
use crate::stream::{RowStream, Streamed};
use crate::transaction::{IsolationLevel, Transaction};
use crate::value::{DataValue, FromSql, SqlType, ToSql};

type Driver = tiberius::Client<Compat<TcpStream>>;

// ---------------------------------------------------------------------------
// Binding: DataValue -> the driver's wire types.
// ---------------------------------------------------------------------------

impl tiberius::ToSql for DataValue {
    fn to_sql(&self) -> tiberius::ColumnData<'_> {
        use tiberius::ColumnData as C;
        match self {
            DataValue::TinyInt(v) => C::U8(Some(*v)),
            DataValue::SmallInt(v) => C::I16(Some(*v)),
            DataValue::Int(v) => C::I32(Some(*v)),
            DataValue::BigInt(v) => C::I64(Some(*v)),
            DataValue::Float(v) => C::F64(Some(*v)),
            DataValue::Decimal(v) => v.to_sql(),
            DataValue::Bool(v) => C::Bit(Some(*v)),
            DataValue::Text(v) => C::String(Some(Cow::Borrowed(v.as_str()))),
            DataValue::Binary(v) => C::Binary(Some(Cow::Borrowed(v.as_slice()))),
            DataValue::Guid(v) => C::Guid(Some(*v)),
            DataValue::Date(v) => v.to_sql(),
            DataValue::Time(v) => v.to_sql(),
            DataValue::DateTime(v) => v.to_sql(),
            DataValue::DateTimeOffset(v) => v.to_sql(),
            // An untyped NULL. SQL Server coerces this for most targets; bind a
            // typed `Option<T>` when the column needs a specific type.
            DataValue::Null => C::I32(None),
        }
    }
}

impl<'a> tiberius::IntoSql<'a> for DataValue {
    fn into_sql(self) -> tiberius::ColumnData<'a> {
        use tiberius::ColumnData as C;
        match self {
            DataValue::Text(v) => C::String(Some(Cow::Owned(v))),
            DataValue::Binary(v) => C::Binary(Some(Cow::Owned(v))),
            // Every other variant borrows nothing.
            other => match tiberius::ToSql::to_sql(&other) {
                C::U8(v) => C::U8(v),
                C::I16(v) => C::I16(v),
                C::I32(v) => C::I32(v),
                C::I64(v) => C::I64(v),
                C::F32(v) => C::F32(v),
                C::F64(v) => C::F64(v),
                C::Bit(v) => C::Bit(v),
                C::Guid(v) => C::Guid(v),
                C::Numeric(v) => C::Numeric(v),
                C::DateTime(v) => C::DateTime(v),
                C::SmallDateTime(v) => C::SmallDateTime(v),
                C::Time(v) => C::Time(v),
                C::Date(v) => C::Date(v),
                C::DateTime2(v) => C::DateTime2(v),
                C::DateTimeOffset(v) => C::DateTimeOffset(v),
                C::String(v) => C::String(v.map(|s| Cow::Owned(s.into_owned()))),
                C::Binary(v) => C::Binary(v.map(|b| Cow::Owned(b.into_owned()))),
                C::Xml(v) => C::Xml(v.map(|x| Cow::Owned(x.into_owned()))),
            },
        }
    }
}

// ---------------------------------------------------------------------------
// Reading: the driver's wire types -> DataValue.
// ---------------------------------------------------------------------------

/// Convert one column value, naming `column` in any error.
///
/// This is the crate's single mapping path: every query goes through it, so
/// the row-returning and scalar-returning paths can never drift.
pub(crate) fn map_column_data(cd: tiberius::ColumnData<'_>, column: &str) -> Result<DataValue> {
    use tiberius::ColumnData as C;

    // A temporal value the server sent but `chrono` cannot represent.
    fn temporal(column: &str, target: &'static str) -> impl Fn(tiberius::error::Error) -> Error {
        let column = column.to_string();
        move |_| Error::Conversion {
            column: column.clone(),
            actual: "temporal",
            target,
        }
    }

    Ok(match cd {
        C::U8(v) => v.map(DataValue::TinyInt).unwrap_or_default(),
        C::I16(v) => v.map(DataValue::SmallInt).unwrap_or_default(),
        C::I32(v) => v.map(DataValue::Int).unwrap_or_default(),
        C::I64(v) => v.map(DataValue::BigInt).unwrap_or_default(),
        C::F32(v) => v.map(|f| DataValue::Float(f as f64)).unwrap_or_default(),
        C::F64(v) => v.map(DataValue::Float).unwrap_or_default(),
        C::Bit(v) => v.map(DataValue::Bool).unwrap_or_default(),
        C::String(v) => v
            .map(|s| DataValue::Text(s.into_owned()))
            .unwrap_or_default(),
        C::Guid(v) => v.map(DataValue::Guid).unwrap_or_default(),
        C::Binary(v) => v
            .map(|b| DataValue::Binary(b.into_owned()))
            .unwrap_or_default(),
        C::Xml(v) => v
            .map(|x| DataValue::Text(x.as_ref().to_string()))
            .unwrap_or_default(),

        // A numeric out of `Decimal`'s range used to become NULL silently.
        C::Numeric(v) => match v {
            None => DataValue::Null,
            Some(n) => rust_decimal::Decimal::try_from_i128_with_scale(n.value(), n.scale() as u32)
                .map(DataValue::Decimal)
                .map_err(|_| Error::Decimal {
                    value: n.value(),
                    scale: n.scale(),
                })?,
        },

        // These used to `.unwrap()`, panicking the caller's task on bad data.
        C::DateTime(v) => {
            <chrono::NaiveDateTime as tiberius::FromSqlOwned>::from_sql_owned(C::DateTime(v))
                .map_err(temporal(column, "NaiveDateTime"))?
                .map(DataValue::DateTime)
                .unwrap_or_default()
        }

        C::SmallDateTime(v) => {
            <chrono::NaiveDateTime as tiberius::FromSqlOwned>::from_sql_owned(C::SmallDateTime(v))
                .map_err(temporal(column, "NaiveDateTime"))?
                .map(DataValue::DateTime)
                .unwrap_or_default()
        }

        C::DateTime2(v) => {
            <chrono::NaiveDateTime as tiberius::FromSqlOwned>::from_sql_owned(C::DateTime2(v))
                .map_err(temporal(column, "NaiveDateTime"))?
                .map(DataValue::DateTime)
                .unwrap_or_default()
        }

        C::Time(v) => <chrono::NaiveTime as tiberius::FromSqlOwned>::from_sql_owned(C::Time(v))
            .map_err(temporal(column, "NaiveTime"))?
            .map(DataValue::Time)
            .unwrap_or_default(),

        C::Date(v) => <chrono::NaiveDate as tiberius::FromSqlOwned>::from_sql_owned(C::Date(v))
            .map_err(temporal(column, "NaiveDate"))?
            .map(DataValue::Date)
            .unwrap_or_default(),

        C::DateTimeOffset(v) => {
            <chrono::DateTime<chrono::FixedOffset> as tiberius::FromSqlOwned>::from_sql_owned(
                C::DateTimeOffset(v),
            )
            .map_err(temporal(column, "DateTime<FixedOffset>"))?
            .map(DataValue::DateTimeOffset)
            .unwrap_or_default()
        }
    })
}

fn columns_of(meta: &tiberius::ResultMetadata) -> Arc<[Column]> {
    meta.columns()
        .iter()
        .map(|c| Column::new(c.name(), SqlType::from(c.column_type())))
        .collect()
}

fn row_of(row: tiberius::Row, columns: &Arc<[Column]>) -> Result<Row> {
    let mut values = Vec::with_capacity(columns.len());
    for (i, cd) in row.into_iter().enumerate() {
        let name = columns.get(i).map(Column::name).unwrap_or("?");
        values.push(map_column_data(cd, name)?);
    }
    Ok(Row::new(Arc::clone(columns), values))
}

/// Collect a query stream into its result sets, in server order.
///
/// `None` if `abandoned` reports the caller gone partway through, in which
/// case the rest of the response is left unread for [`cancel`].
async fn collect_query_stream(
    mut stream: tiberius::QueryStream<'_>,
    abandoned: impl Fn() -> bool,
) -> Result<Option<DataSet>> {
    let mut dataset = DataSet::new();
    let mut current: Option<(DataTable, Arc<[Column]>)> = None;

    while let Some(item) = stream.next().await {
        if abandoned() {
            return Ok(None);
        }
        match item? {
            tiberius::QueryItem::Metadata(meta) => {
                if let Some((table, _)) = current.take() {
                    dataset.push(table);
                }
                let columns = columns_of(&meta);
                let name = format!("table{}", meta.result_index());
                current = Some((DataTable::new(name, Arc::clone(&columns)), columns));
            }
            tiberius::QueryItem::Row(row) => {
                let (table, columns) = current.get_or_insert_with(|| {
                    let columns: Arc<[Column]> = Arc::from(Vec::new());
                    (DataTable::new("table0", Arc::clone(&columns)), columns)
                });
                table.push(row_of(row, columns)?);
            }
        }
    }

    if let Some((table, _)) = current.take() {
        dataset.push(table);
    }
    Ok(Some(dataset))
}

fn as_tiberius_refs(values: &[DataValue]) -> Vec<&dyn tiberius::ToSql> {
    values.iter().map(|v| v as &dyn tiberius::ToSql).collect()
}

pub(crate) fn to_values(params: &[&dyn ToSql]) -> Vec<DataValue> {
    params.iter().map(|p| p.to_value()).collect()
}

// ---------------------------------------------------------------------------
// The connection task
// ---------------------------------------------------------------------------

type Reply<T> = oneshot::Sender<Result<T>>;

/// One unit of work for the connection task.
enum Request {
    Query {
        sql: String,
        params: Vec<DataValue>,
        reply: Reply<DataSet>,
    },
    /// `header` gets the columns, then `rows` gets the rows as they arrive.
    QueryStream {
        sql: String,
        params: Vec<DataValue>,
        header: Reply<Arc<[Column]>>,
        rows: mpsc::Sender<Result<Streamed>>,
    },
    Execute {
        sql: String,
        params: Vec<DataValue>,
        reply: Reply<u64>,
    },
    Batch {
        sql: String,
        reply: Reply<DataSet>,
    },
    ProcQuery {
        name: String,
        params: Vec<(String, DataValue)>,
        reply: Reply<DataSet>,
    },
    ProcExecute {
        name: String,
        params: Vec<(String, DataValue)>,
        reply: Reply<u64>,
    },
    Begin {
        level: IsolationLevel,
        reply: Reply<()>,
    },
    Commit {
        reply: Reply<()>,
    },
    Savepoint {
        name: String,
        reply: Reply<()>,
    },
    /// `reply` is `None` when enqueued from `Transaction::drop`, which has no
    /// caller left to report to.
    Rollback {
        reply: Option<Reply<()>>,
    },
    RollbackTo {
        name: String,
        reply: Option<Reply<()>>,
    },
}

impl Request {
    /// Whether the caller has stopped waiting for this request.
    ///
    /// Transaction control never counts as abandoned: [`Transaction`] relies
    /// on its begin, commit and rollbacks reaching the server whatever became
    /// of the futures that sent them.
    fn abandoned(&self) -> bool {
        match self {
            Request::Query { reply, .. }
            | Request::Batch { reply, .. }
            | Request::ProcQuery { reply, .. } => reply.is_closed(),
            Request::Execute { reply, .. } | Request::ProcExecute { reply, .. } => {
                reply.is_closed()
            }
            Request::QueryStream { header, .. } => header.is_closed(),
            Request::Begin { .. }
            | Request::Commit { .. }
            | Request::Savepoint { .. }
            | Request::Rollback { .. }
            | Request::RollbackTo { .. } => false,
        }
    }
}

/// How long the server gets to acknowledge a cancel. Past this the state of
/// the connection is unknown, so it is closed rather than reused.
const CANCEL_TIMEOUT: Duration = Duration::from_secs(30);

/// Stop the server working on a request whose caller went away partway
/// through, and read up to its acknowledgement so the next request starts on a
/// clean connection.
///
/// Only call this once a response stream has been dropped *between* items.
/// The driver reads a token a field at a time, so a stream dropped while it
/// was waiting for the rest of one would leave half a token behind; between
/// items it has read whole tokens and nothing more.
///
/// `false` if the connection can no longer be trusted.
async fn cancel(driver: &mut Driver) -> bool {
    matches!(
        tokio::time::timeout(CANCEL_TIMEOUT, driver.cancel_query()).await,
        Ok(Ok(()))
    )
}

/// Hand a request's result to its caller, or cancel the rest of the request
/// if the caller went away partway through (`Ok(None)`).
///
/// `false` if the connection can no longer be trusted.
async fn finish<T>(driver: &mut Driver, result: Result<Option<T>>, reply: Reply<T>) -> bool {
    match result.transpose() {
        Some(result) => {
            let _ = reply.send(result);
            true
        }
        None => cancel(driver).await,
    }
}

/// Own the driver and serve requests until every handle is gone.
///
/// Requests are served strictly in order, which is also what TDS requires:
/// there is no MARS, so a connection carries one in-flight request at a time.
///
/// A request whose caller has gone is cancelled: skipped if it has not been
/// sent yet, or stopped between items of its response if it has. If the
/// server does not acknowledge a cancel, the task ends and takes the
/// connection with it, so every request still queued, and every later one,
/// fails with [`Error::ConnectionClosed`].
async fn connection_task(mut driver: Driver, mut rx: mpsc::UnboundedReceiver<Request>) {
    while let Some(request) = rx.recv().await {
        // Not sent yet, so cancelling it is as simple as not sending it.
        if request.abandoned() {
            continue;
        }
        match request {
            Request::Query { sql, params, reply } => {
                let result = run_query(&mut driver, &sql, &params, || reply.is_closed())
                    .await
                    .map_err(|e| e.with_statement(StatementKind::Query, &sql, params.len()));
                if !finish(&mut driver, result, reply).await {
                    return;
                }
            }
            Request::QueryStream {
                sql,
                params,
                header,
                rows,
            } => {
                let abandoned = stream_query(&mut driver, &sql, &params, header, rows).await;
                if abandoned && !cancel(&mut driver).await {
                    return;
                }
            }
            // The driver reads an `execute` response in one call, with no point
            // between items to stop at, so once sent it runs to completion
            // whether or not anyone is still waiting.
            Request::Execute { sql, params, reply } => {
                let refs = as_tiberius_refs(&params);
                let result = driver
                    .execute(&sql, &refs)
                    .await
                    .map(|r| r.total())
                    .map_err(Error::from);
                drop(refs);
                let result =
                    result.map_err(|e| e.with_statement(StatementKind::Query, &sql, params.len()));
                let _ = reply.send(result);
            }
            Request::Batch { sql, reply } => {
                let result = match driver.simple_query(&sql).await {
                    Ok(stream) => collect_query_stream(stream, || reply.is_closed()).await,
                    Err(e) => Err(Error::from(e)),
                }
                .map_err(|e| e.with_statement(StatementKind::Batch, &sql, 0));
                if !finish(&mut driver, result, reply).await {
                    return;
                }
            }
            Request::ProcQuery {
                name,
                params,
                reply,
            } => {
                let count = params.len();
                let result = run_proc_query(&mut driver, &name, params, || reply.is_closed())
                    .await
                    .map_err(|e| e.with_statement(StatementKind::StoredProcedure, &name, count));
                if !finish(&mut driver, result, reply).await {
                    return;
                }
            }
            Request::ProcExecute {
                name,
                params,
                reply,
            } => {
                let count = params.len();
                let result = run_proc_execute(&mut driver, &name, params, || reply.is_closed())
                    .await
                    .map_err(|e| e.with_statement(StatementKind::StoredProcedure, &name, count));
                if !finish(&mut driver, result, reply).await {
                    return;
                }
            }
            Request::Begin { level, reply } => {
                let result = driver
                    .begin_transaction_with_isolation(level.into())
                    .await
                    .map_err(Error::from)
                    .map_err(|e| {
                        e.with_statement(StatementKind::TransactionControl, "BEGIN TRANSACTION", 0)
                    });
                let _ = reply.send(result);
            }
            Request::Commit { reply } => {
                let result = driver
                    .commit_transaction()
                    .await
                    .map_err(Error::from)
                    .map_err(|e| {
                        e.with_statement(StatementKind::TransactionControl, "COMMIT TRANSACTION", 0)
                    });
                let _ = reply.send(result);
            }
            Request::Savepoint { name, reply } => {
                // Not `driver.save_transaction`: tiberius-ng 0.13.1 encodes the
                // TM_SAVE_XACT name length as a count of UTF-16 code units
                // where the server reads a byte count, so any non-empty name is
                // rejected ("odd length N" / "TM request is longer than
                // expected"). T-SQL is the documented alternative and is what
                // the matching rollback already uses. The name is validated as
                // an identifier before it reaches here.
                // Bracket-quoted so reserved words (`inner`, `key`, ...) work.
                // Validation guarantees the name cannot contain `]`.
                let sql = format!("SAVE TRANSACTION [{name}]");
                let result = match driver.simple_query(&sql).await {
                    // Transaction control is never cancelled; see `abandoned`.
                    Ok(stream) => collect_query_stream(stream, || false).await.map(|_| ()),
                    Err(e) => Err(Error::from(e)),
                };
                let _ = reply.send(
                    result
                        .map_err(|e| e.with_statement(StatementKind::TransactionControl, &sql, 0)),
                );
            }
            Request::Rollback { reply } => {
                let result = driver
                    .rollback_transaction()
                    .await
                    .map_err(Error::from)
                    .map_err(|e| {
                        e.with_statement(
                            StatementKind::TransactionControl,
                            "ROLLBACK TRANSACTION",
                            0,
                        )
                    });
                if let Some(reply) = reply {
                    let _ = reply.send(result);
                }
            }
            Request::RollbackTo { name, reply } => {
                // TDS has no "rollback to savepoint" manager request; T-SQL is
                // the documented route. The name is validated as an identifier
                // before it ever reaches here.
                let sql = format!("ROLLBACK TRANSACTION [{name}]");
                let result = match driver.simple_query(&sql).await {
                    Ok(stream) => collect_query_stream(stream, || false).await.map(|_| ()),
                    Err(e) => Err(Error::from(e)),
                }
                .map_err(|e| e.with_statement(StatementKind::TransactionControl, &sql, 0));
                if let Some(reply) = reply {
                    let _ = reply.send(result);
                }
            }
        }
    }

    // Every handle is gone. Close politely; an error here has nobody to tell.
    let _ = driver.close().await;
}

/// `None` if the caller went away partway through; see [`collect_query_stream`].
async fn run_query(
    driver: &mut Driver,
    sql: &str,
    params: &[DataValue],
    abandoned: impl Fn() -> bool,
) -> Result<Option<DataSet>> {
    let refs = as_tiberius_refs(params);
    let stream = driver.query(sql, &refs).await?;
    collect_query_stream(stream, abandoned).await
}

/// How many rows the connection task may read ahead of a [`RowStream`]'s
/// consumer. Past this it stops reading, the server's send buffer fills, and
/// the query waits -- so a slow consumer slows the query down instead of
/// piling rows up in memory. Kept small because a row can be large.
const STREAM_BUFFER: usize = 64;

/// Run a query, handing the rows of its first result set to `rows` as they
/// arrive rather than collecting them.
///
/// `header` gets the result set's columns before any row is sent, or the error
/// if the statement failed before producing one. Everything after that goes to
/// `rows`, which always ends with [`Streamed::End`] or an error.
///
/// Later result sets are read and dropped, but an error in one still fails the
/// stream, just as it fails `query`.
///
/// Returns `true` if the consumer went away partway through -- dropped the
/// [`RowStream`], or the `query_stream` call before it returned -- leaving the
/// rest of the response unread for [`cancel`].
async fn stream_query(
    driver: &mut Driver,
    sql: &str,
    params: &[DataValue],
    header: Reply<Arc<[Column]>>,
    rows: mpsc::Sender<Result<Streamed>>,
) -> bool {
    let mut header = Some(header);
    let mut columns: Arc<[Column]> = Arc::from(Vec::new());
    // False once the first result set is over or the consumer has gone.
    let mut forwarding = true;
    // The first failure, held back until the response has been fully read so
    // the connection is clean for the next request.
    let mut failure: Option<Error> = None;
    // The consumer holds `header`'s receiver until the columns arrive, and
    // `rows`' after.
    let gone = |header: &Option<Reply<Arc<[Column]>>>| match header {
        Some(header) => header.is_closed(),
        None => rows.is_closed(),
    };

    let refs = as_tiberius_refs(params);
    match driver.query(sql, &refs).await {
        Err(e) => failure = Some(e.into()),
        Ok(mut stream) => {
            while let Some(item) = stream.next().await {
                if gone(&header) {
                    return true;
                }
                match item {
                    Ok(tiberius::QueryItem::Metadata(meta)) => match header.take() {
                        Some(header) => {
                            columns = columns_of(&meta);
                            forwarding = header.send(Ok(Arc::clone(&columns))).is_ok();
                        }
                        None => forwarding = false,
                    },
                    Ok(tiberius::QueryItem::Row(row)) if forwarding => {
                        match row_of(row, &columns) {
                            Ok(row) => {
                                forwarding = rows.send(Ok(Streamed::Row(row))).await.is_ok();
                            }
                            Err(e) => {
                                failure = Some(e);
                                forwarding = false;
                            }
                        }
                    }
                    Ok(tiberius::QueryItem::Row(_)) => {}
                    // The driver holds a server error back until it has read
                    // to the end of the response, so nothing follows this.
                    Err(e) => {
                        failure.get_or_insert(e.into());
                        break;
                    }
                }
            }
        }
    }

    let last = match failure {
        Some(e) => Err(e.with_statement(StatementKind::Query, sql, params.len())),
        None => Ok(Streamed::End),
    };
    match header {
        // Still unsent: the statement failed, or produced no result set.
        Some(header) => match last {
            Err(e) => {
                let _ = header.send(Err(e));
            }
            Ok(end) => {
                if header.send(Ok(columns)).is_ok() {
                    let _ = rows.send(Ok(end)).await;
                }
            }
        },
        None => {
            let _ = rows.send(last).await;
        }
    }
    false
}

fn build_proc<'a>(name: &str, params: Vec<(String, DataValue)>) -> tiberius::Command<'a> {
    let mut cmd = tiberius::Command::new(name.to_string());
    for (pname, value) in params {
        cmd.bind_param(pname, value);
    }
    cmd
}

/// `None` if the caller went away partway through; see [`collect_query_stream`].
async fn run_proc_query(
    driver: &mut Driver,
    name: &str,
    params: Vec<(String, DataValue)>,
    abandoned: impl Fn() -> bool,
) -> Result<Option<DataSet>> {
    let mut stream = build_proc(name, params).exec(driver).await?;

    let mut dataset = DataSet::new();
    let mut current: Option<(DataTable, Arc<[Column]>)> = None;

    while let Some(item) = stream.next().await {
        if abandoned() {
            return Ok(None);
        }
        match item? {
            tiberius::CommandItem::Metadata(meta) => {
                if let Some((table, _)) = current.take() {
                    dataset.push(table);
                }
                let columns = columns_of(&meta);
                let tname = format!("table{}", meta.result_index());
                current = Some((DataTable::new(tname, Arc::clone(&columns)), columns));
            }
            tiberius::CommandItem::Row(row) => {
                let (table, columns) = current.get_or_insert_with(|| {
                    let columns: Arc<[Column]> = Arc::from(Vec::new());
                    (DataTable::new("table0", Arc::clone(&columns)), columns)
                });
                table.push(row_of(row, columns)?);
            }
            _ => {}
        }
    }

    if let Some((table, _)) = current.take() {
        dataset.push(table);
    }
    Ok(Some(dataset))
}

/// `None` if the caller went away partway through; see [`collect_query_stream`].
async fn run_proc_execute(
    driver: &mut Driver,
    name: &str,
    params: Vec<(String, DataValue)>,
    abandoned: impl Fn() -> bool,
) -> Result<Option<u64>> {
    let mut stream = build_proc(name, params).exec(driver).await?;
    let mut total = 0u64;
    while let Some(item) = stream.next().await {
        if abandoned() {
            return Ok(None);
        }
        if let tiberius::CommandItem::RowsAffected(n) = item? {
            total += n;
        }
    }
    Ok(Some(total))
}

// ---------------------------------------------------------------------------
// Client
// ---------------------------------------------------------------------------

/// A connection to SQL Server.
///
/// Connect once and reuse it; every statement travels over the same
/// connection.
///
/// ```no_run
/// use tdsql::{Client, Config};
///
/// # async fn f() -> tdsql::Result<()> {
/// let mut client = Client::connect(&Config::new()
///     .host("localhost")
///     .database("master")
///     .auth("sa", "YourStrong!Passw0rd")
///     .trust_cert()).await?;
///
/// let rows = client.query("SELECT id, name FROM users WHERE id > @P1", &[&10i32]).await?;
/// for row in &rows {
///     let id: i32 = row.get("id");
///     let name: String = row.get("name");
///     println!("{id}: {name}");
/// }
/// # Ok(())
/// # }
/// ```
///
/// # Cancelling
///
/// Stop waiting for a request -- drop its future, let a `timeout` or
/// `select!` give up on it, abort the task running it, drop a [`RowStream`]
/// early -- and it is cancelled on the server too, so it stops holding the
/// connection, the server's time and any locks it has taken. The connection
/// stays usable. A request still queued behind another never runs; one
/// already running is stopped at the next row or result set to arrive, and
/// the statements after it in the same batch don't run.
///
/// Two things can't be stopped partway:
///
/// - A statement that hasn't sent anything back yet -- still sorting,
///   aggregating or waiting on a lock -- runs until it does, and is cancelled
///   then.
/// - [`execute`](Self::execute), and [`execute_command`](Self::execute_command)
///   with a text command: the driver reads their whole response in one go, so
///   once sent they run to completion.
///
/// Transaction control -- begin, commit, rollback, savepoints -- always runs,
/// whatever became of its future. And cancelling stops a statement, not the
/// transaction around it: work finished before the cancel stays done, as it
/// would if the statement had failed, so put statements that must stand or
/// fall together in a [`Transaction`].
pub struct Client {
    tx: mpsc::UnboundedSender<Request>,
    savepoints: AtomicU64,
}

impl std::fmt::Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("open", &!self.tx.is_closed())
            .finish()
    }
}

impl Client {
    /// Connect using this configuration.
    pub async fn connect(config: &Config) -> Result<Self> {
        let cfg = config.to_tiberius()?;
        let addr = cfg.get_addr();
        let tcp = TcpStream::connect(&addr).await?;
        tcp.set_nodelay(true)?;
        let driver = tiberius::Client::connect(cfg, tcp.compat_write())
            .await
            .map_err(|e| Error::Connect {
                server: addr,
                source: Box::new(e),
            })?;

        let (tx, rx) = mpsc::unbounded_channel();
        tokio::spawn(connection_task(driver, rx));

        Ok(Self {
            tx,
            savepoints: AtomicU64::new(0),
        })
    }

    /// Connect using an ADO.NET-style connection string.
    pub async fn connect_str(connection_string: &str) -> Result<Self> {
        Self::connect(&Config::from_ado_string(connection_string)).await
    }

    /// Send a request and await its reply.
    async fn call<T>(&self, build: impl FnOnce(Reply<T>) -> Request) -> Result<T> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.tx
            .send(build(reply_tx))
            .map_err(|_| Error::ConnectionClosed)?;
        reply_rx.await.map_err(|_| Error::ConnectionClosed)?
    }

    /// Enqueue work without waiting for it. Used by `Transaction::drop`, which
    /// cannot await. A closed channel means the connection is already gone, in
    /// which case the server has rolled the transaction back for us.
    fn enqueue(&self, request: Request) {
        let _ = self.tx.send(request);
    }

    /// Run a query and collect every row of the first result set.
    ///
    /// Parameters are positional: `@P1` is `params[0]`.
    pub async fn query(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<Vec<Row>> {
        Ok(self
            .query_dataset_values(sql, to_values(params))
            .await?
            .into_rows())
    }

    /// Run a query and receive the rows of its first result set as the server
    /// sends them, rather than all at once after the last one arrives.
    ///
    /// Returns as soon as the server has described the columns, so a statement
    /// the server rejects fails here, just as it would from
    /// [`query`](Self::query). A failure partway through arrives as the
    /// stream's last item instead, after the rows that came before it.
    ///
    /// The stream borrows the client until it is dropped; see [`RowStream`] for
    /// what that holds up, and what happens if it is dropped early.
    ///
    /// ```no_run
    /// # use tdsql::Client;
    /// # async fn f(client: &mut Client) -> tdsql::Result<()> {
    /// let mut rows = client
    ///     .query_stream("SELECT id, name FROM users WHERE id > @P1", &[&10i32])
    ///     .await?;
    ///
    /// while let Some(row) = rows.try_next().await? {
    ///     let name: String = row.get("name");
    ///     println!("{name}");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query_stream(
        &mut self,
        sql: &str,
        params: &[&dyn ToSql],
    ) -> Result<RowStream<'_>> {
        let (sql, params) = (sql.to_string(), to_values(params));
        let (rows, receiver) = mpsc::channel(STREAM_BUFFER);
        let columns = self
            .call(|header| Request::QueryStream {
                sql,
                params,
                header,
                rows,
            })
            .await?;
        Ok(RowStream::new(columns, receiver))
    }

    /// Run a query expecting exactly one row.
    ///
    /// Fails with [`Error::UnexpectedRowCount`] for any other count.
    pub async fn query_one(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<Row> {
        let mut rows = self.query(sql, params).await?;
        match rows.len() {
            1 => Ok(rows.pop().expect("length checked")),
            found => Err(Error::UnexpectedRowCount { found }),
        }
    }

    /// Run a query expecting at most one row.
    ///
    /// Fails with [`Error::UnexpectedRowCount`] if more than one comes back.
    pub async fn query_opt(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<Option<Row>> {
        let mut rows = self.query(sql, params).await?;
        match rows.len() {
            0 => Ok(None),
            1 => Ok(Some(rows.pop().expect("length checked"))),
            found => Err(Error::UnexpectedRowCount { found }),
        }
    }

    /// Run a query and read the first column of the first row.
    ///
    /// Use `Option<T>` as `T` to accept a `NULL`. Fails with
    /// [`Error::UnexpectedRowCount`] if the query returns no rows.
    pub async fn query_scalar<T: FromSql>(
        &mut self,
        sql: &str,
        params: &[&dyn ToSql],
    ) -> Result<T> {
        let row = self.query_one(sql, params).await?;
        if row.is_empty() {
            return Err(Error::ColumnIndexOutOfRange { index: 0, len: 0 });
        }
        row.try_get(0usize)
    }

    /// Run a statement that returns no rows, and report how many rows it
    /// affected.
    ///
    /// If the batch holds several statements, the counts are summed.
    pub async fn execute(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<u64> {
        let (sql, params) = (sql.to_string(), to_values(params));
        self.call(|reply| Request::Execute { sql, params, reply })
            .await
    }

    /// Run a SQL batch verbatim, with no parameters, and collect its rows.
    ///
    /// Parameterised statements are sent as an RPC, and some statements —
    /// `CREATE PROCEDURE`, `CREATE VIEW`, `CREATE TRIGGER` — must be the first
    /// statement of their own batch, so they cannot go that route. Use this for
    /// DDL, and [`execute`](Self::execute) or [`query`](Self::query) for
    /// everything else.
    ///
    /// Because the SQL is sent verbatim, never interpolate untrusted input into
    /// it; use parameters instead.
    ///
    /// ```no_run
    /// # use tdsql::Client;
    /// # async fn f(client: &mut Client) -> tdsql::Result<()> {
    /// client
    ///     .batch("CREATE PROCEDURE dbo.sp_demo AS BEGIN SELECT 1; END")
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn batch(&mut self, sql: &str) -> Result<Vec<Row>> {
        Ok(self.batch_dataset(sql).await?.into_rows())
    }

    /// Run a SQL batch verbatim and collect every result set it produces.
    pub async fn batch_dataset(&mut self, sql: &str) -> Result<DataSet> {
        let sql = sql.to_string();
        self.call(|reply| Request::Batch { sql, reply }).await
    }

    /// Run a [`Command`] and collect every result set it produces.
    ///
    /// This is the path for stored procedures and for batches that return more
    /// than one result set.
    pub async fn query_dataset(&mut self, command: &Command) -> Result<DataSet> {
        match command.prepare() {
            Prepared::Text { sql, params } => self.query_dataset_values(&sql, params).await,
            Prepared::Proc { name, params } => {
                self.call(|reply| Request::ProcQuery {
                    name,
                    params,
                    reply,
                })
                .await
            }
        }
    }

    /// Run a [`Command`] for its row count rather than its rows.
    pub async fn execute_command(&mut self, command: &Command) -> Result<u64> {
        match command.prepare() {
            Prepared::Text { sql, params } => {
                self.call(|reply| Request::Execute { sql, params, reply })
                    .await
            }
            Prepared::Proc { name, params } => {
                self.call(|reply| Request::ProcExecute {
                    name,
                    params,
                    reply,
                })
                .await
            }
        }
    }

    async fn query_dataset_values(&self, sql: &str, params: Vec<DataValue>) -> Result<DataSet> {
        let sql = sql.to_string();
        self.call(|reply| Request::Query { sql, params, reply })
            .await
    }

    /// Begin a transaction at the server's default isolation level.
    ///
    /// The returned [`Transaction`] borrows this client, so nothing else can
    /// use the connection until it is committed, rolled back, or dropped.
    ///
    /// ```no_run
    /// # use tdsql::Client;
    /// # async fn f(client: &mut Client) -> tdsql::Result<()> {
    /// let mut tx = client.transaction().await?;
    /// tx.execute("INSERT INTO orders (id) VALUES (@P1)", &[&1i32]).await?;
    /// tx.commit().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn transaction(&mut self) -> Result<Transaction<'_>> {
        Transaction::begin(self, IsolationLevel::Unspecified).await
    }

    /// Begin a transaction at an explicit isolation level.
    pub async fn transaction_with_isolation(
        &mut self,
        level: IsolationLevel,
    ) -> Result<Transaction<'_>> {
        Transaction::begin(self, level).await
    }

    /// Close the connection, sending a graceful logout.
    pub async fn close(self) -> Result<()> {
        // Dropping the last sender ends the task, which closes the driver.
        drop(self);
        Ok(())
    }

    // -- used by Transaction ------------------------------------------------

    pub(crate) async fn begin_transaction(&mut self, level: IsolationLevel) -> Result<()> {
        self.call(|reply| Request::Begin { level, reply }).await
    }

    pub(crate) async fn commit_transaction(&mut self) -> Result<()> {
        self.call(|reply| Request::Commit { reply }).await
    }

    pub(crate) async fn rollback_transaction(&mut self) -> Result<()> {
        self.call(|reply| Request::Rollback { reply: Some(reply) })
            .await
    }

    pub(crate) async fn save_transaction(&mut self, name: &str) -> Result<()> {
        let name = name.to_string();
        self.call(|reply| Request::Savepoint { name, reply }).await
    }

    pub(crate) async fn rollback_to_savepoint(&mut self, name: &str) -> Result<()> {
        let name = name.to_string();
        self.call(|reply| Request::RollbackTo {
            name,
            reply: Some(reply),
        })
        .await
    }

    /// Queue a rollback without awaiting it. Called from `Transaction::drop`.
    pub(crate) fn enqueue_rollback(&self) {
        self.enqueue(Request::Rollback { reply: None });
    }

    /// Queue a rollback to a savepoint without awaiting it.
    pub(crate) fn enqueue_rollback_to(&self, name: String) {
        self.enqueue(Request::RollbackTo { name, reply: None });
    }

    pub(crate) fn next_savepoint_id(&self) -> u64 {
        self.savepoints.fetch_add(1, Ordering::Relaxed)
    }
}

#[cfg(test)]
mod tests;