Fix load_hashes.

This commit is contained in:
Syfaro 2020-09-23 15:49:55 -04:00
parent 2a2c3b88fc
commit 9f70ae0aca

View File

@ -1,6 +1,6 @@
use bb8::Pool; use bb8::Pool;
use bb8_postgres::PostgresConnectionManager; use bb8_postgres::PostgresConnectionManager;
use futures::future::FutureExt; use futures::StreamExt;
struct NeededPost { struct NeededPost {
id: i32, id: i32,
@ -8,9 +8,10 @@ struct NeededPost {
} }
async fn hash_url( async fn hash_url(
id: i32,
client: std::sync::Arc<reqwest::Client>, client: std::sync::Arc<reqwest::Client>,
url: String, url: String,
) -> Result<(img_hash::ImageHash<[u8; 8]>, i64), image::ImageError> { ) -> (i32, Result<i64, image::ImageError>) {
let data = client let data = client
.get(&url) .get(&url)
.send() .send()
@ -25,7 +26,7 @@ async fn hash_url(
Ok(image) => image, Ok(image) => image,
Err(e) => { Err(e) => {
println!("{:?}", &data[0..50]); println!("{:?}", &data[0..50]);
return Err(e); return (id, Err(e));
} }
}; };
@ -37,7 +38,7 @@ async fn hash_url(
println!("{} - {}", url, num); println!("{} - {}", url, num);
Ok((hash, num)) (id, Ok(num))
} }
async fn load_next_posts( async fn load_next_posts(
@ -103,16 +104,18 @@ async fn main() {
continue; continue;
} }
for chunk in needed_posts.chunks(8) { futures::stream::iter(
let futs = chunk.iter().map(|post| { needed_posts
let db = pool.clone(); .into_iter()
let client = client.clone(); .map(|post| hash_url(post.id, client.clone(), post.full_url)),
let id = post.id; )
.buffer_unordered(8)
.for_each(|res: (i32, Result<i64, image::ImageError>)| async {
let db = pool.get().await.expect("unable to get from pool");
hash_url(client, post.full_url.clone()).then(move |res| async move {
match res { match res {
Ok((_hash, num)) => { (id, Ok(num)) => {
let mut conn = db.get().await.unwrap(); let mut conn = pool.get().await.unwrap();
let tx = conn let tx = conn
.transaction() .transaction()
@ -131,31 +134,20 @@ async fn main() {
.expect("Unable to insert hash to hashes table"); .expect("Unable to insert hash to hashes table");
tx.commit().await.expect("Unable to commit tx"); tx.commit().await.expect("Unable to commit tx");
drop(conn);
} }
Err(e) => { (id, Err(e)) => {
let desc = e.to_string(); let desc = e.to_string();
println!("[{}] hashing error - {}", id, desc); println!("[{}] hashing error - {}", id, desc);
db.get() db.execute(
.await
.unwrap()
.execute(
"UPDATE e621 SET hash_error = $2 WHERE id = $1", "UPDATE e621 SET hash_error = $2 WHERE id = $1",
&[&id, &desc], &[&id, &desc],
) )
.await .await
.expect("Unable to update hash error in database"); .expect("Unable to update hash error in database");
} }
};
})
});
println!("joining futs");
futures::future::join_all(futs).await;
println!("futs completed");
} }
()
})
.await;
} }
} }