diff --git a/src/options.rs b/src/options.rs index bfccb19..ae34a6d 100644 --- a/src/options.rs +++ b/src/options.rs @@ -1,6 +1,6 @@ use nom::branch::alt; use nom::bytes::complete::{escaped, tag, tag_no_case}; -use nom::character::complete::{alpha1, digit0, multispace0, multispace1, none_of}; +use nom::character::complete::{alpha1, digit0, digit1, multispace0, multispace1, none_of}; use nom::combinator::{map, recognize}; use nom::{IResult, Parser}; use nom::sequence::{delimited, preceded, separated_pair, terminated, tuple}; @@ -9,6 +9,7 @@ pub enum UsingOption { File(String), Worksheet(String), Range(String), + ColNames(String), } pub fn parse_option(input: &str) -> IResult<&str, UsingOption> { @@ -16,6 +17,7 @@ pub fn parse_option(input: &str) -> IResult<&str, UsingOption> { parse_filename_option, parse_worksheet_option, parse_range_option, + parse_colnames_option, ))).parse(input) } @@ -49,6 +51,16 @@ fn parse_range_option(input: &str) -> IResult<&str, UsingOption> { |t: (&str, &str)| UsingOption::Range(t.1.to_string()))(input) } +fn parse_colnames_option(input: &str) -> IResult<&str, UsingOption> { + let option = tag_no_case("COLNAMES"); + + let value = preceded( + tag("'"), terminated(digit1, tag("'"))); + + map(separated_pair(option, multispace1, value), + |t: (&str, &str)| UsingOption::ColNames(t.1.to_string()))(input) +} + fn parse_with_spaces<'a, T>(parser: impl Parser<&'a str, T, nom::error::Error<&'a str>>) -> impl Parser<&'a str, T, nom::error::Error<&'a str>> { preceded(multispace0, terminated(parser, multispace0)) @@ -120,4 +132,17 @@ mod tests { _ => panic!("Expected range option") } } + + #[test] + fn parse_colnames_option_produces_colname() { + let (output, option) = parse_colnames_option("COLNAMES '152'").unwrap(); + + assert_eq!(output, ""); + match option { + UsingOption::ColNames(colname) => { + assert_eq!(colname, "152"); + }, + _ => panic!("Expected colnames option") + } + } } diff --git a/src/spreadsheet/manager.rs b/src/spreadsheet/manager.rs index 15a6c76..ce19387 100644 --- a/src/spreadsheet/manager.rs +++ b/src/spreadsheet/manager.rs @@ -7,11 +7,13 @@ use calamine::{open_workbook_auto, DataType, Range, Reader, Sheets}; use std::fs::File; use std::io::BufReader; use std::path::Path; +use std::str::FromStr; pub struct DataManager { sheets: Sheets>, worksheet: String, range: Option, + colnames_row: Option, } pub enum DataManagerError { @@ -49,9 +51,17 @@ impl DataManager { pub fn get_columns(&mut self) -> Vec { let range = self.get_effective_range(); if range.get_size().1 > 0 { + let row_workspace_sheet = self.colnames_row + .and_then(|v| Some((v, self.sheets.worksheet_range(self.worksheet.as_str())))) + .and_then(|(row, sheet)| Some((row, sheet?.ok()?))); (range.start().unwrap().1..=range.end().unwrap().1) .into_iter() - .map(|n| CellIndex::new(n + 1, 1).get_x_as_string()) + .map(|n| { + row_workspace_sheet + .as_ref() + .and_then(|(row, sheet)| sheet.get_value((*row, n)).map(|v| v.to_string())) + .unwrap_or_else(|| CellIndex::new(n + 1, 1).get_x_as_string()) + }) .collect() } else { Vec::new() @@ -70,6 +80,7 @@ pub struct DataManagerBuilder { file: Option, worksheet: Option, range: Option, + colnames_row: Option, } impl DataManagerBuilder { @@ -91,6 +102,11 @@ impl DataManagerBuilder { UsingOption::Range(range) => { builder = builder.range(CellRange::try_parse(range.as_str()).unwrap()); } + UsingOption::ColNames(colnames) => { + // We substract 1 to go from excel indexing (which starts at 1) to 0-based + // indexing of the row. + builder = builder.colnames_row(u32::from_str(colnames.as_str()).unwrap().saturating_sub(1)); + }, } } @@ -112,6 +128,11 @@ impl DataManagerBuilder { self } + pub fn colnames_row(mut self, row: u32) -> Self { + self.colnames_row = Some(row); + self + } + pub fn open(self) -> Result { if let Some(file) = self.file { if let Some(worksheet) = self.worksheet { @@ -120,6 +141,7 @@ impl DataManagerBuilder { sheets, worksheet, range: self.range, + colnames_row: self.colnames_row, }), Err(err) => Err(DataManagerError::Calamine(err)), } diff --git a/tests/abcdef_colnames.xlsx b/tests/abcdef_colnames.xlsx new file mode 100644 index 0000000..477c08c Binary files /dev/null and b/tests/abcdef_colnames.xlsx differ diff --git a/tests/lib.rs b/tests/lib.rs index c339b38..7149050 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -141,3 +141,40 @@ fn test_abcdef_file_aggregate() { AbcdAgg { kind: "odd".to_string(), count: 3 }, ]); } + +#[test] +fn test_abcdef_file_with_colnames() { + let connection = init_connection(); + connection.execute("\ + CREATE VIRTUAL TABLE test_data USING xlite(\ + FILENAME './tests/abcdef_colnames.xlsx',\ + WORKSHEET 'Sheet1',\ + RANGE 'A2:D7', + COLNAMES '1' + );\ + ", params![]).unwrap(); + + let mut query = connection.prepare("\ + SELECT alpha, number, word, kind FROM test_data;\ + ").unwrap(); + + let rows = query.query_map(params![], |row| Ok(Abcd { + alpha: row.get(0).unwrap(), + number: row.get(1).unwrap(), + word: row.get(2).unwrap(), + kind: row.get(3).unwrap(), + })).unwrap(); + + let data = rows + .map(|r| r.unwrap()) + .collect::>(); + + assert_eq!(data, vec![ + Abcd { alpha: "A".to_string(), number: 10.0, word: "ten".to_string(), kind: "even".to_string() }, + Abcd { alpha: "B".to_string(), number: 11.0, word: "eleven".to_string(), kind: "odd".to_string() }, + Abcd { alpha: "C".to_string(), number: 12.0, word: "twelve".to_string(), kind: "even".to_string() }, + Abcd { alpha: "D".to_string(), number: 13.0, word: "thirteen".to_string(), kind: "odd".to_string() }, + Abcd { alpha: "E".to_string(), number: 14.0, word: "fourteen".to_string(), kind: "even".to_string() }, + Abcd { alpha: "F".to_string(), number: 15.0, word: "fifteen".to_string(), kind: "odd".to_string() }, + ]); +}