Skip to content

Commit 646fd2b

Browse files
LeonardBessonLeonardBesson
committed
Refactor state refresh to use reqwest and use thread instead of Arbiter
1 parent 5519c40 commit 646fd2b

3 files changed

Lines changed: 89 additions & 277 deletions

File tree

‎src/data.rs‎

Lines changed: 66 additions & 247 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use std::collections::HashMap;
22
use std::fmt::Formatter;
33
use std::fs::File;
4-
use std::io::Read;
4+
use std::io::{Read, BufReader};
55

66
use actix::{System, SystemRunner};
77
use actix_web::client::Client;
@@ -22,9 +22,9 @@ use crate::db::Pool;
2222
use crate::models::{Address, NewAddress, NewState, State};
2323
use crate::postcode::AddressRecord;
2424
use reqwest::Response;
25+
use actix::fut::err;
2526

26-
const STATE_BODY_LIMIT_BYTES: usize = 2_097_152; // 2MB
27-
const ZIP_BODY_LIMIT_BYTES: usize = 1_074_000_000; // 1GB
27+
const APPROXIMATE_ZIP_SIZE_BYTES: usize = 200_097_152; // 200 MB
2828

2929
const BATCH_SIZE: usize = 2500;
3030

@@ -50,20 +50,6 @@ impl std::fmt::Display for RefreshError {
5050
}
5151
}
5252

53-
// TODO
54-
impl From<actix_web::client::SendRequestError> for RefreshError {
55-
fn from(error: actix_web::client::SendRequestError) -> Self {
56-
OldData
57-
}
58-
}
59-
60-
// TODO
61-
impl From<actix_web::client::PayloadError> for RefreshError {
62-
fn from(error: actix_web::client::PayloadError) -> Self {
63-
OldData
64-
}
65-
}
66-
6753
// TODO
6854
impl From<diesel::result::Error> for RefreshError {
6955
fn from(error: diesel::result::Error) -> Self {
@@ -85,6 +71,12 @@ impl From<diesel::result::Error> for RefreshError {
8571
}
8672
}
8773

74+
impl From<reqwest::Error> for RefreshError {
75+
fn from(error: reqwest::Error) -> Self {
76+
OldData
77+
}
78+
}
79+
8880
#[derive(Debug)]
8981
pub struct StateInfo {
9082
pub url: String,
@@ -99,11 +91,8 @@ pub struct StateRefresh {
9991
pub current_state: Option<State>
10092
}
10193

102-
pub fn refresh_state(
103-
system: &mut SystemRunner,
104-
pool: &Pool
105-
) {
106-
let status = system.block_on(futures::lazy(|| { get_state_refresh(pool) }));
94+
pub fn refresh_state(pool: &PgConnection) {
95+
let status = get_state_refresh(pool);
10796
match status {
10897
Ok(state_refresh) => {
10998
match state_refresh.state_info {
@@ -117,7 +106,7 @@ pub fn refresh_state(
117106
info!("Data already up to date (state: {})", state_info.version);
118107
} else {
119108
info!("Updating data...");
120-
match system.block_on(futures::lazy(|| { update_state(pool, state_info) })) {
109+
match update_state(pool, state_info) {
121110
Ok(_) => { info!("Successfuly updated data"); },
122111
Err(err) => { error!("Error while updating state: {}", err); },
123112
};
@@ -137,28 +126,10 @@ pub fn refresh_state(
137126
};
138127
}
139128

140-
fn get_state_refresh<'a>(pool: &'a Pool) -> impl Future<Item = StateRefresh, Error = RefreshError> + 'a {
141-
Client::default()
142-
.get("http://results.openaddresses.io/state.txt")
143-
.send()
144-
.from_err()
145-
.and_then(move |mut resp| {
146-
resp.body()
147-
.limit(STATE_BODY_LIMIT_BYTES)
148-
.from_err()
149-
.map(move |body| {
150-
let state_info = get_info(body);
151-
let connection = pool.get().unwrap();
152-
let current_state = current_state(&connection);
153-
StateRefresh { state_info, current_state }
154-
})
155-
})
156-
}
157-
158-
fn get_info(body: Bytes) -> Option<StateInfo> {
129+
fn get_info<R: std::io::Read>(response: R) -> Option<StateInfo> {
159130
let mut reader = csv::ReaderBuilder::new()
160131
.delimiter(b'\t')
161-
.from_reader(&body[..]);
132+
.from_reader(response);
162133

163134
reader
164135
.records()
@@ -181,6 +152,57 @@ fn get_info(body: Bytes) -> Option<StateInfo> {
181152
})
182153
}
183154

155+
pub fn get_state_refresh(conn: &PgConnection) -> Result<StateRefresh, RefreshError> {
156+
let response = reqwest::get("http://results.openaddresses.io/state.txt")?;
157+
let state_info = get_info(response);
158+
let current_state = current_state(&conn);
159+
Ok(StateRefresh { state_info, current_state })
160+
}
161+
162+
fn update_state(
163+
conn: &PgConnection,
164+
state_info: StateInfo
165+
) -> Result<(), RefreshError> {
166+
info!("Downloading state version {} from {}", state_info.version, state_info.url);
167+
168+
let mut resp = reqwest::get(&state_info.url)?;
169+
let mut buf: Vec<u8> = Vec::with_capacity(APPROXIMATE_ZIP_SIZE_BYTES);
170+
resp.copy_to(&mut buf)?;
171+
info!("Downloaded zip, size: {} MB", buf.len() / 1_000_000);
172+
info!("Searching for csv file");
173+
174+
let mut reader = std::io::Cursor::new(&buf);
175+
let mut zip = ZipArchive::new(reader).expect("Could not create zip archive");
176+
let re = Regex::new(r"nl.*\.csv").expect("Could not create regex");
177+
for i in 0..zip.len() {
178+
let file = zip.by_index(i).unwrap();
179+
info!("File: {}", file.name());
180+
if re.is_match(file.name()) {
181+
info!("Found csv file");
182+
info!("Updating database records...");
183+
let mut reader = csv::Reader::from_reader(file);
184+
185+
let mut batch = Vec::<AddressRecord>::with_capacity(BATCH_SIZE);
186+
let progress_bar = ProgressBar::new(state_info.address_count as u64);
187+
for record in reader.deserialize() {
188+
let address_record: AddressRecord = record.expect("Could not deserialize post code record");
189+
batch.push(address_record);
190+
if batch.len() == BATCH_SIZE {
191+
process_batch(&conn, &mut batch, &progress_bar);
192+
}
193+
};
194+
process_batch(&conn, &mut batch, &progress_bar);
195+
progress_bar.finish();
196+
197+
create_new_state(&conn, &state_info);
198+
info!("Done");
199+
break;
200+
}
201+
}
202+
203+
Ok(())
204+
}
205+
184206
fn current_state(connection: &PgConnection) -> Option<State> {
185207
use crate::schema::states::dsl::*;
186208

@@ -192,60 +214,6 @@ fn current_state(connection: &PgConnection) -> Option<State> {
192214
.unwrap_or(None)
193215
}
194216

195-
fn update_state<'a>(
196-
pool: &'a Pool,
197-
state_info: StateInfo
198-
) -> impl Future<Item = (), Error = RefreshError> + 'a {
199-
info!("Downloading state version {} from {}", state_info.version, state_info.url);
200-
201-
Client::default()
202-
.get(&state_info.url)
203-
.send()
204-
.map_err(RefreshError::from)
205-
.and_then(move |mut resp| {
206-
resp.body()
207-
.limit(ZIP_BODY_LIMIT_BYTES)
208-
.from_err()
209-
.map(move |body| {
210-
info!("Downloaded zip, size: {} MB", body.len() / 1_000_000);
211-
info!("Searching for csv file");
212-
213-
let mut reader = std::io::Cursor::new(body);
214-
let mut zip = ZipArchive::new(reader).expect("Could not create zip archive");
215-
let re = Regex::new(r"nl.*\.csv").expect("Could not create regex");
216-
for i in 0..zip.len()
217-
{
218-
let file = zip.by_index(i).unwrap();
219-
info!("Filename: {}", file.name());
220-
if re.is_match(file.name()) {
221-
info!("Found csv file");
222-
info!("Updating database records...");
223-
let conn = pool.get().unwrap();
224-
let mut reader = csv::Reader::from_reader(file);
225-
226-
let mut batch = Vec::<AddressRecord>::with_capacity(BATCH_SIZE);
227-
let progress_bar = ProgressBar::new(state_info.address_count as u64);
228-
for record in reader.deserialize() {
229-
let address_record: AddressRecord = record.expect("Could not deserialize post code record");
230-
batch.push(address_record);
231-
if batch.len() == BATCH_SIZE {
232-
process_batch(&conn, &mut batch, &progress_bar);
233-
}
234-
};
235-
process_batch(&conn, &mut batch, &progress_bar);
236-
progress_bar.finish();
237-
238-
create_new_state(&conn, &state_info);
239-
info!("Done");
240-
break;
241-
}
242-
}
243-
244-
()
245-
})
246-
})
247-
}
248-
249217
fn process_batch(
250218
conn: &PgConnection,
251219
batch: &mut Vec<AddressRecord>,
@@ -336,153 +304,4 @@ pub fn get_addresses(
336304
.order(number.asc())
337305
.limit(ADDRESSES_RESULT_LIMIT)
338306
.load(&pool.get().unwrap())
339-
}
340-
341-
// TODO: try to compose get_state_info and update_state
342-
// Chaining futures with different return type is a pain
343-
// Apparently Box<dyn Future<blabla>> can fix the issue
344-
// But I couldn't get it to compile.
345-
// Check "Returning from multiple branches" from https://tokio.rs/docs/futures/combinators/#use-impl-future
346-
347-
348-
349-
350-
351-
352-
353-
354-
355-
356-
357-
358-
359-
/// Reqwest
360-
361-
impl From<reqwest::Error> for RefreshError {
362-
fn from(error: reqwest::Error) -> Self {
363-
OldData
364-
}
365-
}
366-
367-
//pub fn refresh_state2(pool: &Pool) {
368-
// let status = get_state_refresh2(pool);
369-
// match status {
370-
// Ok(state_refresh) => {
371-
// match state_refresh.state_info {
372-
// Some(state_info) => {
373-
// let up_to_date = state_refresh
374-
// .current_state
375-
// .filter(|s| s.version == state_info.version)
376-
// .is_some();
377-
//
378-
// if up_to_date {
379-
// info!("Data already up to date (state: {})", state_info.version);
380-
// } else {
381-
// info!("Updating data...");
382-
// match system.block_on(futures::lazy(|| { update_state(pool, state_info) })) {
383-
// Ok(_) => { info!("Successfuly updated data"); },
384-
// Err(err) => { error!("Error while updating state: {}", err); },
385-
// };
386-
// }
387-
// },
388-
// None => {
389-
// if state_refresh.current_state.is_none() {
390-
// panic!("Couldn't fetch data and no fallback");
391-
// } else {
392-
// info!("Falling back");
393-
// };
394-
// }
395-
// }
396-
// },
397-
// Err(RefreshError::NoData) => { panic!("Couldn't fetch data and no fallback"); },
398-
// Err(RefreshError::OldData) => { error!("Falling back"); }
399-
// };
400-
//}
401-
402-
fn get_info2(response: Response) -> Option<StateInfo> {
403-
let mut reader = csv::ReaderBuilder::new()
404-
.delimiter(b'\t')
405-
.from_reader(response);
406-
407-
reader
408-
.records()
409-
.find(|record| {
410-
record.is_ok() &&
411-
record
412-
.as_ref()
413-
.unwrap()
414-
.as_slice()
415-
.starts_with("nl/countrywide.json")
416-
})
417-
.map(|record| {
418-
let r = record.unwrap();
419-
StateInfo {
420-
address_count: r[4].parse::<usize>().unwrap(),
421-
url: r[8].to_owned(),
422-
hash: r[10].to_owned(),
423-
version: r[15].to_owned()
424-
}
425-
})
426-
}
427-
428-
pub fn get_state_refresh2(pool: & Pool) -> Result<StateRefresh, RefreshError> {
429-
let response = reqwest::get("http://results.openaddresses.io/state.txt")?;
430-
let state_info = get_info2(response);
431-
let connection = pool.get().unwrap();
432-
let current_state = current_state(&connection);
433-
Ok(StateRefresh { state_info, current_state })
434-
}
435-
436-
//fn update_state2<'a>(
437-
// pool: &'a Pool,
438-
// state_info: StateInfo
439-
//) -> impl Future<Item = (), Error = RefreshError> + 'a {
440-
// info!("Downloading state version {} from {}", state_info.version, state_info.url);
441-
//
442-
// Client::default()
443-
// .get(&state_info.url)
444-
// .send()
445-
// .map_err(RefreshError::from)
446-
// .and_then(move |mut resp| {
447-
// resp.body()
448-
// .limit(ZIP_BODY_LIMIT_BYTES)
449-
// .from_err()
450-
// .map(move |body| {
451-
// info!("Downloaded zip, size: {} MB", body.len() / 1_000_000);
452-
// info!("Searching for csv file");
453-
//
454-
// let mut reader = std::io::Cursor::new(body);
455-
// let mut zip = ZipArchive::new(reader).expect("Could not create zip archive");
456-
// let re = Regex::new(r"nl.*\.csv").expect("Could not create regex");
457-
// for i in 0..zip.len()
458-
// {
459-
// let file = zip.by_index(i).unwrap();
460-
// info!("Filename: {}", file.name());
461-
// if re.is_match(file.name()) {
462-
// info!("Found csv file");
463-
// info!("Updating database records...");
464-
// let conn = pool.get().unwrap();
465-
// let mut reader = csv::Reader::from_reader(file);
466-
//
467-
// let mut batch = Vec::<AddressRecord>::with_capacity(BATCH_SIZE);
468-
// let progress_bar = ProgressBar::new(state_info.address_count as u64);
469-
// for record in reader.deserialize() {
470-
// let address_record: AddressRecord = record.expect("Could not deserialize post code record");
471-
// batch.push(address_record);
472-
// if batch.len() == BATCH_SIZE {
473-
// process_batch(&conn, &mut batch, &progress_bar);
474-
// }
475-
// };
476-
// process_batch(&conn, &mut batch, &progress_bar);
477-
// progress_bar.finish();
478-
//
479-
// create_new_state(&conn, &state_info);
480-
// info!("Done");
481-
// break;
482-
// }
483-
// }
484-
//
485-
// ()
486-
// })
487-
// })
488-
//}
307+
}

0 commit comments

Comments
 (0)