This commit is contained in:
Syfaro 2020-03-01 19:39:05 -06:00
parent 057a91e498
commit eee2f7d65d
3 changed files with 689 additions and 484 deletions

1100
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -5,9 +5,9 @@ authors = ["Syfaro <syfaro@huefox.com>"]
edition = "2018" edition = "2018"
[dependencies] [dependencies]
reqwest = "0.10.1" reqwest = "0.10"
postgres = { version = "0.17.0", features = ["with-chrono-0_4"] } postgres = { version = "0.17.0", features = ["with-chrono-0_4"] }
tokio = { version = "0.2", features = ["macros", "time"] } tokio = { version = "0.2", features = ["full"] }
tokio-postgres = { version = "0.5.0" } tokio-postgres = { version = "0.5.0" }
r2d2_postgres = " 0.16.0" r2d2_postgres = " 0.16.0"
r2d2 = "0.8" r2d2 = "0.8"
@ -15,4 +15,4 @@ chrono = "0.4"
[dependencies.furaffinity-rs] [dependencies.furaffinity-rs]
git = "https://git.huefox.com/syfaro/furaffinity-rs" git = "https://git.huefox.com/syfaro/furaffinity-rs"
rev = "1548d7f0132e65b360af8bef16f864a1c3b3df02" features = ["cloudflare-bypass"]

View File

@ -1,9 +1,9 @@
type Client = use tokio_postgres::Client;
r2d2::PooledConnection<r2d2_postgres::PostgresConnectionManager<tokio_postgres::tls::NoTls>>;
fn lookup_tag(client: &mut Client, tag: &str) -> i32 { async fn lookup_tag(client: &Client, tag: &str) -> i32 {
if let Some(row) = client if let Some(row) = client
.query("SELECT id FROM tag WHERE name = $1", &[&tag]) .query("SELECT id FROM tag WHERE name = $1", &[&tag])
.await
.unwrap() .unwrap()
.into_iter() .into_iter()
.next() .next()
@ -13,6 +13,7 @@ fn lookup_tag(client: &mut Client, tag: &str) -> i32 {
client client
.query("INSERT INTO tag (name) VALUES ($1) RETURNING id", &[&tag]) .query("INSERT INTO tag (name) VALUES ($1) RETURNING id", &[&tag])
.await
.unwrap() .unwrap()
.into_iter() .into_iter()
.next() .next()
@ -20,9 +21,10 @@ fn lookup_tag(client: &mut Client, tag: &str) -> i32 {
.get("id") .get("id")
} }
fn lookup_artist(client: &mut Client, artist: &str) -> i32 { async fn lookup_artist(client: &Client, artist: &str) -> i32 {
if let Some(row) = client if let Some(row) = client
.query("SELECT id FROM artist WHERE name = $1", &[&artist]) .query("SELECT id FROM artist WHERE name = $1", &[&artist])
.await
.unwrap() .unwrap()
.into_iter() .into_iter()
.next() .next()
@ -35,6 +37,7 @@ fn lookup_artist(client: &mut Client, artist: &str) -> i32 {
"INSERT INTO artist (name) VALUES ($1) RETURNING id", "INSERT INTO artist (name) VALUES ($1) RETURNING id",
&[&artist], &[&artist],
) )
.await
.unwrap() .unwrap()
.into_iter() .into_iter()
.next() .next()
@ -42,54 +45,52 @@ fn lookup_artist(client: &mut Client, artist: &str) -> i32 {
.get("id") .get("id")
} }
fn has_submission(client: &mut Client, id: i32) -> bool { async fn has_submission(client: &Client, id: i32) -> bool {
client client
.query("SELECT id FROM submission WHERE id = $1", &[&id]) .query("SELECT id FROM submission WHERE id = $1", &[&id])
.await
.expect("unable to run query") .expect("unable to run query")
.into_iter() .into_iter()
.next() .next()
.is_some() .is_some()
} }
fn ids_to_check(client: &mut Client, max: i32) -> Vec<i32> { async fn ids_to_check(client: &Client, max: i32) -> Vec<i32> {
let min = max - 100; let rows = client.query("SELECT sid FROM generate_series((SELECT max(id) FROM submission), $1::int) sid WHERE sid NOT IN (SELECT id FROM submission where id = sid)", &[&max]).await.unwrap();
let rows = client.query("SELECT sid FROM generate_series(LEAST($1::int, (SELECT MIN(id) FROM SUBMISSION)), $2::int) sid WHERE sid NOT IN (SELECT id FROM submission where id = sid)", &[&min, &max]).unwrap();
rows.iter().map(|row| row.get("sid")).collect() rows.iter().map(|row| row.get("sid")).collect()
} }
fn insert_submission( async fn insert_submission(
mut client: &mut Client, mut client: &Client,
sub: &furaffinity_rs::Submission, sub: &furaffinity_rs::Submission,
) -> Result<(), postgres::Error> { ) -> Result<(), postgres::Error> {
let artist_id = lookup_artist(&mut client, &sub.artist); let artist_id = lookup_artist(&mut client, &sub.artist).await;
let tag_ids: Vec<i32> = sub let mut tag_ids = Vec::with_capacity(sub.tags.len());
.tags for tag in &sub.tags {
.iter() tag_ids.push(lookup_tag(&client, &tag).await);
.map(|tag| lookup_tag(&mut client, &tag)) }
.collect();
let hash = sub.hash.clone(); let hash = sub.hash.clone();
let url = sub.content.url(); let url = sub.content.url();
client.execute("INSERT INTO submission (id, artist_id, url, filename, hash, rating, posted_at, description, hash_int, file_id) VALUES ($1, $2, $3, $4, decode($5, 'base64'), $6, $7, $8, $9, CASE WHEN isnumeric(split_part($4, '.', 1)) THEN split_part($4, '.', 1)::int ELSE null END)", &[ client.execute("INSERT INTO submission (id, artist_id, url, filename, hash, rating, posted_at, description, hash_int, file_id) VALUES ($1, $2, $3, $4, decode($5, 'base64'), $6, $7, $8, $9, CASE WHEN isnumeric(split_part($4, '.', 1)) THEN split_part($4, '.', 1)::int ELSE null END)", &[
&sub.id, &artist_id, &url, &sub.filename, &hash, &sub.rating.serialize(), &sub.posted_at, &sub.description, &sub.hash_num, &sub.id, &artist_id, &url, &sub.filename, &hash, &sub.rating.serialize(), &sub.posted_at, &sub.description, &sub.hash_num,
])?; ]).await?;
let stmt = client.prepare( let stmt = client.prepare(
"INSERT INTO tag_to_post (tag_id, post_id) VALUES ($1, $2) ON CONFLICT DO NOTHING", "INSERT INTO tag_to_post (tag_id, post_id) VALUES ($1, $2) ON CONFLICT DO NOTHING",
)?; ).await?;
for tag_id in tag_ids { for tag_id in tag_ids {
client.execute(&stmt, &[&tag_id, &sub.id])?; client.execute(&stmt, &[&tag_id, &sub.id]).await?;
} }
Ok(()) Ok(())
} }
fn insert_null_submission(client: &mut Client, id: i32) -> Result<u64, postgres::Error> { async fn insert_null_submission(client: &Client, id: i32) -> Result<u64, postgres::Error> {
client.execute("INSERT INTO SUBMISSION (id) VALUES ($1)", &[&id]) client.execute("INSERT INTO SUBMISSION (id) VALUES ($1)", &[&id]).await
} }
#[tokio::main] #[tokio::main]
@ -105,19 +106,23 @@ async fn main() {
let dsn = std::env::var("POSTGRES_DSN").expect("missing postgres dsn"); let dsn = std::env::var("POSTGRES_DSN").expect("missing postgres dsn");
let manager = let (client, connection) = tokio_postgres::connect(&dsn, tokio_postgres::NoTls).await.unwrap();
r2d2_postgres::PostgresConnectionManager::new(dsn.parse().unwrap(), postgres::NoTls);
let pool = r2d2::Pool::new(manager).unwrap(); tokio::spawn(async move {
if let Err(e) = connection.await {
panic!(e);
}
});
println!("Started");
'main: loop { 'main: loop {
let mut client = pool.get().unwrap(); println!("Fetching latest ID");
let latest_id = fa.latest_id().await.expect("unable to get latest id"); let latest_id = fa.latest_id().await.expect("unable to get latest id");
for id in ids_to_check(&mut client, latest_id) { for id in ids_to_check(&client, latest_id).await {
'attempt: for attempt in 0..3 { 'attempt: for attempt in 0..3 {
if !has_submission(&mut client, id) { if !has_submission(&client, id).await {
println!("loading submission {}", id); println!("loading submission {}", id);
let sub = match fa.get_submission(id).await { let sub = match fa.get_submission(id).await {
@ -139,7 +144,7 @@ async fn main() {
Some(sub) => sub, Some(sub) => sub,
None => { None => {
println!("did not exist"); println!("did not exist");
insert_null_submission(&mut client, id).unwrap(); insert_null_submission(&client, id).await.unwrap();
break 'attempt; break 'attempt;
} }
}; };
@ -152,7 +157,7 @@ async fn main() {
} }
}; };
insert_submission(&mut client, &sub).unwrap(); insert_submission(&client, &sub).await.unwrap();
break 'attempt; break 'attempt;
} }