mod.rs

// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

pub use mysql_common::named_params;

use mysql_common::{
    constants::DEFAULT_MAX_ALLOWED_PACKET,
    crypto,
    packets::{
        parse_auth_switch_request, parse_handshake_packet, AuthPlugin, AuthSwitchRequest,
        HandshakeResponse, SslRequest,
    },
};

use std::{
    fmt,
    future::Future,
    mem,
    pin::Pin,
    str::FromStr,
    sync::Arc,
    time::{Duration, Instant},
};

use crate::{
    conn::{pool::Pool, stmt_cache::StmtCache},
    connection_like::{streamless::Streamless, ConnectionLike, StmtCacheResult},
    consts::{self, CapabilityFlags},
    error::*,
    io::Stream,
    local_infile_handler::LocalInfileHandler,
    opts::Opts,
    queryable::{query_result, BinaryProtocol, Queryable, TextProtocol},
    Column, OptsBuilder,
};

pub mod pool;
pub mod stmt_cache;

/// Helper that asynchronously disconnects connection on the default tokio executor.
fn disconnect(mut conn: Conn) {
    let disconnected = conn.inner.disconnected;

    // Mark conn as disconnected.
    conn.inner.disconnected = true;

    if !disconnected {
        // We shouldn't call tokio::spawn if unwinding
        if std::thread::panicking() {
            return;
        }

        // Server will report broken connection if spawn fails.
        // this might fail if, say, the runtime is shutting down, but we've done what we could
        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn(async move {
                if let Ok(conn) = conn.cleanup().await {
                    let _ = conn.disconnect().await;
                }
            });
        }
    }
}

/// Mysql connection
struct ConnInner {
    stream: Option<Stream>,
    id: u32,
    version: (u16, u16, u16),
    max_allowed_packet: usize,
    socket: Option<String>,
    capabilities: consts::CapabilityFlags,
    status: consts::StatusFlags,
    last_insert_id: u64,
    affected_rows: u64,
    warnings: u16,
    pool: Option<Pool>,
    has_result: Option<(Arc<Vec<Column>>, Option<StmtCacheResult>)>,
    in_transaction: bool,
    opts: Opts,
    last_io: Instant,
    wait_timeout: Duration,
    stmt_cache: StmtCache,
    nonce: Vec<u8>,
    auth_plugin: AuthPlugin<'static>,
    auth_switched: bool,
    /// Connection is already disconnected.
    disconnected: bool,
}

impl fmt::Debug for ConnInner {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Conn")
            .field("connection id", &self.id)
            .field("server version", &self.version)
            .field("pool", &self.pool)
            .field("has result", &self.has_result.is_some())
            .field("in transaction", &self.in_transaction)
            .field("stream", &self.stream)
            .field("options", &self.opts)
            .finish()
    }
}

impl ConnInner {
    /// Constructs an empty connection.
    fn empty(opts: Opts) -> ConnInner {
        ConnInner {
            capabilities: opts.get_capabilities(),
            status: consts::StatusFlags::empty(),
            last_insert_id: 0,
            affected_rows: 0,
            stream: None,
            max_allowed_packet: DEFAULT_MAX_ALLOWED_PACKET,
            warnings: 0,
            version: (0, 0, 0),
            id: 0,
            has_result: None,
            pool: None,
            in_transaction: false,
            last_io: Instant::now(),
            wait_timeout: Duration::from_secs(0),
            stmt_cache: StmtCache::new(opts.get_stmt_cache_size()),
            socket: opts.get_socket().map(Into::into),
            opts,
            nonce: Vec::default(),
            auth_plugin: AuthPlugin::MysqlNativePassword,
            auth_switched: false,
            disconnected: false,
        }
    }
}

#[derive(Debug)]
pub struct Conn {
    inner: Box<ConnInner>,
}

impl Conn {
    /// Returns the ID generated by a query (usually `INSERT`) on a table with a column having the
    /// `AUTO_INCREMENT` attribute. Returns `None` if there was no previous query on the connection
    /// or if the query did not update an AUTO_INCREMENT value.
    pub fn last_insert_id(&self) -> Option<u64> {
        self.get_last_insert_id()
    }

    /// Returns the number of rows affected by the last `INSERT`, `UPDATE`, `REPLACE` or `DELETE`
    /// query.
    pub fn affected_rows(&self) -> u64 {
        self.get_affected_rows()
    }

    async fn close(mut self) -> Result<()> {
        self.inner.disconnected = true;
        self.cleanup().await?.disconnect().await
    }

    fn is_secure(&self) -> bool {
        if let Some(ref stream) = self.inner.stream {
            stream.is_secure()
        } else {
            false
        }
    }

    /// Hacky way to move connection through &mut. `self` becomes unusable.
    fn take(&mut self) -> Conn {
        let inner = mem::replace(&mut *self.inner, ConnInner::empty(Default::default()));
        Conn {
            inner: Box::new(inner),
        }
    }

    fn empty(opts: Opts) -> Self {
        Self {
            inner: Box::new(ConnInner::empty(opts)),
        }
    }

    fn setup_stream(mut self) -> Result<Conn> {
        if let Some(stream) = self.inner.stream.take() {
            stream.set_keepalive_ms(self.inner.opts.get_tcp_keepalive())?;
            stream.set_tcp_nodelay(self.inner.opts.get_tcp_nodelay())?;
            self.inner.stream = Some(stream);
            Ok(self)
        } else {
            unreachable!();
        }
    }

    async fn handle_handshake(self) -> Result<Conn> {
        let (mut conn, packet) = self.read_packet().await?;
        let handshake = parse_handshake_packet(&*packet)?;
        conn.inner.nonce = {
            let mut nonce = Vec::from(handshake.scramble_1_ref());
            nonce.extend_from_slice(handshake.scramble_2_ref().unwrap_or(&[][..]));
            nonce
        };

        conn.inner.capabilities = handshake.capabilities() & conn.inner.opts.get_capabilities();
        conn.inner.version = handshake.server_version_parsed().unwrap_or((0, 0, 0));
        conn.inner.id = handshake.connection_id();
        conn.inner.status = handshake.status_flags();
        conn.inner.auth_plugin = match handshake.auth_plugin() {
            Some(AuthPlugin::MysqlNativePassword) => AuthPlugin::MysqlNativePassword,
            Some(AuthPlugin::CachingSha2Password) => AuthPlugin::CachingSha2Password,
            Some(AuthPlugin::Other(ref name)) => {
                let name = String::from_utf8_lossy(name).into();
                return Err(DriverError::UnknownAuthPlugin { name }.into());
            }
            None => AuthPlugin::MysqlNativePassword,
        };
        Ok(conn)
    }

    async fn switch_to_ssl_if_needed(self) -> Result<Conn> {
        if self
            .inner
            .opts
            .get_capabilities()
            .contains(CapabilityFlags::CLIENT_SSL)
        {
            let ssl_request = SslRequest::new(self.inner.capabilities);
            let conn = self.write_packet(ssl_request.as_ref()).await?;
            let ssl_opts = conn
                .get_opts()
                .get_ssl_opts()
                .cloned()
                .expect("unreachable");
            let domain = conn.get_opts().get_ip_or_hostname().into();
            let (streamless, stream) = conn.take_stream();
            let stream = stream.make_secure(domain, ssl_opts).await?;
            Ok(streamless.return_stream(stream))
        } else {
            Ok(self)
        }
    }

    async fn do_handshake_response(self) -> Result<Conn> {
        let auth_data = self
            .inner
            .auth_plugin
            .gen_data(self.inner.opts.get_pass(), &*self.inner.nonce);

        let handshake_response = HandshakeResponse::new(
            &auth_data,
            self.inner.version,
            self.inner.opts.get_user(),
            self.inner.opts.get_db_name(),
            &self.inner.auth_plugin,
            self.get_capabilities(),
            &Default::default(), // TODO: Add support
        );

        self.write_packet(handshake_response.as_ref()).await
    }

    async fn perform_auth_switch(
        mut self,
        auth_switch_request: AuthSwitchRequest<'_>,
    ) -> Result<Conn> {
        if !self.inner.auth_switched {
            self.inner.auth_switched = true;
            self.inner.nonce = auth_switch_request.plugin_data().into();
            self.inner.auth_plugin = auth_switch_request.auth_plugin().clone().into_owned();
            let plugin_data = self
                .inner
                .auth_plugin
                .gen_data(self.inner.opts.get_pass(), &*self.inner.nonce)
                .unwrap_or_else(Vec::new);
            self.write_packet(plugin_data).await?.continue_auth().await
        } else {
            unreachable!("auth_switched flag should be checked by caller")
        }
    }

    fn continue_auth(self) -> Pin<Box<dyn Future<Output = Result<Conn>> + Send>> {
        // NOTE: we need to box this since it may recurse
        // see https://github.com/rust-lang/rust/issues/46415#issuecomment-528099782
        Box::pin(async move {
            match self.inner.auth_plugin {
                AuthPlugin::MysqlNativePassword => self.continue_mysql_native_password_auth().await,
                AuthPlugin::CachingSha2Password => self.continue_caching_sha2_password_auth().await,
                AuthPlugin::Other(ref name) => Err(DriverError::UnknownAuthPlugin {
                    name: String::from_utf8_lossy(name.as_ref()).to_string(),
                })?,
            }
        })
    }

    fn switch_to_compression(mut self) -> Result<Conn> {
        if self
            .get_capabilities()
            .contains(CapabilityFlags::CLIENT_COMPRESS)
        {
            if let Some(compression) = self.inner.opts.get_compression() {
                if let Some(stream) = self.inner.stream.as_mut() {
                    stream.compress(compression);
                }
            }
        }
        Ok(self)
    }

    async fn continue_caching_sha2_password_auth(self) -> Result<Conn> {
        let (conn, packet) = self.read_packet().await?;
        match packet.get(0) {
            Some(0x00) => {
                // ok packet for empty password
                Ok(conn)
            }
            Some(0x01) => match packet.get(1) {
                Some(0x03) => {
                    // auth ok
                    conn.drop_packet().await
                }
                Some(0x04) => {
                    let mut pass = conn
                        .inner
                        .opts
                        .get_pass()
                        .map(Vec::from)
                        .unwrap_or_default();
                    pass.push(0);

                    let conn = if conn.is_secure() {
                        conn.write_packet(&*pass).await?
                    } else {
                        let conn = conn.write_packet(&[0x02][..]).await?;
                        let (conn, packet) = conn.read_packet().await?;
                        let key = &packet[1..];
                        for (i, byte) in pass.iter_mut().enumerate() {
                            *byte ^= conn.inner.nonce[i % conn.inner.nonce.len()];
                        }
                        let encrypted_pass = crypto::encrypt(&*pass, key);
                        conn.write_packet(&*encrypted_pass).await?
                    };
                    conn.drop_packet().await
                }
                _ => Err(DriverError::UnexpectedPacket {
                    payload: packet.into(),
                }
                .into()),
            },
            Some(0xfe) if !conn.inner.auth_switched => {
                let auth_switch_request = parse_auth_switch_request(&*packet)?.into_owned();
                conn.perform_auth_switch(auth_switch_request).await
            }
            _ => Err(DriverError::UnexpectedPacket {
                payload: packet.into(),
            }
            .into()),
        }
    }

    async fn continue_mysql_native_password_auth(self) -> Result<Conn> {
        let (this, packet) = self.read_packet().await?;
        match packet.get(0) {
            Some(0x00) => Ok(this),
            Some(0xfe) if !this.inner.auth_switched => {
                let auth_switch_request = parse_auth_switch_request(packet.as_ref())?.into_owned();
                this.perform_auth_switch(auth_switch_request).await
            }
            _ => Err(DriverError::UnexpectedPacket { payload: packet }.into()),
        }
    }

    async fn drop_packet(self) -> Result<Conn> {
        Ok(self.read_packet().await?.0)
    }

    async fn run_init_commands(self) -> Result<Conn> {
        let mut init: Vec<_> = self.inner.opts.get_init().iter().cloned().collect();

        let mut conn = self;
        while let Some(query) = init.pop() {
            conn = conn.drop_query(query).await?;
        }
        Ok(conn)
    }

    /// Returns future that resolves to [`Conn`].
    pub async fn new<T: Into<Opts>>(opts: T) -> Result<Conn> {
        let opts = opts.into();
        let mut conn = Conn::empty(opts.clone());

        let stream = if let Some(path) = opts.get_socket() {
            Stream::connect_socket(path.to_owned()).await?
        } else {
            Stream::connect_tcp((opts.get_ip_or_hostname(), opts.get_tcp_port())).await?
        };

        conn.inner.stream = Some(stream);
        conn.setup_stream()?
            .handle_handshake()
            .await?
            .switch_to_ssl_if_needed()
            .await?
            .do_handshake_response()
            .await?
            .continue_auth()
            .await?
            .switch_to_compression()?
            .read_socket()
            .await?
            .reconnect_via_socket_if_needed()
            .await?
            .read_max_allowed_packet()
            .await?
            .read_wait_timeout()
            .await?
            .run_init_commands()
            .await
    }

    /// Returns future that resolves to [`Conn`].
    pub async fn from_url<T: AsRef<str>>(url: T) -> Result<Conn> {
        Conn::new(Opts::from_str(url.as_ref())?).await
    }

    /// Will try to connect via socket using socket address in `self.inner.socket`.
    ///
    /// Returns new connection on success or self on error.
    ///
    /// Won't try to reconnect if socket connection is already enforced in [`Opts`].
    fn reconnect_via_socket_if_needed(self) -> Pin<Box<dyn Future<Output = Result<Conn>> + Send>> {
        // NOTE: we need to box this since it may recurse
        // see https://github.com/rust-lang/rust/issues/46415#issuecomment-528099782
        Box::pin(async move {
            if let Some(socket) = self.inner.socket.as_ref() {
                let opts = self.inner.opts.clone();
                if opts.get_socket().is_none() {
                    let mut builder = OptsBuilder::from_opts(opts);
                    builder.socket(Some(&**socket));
                    match Conn::new(builder).await {
                        Ok(conn) => return Ok(conn),
                        Err(_) => return Ok(self),
                    }
                }
            }
            Ok(self)
        })
    }

    /// Returns future that resolves to [`Conn`] with socket address stored in it.
    ///
    /// Do nothing if socket address is already in [`Opts`] or if `prefer_socket` is `false`.
    async fn read_socket(self) -> Result<Self> {
        if self.inner.opts.get_prefer_socket() && self.inner.socket.is_none() {
            let (mut this, row_opt) = self.first("SELECT @@socket").await?;
            this.inner.socket = row_opt.unwrap_or((None,)).0;
            Ok(this)
        } else {
            Ok(self)
        }
    }

    /// Returns future that resolves to [`Conn`] with `max_allowed_packet` stored in it.
    async fn read_max_allowed_packet(self) -> Result<Self> {
        let (mut this, row_opt): (Self, _) = self.first("SELECT @@max_allowed_packet").await?;
        if let Some(stream) = this.inner.stream.as_mut() {
            stream.set_max_allowed_packet(row_opt.unwrap_or((DEFAULT_MAX_ALLOWED_PACKET,)).0);
        }
        Ok(this)
    }

    /// Returns future that resolves to [`Conn`] with `wait_timeout` stored in it.
    async fn read_wait_timeout(self) -> Result<Self> {
        let (mut this, row_opt) = self.first("SELECT @@wait_timeout").await?;
        let wait_timeout_secs = row_opt.unwrap_or((28800,)).0;
        this.inner.wait_timeout = Duration::from_secs(wait_timeout_secs);
        Ok(this)
    }

    /// Returns true if time since last io exceeds wait_timeout (or conn_ttl if specified in opts).
    fn expired(&self) -> bool {
        let ttl = self
            .inner
            .opts
            .get_conn_ttl()
            .unwrap_or(self.inner.wait_timeout);
        self.idling() > ttl
    }

    /// Returns duration since last io.
    fn idling(&self) -> Duration {
        self.inner.last_io.elapsed()
    }

    /// Returns future that resolves to a [`Conn`] with `COM_RESET_CONNECTION` executed on it.
    pub async fn reset(self) -> Result<Conn> {
        let pool = self.inner.pool.clone();
        let mut conn = if self.inner.version > (5, 7, 2) {
            self.write_command_data(consts::Command::COM_RESET_CONNECTION, &[])
                .await?
                .read_packet()
                .await?
                .0
        } else {
            Conn::new(self.inner.opts.clone()).await?
        };

        conn.inner.stmt_cache.clear();
        conn.inner.pool = pool;
        Ok(conn)
    }

    async fn rollback_transaction(mut self) -> Result<Self> {
        assert!(self.inner.in_transaction);
        self.inner.in_transaction = false;
        self.drop_query("ROLLBACK").await
    }

    async fn drop_result(mut self) -> Result<Conn> {
        match self.inner.has_result.take() {
            Some((columns, None)) => {
                query_result::assemble::<_, TextProtocol>(self, Some(columns), None)
                    .drop_result()
                    .await
            }
            Some((columns, cached)) => {
                query_result::assemble::<_, BinaryProtocol>(self, Some(columns), cached)
                    .drop_result()
                    .await
            }
            None => Ok(self),
        }
    }

    fn cleanup(self) -> Pin<Box<dyn Future<Output = Result<Conn>> + Send>> {
        // NOTE: we need to box this since it may recurse
        // see https://github.com/rust-lang/rust/issues/46415#issuecomment-528099782
        Box::pin(async move {
            if self.inner.has_result.is_some() {
                self.drop_result().await?.cleanup().await
            } else if self.inner.in_transaction {
                self.rollback_transaction().await?.cleanup().await
            } else {
                Ok(self)
            }
        })
    }
}

impl ConnectionLike for Conn {
    fn take_stream(mut self) -> (Streamless<Self>, Stream) {
        let stream = self.inner.stream.take().expect("Logic error: stream taken");
        (Streamless::new(self), stream)
    }

    fn return_stream(&mut self, stream: Stream) {
        self.inner.stream = Some(stream);
    }

    fn stmt_cache_ref(&self) -> &StmtCache {
        &self.inner.stmt_cache
    }

    fn stmt_cache_mut(&mut self) -> &mut StmtCache {
        &mut self.inner.stmt_cache
    }

    fn get_affected_rows(&self) -> u64 {
        self.inner.affected_rows
    }

    fn get_capabilities(&self) -> consts::CapabilityFlags {
        self.inner.capabilities
    }

    fn get_in_transaction(&self) -> bool {
        self.inner.in_transaction
    }

    fn get_last_insert_id(&self) -> Option<u64> {
        match self.inner.last_insert_id {
            0 => None,
            x => Some(x),
        }
    }

    fn get_local_infile_handler(&self) -> Option<Arc<dyn LocalInfileHandler>> {
        self.inner.opts.get_local_infile_handler()
    }

    fn get_max_allowed_packet(&self) -> usize {
        self.inner.max_allowed_packet
    }

    fn get_opts(&self) -> &Opts {
        &self.inner.opts
    }

    fn get_pending_result(&self) -> Option<&(Arc<Vec<Column>>, Option<StmtCacheResult>)> {
        self.inner.has_result.as_ref()
    }

    fn get_server_version(&self) -> (u16, u16, u16) {
        self.inner.version
    }

    fn get_status(&self) -> consts::StatusFlags {
        self.inner.status
    }

    fn set_affected_rows(&mut self, affected_rows: u64) {
        self.inner.affected_rows = affected_rows;
    }

    fn set_in_transaction(&mut self, in_transaction: bool) {
        self.inner.in_transaction = in_transaction;
    }

    fn set_last_insert_id(&mut self, last_insert_id: u64) {
        self.inner.last_insert_id = last_insert_id;
    }

    fn set_pending_result(&mut self, meta: Option<(Arc<Vec<Column>>, Option<StmtCacheResult>)>) {
        self.inner.has_result = meta;
    }

    fn set_status(&mut self, status: consts::StatusFlags) {
        self.inner.status = status;
    }

    fn set_warnings(&mut self, warnings: u16) {
        self.inner.warnings = warnings;
    }

    fn reset_seq_id(&mut self) {
        if let Some(stream) = self.inner.stream.as_mut() {
            stream.reset_seq_id();
        }
    }

    fn sync_seq_id(&mut self) {
        if let Some(stream) = self.inner.stream.as_mut() {
            stream.sync_seq_id();
        }
    }

    fn touch(&mut self) {
        self.inner.last_io = Instant::now();
    }

    fn on_disconnect(&mut self) {
        self.inner.disconnected = true;
    }
}

#[cfg(test)]
mod test {
    use crate::{
        from_row, params, prelude::*, test_misc::get_opts, Conn, OptsBuilder, TransactionOptions,
        WhiteListFsLocalInfileHandler,
    };

    #[test]
    fn opts_should_satisfy_send_and_sync() {
        struct A<T: Sync + Send>(T);
        A(get_opts());
    }

    #[tokio::test]
    async fn should_connect_without_database() -> super::Result<()> {
        let mut opts = get_opts();

        // no database name
        opts.db_name(None::<String>);
        let conn: Conn = Conn::new(opts.clone()).await?.ping().await?;
        conn.disconnect().await?;

        // empty database name
        opts.db_name(Some(""));
        let conn: Conn = Conn::new(opts).await?.ping().await?;
        conn.disconnect().await?;

        Ok(())
    }

    #[tokio::test]
    async fn should_connect() -> super::Result<()> {
        let conn: Conn = Conn::new(get_opts()).await?.ping().await?;

        let (mut conn, plugins): (Conn, _) = conn
            .query("SHOW PLUGINS")
            .await?
            .map_and_drop(|mut row| row.take::<String, _>("Name").unwrap())
            .await?;

        // Should connect with any combination of supported plugin and empty-nonempty password.
        let variants = vec![
            ("caching_sha2_password", 2, "non-empty"),
            ("caching_sha2_password", 2, ""),
            ("mysql_native_password", 0, "non-empty"),
            ("mysql_native_password", 0, ""),
        ]
        .into_iter()
        .filter(|variant| plugins.iter().any(|p| p == variant.0));

        for (plug, val, pass) in variants {
            let query = format!("CREATE USER 'test_user'@'%' IDENTIFIED WITH {}", plug);
            conn = conn.drop_query(query).await.unwrap();

            conn = if (8, 0, 11) <= conn.inner.version && conn.inner.version <= (9, 0, 0) {
                conn.drop_query(format!("SET PASSWORD FOR 'test_user'@'%' = '{}'", pass))
                    .await
                    .unwrap()
            } else {
                conn = conn
                    .drop_query(format!("SET old_passwords = {}", val))
                    .await
                    .unwrap();
                conn.drop_query(format!(
                    "SET PASSWORD FOR 'test_user'@'%' = PASSWORD('{}')",
                    pass
                ))
                .await
                .unwrap()
            };

            let mut opts = get_opts();
            opts.user(Some("test_user"))
                .pass(Some(pass))
                .db_name(None::<String>);
            let result = Conn::new(opts).await;

            conn = conn.drop_query("DROP USER 'test_user'@'%'").await.unwrap();

            result?.disconnect().await?;
        }

        if crate::test_misc::test_compression() {
            assert!(format!("{:?}", conn).contains("Compression"));
        }

        if crate::test_misc::test_ssl() {
            assert!(format!("{:?}", conn).contains("Tls"));
        }

        conn.disconnect().await?;
        Ok(())
    }

    #[test]
    fn should_not_panic_if_dropped_without_tokio_runtime() {
        let fut = Conn::new(get_opts());
        let mut runtime = tokio::runtime::Runtime::new().unwrap();
        runtime.block_on(async {
            fut.await.unwrap();
        });
        // connection will drop here
    }

    #[tokio::test]
    async fn should_execute_init_queries_on_new_connection() -> super::Result<()> {
        let mut opts_builder = OptsBuilder::from_opts(get_opts());
        opts_builder.init(vec!["SET @a = 42", "SET @b = 'foo'"]);
        let (conn, result) = Conn::new(opts_builder)
            .await?
            .query("SELECT @a, @b")
            .await?
            .collect_and_drop::<(u8, String)>()
            .await?;
        conn.disconnect().await?;
        assert_eq!(result, vec![(42, "foo".into())]);
        Ok(())
    }

    #[tokio::test]
    async fn should_reset_the_connection() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let conn = conn.drop_exec("SELECT ?", (1,)).await?;
        let conn = conn.reset().await?;
        let conn = conn.drop_exec("SELECT ?", (1,)).await?;
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_not_cache_statements_if_stmt_cache_size_is_zero() -> super::Result<()> {
        let mut opts = OptsBuilder::from_opts(get_opts());
        opts.stmt_cache_size(0);
        let conn = Conn::new(opts).await?;
        let conn = conn.drop_exec("DO ?", (1,)).await?;
        let stmt = conn.prepare("DO 2").await?;
        let (stmt, _) = stmt.first::<_, (crate::Value,)>(()).await?;
        let (stmt, _) = stmt.first::<_, (crate::Value,)>(()).await?;
        let conn = stmt.close().await?;
        let conn = conn.prep_exec("DO 3", ()).await?.drop_result().await?;
        let conn = conn.batch_exec("DO 4", vec![(), ()]).await?;
        let (conn, _) = conn.first_exec::<_, _, (u8,)>("DO 5", ()).await?;
        let (conn, row) = conn
            .first("SHOW SESSION STATUS LIKE 'Com_stmt_close';")
            .await?;
        assert_eq!(from_row::<(String, usize)>(row.unwrap()).1, 5);
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_hold_stmt_cache_size_bound() -> super::Result<()> {
        use crate::connection_like::ConnectionLike;

        let mut opts = OptsBuilder::from_opts(get_opts());
        opts.stmt_cache_size(3);
        let conn = Conn::new(opts)
            .await?
            .drop_exec("DO 1", ())
            .await?
            .drop_exec("DO 2", ())
            .await?
            .drop_exec("DO 3", ())
            .await?
            .drop_exec("DO 1", ())
            .await?
            .drop_exec("DO 4", ())
            .await?
            .drop_exec("DO 3", ())
            .await?
            .drop_exec("DO 5", ())
            .await?
            .drop_exec("DO 6", ())
            .await?;
        let (conn, row_opt) = conn
            .first("SHOW SESSION STATUS LIKE 'Com_stmt_close';")
            .await?;
        let (_, count): (String, usize) = row_opt.unwrap();
        assert_eq!(count, 3);
        let order = conn
            .stmt_cache_ref()
            .iter()
            .map(Clone::clone)
            .collect::<Vec<String>>();
        assert_eq!(order, &["DO 3", "DO 5", "DO 6"]);
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_perform_queries() -> super::Result<()> {
        let long_string = ::std::iter::repeat('A')
            .take(18 * 1024 * 1024)
            .collect::<String>();
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(format!(r"SELECT '{}', 231", long_string))
            .await?;
        let (conn, result) = result
            .reduce_and_drop(vec![], move |mut acc, row| {
                acc.push(from_row(row));
                acc
            })
            .await?;
        conn.disconnect().await?;
        assert_eq!((long_string, 231), result[0]);
        Ok(())
    }

    #[tokio::test]
    async fn should_drop_query() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let (conn, result) = conn
            .drop_query("CREATE TEMPORARY TABLE tmp (id int DEFAULT 10, name text)")
            .await?
            .drop_query("INSERT INTO tmp VALUES (1, 'foo')")
            .await?
            .first::<_, (u8,)>("SELECT COUNT(*) FROM tmp")
            .await?;
        conn.disconnect().await?;
        assert_eq!(result, Some((1,)));
        Ok(())
    }

    #[tokio::test]
    async fn should_try_collect() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(
                r"SELECT 'hello', 123
                    UNION ALL
                    SELECT 'world', 'bar'
                    UNION ALL
                    SELECT 'hello', 123
                ",
            )
            .await?;
        let (result, mut rows) = result.try_collect::<(String, u8)>().await?;
        assert!(rows.pop().unwrap().is_ok());
        assert!(rows.pop().unwrap().is_err());
        assert!(rows.pop().unwrap().is_ok());
        let conn = result.drop_result().await?;
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_try_collect_and_drop() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let (conn, mut rows) = conn
            .query(
                r"SELECT 'hello', 123
                    UNION ALL
                    SELECT 'world', 'bar'
                    UNION ALL
                    SELECT 'hello', 123;
                    SELECT 'foo', 255;
                ",
            )
            .await?
            .try_collect_and_drop::<(String, u8)>()
            .await?;
        assert!(rows.pop().unwrap().is_ok());
        assert!(rows.pop().unwrap().is_err());
        assert!(rows.pop().unwrap().is_ok());
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_handle_mutliresult_set() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(
                r"SELECT 'hello', 123
                    UNION ALL
                    SELECT 'world', 231;
                    SELECT 'foo', 255;
                ",
            )
            .await?;
        let (result, rows_1) = result.collect::<(String, u8)>().await?;
        let (conn, rows_2) = result.collect_and_drop().await?;
        conn.disconnect().await?;

        assert_eq!((String::from("hello"), 123), rows_1[0]);
        assert_eq!((String::from("world"), 231), rows_1[1]);
        assert_eq!((String::from("foo"), 255), rows_2[0]);
        Ok(())
    }

    #[tokio::test]
    async fn should_map_resultset() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(
                r"
                    SELECT 'hello', 123
                    UNION ALL
                    SELECT 'world', 231;
                    SELECT 'foo', 255;
                ",
            )
            .await?;

        let (result, rows_1) = result.map(|row| from_row::<(String, u8)>(row)).await?;
        let (conn, rows_2) = result.map_and_drop(from_row).await?;
        conn.disconnect().await?;

        assert_eq!((String::from("hello"), 123), rows_1[0]);
        assert_eq!((String::from("world"), 231), rows_1[1]);
        assert_eq!((String::from("foo"), 255), rows_2[0]);
        Ok(())
    }

    #[tokio::test]
    async fn should_reduce_resultset() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(
                r"SELECT 5
                    UNION ALL
                    SELECT 6;
                    SELECT 7;",
            )
            .await?;
        let (result, reduced) = result
            .reduce(0, |mut acc, row| {
                acc += from_row::<i32>(row);
                acc
            })
            .await?;
        let (conn, rows_2) = result.collect_and_drop::<i32>().await?;
        conn.disconnect().await?;
        assert_eq!(11, reduced);
        assert_eq!(7, rows_2[0]);
        Ok(())
    }

    #[tokio::test]
    async fn should_handle_multi_result_sets_where_some_results_have_no_output() -> super::Result<()>
    {
        const QUERY: &str = r"SELECT 1;
            UPDATE time_zone SET Time_zone_id = 1 WHERE Time_zone_id = 1;
            SELECT 2;
            SELECT 3;
            UPDATE time_zone SET Time_zone_id = 1 WHERE Time_zone_id = 1;
            UPDATE time_zone SET Time_zone_id = 1 WHERE Time_zone_id = 1;
            SELECT 4;";

        let c = Conn::new(get_opts()).await?;
        let c = c
            .drop_query("CREATE TEMPORARY TABLE time_zone (Time_zone_id INT)")
            .await
            .unwrap();
        let t = c.start_transaction(TransactionOptions::new()).await?;
        let t = t.drop_query(QUERY).await?;
        let r = t.query(QUERY).await?;
        let (t, out) = r.collect_and_drop::<u8>().await?;
        assert_eq!(vec![1], out);
        let r = t.query(QUERY).await?;
        let t = r
            .for_each_and_drop(|x| assert_eq!(from_row::<u8>(x), 1))
            .await?;
        let r = t.query(QUERY).await?;
        let (t, out) = r.map_and_drop(|row| from_row::<u8>(row)).await?;
        assert_eq!(vec![1], out);
        let r = t.query(QUERY).await?;
        let (t, out) = r
            .reduce_and_drop(0u8, |acc, x| acc + from_row::<u8>(x))
            .await?;
        assert_eq!(1, out);
        let t = t.query(QUERY).await?.drop_result().await?;
        let c = t.commit().await?;
        let (c, result) = c.first_exec::<_, _, u8>("SELECT 1", ()).await?;
        c.disconnect().await?;
        assert_eq!(result, Some(1));
        Ok(())
    }

    #[tokio::test]
    async fn should_iterate_over_resultset() -> super::Result<()> {
        use std::sync::{
            atomic::{AtomicUsize, Ordering},
            Arc,
        };

        let acc = Arc::new(AtomicUsize::new(0));

        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .query(
                r"SELECT 2
                    UNION ALL
                    SELECT 3;
                    SELECT 5;",
            )
            .await?;
        let result = result
            .for_each({
                let acc = acc.clone();
                move |row| {
                    acc.fetch_add(from_row::<usize>(row), Ordering::SeqCst);
                }
            })
            .await?;
        let conn = result
            .for_each_and_drop({
                let acc = acc.clone();
                move |row| {
                    acc.fetch_add(from_row::<usize>(row), Ordering::SeqCst);
                }
            })
            .await?;
        conn.disconnect().await?;
        assert_eq!(acc.load(Ordering::SeqCst), 10);
        Ok(())
    }

    #[tokio::test]
    async fn should_prepare_statement() -> super::Result<()> {
        Conn::new(get_opts())
            .await?
            .prepare(r"SELECT ?")
            .await?
            .close()
            .await?
            .disconnect()
            .await?;

        Conn::new(get_opts())
            .await?
            .prepare(r"SELECT :foo")
            .await?
            .close()
            .await?
            .disconnect()
            .await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_execute_statement() -> super::Result<()> {
        let long_string = ::std::iter::repeat('A')
            .take(18 * 1024 * 1024)
            .collect::<String>();
        let conn = Conn::new(get_opts()).await?;
        let stmt = conn.prepare(r"SELECT ?").await?;
        let result = stmt.execute((&long_string,)).await?;
        let (stmt, mut mapped) = result
            .map_and_drop(|row| from_row::<(String,)>(row))
            .await?;
        assert_eq!(mapped.len(), 1);
        assert_eq!(mapped.pop(), Some((long_string,)));
        let result = stmt.execute((42,)).await?;
        let (stmt, collected) = result.collect_and_drop::<(u8,)>().await?;
        assert_eq!(collected, vec![(42u8,)]);
        let result = stmt.execute((8,)).await?;
        let (stmt, reduced) = result
            .reduce_and_drop(2, |mut acc, row| {
                acc += from_row::<i32>(row);
                acc
            })
            .await?;
        stmt.close().await?.disconnect().await?;
        assert_eq!(reduced, 10);

        let conn = Conn::new(get_opts()).await?;
        let stmt = conn.prepare(r"SELECT :foo, :bar, :foo, 3").await?;
        let result = stmt
            .execute(params! { "foo" => "quux", "bar" => "baz" })
            .await?;
        let (stmt, mut mapped) = result
            .map_and_drop(|row| from_row::<(String, String, String, u8)>(row))
            .await?;
        assert_eq!(mapped.len(), 1);
        assert_eq!(
            mapped.pop(),
            Some(("quux".into(), "baz".into(), "quux".into(), 3))
        );
        let result = stmt.execute(params! { "foo" => 2, "bar" => 3 }).await?;
        let (stmt, collected) = result.collect_and_drop::<(u8, u8, u8, u8)>().await?;
        assert_eq!(collected, vec![(2, 3, 2, 3)]);
        let result = stmt.execute(params! { "foo" => 2, "bar" => 3 }).await?;
        let (stmt, reduced) = result
            .reduce_and_drop(0, |acc, row| {
                let (a, b, c, d): (u8, u8, u8, u8) = from_row(row);
                acc + a + b + c + d
            })
            .await?;
        stmt.close().await?.disconnect().await?;
        assert_eq!(reduced, 10);
        Ok(())
    }

    #[tokio::test]
    async fn should_prep_exec_statement() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let result = conn
            .prep_exec(r"SELECT :a, :b, :a", params! { "a" => 2, "b" => 3 })
            .await?;
        let (conn, output) = result
            .map_and_drop(|row| {
                let (a, b, c): (u8, u8, u8) = from_row(row);
                a * b * c
            })
            .await?;
        conn.disconnect().await?;
        assert_eq!(output[0], 12u8);
        Ok(())
    }

    #[tokio::test]
    async fn should_first_exec_statement() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let (conn, output): (_, Option<(u8,)>) = conn
            .first_exec(
                r"SELECT :a UNION ALL SELECT :b",
                params! { "a" => 2, "b" => 3 },
            )
            .await?;
        conn.disconnect().await?;
        assert_eq!(output.unwrap(), (2u8,));
        Ok(())
    }

    #[tokio::test]
    async fn should_run_transactions() -> super::Result<()> {
        let conn = Conn::new(get_opts()).await?;
        let conn = conn
            .drop_query("CREATE TEMPORARY TABLE tmp (id INT, name TEXT)")
            .await?;
        let transaction = conn.start_transaction(Default::default()).await?;
        let conn = transaction
            .drop_query("INSERT INTO tmp VALUES (1, 'foo'), (2, 'bar')")
            .await?
            .commit()
            .await?;
        let (conn, output_opt) = conn.first("SELECT COUNT(*) FROM tmp").await?;
        assert_eq!(output_opt, Some((2u8,)));
        let transaction = conn.start_transaction(Default::default()).await?;
        let transaction = transaction
            .drop_query("INSERT INTO tmp VALUES (3, 'baz'), (4, 'quux')")
            .await?;
        let (t, output_opt) = transaction
            .first_exec("SELECT COUNT(*) FROM tmp", ())
            .await?;
        assert_eq!(output_opt, Some((4u8,)));
        let conn = t.rollback().await?;
        let (conn, output_opt) = conn.first("SELECT COUNT(*) FROM tmp").await?;
        assert_eq!(output_opt, Some((2u8,)));
        conn.disconnect().await?;
        Ok(())
    }

    #[tokio::test]
    async fn should_handle_local_infile() -> super::Result<()> {
        use std::{io::Write, path::Path};

        let file_path = tempfile::Builder::new().tempfile_in("").unwrap();
        let file_path = file_path.path();
        let file_name = Path::new(file_path.file_name().unwrap());

        let mut opts = OptsBuilder::from_opts(get_opts());
        opts.local_infile_handler(Some(WhiteListFsLocalInfileHandler::new(&[file_name][..])));

        let conn = Conn::new(opts).await.unwrap();
        let conn = conn
            .drop_query("CREATE TEMPORARY TABLE tmp (a TEXT);")
            .await
            .unwrap();

        let mut file = ::std::fs::File::create(file_name).unwrap();
        let _ = file.write(b"AAAAAA\n");
        let _ = file.write(b"BBBBBB\n");
        let _ = file.write(b"CCCCCC\n");
        let conn = match conn
            .drop_query(dbg!(format!(
                r#"LOAD DATA LOCAL INFILE "{}" INTO TABLE tmp;"#,
                file_name.display()
            )))
            .await
        {
            Ok(conn) => conn,
            Err(super::Error::Server(ref err)) if err.code == 1148 => {
                // The used command is not allowed with this MySQL version
                return Ok(());
            }
            Err(super::Error::Server(ref err)) if err.code == 3948 => {
                // Loading local data is disabled;
                // this must be enabled on both the client and server sides
                return Ok(());
            }
            e @ Err(_) => e.unwrap(),
        };
        let (conn, result) = conn
            .prep_exec("SELECT * FROM tmp;", ())
            .await
            .unwrap()
            .map_and_drop(|row| from_row::<(String,)>(row).0)
            .await
            .unwrap();
        assert_eq!(result.len(), 3);
        assert_eq!(result[0], "AAAAAA");
        assert_eq!(result[1], "BBBBBB");
        assert_eq!(result[2], "CCCCCC");
        let result = conn.disconnect().await;

        if let Err(crate::error::Error::Server(ref err)) = result {
            if err.code == 1148 {
                // The used command is not allowed with this MySQL version
                return Ok(());
            }
        }
        result.unwrap();
        Ok(())
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use futures_util::try_future::TryFutureExt;

        use crate::{conn::Conn, queryable::Queryable, test_misc::get_opts};

        #[bench]
        fn simple_exec(bencher: &mut test::Bencher) {
            let runtime = tokio::runtime::Runtime::new().unwrap();
            let mut conn_opt = Some(runtime.block_on(Conn::new(get_opts())).unwrap());

            bencher.iter(|| {
                let conn = conn_opt.take().unwrap();
                conn_opt = Some(runtime.block_on(conn.drop_query("DO 1")).unwrap());
            });

            runtime
                .block_on(conn_opt.take().unwrap().disconnect())
                .unwrap();
            runtime.shutdown_on_idle();
        }

        #[bench]
        fn select_large_string(bencher: &mut test::Bencher) {
            let runtime = tokio::runtime::Runtime::new().unwrap();
            let mut conn_opt = Some(runtime.block_on(Conn::new(get_opts())).unwrap());

            bencher.iter(|| {
                let conn = conn_opt.take().unwrap();
                conn_opt = Some(
                    runtime
                        .block_on(conn.drop_query("SELECT REPEAT('A', 10000)"))
                        .unwrap(),
                );
            });

            runtime
                .block_on(conn_opt.take().unwrap().disconnect())
                .unwrap();
            runtime.shutdown_on_idle();
        }

        #[bench]
        fn prepared_exec(bencher: &mut test::Bencher) {
            let runtime = tokio::runtime::Runtime::new().unwrap();
            let mut stmt_opt = Some(
                runtime
                    .block_on(Conn::new(get_opts()).and_then(|conn| conn.prepare("DO 1")))
                    .unwrap(),
            );

            bencher.iter(|| {
                let stmt = stmt_opt.take().unwrap();
                stmt_opt = Some(
                    runtime
                        .block_on(stmt.execute(()).and_then(|result| result.drop_result()))
                        .unwrap(),
                );
            });

            runtime
                .block_on(
                    stmt_opt
                        .take()
                        .unwrap()
                        .close()
                        .and_then(|conn| conn.disconnect()),
                )
                .unwrap();
            runtime.shutdown_on_idle();
        }

        #[bench]
        fn prepare_and_exec(bencher: &mut test::Bencher) {
            let runtime = tokio::runtime::Runtime::new().unwrap();
            let mut conn_opt = Some(runtime.block_on(Conn::new(get_opts())).unwrap());

            bencher.iter(|| {
                let conn = conn_opt.take().unwrap();
                conn_opt = Some(
                    runtime
                        .block_on(
                            conn.prepare("SELECT ?")
                                .and_then(|stmt| stmt.execute((0,)))
                                .and_then(|result| result.drop_result())
                                .and_then(|stmt| stmt.close()),
                        )
                        .unwrap(),
                );
            });

            runtime
                .block_on(conn_opt.take().unwrap().disconnect())
                .unwrap();
            runtime.shutdown_on_idle();
        }
    }
}