main.rs (7041B)
1 use std::{fs::File, time::Duration}; 2 3 use clap::Parser; 4 use color_eyre::eyre::Result; 5 use git2::Signature; 6 use memmap2::Mmap; 7 use tracing::{info, warn}; 8 9 use crate::{git::init_git_repository, osm::osm_data::convert_objects_to_git}; 10 11 mod git; 12 mod osm; 13 14 #[derive(Parser)] 15 #[command(author, version, about, long_about = None)] 16 struct Cli { 17 /// Path to the git repo to replay changesets to 18 #[arg(short, long, default_value = "./osm-git-repo")] 19 git_repo_path: String, 20 /// The server to get day replication files from 21 #[arg( 22 short, 23 long, 24 default_value = "https://planet.openstreetmap.org/replication/day" 25 )] 26 replication_server: String, 27 /// Where to write cache files 28 #[arg(long, default_value = "./cache")] 29 cache_path: String, 30 /// If the git repo should be removed and recreated 31 #[arg(short, long)] 32 clean: bool, 33 /// Where to start downloading data from 34 #[arg(long, default_value = "000/000/000")] 35 start_data: String, 36 /// The time to wait between downloading data 37 /// This is to avoid causing a lot of load on the OSM servers 38 #[arg(long, default_value = "500")] 39 wait_time: u64, 40 } 41 42 #[tokio::main] 43 async fn main() -> Result<()> { 44 color_eyre::install()?; 45 tracing_subscriber::fmt::init(); 46 let cli = Cli::parse(); 47 48 info!( 49 "Starting to replay osm changesets to git repo at {}", 50 cli.git_repo_path 51 ); 52 53 let client = reqwest::Client::builder() 54 .user_agent("osm-git-replay/0.1.0") 55 .gzip(true) 56 .timeout(Duration::from_secs(60)) 57 .build()?; 58 59 if cli.clean { 60 info!("Cleaning git repo at {}", cli.git_repo_path); 61 if std::path::Path::new(&cli.git_repo_path).exists() { 62 std::fs::remove_dir_all(&cli.git_repo_path)?; 63 } 64 } 65 66 let author = Signature::now("osm-git-replay", "osm-git-replay@localhost")?; 67 68 let repository = init_git_repository(&cli.git_repo_path, &cli.replication_server, &author)?; 69 info!("Git repository initialized"); 70 71 // Data download metadata 72 // TODO: We should probably detect where to resume from 73 let mut data_position_top = cli.start_data[0..3].parse::<u16>()?; 74 let mut data_position_middle = cli.start_data[4..7].parse::<u16>()?; 75 let mut data_position_bottom = cli.start_data[8..11].parse::<u16>()?; 76 77 // Parse the changesets and convert them to git objects 78 loop { 79 // Check for cache and use it if it exists 80 let cache_file_path = format!( 81 "{}/replication/{:03}/{:03}/{:03}.osm.gz", 82 cli.cache_path, data_position_top, data_position_middle, data_position_bottom 83 ); 84 85 if std::path::Path::new(&cache_file_path).exists() { 86 info!("Using cached data file at {}", cache_file_path); 87 let file = File::open(&cache_file_path)?; 88 let data = unsafe { Mmap::map(&file)? }; 89 let changeset_location = format!("{}/changesets/torrents", cli.cache_path); 90 convert_objects_to_git(&repository, &author, &data, &changeset_location)?; 91 info!("Data file parsed"); 92 93 // Increment the data position 94 if data_position_top == 999 95 && data_position_middle == 999 96 && data_position_bottom == 999 97 { 98 // Uhhhhhh?! 99 break; 100 } 101 102 if data_position_middle == 999 && data_position_bottom == 999 { 103 data_position_middle = 0; 104 data_position_bottom = 0; 105 data_position_top += 1; 106 } 107 108 if data_position_bottom == 999 { 109 data_position_bottom = 0; 110 data_position_middle += 1; 111 } 112 } else { 113 { 114 // Download minute replication files and find the changesets that were modified in that minute 115 let data_url = format!( 116 "{}/{:03}/{:03}/{:03}.osc.gz", 117 cli.replication_server, 118 data_position_top, 119 data_position_middle, 120 data_position_bottom 121 ); 122 info!("Downloading data file from {}", data_url); 123 let data_response: reqwest::Response = client.get(&data_url).send().await?; 124 125 if data_response.status() == reqwest::StatusCode::NOT_FOUND { 126 warn!("data file not found at {}", data_url); 127 // Increment the data position 128 if data_position_top == 999 129 && data_position_middle == 999 130 && data_position_bottom == 999 131 { 132 // Uhhhhhh?! 133 break; 134 } 135 136 if data_position_middle == 999 && data_position_bottom == 999 { 137 data_position_middle = 0; 138 data_position_bottom = 0; 139 data_position_top += 1; 140 } 141 142 if data_position_bottom == 999 { 143 data_position_bottom = 0; 144 data_position_middle += 1; 145 } 146 147 if data_position_bottom < 999 { 148 data_position_bottom += 1; 149 } 150 151 continue; 152 } 153 154 let data = data_response.bytes().await?; 155 info!("Caching Data file to disk"); 156 std::fs::create_dir_all(std::path::Path::new(&cache_file_path).parent().unwrap())?; 157 std::fs::write(&cache_file_path, &data)?; 158 info!("Data file downloaded"); 159 }; 160 161 let file = File::open(cache_file_path)?; 162 let data = unsafe { Mmap::map(&file)? }; 163 164 let changeset_location = format!("{}/changesets/torrents", cli.cache_path); 165 convert_objects_to_git(&repository, &author, &data, &changeset_location)?; 166 167 // Increment the data position 168 if data_position_top == 999 169 && data_position_middle == 999 170 && data_position_bottom == 999 171 { 172 // Uhhhhhh?! 173 break; 174 } 175 176 if data_position_middle == 999 && data_position_bottom == 999 { 177 data_position_middle = 0; 178 data_position_bottom = 0; 179 data_position_top += 1; 180 } 181 182 if data_position_bottom == 999 { 183 data_position_bottom = 0; 184 data_position_middle += 1; 185 } 186 187 if data_position_bottom < 999 { 188 data_position_bottom += 1; 189 } 190 191 // Wait a few seconds before downloading the next data file 192 tokio::time::sleep(Duration::from_millis(cli.wait_time)).await; 193 } 194 } 195 196 info!( 197 "Downloaded data until {} {} {}", 198 data_position_top, 199 data_position_middle, 200 data_position_bottom - 1 201 ); 202 203 Ok(()) 204 }