more error handling and overwrite option

src/main.rs

@@ -1,8 +1,9 @@
 use anyhow::{Error, Ok, Result};
 use indicatif::{ProgressBar, ProgressState, ProgressStyle};
 use lazy_static::lazy_static;
-use std::{env, fmt::Write, io::BufRead, thread, time::Duration};
+use std::{env, fmt::Write, io::BufRead};
 use surrealdb::{engine::remote::ws::Ws, opt::auth::Root, Surreal};
+use tokio::time::{sleep, Duration};
 mod utils;
 use utils::*;
 
@@ -18,7 +19,7 @@ lazy_static! {
 
 #[tokio::main]
 async fn main() -> Result<(), Error> {
-    thread::sleep(Duration::from_secs(10));
+    sleep(Duration::from_secs(10)).await;
     let total_size = 113_000_000;
 
     let pb = ProgressBar::new(total_size);
@@ -48,14 +49,30 @@ async fn main() -> Result<(), Error> {
     if !*THREADED_REQUESTS {
         let mut counter = 0;
         for line in reader.lines() {
-            create_db_entity(&db, line?).await?;
+            let mut retries = 0;
+            let line = line?;
+
+            loop {
+                if create_db_entity(&db, &line).await.is_ok() {
+                    break;
+                }
+                if retries >= 60 * 10 {
+                    panic!("Failed to create entities, too many retries");
+                }
+                retries += 1;
+                sleep(Duration::from_secs(1)).await;
+                if db.use_ns("wikidata").use_db("wikidata").await.is_err() {
+                    continue;
+                };
+            }
+
             counter += 1;
             if counter % 100 == 0 {
                 pb.inc(100);
             }
         }
     } else {
-        create_db_entities_threaded(&db, reader, Some(pb.clone()), 1000, 100).await?;
+        create_db_entities_threaded(&db, reader, Some(pb.clone()), 2500, 100).await?;
     }
 
     pb.finish();

src/utils.rs

@@ -2,17 +2,27 @@ use anyhow::{Error, Ok, Result};
 use bzip2::read::MultiBzDecoder;
 use futures::future::join_all;
 use indicatif::ProgressBar;
+use lazy_static::lazy_static;
 use serde_json::{from_str, Value};
 use std::{
+    env,
     fs::File,
     io::{BufRead, BufReader},
 };
 use surrealdb::{Connection, Surreal};
+use tokio::time::{sleep, Duration};
 use wikidata::Entity;
 
 mod tables;
 use tables::*;
 
+lazy_static! {
+    static ref OVERWRITE_DB: bool = env::var("OVERWRITE_DB")
+        .expect("OVERWRITE_DB not set")
+        .parse()
+        .expect("Failed to parse OVERWRITE_DB");
+}
+
 #[allow(non_camel_case_types)]
 pub enum File_Format {
     json,
@@ -35,7 +45,7 @@ impl File_Format {
     }
 }
 
-pub async fn create_db_entity<C: Connection>(db: &Surreal<C>, line: String) -> Result<(), Error> {
+pub async fn create_db_entity<C: Connection>(db: &Surreal<C>, line: &str) -> Result<(), Error> {
     let line = line.trim().trim_end_matches(',').to_string();
     if line == "[" || line == "]" {
         return Ok(());
@@ -48,15 +58,13 @@ pub async fn create_db_entity<C: Connection>(db: &Surreal<C>, line: String) -> R
 
     let id = data.id.clone().expect("No ID");
     data.id = None;
-    let _ = db.create::<Option<EntityMini>>(&id).await.is_err();
-    {
+    if db.create::<Option<EntityMini>>(&id).await.is_err() && *OVERWRITE_DB {
         db.update::<Option<EntityMini>>(&id).content(data).await?;
-    };
+    }
 
     let id = claims.id.clone().expect("No ID");
     claims.id = None;
-    let _ = db.create::<Option<Claims>>(&id).await.is_err();
-    {
+    if db.create::<Option<Claims>>(&id).await.is_err() && *OVERWRITE_DB {
         db.update::<Option<Claims>>(&id).content(claims).await?;
     }
     Ok(())
@@ -64,12 +72,12 @@ pub async fn create_db_entity<C: Connection>(db: &Surreal<C>, line: String) -> R
 
 pub async fn create_db_entities<C: Connection>(
     db: &Surreal<C>,
-    lines: Vec<String>,
-    pb: Option<ProgressBar>,
+    lines: &Vec<String>,
+    pb: &Option<ProgressBar>,
 ) -> Result<(), Error> {
     let mut counter = 0;
     for line in lines {
-        create_db_entity(db, line.to_string()).await?;
+        create_db_entity(db, line).await?;
         counter += 1;
         if counter % 100 == 0 {
             if let Some(ref p) = pb {
@@ -92,7 +100,7 @@ pub async fn create_db_entities_threaded<C: Connection>(
     let mut chunk_counter = 0;
 
     for line in reader.lines() {
-        chunk.push(line.unwrap());
+        chunk.push(line?);
 
         if chunk.len() >= batch_size {
             let db = db.clone();
@@ -100,7 +108,20 @@ pub async fn create_db_entities_threaded<C: Connection>(
             let pb = pb.clone();
 
             futures.push(tokio::spawn(async move {
-                create_db_entities(&db, lines, pb).await.unwrap();
+                let mut retries = 0;
+                loop {
+                    if create_db_entities(&db, &lines, &pb).await.is_ok() {
+                        break;
+                    }
+                    if retries >= 60 * 10 {
+                        panic!("Failed to create entities, too many retries");
+                    }
+                    retries += 1;
+                    sleep(Duration::from_secs(1)).await;
+                    if db.use_ns("wikidata").use_db("wikidata").await.is_err() {
+                        continue;
+                    };
+                }
             }));
             chunk_counter += 1;
             chunk.clear();
@@ -113,7 +134,7 @@ pub async fn create_db_entities_threaded<C: Connection>(
         }
     }
 
-    create_db_entities(db, chunk, pb).await.unwrap();
+    create_db_entities(db, &chunk, &pb).await?;
     join_all(futures).await;
     Ok(())
 }