Select column value

src/args.rs

@@ -1,29 +1,55 @@
-use anyhow::{Context, Error, Result};
+use anyhow::{bail, Context, Error, Result};
 use regex::Regex;
 use std::env;
 
 #[derive(Debug)]
 pub enum Command {
     Info,
     Tables,
-    Count(String),
+    Select(Select),
 }
 
-impl TryFrom<String> for Command {
+#[derive(Debug)]
+pub enum Select {
+    Count { table: String },
+    Column { table: String, column: String },
+}
+
+impl TryFrom<String> for Select {
     type Error = Error;
 
     fn try_from(value: String) -> Result<Self> {
-        let rg_select = Regex::new(r"select count\(\*\) from (?P<table_name>[A-Za-z]+)")?;
+        let rg_count = Regex::new(r"select count\(\*\) from (?P<table>[A-Za-z]+)")?;
+        let rg_col = Regex::new(r"select (?P<column>[A-Za-z]+) from (?P<table>[A-Za-z]+)")?;
+        match value.as_str() {
+            s if rg_count.is_match(s) => {
+                let caps = rg_count.captures(s).context("select count regex")?;
+                let table = (&caps["table"]).to_string();
+                Ok(Select::Count { table })
+            }
+            s if rg_col.is_match(s) => {
+                let caps = rg_col.captures(s).context("select col regex")?;
+                let table = (&caps["table"]).to_string();
+                let column = (&caps["column"]).to_string();
+                Ok(Select::Column { table, column })
+            }
+            e => bail!("Not supported: {e}"),
+        }
+    }
+}
 
+impl TryFrom<String> for Command {
+    type Error = Error;
+
+    fn try_from(value: String) -> Result<Self> {
         match value.as_str() {
             ".dbinfo" => Ok(Command::Info),
             ".tables" => Ok(Command::Tables),
-            count if (rg_select.is_match(count)) => {
-                let caps = rg_select.captures(count).context("regex")?;
-                let name = (&caps["table_name"]).to_string();
-                Ok(Command::Count(name))
+            s if s.starts_with("select") => {
+                let sel = Select::try_from(value)?;
+                Ok(Command::Select(sel))
             }
-            e => Err(anyhow::anyhow!("Not a command: {e}")),
+            e => bail!("Not a command: {e}"),
         }
     }
 }
@@ -37,6 +63,7 @@ pub struct Args {
 pub fn build() -> Result<Args> {
     let mut args = env::args().skip(1);
     let filename = args.next().context("Missing filename")?;
-    let cmds = args.flat_map(Command::try_from).collect();
+    let cmds = args.map(Command::try_from).collect::<Result<_>>();
+    let cmds = cmds?;
     Ok(Args { filename, cmds })
 }

src/codec.rs

@@ -85,3 +85,52 @@ pub mod float {
         Ok(f64::from_be_bytes(buf))
     }
 }
+
+pub mod sql {
+    use nom::bytes::complete::{tag, take_until, take_while};
+    use nom::character::complete::{char, multispace0};
+    use nom::multi::separated_list0;
+    use nom::sequence::preceded;
+    use nom::{IResult, Parser};
+
+    fn extract_name(io: &str) -> IResult<&str, &str> {
+        take_while(|c| c != ',')
+            .and_then(preceded(multispace0, take_while(char::is_alphanumeric)))
+            .parse(io)
+    }
+
+    fn parse_between(io: &str) -> IResult<&str, Vec<&str>> {
+        let mut drop_statement = take_until("(");
+        let (io, _) = drop_statement.parse(io)?;
+
+        preceded(tag("("), take_until(")"))
+            .and_then(separated_list0(char(','), extract_name))
+            .parse(io)
+    }
+
+    pub fn column_names(io: &str) -> Vec<&str> {
+        parse_between(io).map(|r| r.1).unwrap_or_default()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    const SQL_CREATE: &str = "CREATE TABLE butterscotch (id integer primary key, grape text,chocolate text,coconut text,coffee text,butterscotch text)";
+
+    #[test]
+    fn test_sql_column_names() {
+        let names = sql::column_names(SQL_CREATE);
+        assert_eq!(
+            &names,
+            &[
+                "id",
+                "grape",
+                "chocolate",
+                "coconut",
+                "coffee",
+                "butterscotch"
+            ]
+        );
+    }
+}

src/file.rs

@@ -1,20 +1,19 @@
 mod header;
 
-use crate::offset::Offset;
-use crate::page::{self, Cell, Page, TableLeafCell};
-use anyhow::{Context, Result};
+use crate::schema::Schema;
+use crate::{
+    offset::Offset,
+    page::{self, Page},
+};
+use anyhow::Result;
 use header::Header;
-use itertools::Itertools;
-use std::io::Read;
 use std::{
     fs::File,
-    io::{BufReader, Seek},
+    io::{BufReader, Read, Seek},
 };
 
 pub trait SQL {
-    fn schema(&mut self) -> Result<Page>;
-    fn tables(&mut self) -> Result<Vec<TableLeafCell>>;
-    fn table_named(&mut self, name: &str) -> Result<Page>;
+    fn schema(&mut self) -> Result<Schema>;
     fn page_at(&mut self, idx: u64) -> Result<Page>;
 }
 
@@ -33,27 +32,8 @@ impl SQLiteFile {
 }
 
 impl SQL for SQLiteFile {
-    fn schema(&mut self) -> Result<Page> {
-        parser::read_page(self, 0u64.into(), Header::size())
-    }
-
-    fn tables(&mut self) -> Result<Vec<TableLeafCell>> {
-        let schema = self.schema()?;
-        let cells = schema.cells();
-        Ok(cells.iter().filter(|c| c.is_table()).cloned().collect_vec())
-    }
-
-    fn table_named(&mut self, name: &str) -> Result<Page> {
-        let schema = self.schema()?;
-        let cells = schema.cells();
-        let cell = cells
-            .iter()
-            .filter(|c| c.is_table())
-            .find(|c| c.record.values.get(1).map(|s| s.as_str()).flatten() == Some(name))
-            .context("table cell")?;
-
-        let page_id = cell.record.values[3].as_int().unwrap() as u64 - 1;
-        self.page_at(page_id)
+    fn schema(&mut self) -> Result<Schema> {
+        parser::read_page(self, 0u64.into(), Header::size()).and_then(Schema::try_from)
     }
 
     fn page_at(&mut self, idx: u64) -> Result<Page> {

src/file/header.rs

@@ -4,7 +4,6 @@ const HEADER_SIZE: usize = 100;
 
 pub struct Header([u8; HEADER_SIZE]);
 
-#[allow(dead_code)]
 impl Header {
     pub const fn size() -> usize {
         HEADER_SIZE
@@ -25,6 +24,7 @@ impl Header {
         }
     }
 
+    #[allow(dead_code)]
     /// Reserved chunk at the end of each page.
     pub fn reserved_page_size(&self) -> u8 {
         self.0[20]

src/main.rs

@@ -3,11 +3,11 @@ mod codec;
 mod file;
 mod offset;
 mod page;
+mod schema;
 mod value;
 
-use crate::page::Cell;
-use anyhow::Result;
-use args::Command;
+use anyhow::{Context, Result};
+use args::{Command, Select};
 use file::{SQLiteFile, SQL};
 use itertools::Itertools;
 
@@ -24,18 +24,39 @@ fn main() -> Result<()> {
                 println!("number of tables: {}", file.schema()?.head.cell_count);
             }
             Command::Tables => {
-                let tables = file.tables()?;
+                let schema = file.schema()?;
+                let tables = schema.tables();
                 let msg = tables
-                    .iter()
-                    .filter(|c| !c.is_internal())
-                    .flat_map(|t| t.record.values[2].as_str())
+                    .filter(|c| !c.internal)
+                    .map(|t| t.name.as_str())
                     .join(" ");
                 println!("{msg}");
             }
-            Command::Count(name) => {
-                let page = file.table_named(&name)?;
+            Command::Select(Select::Count { table }) => {
+                let schema = file.schema()?;
+                let table = schema.table_named(&table)?;
+                let page = file.page_at(table.id)?;
                 println!("{}", page.head.cell_count);
             }
+            Command::Select(Select::Column { table, column }) => {
+                let schema = file.schema()?;
+                let table = schema.table_named(&table)?;
+                let page = file.page_at(table.id)?;
+                let cols = table.column_names();
+                let col_idx = cols
+                    .iter()
+                    .find_position(|c| *c == &column.as_str())
+                    .context("column not found")?
+                    .0;
+
+                let cells = page.cells();
+                let names = cells
+                    .iter()
+                    .flat_map(|cell| cell.record.values.get(col_idx))
+                    .map(|v| v.to_string());
+
+                names.for_each(|name| println!("{name}"));
+            }
         }
     }
 

src/page.rs

@@ -4,8 +4,8 @@ mod kind;
 
 use crate::offset::{self, Offset};
 use anyhow::{anyhow, Result};
-pub use cell::{Cell, TableLeafCell};
-use header::Header;
+pub use cell::TableLeafCell;
+pub use header::Header;
 
 #[derive(Debug)]
 pub struct Page {

src/page/cell.rs

@@ -4,38 +4,13 @@ use crate::{
     Result,
 };
 
-const CELL_TYPE_TABLE: &str = "table";
-const NAME_PREFIX_SQLITE: &str = "sqlite_";
-
-pub trait Cell {
-    fn is_table(&self) -> bool;
-    fn is_internal(&self) -> bool;
-}
-
-#[allow(dead_code)]
 #[derive(Debug, Clone)]
 pub struct TableLeafCell {
     pub id: u64,
     pub len: u64,
     pub record: Record,
 }
 
-impl Cell for TableLeafCell {
-    fn is_table(&self) -> bool {
-        self.record.values.first().map(|v| v.as_str()).flatten() == Some(CELL_TYPE_TABLE)
-    }
-
-    fn is_internal(&self) -> bool {
-        self.record
-            .values
-            .get(1)
-            .map(|s| s.as_str())
-            .flatten()
-            .map(|s| s.starts_with(NAME_PREFIX_SQLITE))
-            .unwrap_or(false)
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct Record {
     pub types: Vec<Type>,