Merge pull request #11 from roblabla/colnames

Add COLNAMES option to give names to the columns

Author: x2bool

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,13 +9,15 @@ pub enum UsingOption {
     File(String),
     Worksheet(String),
     Range(String),
+    ColNames(String),
 }
 
 pub fn parse_option(input: &str) -> IResult<&str, UsingOption> {
     parse_with_spaces(alt((
         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")
+        }
+    }
 }

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<BufReader<File>>,
     worksheet: String,
     range: Option<CellRange>,
+    colnames_row: Option<u32>,
 }
 
 pub enum DataManagerError {
@@ -49,9 +51,17 @@ impl DataManager {
     pub fn get_columns(&mut self) -> Vec<String> {
         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<String>,
     worksheet: Option<String>,
     range: Option<CellRange>,
+    colnames_row: Option<u32>,
 }
 
 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<DataManager, DataManagerError> {
         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)),
                 }

tests/abcdef_colnames.xlsx

@@ -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::<Vec<Abcd>>();
+
+    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() },
+    ]);
+}