feat(mbtiles): add Mbtiles::stream_coords and Mbtiles::stream_tiles

martin-tile-utils/src/lib.rs

@@ -23,6 +23,9 @@ pub struct TileCoord {
     pub y: u32,
 }
 
+pub type TileData = Vec<u8>;
+pub type Tile = (TileCoord, Option<TileData>);
+
 impl Display for TileCoord {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         if f.alternate() {
@@ -33,6 +36,36 @@ impl Display for TileCoord {
     }
 }
 
+impl TileCoord {
+    /// Checks provided coordinates for validity
+    /// before constructing [`TileCoord`] instance.
+    ///
+    /// Check [`Self::new_unchecked`] if you are sure that your inputs are possible.
+    #[must_use]
+    pub fn new_checked(z: u8, x: u32, y: u32) -> Option<TileCoord> {
+        Self::is_possible_on_zoom_level(z, x, y).then_some(Self { z, x, y })
+    }
+
+    /// Constructs [`TileCoord`] instance from arguments without checking that the tiles can exist.
+    ///
+    /// Check [`Self::new_checked`] if you are unsure if your inputs are possible.
+    #[must_use]
+    pub fn new_unchecked(z: u8, x: u32, y: u32) -> TileCoord {
+        Self { z, x, y }
+    }
+
+    /// Checks that zoom `z` is plausibily small and `x`/`y` is possible on said zoom level
+    #[must_use]
+    pub fn is_possible_on_zoom_level(z: u8, x: u32, y: u32) -> bool {
+        if z > MAX_ZOOM {
+            return false;
+        }
+
+        let side_len = 1_u32 << z;
+        x < side_len && y < side_len
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Format {
     Gif,
@@ -484,4 +517,50 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_tile_coord_zoom_range() {
+        for z in 0..=MAX_ZOOM {
+            assert!(TileCoord::is_possible_on_zoom_level(z, 0, 0));
+            assert_eq!(
+                TileCoord::new_checked(z, 0, 0),
+                Some(TileCoord { z, x: 0, y: 0 })
+            );
+        }
+        assert!(!TileCoord::is_possible_on_zoom_level(MAX_ZOOM + 1, 0, 0));
+        assert_eq!(TileCoord::new_checked(MAX_ZOOM + 1, 0, 0), None);
+    }
+
+    #[test]
+    fn test_tile_coord_new_checked_xy_for_zoom() {
+        assert!(TileCoord::is_possible_on_zoom_level(5, 0, 0));
+        assert_eq!(
+            TileCoord::new_checked(5, 0, 0),
+            Some(TileCoord { z: 5, x: 0, y: 0 })
+        );
+        assert!(TileCoord::is_possible_on_zoom_level(5, 31, 31));
+        assert_eq!(
+            TileCoord::new_checked(5, 31, 31),
+            Some(TileCoord { z: 5, x: 31, y: 31 })
+        );
+        assert!(!TileCoord::is_possible_on_zoom_level(5, 31, 32));
+        assert_eq!(TileCoord::new_checked(5, 31, 32), None);
+        assert!(!TileCoord::is_possible_on_zoom_level(5, 32, 31));
+        assert_eq!(TileCoord::new_checked(5, 32, 31), None);
+    }
+
+    #[test]
+    /// Any (u8, u32, u32) values can be put inside [`TileCoord`], of course, but some
+    /// functions may panic at runtime (e.g. [`mbtiles::invert_y_value`]) if they are impossible,
+    /// so let's not do that.
+    fn test_tile_coord_new_unchecked() {
+        assert_eq!(
+            TileCoord::new_unchecked(u8::MAX, u32::MAX, u32::MAX),
+            TileCoord {
+                z: u8::MAX,
+                x: u32::MAX,
+                y: u32::MAX
+            }
+        );
+    }
 }

mbtiles/src/mbtiles.rs

@@ -1,9 +1,12 @@
 use std::ffi::OsStr;
 use std::fmt::{Display, Formatter};
 use std::path::Path;
+use std::pin::Pin;
 
 use enum_display::EnumDisplay;
+use futures::Stream;
 use log::debug;
+use martin_tile_utils::{Tile, TileCoord};
 use serde::{Deserialize, Serialize};
 use sqlite_compressions::{register_bsdiffraw_functions, register_gzip_functions};
 use sqlite_hashes::register_md5_functions;
@@ -131,6 +134,87 @@ impl Mbtiles {
         Ok(())
     }
 
+    /// Stream over coordinates of all tiles in the database.
+    ///
+    /// No particular order is guaranteed.
+    ///
+    /// Note that returned [Stream] holds a mutable reference to the given
+    /// connection, making it unusable for anything else until the stream
+    /// is dropped.
+    pub fn stream_coords<'e, T>(
+        &self,
+        conn: &'e mut T,
+    ) -> Pin<Box<dyn Stream<Item = MbtResult<TileCoord>> + Send + 'e>>
+    where
+        &'e mut T: SqliteExecutor<'e>,
+    {
+        use futures::StreamExt;
+
+        let query = query! {"SELECT zoom_level, tile_column, tile_row FROM tiles"};
+        let stream = query.fetch(conn);
+
+        // We only need `&self` for `self.filepath`, which in turn we only
+        // need to create proper `MbtError::InvalidTileIndex`es.
+        // Cloning the filepath allows us to drop [Mbtiles] instance while returned
+        // stream is still alive.
+        let filepath = self.filepath.clone();
+
+        Box::pin(stream.map(move |result| {
+            result.map_err(MbtError::from).and_then(|row| {
+                let z = row.zoom_level;
+                let x = row.tile_column;
+                let y = row.tile_row;
+                let coord = parse_tile_index(z, x, y).ok_or_else(|| {
+                    MbtError::InvalidTileIndex(
+                        filepath.to_string(),
+                        format!("{z:?}"),
+                        format!("{x:?}"),
+                        format!("{y:?}"),
+                    )
+                })?;
+                Ok(coord)
+            })
+        }))
+    }
+
+    /// Returns a stream over all tiles in the database.
+    ///
+    /// No particular order is guaranteed.
+    ///
+    /// Note that returned [Stream] holds a mutable reference to the given
+    /// connection, making it unusable for anything else until the stream
+    /// is dropped.
+    pub fn stream_tiles<'e, T>(
+        &self,
+        conn: &'e mut T,
+    ) -> Pin<Box<dyn Stream<Item = MbtResult<Tile>> + Send + 'e>>
+    where
+        &'e mut T: SqliteExecutor<'e>,
+    {
+        use futures::StreamExt;
+
+        let query = query! {"SELECT zoom_level, tile_column, tile_row, tile_data FROM tiles"};
+        let stream = query.fetch(conn);
+        let filepath = self.filepath.clone();
+
+        Box::pin(stream.map(move |result| {
+            result.map_err(MbtError::from).and_then(|row| {
+                let z = row.zoom_level;
+                let x = row.tile_column;
+                let y = row.tile_row;
+                let coord = parse_tile_index(z, x, y).ok_or_else(|| {
+                    MbtError::InvalidTileIndex(
+                        filepath.to_string(),
+                        format!("{z:?}"),
+                        format!("{x:?}"),
+                        format!("{y:?}"),
+                    )
+                })?;
+                Ok((coord, row.tile_data))
+            })
+        }))
+    }
+
     /// Get a tile from the database
     ///
     /// See [`Mbtiles::get_tile_and_hash`] if you do need the hash
@@ -283,6 +367,17 @@ pub async fn attach_sqlite_fn(conn: &mut SqliteConnection) -> MbtResult<()> {
     Ok(())
 }
 
+fn parse_tile_index(z: Option<i64>, x: Option<i64>, y: Option<i64>) -> Option<TileCoord> {
+    let z: u8 = z?.try_into().ok()?;
+    let x: u32 = x?.try_into().ok()?;
+    let y: u32 = y?.try_into().ok()?;
+
+    // Inverting `y` value can panic if it is greater than `(1 << z) - 1`,
+    // so we must ensure that it is vald first.
+    TileCoord::is_possible_on_zoom_level(z, x, y)
+        .then(|| TileCoord::new_unchecked(z, x, invert_y_value(z, y)))
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;

mbtiles/tests/streams.rs

@@ -0,0 +1,106 @@
+use futures::{StreamExt, TryStreamExt};
+use martin_tile_utils::{Tile, TileCoord};
+use mbtiles::{Mbtiles, create_metadata_table};
+use sqlx::{Executor as _, SqliteConnection, query};
+
+fn coord_key(coord: &TileCoord) -> (u8, u32, u32) {
+    let TileCoord { z, x, y } = *coord;
+    (z, x, y)
+}
+
+fn tile_key(tile: &Tile) -> (u8, u32, u32) {
+    coord_key(&tile.0)
+}
+
+async fn new(rows: &[&str]) -> (Mbtiles, SqliteConnection) {
+    let mbtiles = Mbtiles::new(":memory:").unwrap();
+    let mut conn = mbtiles.open().await.unwrap();
+    create_metadata_table(&mut conn).await.unwrap();
+
+    conn.execute(
+        "CREATE TABLE tiles (
+             zoom_level integer,
+             tile_column integer,
+             tile_row integer,
+             tile_data blob,
+             PRIMARY KEY(zoom_level, tile_column, tile_row));",
+    )
+    .await
+    .unwrap();
+
+    for row in rows {
+        let sql = format!(
+            "INSERT INTO tiles (zoom_level, tile_column, tile_row, tile_data)
+            VALUES ({row});"
+        );
+        query(&sql).execute(&mut conn).await.expect(&sql);
+    }
+
+    (mbtiles, conn)
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn mbtiles_stream_tiles() {
+    let (mbtiles, mut conn) = new(&[
+        // Note that `y`-coordinates are inverted.
+        "1, 0, 1, CAST('tl' AS BLOB)",
+        "1, 1, 0, CAST('br' AS BLOB)",
+        "2, 0, 0, NULL",
+    ])
+    .await;
+
+    {
+        let mut coords: Vec<TileCoord> = mbtiles
+            .stream_coords(&mut conn)
+            .try_collect()
+            .await
+            .expect("Failed to colled tile coords");
+
+        // Iteration order is not guaranteed.
+        coords.sort_by_key(coord_key);
+
+        assert_eq!(
+            coords,
+            [
+                TileCoord { z: 1, x: 0, y: 0 },
+                TileCoord { z: 1, x: 1, y: 1 },
+                TileCoord { z: 2, x: 0, y: 3 },
+            ]
+        );
+    }
+
+    {
+        let mut tiles: Vec<Tile> = mbtiles
+            .stream_tiles(&mut conn)
+            .try_collect()
+            .await
+            .expect("Failed to collect tiles");
+
+        tiles.sort_by_key(tile_key);
+
+        assert_eq!(
+            tiles,
+            [
+                (TileCoord { z: 1, x: 0, y: 0 }, Some(b"tl".to_vec())),
+                (TileCoord { z: 1, x: 1, y: 1 }, Some(b"br".to_vec())),
+                (TileCoord { z: 2, x: 0, y: 3 }, None),
+            ]
+        );
+    }
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn mbtiles_stream_errors() {
+    let (mbtiles, mut conn) = new(&[
+        // Note that `y`-coordinates are inverted.
+        // `4` is an invalid value for `x` at `z = 2`. Valid range is `0..=3`.
+        "2, 4, 0, NULL",
+    ])
+    .await;
+
+    let mut stream = mbtiles.stream_coords(&mut conn);
+    match stream.next().await {
+        Some(Err(mbtiles::MbtError::InvalidTileIndex(_filename, _z, _x, _y))) => {}
+        _ => panic!("Unexpected value returned from stream!"),
+    }
+}