add support for SQLite output

3 files changed, 73 insertions, 7 deletions

diff --git a/src/main.rs b/src/main.rs
index 7230dce..e5b835e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,7 +8,7 @@ use color_eyre::{
 use hittekaart::{
     gpx::{self, Compression},
     renderer,
-    storage::{Folder, Storage},
+    storage::{Folder, Sqlite, Storage},
 };
 use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
 use is_terminal::IsTerminal;
@@ -36,9 +36,15 @@ struct Args {
     #[arg(long, short, default_value_t = 0)]
     threads: usize,
 
-    /// The output directory. Will be created if it does not exist.
-    #[arg(long, short, default_value = "tiles")]
-    output_directory: PathBuf,
+    /// The output directory. Will be created if it does not exist. Defaults to "tiles" for the
+    /// folder-based storage, and "tiles.sqlite" for the SQLite-based storage.
+    #[arg(long, short)]
+    output: Option<PathBuf>,
+
+    /// Store the tiles in a SQLite database. If given, `--output` will determine the SQLite
+    /// filename.
+    #[arg(long)]
+    sqlite: bool,
 }
 
 fn main() -> Result<()> {
@@ -85,7 +91,13 @@ fn main() -> Result<()> {
     let tracks = tracks.into_iter().collect::<Result<Vec<_>>>()?;
     bar.finish();
 
-    let mut storage = Folder::new(args.output_directory.clone());
+    let mut storage: Box<dyn Storage + Send> = if args.sqlite {
+        let output = args.output.unwrap_or_else(|| "tiles.sqlite".into());
+        Box::new(Sqlite::connect(output)?)
+    } else {
+        let output = args.output.unwrap_or_else(|| "tiles".into());
+        Box::new(Folder::new(output))
+    };
     storage.prepare()?;
 
     let multibar = MultiProgress::new();
diff --git a/src/renderer.rs b/src/renderer.rs
index 4840e73..74a321c 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -139,7 +139,7 @@ fn colorize_tile(tile: &ImageBuffer<Luma<u8>, Vec<u8>>, max: u32) -> RgbaImage {
 /// rendering the next one.
 ///
 /// This has a way lower memory usage than [`colorize_heatcounter`].
-pub fn lazy_colorization<F: FnMut(RenderedTile) -> Result<()> + Send + Sync>(
+pub fn lazy_colorization<F: FnMut(RenderedTile) -> Result<()> + Send>(
     layer: HeatCounter,
     mut save_callback: F,
 ) -> Result<()> {
diff --git a/src/storage.rs b/src/storage.rs
index fef8194..51a418e 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -1,4 +1,8 @@
-use color_eyre::{eyre::{bail, WrapErr}, Result};
+use color_eyre::{
+    eyre::{bail, WrapErr},
+    Result,
+};
+use rusqlite::{params, Connection};
 use std::{
     fs,
     io::ErrorKind,
@@ -65,3 +69,53 @@ impl Storage for Folder {
         Ok(())
     }
 }
+
+#[derive(Debug)]
+pub struct Sqlite {
+    connection: Connection,
+}
+
+impl Sqlite {
+    pub fn connect<P: AsRef<Path>>(file: P) -> Result<Self> {
+        let path = file.as_ref();
+        if fs::metadata(path).is_ok() {
+            bail!("Path {path:?} already exists, refusing to open")
+        }
+        let connection = Connection::open(path)?;
+        Ok(Sqlite { connection })
+    }
+}
+
+impl Storage for Sqlite {
+    fn prepare(&mut self) -> Result<()> {
+        self.connection.execute(
+            "CREATE TABLE tiles (
+            zoom INTEGER,
+            x INTEGER,
+            y INTEGER,
+            data BLOB,
+            PRIMARY KEY (zoom, x, y)
+        );",
+            (),
+        )?;
+        self.connection.execute("BEGIN;", ())?;
+        Ok(())
+    }
+
+    fn prepare_zoom(&mut self, _zoom: u32) -> Result<()> {
+        Ok(())
+    }
+
+    fn store(&mut self, zoom: u32, x: u64, y: u64, data: &[u8]) -> Result<()> {
+        self.connection.execute(
+            "INSERT INTO tiles (zoom, x, y, data) VALUES (?, ?, ?, ?)",
+            params![zoom, x, y, data],
+        )?;
+        Ok(())
+    }
+
+    fn finish(&mut self) -> Result<()> {
+        self.connection.execute("COMMIT;", ())?;
+        Ok(())
+    }
+}