1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use crate::config::MbtilesStoreCfg;
use crate::mbtiles_ds::{Error as MbtilesDsError, MbtilesDatasource};
use crate::store::{TileReader, TileStoreError, TileWriter};
use async_trait::async_trait;
use bbox_core::{Compression, TileResponse};
use log::info;
use martin_mbtiles::{CopyDuplicateMode, Metadata};
use martin_tile_utils::{Encoding as TileEncoding, Format as TileFormat};
use std::ffi::OsStr;
use std::io::Cursor;
use std::path::Path;
use tile_grid::Xyz;

#[derive(Clone, Debug)]
pub struct MbtilesStore {
    pub(crate) mbt: MbtilesDatasource,
}

impl MbtilesStore {
    pub async fn from_config(cfg: &MbtilesStoreCfg) -> Result<Self, MbtilesDsError> {
        info!("Creating connection pool for {}", &cfg.abs_path().display());
        let mbt = MbtilesDatasource::from_config(cfg, None).await?;
        //let opt = SqliteConnectOptions::new().filename(file).read_only(true);
        Ok(MbtilesStore { mbt })
    }
    pub async fn from_config_writable(
        cfg: &MbtilesStoreCfg,
        metadata: Metadata,
    ) -> Result<Self, MbtilesDsError> {
        info!("Creating connection pool for {}", &cfg.abs_path().display());
        let mbt = MbtilesDatasource::from_config(cfg, Some(metadata)).await?;
        Ok(MbtilesStore { mbt })
    }
    pub fn config_from_cli_arg(file_or_url: &str) -> Option<MbtilesStoreCfg> {
        match Path::new(file_or_url).extension().and_then(OsStr::to_str) {
            Some("mbtiles") => {
                let cfg = MbtilesStoreCfg {
                    path: file_or_url.into(),
                };
                Some(cfg)
            }
            _ => None,
        }
    }
}

#[async_trait]
impl TileWriter for MbtilesStore {
    fn compression(&self) -> Compression {
        match self.mbt.format_info.encoding {
            TileEncoding::Gzip => Compression::Gzip,
            _ => Compression::None,
        }
    }
    async fn exists(&self, xyz: &Xyz) -> bool {
        match self.mbt.get_tile(xyz.z, xyz.x as u32, xyz.y as u32).await {
            Ok(None) | Err(_) => false,
            Ok(_) => true,
        }
    }
    async fn put_tile(&self, xyz: &Xyz, data: Vec<u8>) -> Result<(), TileStoreError> {
        let mut conn = self.mbt.pool.acquire().await?;
        self.mbt
            .mbtiles
            .insert_tiles(
                &mut conn,
                self.mbt.layout,
                CopyDuplicateMode::Override,
                &[(xyz.z, xyz.x as u32, xyz.y as u32, data)],
            )
            .await?;
        Ok(())
    }
    async fn put_tiles(&mut self, tiles: &[(u8, u32, u32, Vec<u8>)]) -> Result<(), TileStoreError> {
        let mut conn = self.mbt.pool.acquire().await?;
        self.mbt
            .mbtiles
            .insert_tiles(
                &mut conn,
                self.mbt.layout,
                CopyDuplicateMode::Override,
                tiles,
            )
            .await?;
        Ok(())
    }
}

#[async_trait]
impl TileReader for MbtilesStore {
    async fn get_tile(&self, xyz: &Xyz) -> Result<Option<TileResponse>, TileStoreError> {
        let resp =
            if let Some(content) = self.mbt.get_tile(xyz.z, xyz.x as u32, xyz.y as u32).await? {
                let mut response = TileResponse::new();
                if self.mbt.format_info.format == TileFormat::Mvt {
                    response.set_content_type("application/x-protobuf");
                }
                if let Some(encoding) = self.mbt.format_info.encoding.content_encoding() {
                    response.insert_header(("Content-Encoding", encoding));
                }
                let body = Box::new(Cursor::new(content));
                Some(response.with_body(body))
            } else {
                None
            };
        Ok(resp)
    }
}