match CREATE_MODE

This commit is contained in:
Elijah McMorris 2024-02-06 10:05:58 -08:00
parent 305bf5273b
commit 38fdee5728
Signed by: NexVeridian
SSH key fingerprint: SHA256:bsA1SKZxuEcEVHAy3gY1HUeM5ykRJl0U0kQHQn0hMg8
7 changed files with 87 additions and 60 deletions

View file

@ -11,14 +11,22 @@ lazy_static! {
env::var("WIKIDATA_FILE_FORMAT").expect("FILE_FORMAT not set");
static ref WIKIDATA_FILE_NAME: String =
env::var("WIKIDATA_FILE_NAME").expect("FILE_NAME not set");
static ref THREADED_REQUESTS: bool = env::var("THREADED_REQUESTS")
.expect("THREADED_REQUESTS not set")
.parse()
.expect("Failed to parse THREADED_REQUESTS");
static ref WIKIDATA_BULK_INSERT: bool = env::var("WIKIDATA_BULK_INSERT")
.expect("WIKIDATA_BULK_INSERT not set")
.parse()
.expect("Failed to parse WIKIDATA_BULK_INSERT");
static ref CREATE_MODE: CreateMode = match env::var("CREATE_MODE")
.expect("CREATE_MODE not set")
.as_str()
{
"Single" => CreateMode::Single,
"ThreadedSingle" => CreateMode::ThreadedSingle,
"ThreadedBulk" => CreateMode::ThreadedBulk,
_ => panic!("Unknown CREATE_MODE"),
};
}
#[derive(Clone, Copy)]
pub enum CreateMode {
Single,
ThreadedSingle,
ThreadedBulk,
}
#[tokio::main]
@ -29,51 +37,55 @@ async fn main() -> Result<(), Error> {
let db = create_db_ws().await?;
let reader = File_Format::new(&WIKIDATA_FILE_FORMAT).reader(&WIKIDATA_FILE_NAME)?;
if !*THREADED_REQUESTS {
let mut counter = 0;
for line in reader.lines() {
let mut retries = 0;
let line = line?;
match *CREATE_MODE {
CreateMode::Single => {
let mut counter = 0;
for line in reader.lines() {
let mut retries = 0;
let line = line?;
loop {
if create_db_entity(&db, &line).await.is_ok() {
break;
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;
};
}
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);
counter += 1;
if counter % 100 == 0 {
pb.inc(100);
}
}
}
} else if *WIKIDATA_BULK_INSERT {
create_db_entities_threaded(
None::<Surreal<Client>>,
reader,
Some(pb.clone()),
2500,
100,
CreateVersion::Bulk,
)
.await?;
} else {
create_db_entities_threaded(
None::<Surreal<Client>>,
reader,
Some(pb.clone()),
2500,
100,
CreateVersion::Single,
)
.await?;
CreateMode::ThreadedSingle => {
create_db_entities_threaded(
None::<Surreal<Client>>,
reader,
Some(pb.clone()),
2_500,
100,
CreateVersion::Single,
)
.await?;
}
CreateMode::ThreadedBulk => {
create_db_entities_threaded(
None::<Surreal<Client>>,
reader,
Some(pb.clone()),
500,
1000,
CreateVersion::Bulk,
)
.await?;
}
}
pb.finish();