refactor init_db and init_progress_bar

src/utils/init_db.rs

@@ -0,0 +1,38 @@
+use anyhow::Error;
+use anyhow::Result;
+use lazy_static::lazy_static;
+use std::env;
+use surrealdb::{
+    engine::{
+        local::{Db, Mem},
+        remote::ws::{Client, Ws},
+    },
+    opt::auth::Root,
+    Surreal,
+};
+
+lazy_static! {
+    static ref DB_USER: String = env::var("DB_USER").expect("DB_USER not set");
+    static ref DB_PASSWORD: String = env::var("DB_PASSWORD").expect("DB_PASSWORD not set");
+    static ref WIKIDATA_DB_PORT: String =
+        env::var("WIKIDATA_DB_PORT").expect("WIKIDATA_DB_PORT not set");
+}
+
+pub async fn create_db_ws() -> Result<Surreal<Client>, Error> {
+    let db = Surreal::new::<Ws>(WIKIDATA_DB_PORT.as_str()).await?;
+
+    db.signin(Root {
+        username: &DB_USER,
+        password: &DB_PASSWORD,
+    })
+    .await?;
+    db.use_ns("wikidata").use_db("wikidata").await?;
+
+    Ok(db)
+}
+
+pub async fn create_db_mem() -> Result<Surreal<Db>, Error> {
+    let db = Surreal::new::<Mem>(()).await?;
+    db.use_ns("wikidata").use_db("wikidata").await?;
+    Ok(db)
+}

src/utils/init_progress_bar.rs

@@ -0,0 +1,22 @@
+use indicatif::{ProgressBar, ProgressState, ProgressStyle};
+
+pub async fn create_pb() -> ProgressBar {
+    let total_size = 110_000_000;
+    let pb = ProgressBar::new(total_size);
+    pb.set_style(
+        ProgressStyle::with_template(
+            "[{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ETA:[{eta}]",
+        )
+        .unwrap()
+        .with_key(
+            "eta",
+            |state: &ProgressState, w: &mut dyn std::fmt::Write| {
+                let sec = state.eta().as_secs();
+                let min = (sec / 60) % 60;
+                let hr = (sec / 60) / 60;
+                write!(w, "{}:{:02}:{:02}", hr, min, sec % 60).unwrap()
+            },
+        ),
+    );
+    pb
+}