use crate::core::model::{Category, State, SubTask, Todo};
use crate::core::store::Store;
use anyhow::{Context, Result};
use arrow::array::{Array, BooleanArray, Int64Array, StringArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_writer::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
use std::path::PathBuf;
use std::sync::Arc;
pub struct TruenoStore {
path: PathBuf,
}
impl TruenoStore {
pub fn new(path: impl Into<PathBuf>) -> Self {
Self { path: path.into() }
}
fn get_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::UInt64, false),
Field::new("parent_id", DataType::UInt64, true),
Field::new("category", DataType::Utf8, false),
Field::new("title", DataType::Utf8, false),
Field::new("done", DataType::Boolean, false),
Field::new("notes", DataType::Utf8, true),
Field::new("created_at", DataType::Int64, false),
Field::new("updated_at", DataType::Int64, false),
Field::new("color", DataType::Utf8, false),
]))
}
fn flatten_state(state: &State) -> Vec<FlatItem> {
let mut flat = Vec::new();
for cat in State::all_categories() {
let todos = state.get_category(cat);
for todo in todos {
flat.push(FlatItem {
id: todo.id,
parent_id: None,
category: cat.as_str().to_string(),
title: todo.title.clone(),
done: cat == Category::Completed,
notes: todo.notes.clone(),
created_at: todo.created_at,
updated_at: todo.updated_at,
color: todo.color.as_str().to_string(),
});
Self::flatten_subs(&todo.subtasks, todo.id, cat.as_str(), &mut flat);
}
}
flat
}
fn flatten_subs(subs: &[SubTask], parent_id: u64, cat_str: &str, flat: &mut Vec<FlatItem>) {
for sub in subs {
flat.push(FlatItem {
id: sub.id,
parent_id: Some(parent_id),
category: cat_str.to_string(),
title: sub.title.clone(),
done: sub.done,
notes: sub.notes.clone(),
created_at: sub.created_at,
updated_at: sub.updated_at,
color: sub.color.as_str().to_string(),
});
Self::flatten_subs(&sub.subtasks, sub.id, cat_str, flat);
}
}
fn reconstruct_state(items: Vec<FlatItem>) -> State {
let mut state = State::new();
let mut max_id = 1;
use std::collections::HashMap;
let mut children_map: HashMap<Option<u64>, Vec<FlatItem>> = HashMap::new();
for item in items {
if item.id >= max_id { max_id = item.id + 1; }
children_map.entry(item.parent_id).or_default().push(item);
}
state.next_id = max_id;
if let Some(l1_items) = children_map.remove(&None) {
for item in l1_items {
let cat = Category::from_str(&item.category).unwrap_or(Category::Short);
let todo = Todo {
id: item.id,
title: item.title,
notes: item.notes,
created_at: item.created_at,
updated_at: item.updated_at,
color: crate::core::model::ColorTag::from_str(&item.color),
subtasks: Self::reconstruct_subs(item.id, &mut children_map),
};
state.get_category_mut(cat).push(todo);
}
}
state
}
fn reconstruct_subs(parent_id: u64, map: &mut std::collections::HashMap<Option<u64>, Vec<FlatItem>>) -> Vec<SubTask> {
let mut subs = Vec::new();
if let Some(children) = map.remove(&Some(parent_id)) {
for child in children {
subs.push(SubTask {
id: child.id,
title: child.title,
done: child.done,
created_at: child.created_at,
updated_at: child.updated_at,
color: crate::core::model::ColorTag::from_str(&child.color),
notes: child.notes,
subtasks: Self::reconstruct_subs(child.id, map),
});
}
}
subs
}
}
struct FlatItem {
id: u64,
parent_id: Option<u64>,
category: String,
title: String,
done: bool,
notes: Option<String>,
created_at: i64,
updated_at: i64,
color: String,
}
impl Store for TruenoStore {
fn load(&self) -> Result<State> {
if !self.path.exists() {
return Ok(State::new());
}
let file = File::open(&self.path)?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
let mut reader = builder.build()?;
let mut items = Vec::new();
while let Some(batch_res) = reader.next() {
let batch = batch_res?;
let ids = batch.column(0).as_any().downcast_ref::<UInt64Array>().context("id column missing")?;
let parent_ids = batch.column(1).as_any().downcast_ref::<UInt64Array>().context("parent_id column missing")?;
let categories = batch.column(2).as_any().downcast_ref::<StringArray>().context("category column missing")?;
let titles = batch.column(3).as_any().downcast_ref::<StringArray>().context("title column missing")?;
let dones = batch.column(4).as_any().downcast_ref::<BooleanArray>().context("done column missing")?;
let notes = batch.column(5).as_any().downcast_ref::<StringArray>().context("notes column missing")?;
let created_ats = batch.column(6).as_any().downcast_ref::<Int64Array>().context("created_at column missing")?;
let updated_ats = batch.column(7).as_any().downcast_ref::<Int64Array>().context("updated_at column missing")?;
let colors = batch.column(8).as_any().downcast_ref::<StringArray>().context("color column missing")?;
for i in 0..batch.num_rows() {
items.push(FlatItem {
id: ids.value(i),
parent_id: if parent_ids.is_null(i) { None } else { Some(parent_ids.value(i)) },
category: categories.value(i).to_string(),
title: titles.value(i).to_string(),
done: dones.value(i),
notes: if notes.is_null(i) { None } else { Some(notes.value(i).to_string()) },
created_at: created_ats.value(i),
updated_at: updated_ats.value(i),
color: colors.value(i).to_string(),
});
}
}
Ok(Self::reconstruct_state(items))
}
fn save(&self, state: &State) -> Result<()> {
let items = Self::flatten_state(state);
let schema = Self::get_schema();
let ids: Vec<u64> = items.iter().map(|i| i.id).collect();
let parent_ids: Vec<Option<u64>> = items.iter().map(|i| i.parent_id).collect();
let categories: Vec<String> = items.iter().map(|i| i.category.clone()).collect();
let titles: Vec<String> = items.iter().map(|i| i.title.clone()).collect();
let dones: Vec<bool> = items.iter().map(|i| i.done).collect();
let notes: Vec<Option<String>> = items.iter().map(|i| i.notes.clone()).collect();
let created_ats: Vec<i64> = items.iter().map(|i| i.created_at).collect();
let updated_ats: Vec<i64> = items.iter().map(|i| i.updated_at).collect();
let colors: Vec<String> = items.iter().map(|i| i.color.clone()).collect();
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(UInt64Array::from(ids)),
Arc::new(UInt64Array::from(parent_ids)),
Arc::new(StringArray::from(categories)),
Arc::new(StringArray::from(titles)),
Arc::new(BooleanArray::from(dones)),
Arc::new(StringArray::from(notes)),
Arc::new(Int64Array::from(created_ats)),
Arc::new(Int64Array::from(updated_ats)),
Arc::new(StringArray::from(colors)),
],
)?;
let file = File::create(&self.path)
.with_context(|| format!("Failed to create backup file at {:?}", self.path))?;
let mut writer = ArrowWriter::try_new(file, schema, None)?;
writer.write(&batch)?;
writer.close()?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::model::{ColorTag, SubTask, Todo};
#[test]
fn test_flatten_state() {
let mut state = State::new();
let sub = SubTask {
id: 2,
title: "Sub".to_string(),
done: false,
created_at: 100,
updated_at: 100,
color: ColorTag::Red,
notes: None,
subtasks: vec![],
};
let todo = Todo {
id: 1,
title: "Todo".to_string(),
notes: Some("Note".to_string()),
created_at: 50,
updated_at: 50,
color: ColorTag::Blue,
subtasks: vec![sub],
};
state.short.push(todo);
let flat = TruenoStore::flatten_state(&state);
assert_eq!(flat.len(), 2);
let f_todo = flat.iter().find(|i| i.id == 1).unwrap();
assert_eq!(f_todo.parent_id, None);
assert_eq!(f_todo.category, "Short");
let f_sub = flat.iter().find(|i| i.id == 2).unwrap();
assert_eq!(f_sub.parent_id, Some(1));
assert_eq!(f_sub.category, "Short");
assert_eq!(f_sub.color, "red");
}
}