NexVeridian/wikidata-to-surrealdb
1
0
Fork 0
mirror of https://github.com/NexVeridian/wikidata-to-surrealdb.git synced 2025-09-02 09:59:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

match CREATE_MODE

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

108
src/main.rs
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();

6
src/utils.rs
View file

@ -80,7 +80,7 @@ pub async fn create_db_entity(db: &Surreal<impl Connection>, line: &str) -> Resu
pub async fn create_db_entities(
db: &Surreal<impl Connection>,
lines: &Vec<String>,
lines: &[String],
pb: &Option<ProgressBar>,
) -> Result<(), Error> {
let mut counter = 0;
@ -168,7 +168,7 @@ impl CreateVersion {
pub async fn run(
self,
db: &Surreal<impl Connection>,
chunk: &Vec<String>,
chunk: &[String],
pb: &Option<ProgressBar>,
batch_size: usize,
) -> bool {
@ -234,7 +234,7 @@ pub async fn create_db_entities_threaded(
panic!("Failed to create entities, too many retries");
}
retries += 1;
sleep(Duration::from_millis(100)).await;
sleep(Duration::from_millis(250)).await;
}
}));
chunk_counter += 1;