diesel/mysql_like/connection/stmt/

iterator.rs

#![allow(unsafe_code)] // module uses ffi
use alloc::rc::Rc;
use core::cell::{Ref, RefCell};
use core::marker::PhantomData;

use super::{OutputBinds, Statement, StatementMetadata, StatementUse};
use crate::backend::Backend;
use crate::connection::statement_cache::MaybeCached;
use crate::mysql_like::{MysqlLikeBackend, MysqlType};
use crate::result::Error::DeserializationError;
use crate::result::QueryResult;
use crate::row::*;

#[allow(missing_debug_implementations)]
pub struct StatementIterator<'a, DB: MysqlLikeBackend> {
    stmt: StatementUse<'a, DB>,
    last_row: Rc<RefCell<PrivateMysqlRow>>,
    metadata: Rc<StatementMetadata>,
    len: usize,
}

impl<'a, DB: MysqlLikeBackend> StatementIterator<'a, DB> {
    pub fn from_stmt(
        stmt: MaybeCached<'a, Statement<DB>>,
        types: &[Option<MysqlType>],
    ) -> QueryResult<Self> {
        let mut stmt = if let Some(metadata) = stmt.metadata()? {
            let mut output_binds =
                OutputBinds::from_output_types(types, &metadata).map_err(DeserializationError)?;
            /*
                This may seem redundant but if we don't do this we will hit
                a memory bug in `libmysqlclient`
            */
            stmt.execute_statement(&mut output_binds)?
        } else {
            //Sometimes we must execute the statement to get its metadata.
            unsafe { stmt.execute()? }
        };

        let metadata = stmt.metadata()?;
        let output_binds =
            OutputBinds::from_output_types(types, &metadata).map_err(DeserializationError)?;

        let size = unsafe { stmt.result_size() }?;

        Ok(StatementIterator {
            metadata: Rc::new(metadata),
            last_row: Rc::new(RefCell::new(PrivateMysqlRow::Direct(output_binds))),
            len: size,
            stmt,
        })
    }
}

impl<DB: MysqlLikeBackend> Iterator for StatementIterator<'_, DB> {
    type Item = QueryResult<MysqlRow<DB>>;

    fn next(&mut self) -> Option<Self::Item> {
        // check if we own the only instance of the bind buffer
        // if that's the case we can reuse the underlying allocations
        // if that's not the case, we need to copy the output bind buffers
        // to somewhere else
        let res = if let Some(binds) = Rc::get_mut(&mut self.last_row) {
            if let PrivateMysqlRow::Direct(binds) = RefCell::get_mut(binds) {
                self.stmt.populate_row_buffers(binds)
            } else {
                // any other state than `PrivateMysqlRow::Direct` is invalid here
                // and should not happen. If this ever happens this is a logic error
                // in the code above
                unreachable!(
                    "You've reached an impossible internal state. \
                     If you ever see this error message please open \
                     an issue at https://github.com/diesel-rs/diesel \
                     providing example code how to trigger this error."
                )
            }
        } else {
            // The shared bind buffer is in use by someone else,
            // this means we copy out the values and replace the used reference
            // by the copied values. After this we can advance the statement
            // another step
            let mut last_row = {
                let mut last_row = match self.last_row.try_borrow_mut() {
                    Ok(o) => o,
                    Err(_e) => {
                        return Some(Err(DeserializationError(
                            "Failed to reborrow row. Try to release any `MysqlField` or `MysqlValue` \
                             that exists at this point"
                                .into(),
                        )));
                    }
                };
                let last_row = &mut *last_row;
                let duplicated = last_row.duplicate();
                core::mem::replace(last_row, duplicated)
            };
            let res = if let PrivateMysqlRow::Direct(ref mut binds) = last_row {
                self.stmt.populate_row_buffers(binds)
            } else {
                // any other state than `PrivateMysqlRow::Direct` is invalid here
                // and should not happen. If this ever happens this is a logic error
                // in the code above
                unreachable!(
                    "You've reached an impossible internal state. \
                     If you ever see this error message please open \
                     an issue at https://github.com/diesel-rs/diesel \
                     providing example code how to trigger this error."
                )
            };
            self.last_row = Rc::new(RefCell::new(last_row));
            res
        };

        match res {
            Ok(Some(())) => {
                self.len = self.len.saturating_sub(1);
                Some(Ok(MysqlRow {
                    metadata: self.metadata.clone(),
                    row: self.last_row.clone(),
                    _phantom: PhantomData,
                }))
            }
            Ok(None) => None,
            Err(e) => {
                self.len = self.len.saturating_sub(1);
                Some(Err(e))
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len(), Some(self.len()))
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        self.len()
    }
}

impl<DB: MysqlLikeBackend> ExactSizeIterator for StatementIterator<'_, DB> {
    fn len(&self) -> usize {
        self.len
    }
}

#[derive(Clone)]
#[allow(missing_debug_implementations)]
pub struct MysqlRow<DB: MysqlLikeBackend> {
    row: Rc<RefCell<PrivateMysqlRow>>,
    metadata: Rc<StatementMetadata>,
    _phantom: PhantomData<DB>,
}

enum PrivateMysqlRow {
    Direct(OutputBinds),
    Copied(OutputBinds),
}

impl PrivateMysqlRow {
    fn duplicate(&self) -> Self {
        match self {
            Self::Copied(b) | Self::Direct(b) => Self::Copied(b.clone()),
        }
    }
}

impl<DB: MysqlLikeBackend> RowSealed for MysqlRow<DB> {}

impl<'a, DB: MysqlLikeBackend> Row<'a, DB> for MysqlRow<DB> {
    type Field<'f>
        = MysqlLikeField<'f, DB>
    where
        'a: 'f,
        Self: 'f;
    type InnerPartialRow = Self;

    fn field_count(&self) -> usize {
        self.metadata.fields().len()
    }

    fn get<'b, I>(&'b self, idx: I) -> Option<Self::Field<'b>>
    where
        'a: 'b,
        Self: RowIndex<I>,
    {
        let idx = self.idx(idx)?;
        Some(MysqlLikeField {
            binds: self.row.borrow(),
            metadata: self.metadata.clone(),
            idx,
            _phantom: PhantomData,
        })
    }

    fn partial_row(&self, range: core::ops::Range<usize>) -> PartialRow<'_, Self::InnerPartialRow> {
        PartialRow::new::<DB>(self, range)
    }
}

impl<DB: MysqlLikeBackend> RowIndex<usize> for MysqlRow<DB>
where
    MysqlRow<DB>: for<'a> Row<'a, DB>,
{
    fn idx(&self, idx: usize) -> Option<usize> {
        if idx < self.field_count() {
            Some(idx)
        } else {
            None
        }
    }
}

impl<'a, DB: MysqlLikeBackend> RowIndex<&'a str> for MysqlRow<DB> {
    fn idx(&self, idx: &'a str) -> Option<usize> {
        self.metadata
            .fields()
            .iter()
            .enumerate()
            .find(|(_, field_meta)| field_meta.field_name() == Some(idx))
            .map(|(idx, _)| idx)
    }
}

#[allow(missing_debug_implementations)]
pub struct MysqlLikeField<'a, DB: MysqlLikeBackend> {
    binds: Ref<'a, PrivateMysqlRow>,
    metadata: Rc<StatementMetadata>,
    idx: usize,
    _phantom: PhantomData<DB>,
}

impl<'a, DB: MysqlLikeBackend> Field<'a, DB> for MysqlLikeField<'a, DB> {
    fn field_name(&self) -> Option<&str> {
        self.metadata.fields()[self.idx].field_name()
    }

    fn is_null(&self) -> bool {
        match &*self.binds {
            PrivateMysqlRow::Copied(b) | PrivateMysqlRow::Direct(b) => b[self.idx].is_null(),
        }
    }

    fn value(&self) -> Option<<DB as Backend>::RawValue<'_>> {
        match &*self.binds {
            PrivateMysqlRow::Copied(b) | PrivateMysqlRow::Direct(b) => b[self.idx].value(),
        }
    }
}

#[cfg(test)]
mod tests {

    use crate::connection::LoadConnection;
    use crate::deserialize::{FromSql, FromSqlRow};
    use crate::prelude::*;
    use crate::row::{Field, Row};
    use crate::sql_types;

    #[diesel_test_helper::test]
    #[allow(clippy::drop_non_drop)] // we want to explicitly extend lifetimes here
    fn fun_with_row_iters() {
        crate::table! {
            users(id) {
                id -> Integer,
                name -> Text,
            }
        }

        #[cfg(feature = "mysql")]
        type DB = crate::mysql::Mysql;
        #[cfg(feature = "mariadb")]
        type DB = crate::mariadb::Mariadb;

        let conn = &mut crate::test_helpers::connection();

        crate::sql_query(
            "CREATE TEMPORARY TABLE IF NOT EXISTS users(id INTEGER PRIMARY KEY, name TEXT NOT NULL);",
        )
        .execute(conn)
        .unwrap();

        crate::insert_into(users::table)
            .values(vec![
                (users::id.eq(1), users::name.eq("Sean")),
                (users::id.eq(2), users::name.eq("Tess")),
            ])
            .execute(conn)
            .unwrap();

        let query = users::table.select((users::id, users::name));

        let expected = vec![(1, String::from("Sean")), (2, String::from("Tess"))];

        {
            let row_iter = conn.load(query).unwrap();
            for (row, expected) in row_iter.zip(&expected) {
                let row = row.unwrap();

                let deserialized = <(i32, String) as FromSqlRow<
                    (sql_types::Integer, sql_types::Text),
                    _,
                >>::build_from_row(&row)
                .unwrap();

                assert_eq!(&deserialized, expected);
            }
        }

        {
            let collected_rows = conn.load(query).unwrap().collect::<Vec<_>>();
            assert_eq!(collected_rows.len(), 2);
            for (row, expected) in collected_rows.iter().zip(&expected) {
                let deserialized = row
                    .as_ref()
                    .map(|row| {
                        <(i32, String) as FromSqlRow<
                            (sql_types::Integer, sql_types::Text),
                        _,
                        >>::build_from_row(row).unwrap()
                    })
                    .unwrap();
                assert_eq!(&deserialized, expected);
            }
        }

        let mut row_iter = conn.load(query).unwrap();

        let first_row = row_iter.next().unwrap().unwrap();
        let first_fields = (
            Row::get(&first_row, 0).unwrap(),
            Row::get(&first_row, 1).unwrap(),
        );
        let first_values = (first_fields.0.value(), first_fields.1.value());

        assert!(row_iter.next().unwrap().is_err());
        std::mem::drop(first_values);
        assert!(row_iter.next().unwrap().is_err());
        std::mem::drop(first_fields);

        let second_row = row_iter.next().unwrap().unwrap();
        let second_fields = (
            Row::get(&second_row, 0).unwrap(),
            Row::get(&second_row, 1).unwrap(),
        );
        let second_values = (second_fields.0.value(), second_fields.1.value());

        assert!(row_iter.next().unwrap().is_err());
        std::mem::drop(second_values);
        assert!(row_iter.next().unwrap().is_err());
        std::mem::drop(second_fields);

        assert!(row_iter.next().is_none());

        let first_fields = (
            Row::get(&first_row, 0).unwrap(),
            Row::get(&first_row, 1).unwrap(),
        );
        let second_fields = (
            Row::get(&second_row, 0).unwrap(),
            Row::get(&second_row, 1).unwrap(),
        );

        let first_values = (first_fields.0.value(), first_fields.1.value());
        let second_values = (second_fields.0.value(), second_fields.1.value());

        assert_eq!(
            <i32 as FromSql<sql_types::Integer, DB>>::from_nullable_sql(first_values.0).unwrap(),
            expected[0].0
        );
        assert_eq!(
            <String as FromSql<sql_types::Text, DB>>::from_nullable_sql(first_values.1).unwrap(),
            expected[0].1
        );

        assert_eq!(
            <i32 as FromSql<sql_types::Integer, DB>>::from_nullable_sql(second_values.0).unwrap(),
            expected[1].0
        );
        assert_eq!(
            <String as FromSql<sql_types::Text, DB>>::from_nullable_sql(second_values.1).unwrap(),
            expected[1].1
        );

        let first_fields = (
            Row::get(&first_row, 0).unwrap(),
            Row::get(&first_row, 1).unwrap(),
        );
        let first_values = (first_fields.0.value(), first_fields.1.value());

        assert_eq!(
            <i32 as FromSql<sql_types::Integer, DB>>::from_nullable_sql(first_values.0).unwrap(),
            expected[0].0
        );
        assert_eq!(
            <String as FromSql<sql_types::Text, DB>>::from_nullable_sql(first_values.1).unwrap(),
            expected[0].1
        );
    }

    #[diesel_test_helper::test]
    fn output_binds_are_rebound_after_column_error() {
        use crate::dsl::sql;
        use crate::mysql_like::connection::DefaultLoadingMode;
        use crate::sql_types::{Integer, Text};

        let conn = &mut crate::test_helpers::connection();

        diesel::sql_query(
            "CREATE TEMPORARY TABLE f02 (id INTEGER PRIMARY KEY, txt TEXT NOT NULL, num VARCHAR(32) NOT NULL)",
        )
        .execute(conn)
        .unwrap();

        diesel::sql_query(
            "INSERT INTO f02 VALUES \
             (1, REPEAT('a', 4096), '1'), \
             (2, REPEAT('b', 8192), '99999999999'), \
             (3, REPEAT('c', 4096), '3')",
        )
        .execute(conn)
        .unwrap();

        let mut iter = sql::<(Text, Integer)>("SELECT txt, num FROM f02 ORDER BY id")
            .load_iter::<(String, i32), DefaultLoadingMode>(conn)
            .unwrap();

        let first = iter.next().unwrap().unwrap();
        assert_eq!(first, ("a".repeat(4096), 1));

        // the overflowing column is reported as an error
        assert!(iter.next().unwrap().is_err());

        // keep iterating: this must neither touch freed memory nor return stale data
        let third = iter.next().unwrap().unwrap();
        assert_eq!(third, ("c".repeat(4096), 3));
        assert!(iter.next().is_none());
    }

    #[test]
    fn row_field_names_after_statement_close() {
        let conn = &mut crate::test_helpers::connection();
        conn.set_prepared_statement_cache_size(crate::connection::CacheSize::Disabled);
        let rows: Vec<_> = conn
            .load(diesel::sql_query("SELECT 1 AS some_column_name"))
            .unwrap()
            .collect::<Result<_, _>>()
            .unwrap();
        // the iterator, and with it the uncached MYSQL_STMT, is gone now
        let f = Row::get(&rows[0], 0usize).unwrap();
        assert_eq!(f.field_name(), Some("some_column_name"));
        assert!(Row::get(&rows[0], "some_column_name").is_some());
    }

    #[test]
    fn row_field_names_after_connection_drop() {
        let rows: Vec<_> = {
            let conn = &mut crate::test_helpers::connection();
            conn.load(diesel::sql_query("SELECT 1 AS some_column_name"))
                .unwrap()
                .collect::<Result<_, _>>()
                .unwrap()
        };
        let f = Row::get(&rows[0], 0usize).unwrap();
        assert_eq!(f.field_name(), Some("some_column_name"));
    }
}