1 files changed, 38 insertions, 17 deletions

diff --git a/hittekaart-py/src/lib.rs b/hittekaart-py/src/lib.rs
index 9f4bd8d..4b67777 100644
--- a/hittekaart-py/src/lib.rs
+++ b/hittekaart-py/src/lib.rs
@@ -7,6 +7,7 @@ use std::error::Error;
 use std::ffi::OsStr;
 use std::fmt::Write as _;
 use std::os::unix::ffi::OsStrExt as _;
+use std::path::PathBuf;
 
 create_exception!(hittekaart_py, HitteError, pyo3::exceptions::PyException);
 
@@ -61,20 +62,45 @@ impl Track {
     }
 }
 
-#[pyclass]
-struct Storage {
-    inner: Box<dyn hittekaart::storage::Storage + Send + Sync>,
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum StorageType {
+    Folder(PathBuf),
+    Sqlite(PathBuf),
 }
 
+#[pyclass]
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct Storage(StorageType);
+
 #[pymethods]
 impl Storage {
     #[staticmethod]
     #[pyo3(name = "Folder")]
     fn folder(path: &[u8]) -> Self {
         let path = OsStr::from_bytes(path);
-        let storage = hittekaart::storage::Folder::new(path.into());
-        Storage {
-            inner: Box::new(storage),
+        Storage(StorageType::Folder(path.into()))
+    }
+
+    #[staticmethod]
+    #[pyo3(name = "Sqlite")]
+    fn sqlite(path: &[u8]) -> Self {
+        let path = OsStr::from_bytes(path);
+        Storage(StorageType::Sqlite(path.into()))
+    }
+}
+
+impl Storage {
+    fn open(&self) -> PyResult<Box<dyn hittekaart::storage::Storage>> {
+        match self.0 {
+            StorageType::Folder(ref path) => {
+                let storage = hittekaart::storage::Folder::new(path.clone());
+                Ok(Box::new(storage))
+            }
+            StorageType::Sqlite(ref path) => {
+                let storage = hittekaart::storage::Sqlite::connect(path.clone())
+                    .map_err(|e| err_to_py(&e))?;
+                Ok(Box::new(storage))
+            }
         }
     }
 }
@@ -108,7 +134,7 @@ fn generate(
     }
 
     if let Ok(r) = renderer.downcast::<HeatmapRenderer>() {
-        do_generate(tracks, &r.borrow().inner, &mut storage.borrow_mut())
+        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
     } else {
         Err(PyTypeError::new_err("Expected a HeatmapRenderer"))
     }
@@ -117,28 +143,23 @@ fn generate(
 fn do_generate<R: Renderer>(
     tracks: Vec<Vec<Coordinates>>,
     renderer: &R,
-    storage: &mut Storage,
+    storage: &mut dyn hittekaart::storage::Storage,
 ) -> PyResult<()> {
-    storage.inner.prepare().map_err(|e| err_to_py(&e))?;
+    storage.prepare().map_err(|e| err_to_py(&e))?;
 
     for zoom in 0..=19 {
         let counter =
             renderer::prepare(renderer, zoom, &tracks, || Ok(())).map_err(|e| err_to_py(&e))?;
 
-        storage
-            .inner
-            .prepare_zoom(zoom)
-            .map_err(|e| err_to_py(&e))?;
+        storage.prepare_zoom(zoom).map_err(|e| err_to_py(&e))?;
 
         renderer::colorize(renderer, counter, |rendered_tile| {
-            storage
-                .inner
-                .store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
+            storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
             Ok(())
         })
         .map_err(|e| err_to_py(&e))?;
     }
-    storage.inner.finish().map_err(|e| err_to_py(&e))?;
+    storage.finish().map_err(|e| err_to_py(&e))?;
 
     Ok(())
 }