Skip to content

Commit

Permalink
Make copy into existing mbtiles file explicit
Browse files Browse the repository at this point in the history
When using `martin-cp` or `mbtiles copy` into an existing file,
require the user to set `--on-duplicate`. This prevents accidental
writing to an existing file.
  • Loading branch information
nyurik committed Dec 14, 2023
1 parent 9792e90 commit c55bb3c
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 25 deletions.
22 changes: 15 additions & 7 deletions martin/src/bin/martin-cp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ use martin::{
use martin_tile_utils::{bbox_to_xyz, TileInfo};
use mbtiles::sqlx::SqliteConnection;
use mbtiles::{
init_mbtiles_schema, is_empty_database, CopyDuplicateMode, MbtType, MbtTypeCli, Mbtiles,
init_mbtiles_schema, is_empty_database, CopyDuplicateMode, MbtError, MbtType, MbtTypeCli,
Mbtiles,
};
use tilejson::Bounds;
use tokio::sync::mpsc::channel;
Expand Down Expand Up @@ -73,9 +74,9 @@ pub struct CopyArgs {
/// Use `identity` to disable compression. Ignored for non-encodable tiles like PNG and JPEG.
#[arg(long, alias = "encodings", default_value = "gzip")]
pub encoding: String,
/// Specify the behaviour when generated tile already exists in the destination file.
#[arg(long, value_enum, default_value_t = CopyDuplicateMode::default())]
pub on_duplicate: CopyDuplicateMode,
/// Allow copying to existing files, and indicate what to do if a tile with the same Z/X/Y already exists
#[arg(long, value_enum)]
pub on_duplicate: Option<CopyDuplicateMode>,
/// Number of concurrent connections to use.
#[arg(long, default_value = "1")]
pub concurrency: Option<usize>,
Expand Down Expand Up @@ -220,7 +221,7 @@ enum MartinCpError {
#[error(transparent)]
Actix(#[from] actix_web::Error),
#[error(transparent)]
Mbt(#[from] mbtiles::MbtError),
Mbt(#[from] MbtError),
}

impl Display for Progress {
Expand Down Expand Up @@ -273,6 +274,13 @@ async fn run_tile_copy(args: CopyArgs, state: ServerState) -> MartinCpResult<()>
let tiles = compute_tile_ranges(&args);
let mbt = Mbtiles::new(output_file)?;
let mut conn = mbt.open_or_new().await?;
let on_duplicate = if let Some(on_duplicate) = args.on_duplicate {
on_duplicate
} else if !is_empty_database(&mut conn).await? {
return Err(MbtError::DestinationFileExists(output_file.clone()).into());
} else {
CopyDuplicateMode::Override
};
let mbt_type = init_schema(&mbt, &mut conn, sources, tile_info, args.mbt_type).await?;
let query = args.url_query.as_deref();
let req = TestRequest::default()
Expand Down Expand Up @@ -317,7 +325,7 @@ async fn run_tile_copy(args: CopyArgs, state: ServerState) -> MartinCpResult<()>
} else {
batch.push((tile.xyz.z, tile.xyz.x, tile.xyz.y, tile.data));
if batch.len() >= BATCH_SIZE || last_saved.elapsed() > SAVE_EVERY {
mbt.insert_tiles(&mut conn, mbt_type, args.on_duplicate, &batch)
mbt.insert_tiles(&mut conn, mbt_type, on_duplicate, &batch)
.await?;
batch.clear();
last_saved = Instant::now();
Expand All @@ -332,7 +340,7 @@ async fn run_tile_copy(args: CopyArgs, state: ServerState) -> MartinCpResult<()>
}
}
if !batch.is_empty() {
mbt.insert_tiles(&mut conn, mbt_type, args.on_duplicate, &batch)
mbt.insert_tiles(&mut conn, mbt_type, on_duplicate, &batch)
.await?;
}
Ok(())
Expand Down
6 changes: 3 additions & 3 deletions mbtiles/src/bin/mbtiles.rs
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,7 @@ mod tests {
#[test]
fn test_copy_diff_with_override_copy_duplicate_mode() {
let mut opt = MbtilesCopier::new(PathBuf::from("src_file"), PathBuf::from("dst_file"));
opt.on_duplicate = CopyDuplicateMode::Override;
opt.on_duplicate = Some(CopyDuplicateMode::Override);
assert_eq!(
Args::parse_from([
"mbtiles",
Expand All @@ -332,7 +332,7 @@ mod tests {
#[test]
fn test_copy_diff_with_ignore_copy_duplicate_mode() {
let mut opt = MbtilesCopier::new(PathBuf::from("src_file"), PathBuf::from("dst_file"));
opt.on_duplicate = CopyDuplicateMode::Ignore;
opt.on_duplicate = Some(CopyDuplicateMode::Ignore);
assert_eq!(
Args::parse_from([
"mbtiles",
Expand All @@ -352,7 +352,7 @@ mod tests {
#[test]
fn test_copy_diff_with_abort_copy_duplicate_mode() {
let mut opt = MbtilesCopier::new(PathBuf::from("src_file"), PathBuf::from("dst_file"));
opt.on_duplicate = CopyDuplicateMode::Abort;
opt.on_duplicate = Some(CopyDuplicateMode::Abort);
assert_eq!(
Args::parse_from([
"mbtiles",
Expand Down
37 changes: 22 additions & 15 deletions mbtiles/src/copier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,10 @@ use crate::{
AGG_TILES_HASH_IN_DIFF,
};

#[derive(PartialEq, Eq, Default, Debug, Clone, Copy, EnumDisplay, Serialize, Deserialize)]
#[derive(PartialEq, Eq, Debug, Clone, Copy, EnumDisplay, Serialize, Deserialize)]
#[enum_display(case = "Kebab")]
#[cfg_attr(feature = "cli", derive(ValueEnum))]
pub enum CopyDuplicateMode {
#[default]
Override,
Ignore,
Abort,
Expand Down Expand Up @@ -62,9 +61,9 @@ pub struct MbtilesCopier {
/// Destination type with options
#[cfg_attr(feature = "cli", arg(skip))]
pub dst_type: Option<MbtType>,
/// Specify copying behaviour when tiles with duplicate (zoom_level, tile_column, tile_row) values are found
#[cfg_attr(feature = "cli", arg(long, value_enum, default_value_t = CopyDuplicateMode::default()))]
pub on_duplicate: CopyDuplicateMode,
/// Allow copying to existing files, and indicate what to do if a tile with the same Z/X/Y already exists
#[cfg_attr(feature = "cli", arg(long, value_enum))]
pub on_duplicate: Option<CopyDuplicateMode>,
/// Minimum zoom level to copy
#[cfg_attr(feature = "cli", arg(long, conflicts_with("zoom_levels")))]
pub min_zoom: Option<u8>,
Expand Down Expand Up @@ -102,7 +101,7 @@ impl MbtilesCopier {
dst_file: dst_filepath,
dst_type_cli: None,
dst_type: None,
on_duplicate: CopyDuplicateMode::Override,
on_duplicate: None,
min_zoom: None,
max_zoom: None,
zoom_levels: Vec::default(),
Expand Down Expand Up @@ -172,6 +171,14 @@ impl MbtileCopierInt {
let src_type = src_mbt.open_and_detect_type().await?;
let mut conn = dst_mbt.open_or_new().await?;
let is_empty_db = is_empty_database(&mut conn).await?;
let on_duplicate = if let Some(on_duplicate) = self.options.on_duplicate {
on_duplicate
} else if !is_empty_database(&mut conn).await? {
return Err(MbtError::DestinationFileExists(self.options.dst_file));
} else {
CopyDuplicateMode::Override
};

src_mbt.attach_to(&mut conn, "sourceDb").await?;

let dst_type: MbtType;
Expand Down Expand Up @@ -211,8 +218,8 @@ impl MbtileCopierInt {

let (where_clause, query_args) = self.get_where_clause();
let select_from = format!("{select_from} {where_clause}");
let on_dupl = self.options.on_duplicate.to_sql();
let sql_cond = self.get_on_duplicate_sql_cond(dst_type);
let on_dupl = on_duplicate.to_sql();
let sql_cond = Self::get_on_duplicate_sql_cond(on_duplicate, dst_type);

debug!("Copying tiles with 'INSERT {on_dupl}' {src_type} -> {dst_type} ({sql_cond})");

Expand Down Expand Up @@ -417,8 +424,8 @@ impl MbtileCopierInt {
}

/// Returns WHERE condition SQL depending on the override and destination type
fn get_on_duplicate_sql_cond(&self, dst_type: MbtType) -> String {
match &self.options.on_duplicate {
fn get_on_duplicate_sql_cond(on_duplicate: CopyDuplicateMode, dst_type: MbtType) -> String {
match on_duplicate {
CopyDuplicateMode::Ignore | CopyDuplicateMode::Override => String::new(),
CopyDuplicateMode::Abort => {
let (main_table, tile_identifier) = match dst_type {
Expand Down Expand Up @@ -820,7 +827,7 @@ mod tests {
let dst = PathBuf::from("../tests/fixtures/mbtiles/world_cities.mbtiles");

let mut opt = MbtilesCopier::new(src.clone(), dst.clone());
opt.on_duplicate = CopyDuplicateMode::Abort;
opt.on_duplicate = Some(CopyDuplicateMode::Abort);

assert!(matches!(
opt.run().await.unwrap_err(),
Expand All @@ -841,9 +848,9 @@ mod tests {
.run()
.await?;

let mut dst_conn = MbtilesCopier::new(src_file.clone(), dst.clone())
.run()
.await?;
let mut opt = MbtilesCopier::new(src_file.clone(), dst.clone());
opt.on_duplicate = Some(CopyDuplicateMode::Override);
let mut dst_conn = opt.run().await?;

// Verify the tiles in the destination file is a superset of the tiles in the source file
Mbtiles::new(src_file)?
Expand Down Expand Up @@ -871,7 +878,7 @@ mod tests {
.await?;

let mut opt = MbtilesCopier::new(src_file.clone(), dst.clone());
opt.on_duplicate = CopyDuplicateMode::Ignore;
opt.on_duplicate = Some(CopyDuplicateMode::Ignore);
let mut dst_conn = opt.run().await?;

// Verify the tiles in the destination file are the same as those in the source file except for those with duplicate (zoom_level, tile_column, tile_row)
Expand Down
3 changes: 3 additions & 0 deletions mbtiles/src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,9 @@ pub enum MbtError {

#[error("The MBTiles file {0} has data of type {1}, but the desired type was set to {2}")]
MismatchedTargetType(PathBuf, MbtType, MbtType),

#[error("Unless --on-duplicate (override|ignore|abort) is set, writing tiles to an existing non-empty MBTiles file is disabled. Either set --on-duplicate flag, or delete {}", .0.display())]
DestinationFileExists(PathBuf),
}

pub type MbtResult<T> = Result<T, MbtError>;

0 comments on commit c55bb3c

Please sign in to comment.