rusqlite/src/column.rs

use std::str;

use crate::{Error, Result, Row, Rows, Statement};

/// Information about a column of a SQLite query.
#[derive(Debug)]
pub struct Column<'stmt> {
    name: &'stmt str,
    decl_type: Option<&'stmt str>,
}

impl Column<'_> {
    /// Returns the name of the column.
    pub fn name(&self) -> &str {
        self.name
    }

    /// Returns the type of the column (`None` for expression).
    pub fn decl_type(&self) -> Option<&str> {
        self.decl_type
    }
}

impl Statement<'_> {
    /// Get all the column names in the result set of the prepared statement.
    pub fn column_names(&self) -> Vec<&str> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let slice = self.stmt.column_name(i);
            let s = str::from_utf8(slice.to_bytes()).unwrap();
            cols.push(s);
        }
        cols
    }

    /// Return the number of columns in the result set returned by the prepared
    /// statement.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns the column index in the result set for a given column name.
    ///
    /// If there is no AS clause then the name of the column is unspecified and
    /// may change from one release of SQLite to the next.
    ///
    /// # Failure
    ///
    /// Will return an `Error::InvalidColumnName` when there is no column with
    /// the specified `name`.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        let bytes = name.as_bytes();
        let n = self.column_count();
        for i in 0..n {
            if bytes.eq_ignore_ascii_case(self.stmt.column_name(i).to_bytes()) {
                return Ok(i);
            }
        }
        Err(Error::InvalidColumnName(String::from(name)))
    }

    /// Returns a slice describing the columns of the result of the query.
    pub fn columns<'stmt>(&'stmt self) -> Vec<Column<'stmt>> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let slice = self.stmt.column_name(i);
            let name = str::from_utf8(slice.to_bytes()).unwrap();
            let slice = self.stmt.column_decltype(i);
            let decl_type = slice.map(|s| str::from_utf8(s.to_bytes()).unwrap());
            cols.push(Column { name, decl_type });
        }
        cols
    }
}

impl<'stmt> Rows<'stmt> {
    /// Get all the column names.
    pub fn column_names(&self) -> Option<Vec<&str>> {
        self.stmt.map(Statement::column_names)
    }

    /// Return the number of columns.
    pub fn column_count(&self) -> Option<usize> {
        self.stmt.map(Statement::column_count)
    }

    /// Returns a slice describing the columns of the Rows.
    pub fn columns(&self) -> Option<Vec<Column<'stmt>>> {
        self.stmt.map(Statement::columns)
    }
}

impl<'stmt> Row<'stmt> {
    /// Return the number of columns in the current row.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns a slice describing the columns of the Row.
    pub fn columns(&self) -> Vec<Column<'stmt>> {
        self.stmt.columns()
    }
}

#[cfg(test)]
mod test {
    use super::Column;
    use crate::Connection;

    #[test]
    fn test_columns() {
        let db = Connection::open_in_memory().unwrap();
        let query = db.prepare("SELECT * FROM sqlite_master").unwrap();
        let columns = query.columns();
        let column_names: Vec<&str> = columns.iter().map(Column::name).collect();
        assert_eq!(
            column_names.as_slice(),
            &["type", "name", "tbl_name", "rootpage", "sql"]
        );
        let column_types: Vec<Option<&str>> = columns.iter().map(Column::decl_type).collect();
        assert_eq!(
            column_types.as_slice(),
            &[
                Some("text"),
                Some("text"),
                Some("text"),
                Some("int"),
                Some("text")
            ]
        );
    }
}