bbox_tile_server/store/
pmtiles.rs

1use crate::config::PmtilesStoreCfg;
2use crate::store::{TileReader, TileStoreError, TileWriter};
3use async_trait::async_trait;
4use bbox_core::{Compression, Format, TileResponse};
5use log::{debug, info};
6use martin_mbtiles::Metadata;
7use pmtiles::async_reader::AsyncPmTilesReader;
8use pmtiles::mmap::MmapBackend;
9use pmtiles2::{util::tile_id, Compression as PmCompression, PMTiles, TileType};
10use serde_json::json;
11use std::ffi::OsStr;
12use std::fs::File;
13use std::io::Cursor;
14use std::path::{Path, PathBuf};
15use tile_grid::Xyz;
16
17pub struct PmtilesStoreReader {
18    pub path: PathBuf,
19    reader: AsyncPmTilesReader<MmapBackend>,
20}
21
22#[derive(Debug)]
23pub struct PmtilesStoreWriter {
24    path: PathBuf,
25    format: Format,
26    metadata: Metadata,
27    // We need an option for consuming PMTiles when finalizing
28    archive: Option<PMTiles<Cursor<Vec<u8>>>>,
29}
30
31// Custom impl because `Clone` is not implemented for `AsyncPmTilesReader`
32impl Clone for PmtilesStoreReader {
33    fn clone(&self) -> Self {
34        futures::executor::block_on(async {
35            Self::create_reader(self.path.clone()).await.expect("clone")
36        })
37    }
38}
39
40// Custom impl because `Clone` is not implemented for `PMTiles`
41impl Clone for PmtilesStoreWriter {
42    fn clone(&self) -> Self {
43        Self::new(self.path.clone(), self.metadata.clone(), &self.format)
44    }
45}
46
47impl PmtilesStoreReader {
48    pub async fn create_reader(path: PathBuf) -> Result<Self, TileStoreError> {
49        let reader = AsyncPmTilesReader::new_with_path(&path).await?;
50        Ok(Self { path, reader })
51    }
52    pub async fn from_config(cfg: &PmtilesStoreCfg) -> Result<Self, TileStoreError> {
53        Self::create_reader(cfg.abs_path()).await
54    }
55    pub fn config_from_cli_arg(file_or_url: &str) -> Option<PmtilesStoreCfg> {
56        match Path::new(file_or_url).extension().and_then(OsStr::to_str) {
57            Some("pmtiles") => {
58                let cfg = PmtilesStoreCfg {
59                    path: file_or_url.into(),
60                };
61                Some(cfg)
62            }
63            _ => None,
64        }
65    }
66    pub async fn get_metadata(&self) -> Result<String, ::pmtiles::error::Error> {
67        self.reader.get_metadata().await
68    }
69}
70
71#[async_trait]
72impl TileReader for PmtilesStoreReader {
73    async fn get_tile(&self, xyz: &Xyz) -> Result<Option<TileResponse>, TileStoreError> {
74        let resp = if let Some(tile) = self.reader.get_tile(xyz.z, xyz.x, xyz.y).await {
75            let mut response = TileResponse::new();
76            response.set_content_type(tile.tile_type.content_type());
77            if let Some(encoding) = tile.tile_compression.content_encoding() {
78                response.insert_header(("Content-Encoding", encoding.to_lowercase()));
79            }
80            Some(response.with_body(Box::new(Cursor::new(tile.data))))
81        } else {
82            None
83        };
84        Ok(resp)
85    }
86}
87
88#[async_trait]
89impl TileWriter for PmtilesStoreWriter {
90    fn compression(&self) -> Compression {
91        match self.archive.as_ref().expect("initialized").tile_compression {
92            PmCompression::GZip => Compression::Gzip,
93            _ => Compression::None,
94        }
95    }
96    async fn exists(&self, _xyz: &Xyz) -> bool {
97        // self.archive.get_tile(xyz.x, xyz.y, xyz.z)?.is_some()
98        false
99    }
100    async fn put_tile(&self, _xyz: &Xyz, _data: Vec<u8>) -> Result<(), TileStoreError> {
101        Err(TileStoreError::ReadOnly)
102    }
103    async fn put_tile_mut(&mut self, xyz: &Xyz, data: Vec<u8>) -> Result<(), TileStoreError> {
104        self.archive
105            .as_mut()
106            .expect("initialized")
107            .add_tile(tile_id(xyz.z, xyz.x, xyz.y), data);
108        debug!(
109            "put_tile_mut - num_tiles: {}",
110            self.archive.as_ref().expect("initialized").num_tiles()
111        );
112        Ok(())
113    }
114    fn finalize(&mut self) -> Result<(), TileStoreError> {
115        info!("Writing {}", self.path.display());
116        let mut file = File::create(&self.path)
117            .map_err(|e| TileStoreError::FileError(self.path.clone(), e))?;
118        if let Some(archive) = self.archive.take() {
119            info!("Number of tiles: {}", archive.num_tiles(),);
120            archive
121                .to_writer(&mut file)
122                .map_err(|e| TileStoreError::FileError(self.path.clone(), e))?;
123        }
124        Ok(())
125    }
126}
127
128impl PmtilesStoreWriter {
129    pub fn new(path: PathBuf, metadata: Metadata, format: &Format) -> Self {
130        let mut archive = PMTiles::default();
131        archive.tile_type = match format {
132            Format::Jpeg => TileType::Jpeg,
133            Format::Mvt => TileType::Mvt,
134            Format::Png => TileType::Png,
135            Format::Webp => TileType::WebP,
136            _ => TileType::Unknown,
137        };
138        archive.tile_compression = if *format == Format::Mvt {
139            PmCompression::GZip
140        } else {
141            PmCompression::None
142        };
143        if let Some(minzoom) = metadata.tilejson.minzoom {
144            archive.min_zoom = minzoom;
145        }
146        if let Some(maxzoom) = metadata.tilejson.maxzoom {
147            archive.max_zoom = maxzoom;
148        }
149        if let Some(bounds) = metadata.tilejson.bounds {
150            archive.min_longitude = bounds.left;
151            archive.min_latitude = bounds.bottom;
152            archive.max_longitude = bounds.right;
153            archive.max_latitude = bounds.top;
154        }
155        if let Some(center) = metadata.tilejson.center {
156            archive.center_longitude = center.longitude;
157            archive.center_latitude = center.latitude;
158            archive.center_zoom = center.zoom;
159        }
160        let mut meta_data = json!({
161            "name": &metadata.id, "description": &metadata.tilejson.description, "attribution": &metadata.tilejson.attribution
162        });
163        if let Some(vector_layers) = &metadata.tilejson.vector_layers {
164            meta_data["vector_layers"] = json!(vector_layers);
165        }
166        archive.meta_data = Some(meta_data);
167        Self {
168            path,
169            metadata,
170            format: *format,
171            archive: Some(archive),
172        }
173    }
174    pub fn from_config(cfg: &PmtilesStoreCfg, metadata: Metadata, format: &Format) -> Self {
175        Self::new(cfg.abs_path(), metadata, format)
176    }
177}