#![allow(proc_macro_derive_resolution_fallback)]
#[macro_use]
extern crate structopt;
#[macro_use]
extern crate diesel;
extern crate dotenv;
#[macro_use]
extern crate diesel_migrations;
#[macro_use]
extern crate fake;
extern crate indicatif;
extern crate rand;
use diesel::pg::PgConnection;
use diesel::prelude::*;
use dotenv::dotenv;
use std::env;
use std::process::Command;
use std::error::Error;
use std::io;
use std::io::ErrorKind::InvalidInput;
use structopt::StructOpt;
use rand::seq::SliceRandom;
use rand::thread_rng;
use rayon::prelude::*;
use indicatif::{ProgressBar, ProgressStyle};
mod models;
mod schema;
mod pg_pool;
pub use pg_pool::DbConn;
pub use pg_pool::Pool;
embed_migrations!("./migrations");
use self::models::{
Choice, NewChoice, NewQuestion, NewSurvey, NewUser, NewVote, Question, Survey, User, Vote, Category, NewCategory
};
#[derive(StructOpt, Debug)]
#[structopt(
name = "birdseed",
about = "The libellis database seeder",
long_about = "You can use birdseed to seed a libellis db with junk data!"
)]
pub enum Birdseed {
#[structopt(name = "feed")]
Feed {
#[structopt(short = "r", long = "rows", default_value = "1000")]
row_count: u32,
},
#[structopt(name = "setup")]
Setup,
#[structopt(name = "migrate")]
Migrate {
#[structopt(short = "d", long = "database", default_value = "all")]
db: String,
},
#[structopt(name = "rebuild")]
Rebuild {
#[structopt(short = "d", long = "database", default_value = "all")]
db: String,
},
#[structopt(name = "clear")]
Clear,
}
pub fn run(config: Birdseed) -> Result<(), Box<dyn Error>> {
match config {
Birdseed::Feed { row_count } => populate_all(row_count),
Birdseed::Rebuild { db } => rebuild(&db),
Birdseed::Setup => setup(),
Birdseed::Migrate { db } => migrate(&db),
Birdseed::Clear => clear_all(),
}
}
fn setup() -> Result<(), Box<dyn Error>> {
drop_database("libellis");
drop_database("libellis_test");
println!("\r\n 🐦 Creating Main Database 🐦\r\n",);
setup_database("libellis");
println!("\r\n 🐦 Creating Test Database 🐦\r\n",);
setup_database("libellis_test");
println!("\r\n 🐦 Running Main DB Migrations 🐦\r\n",);
rebuild("main")?;
println!("\r\n 🐦 Running Test DB Migrations 🐦\r\n",);
rebuild("test")?;
println!("\r\n 🐦 All Done! 🐦\r\n",);
Ok(())
}
fn setup_database(database: &str) {
Command::new("createdb")
.arg(database)
.output()
.expect("failed to create database");
}
fn drop_database(database: &str) {
Command::new("dropdb")
.arg(database)
.output()
.expect("failed to drop database");
}
fn populate_all(row_count: u32) -> Result<(), Box<dyn Error>> {
dotenv().ok();
let base_url = env::var("PSQL_URL")?;
env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));
let pool = generate_pool();
println!("\r\n 🐦 Seeding All Tables 🐦\r\n",);
let bar = ProgressBar::new((row_count * 11) as u64);
bar.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
.progress_chars("##-"),
);
let usernames = populate_users(&pool, row_count, &bar)?;
let _ = populate_categories(&pool, "TestCategory", &bar)?;
let survey_ids = populate_surveys(&pool, &usernames, row_count, &bar)?;
let question_ids = populate_questions(&pool, &survey_ids, row_count, &bar)?;
let choice_ids = populate_choices(&pool, &question_ids, row_count, &bar)?;
populate_votes(&pool, &usernames, &choice_ids, &bar)?;
bar.finish();
println!("\r\n 🐦 Birdseed has Finished Seeding! 🐦\r\n",);
Ok(())
}
fn migrate(database: &str) -> Result<(), Box<dyn Error>> {
dotenv().ok();
let base_url = env::var("PSQL_URL")?;
match database {
"all" | "a" => {
migrate_main(&base_url)?;
migrate_test(&base_url)?;
}
"main" | "m" => migrate_main(&base_url)?,
"test" | "t" => migrate_test(&base_url)?,
_ => {
return Err(io::Error::new(
InvalidInput,
"Invalid Database Type, choose 'main', 'test', or 'all'",
)
.into());
}
};
Ok(())
}
fn migrate_main(base_url: &String) -> Result<(), Box<dyn Error>> {
env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));
let conn = establish_connection();
println!("\r\n 🐦 Running Migrations on Main DB 🐦\r\n");
embedded_migrations::run_with_output(&conn, &mut std::io::stdout())?;
Ok(())
}
fn migrate_test(base_url: &String) -> Result<(), Box<dyn Error>> {
env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis_test"));
let conn = establish_connection();
println!("\r\n 🐦 Running Migrations on Test DB 🐦\r\n");
embedded_migrations::run_with_output(&conn, &mut std::io::stdout())?;
Ok(())
}
fn clear_all() -> Result<(), Box<dyn Error>> {
use self::schema::*;
dotenv().ok();
let base_url = env::var("PSQL_URL")?;
std::env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));
let conn = establish_connection();
println!("\r\n 🐦 Clearing all Tables 🐦\r\n");
let bar = ProgressBar::new(5);
bar.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
.progress_chars("##-"),
);
diesel::delete(votes::table).execute(&conn)?;
bar.inc(1);
diesel::delete(choices::table).execute(&conn)?;
bar.inc(1);
diesel::delete(questions::table).execute(&conn)?;
bar.inc(1);
diesel::delete(surveys::table).execute(&conn)?;
bar.inc(1);
diesel::delete(categories::table).execute(&conn)?;
bar.inc(1);
diesel::delete(users::table).execute(&conn)?;
bar.inc(1);
bar.finish();
println!("\r\n 🐦 All Tables Cleared! 🐦\r\n");
Ok(())
}
fn drop_all(conn: &PgConnection) {
let bar = ProgressBar::new(8);
bar.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
.progress_chars("##-"),
);
conn.execute("DROP VIEW users_votes");
bar.inc(1);
conn.execute("DROP TABLE votes");
bar.inc(1);
conn.execute("DROP TABLE choices");
bar.inc(1);
conn.execute("DROP TABLE questions");
bar.inc(1);
conn.execute("DROP TABLE surveys");
bar.inc(1);
conn.execute("DROP TABLE categories");
bar.inc(1);
conn.execute("DROP TABLE users");
bar.inc(1);
conn.execute("DROP TABLE __diesel_schema_migrations");
bar.inc(1);
bar.finish();
}
fn rebuild(database: &str) -> Result<(), Box<dyn Error>> {
dotenv().ok();
let base_url = env::var("PSQL_URL")?;
match database {
"all" | "a" => {
rebuild_main(&base_url)?;
rebuild_test(&base_url)?;
}
"main" | "m" => rebuild_main(&base_url)?,
"test" | "t" => rebuild_test(&base_url)?,
_ => {
return Err(io::Error::new(
InvalidInput,
"Invalid Database Type, choose 'main', 'test', or 'all'",
)
.into());
}
};
Ok(())
}
fn rebuild_main(base_url: &str) -> Result<(), Box<dyn Error>> {
std::env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));
println!("\r\n 🐦 Connecting to Main DB 🐦\r\n");
let conn = establish_connection();
println!("\r\n 🐦 Dropping all Tables 🐦\r\n");
drop_all(&conn);
println!("\r\n 🐦 Running Migrations 🐦\r\n");
embedded_migrations::run_with_output(&conn, &mut std::io::stdout())?;
println!("\r\n 🐦 Tables Successfully Rebuilt! 🐦\r\n");
Ok(())
}
fn rebuild_test(base_url: &str) -> Result<(), Box<dyn Error>> {
std::env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis_test"));
println!("\r\n 🐦 Connecting to Test DB 🐦\r\n");
let conn = establish_connection();
println!("\r\n 🐦 Dropping all Tables 🐦\r\n");
drop_all(&conn);
println!("\r\n 🐦 Running Migrations 🐦\r\n");
embedded_migrations::run_with_output(&conn, &mut std::io::stdout())?;
println!("\r\n 🐦 Tables Successfully Rebuilt! 🐦\r\n");
Ok(())
}
fn populate_users(
pool: &Pool,
row_count: u32,
bar: &ProgressBar,
) -> Result<Vec<String>, Box<dyn Error>> {
bar.set_message(&format!("Seeding {} users", row_count));
let usernames: Vec<String> = (0..row_count).into_par_iter().map(|_| {
let pool = pool.clone();
let conn = pool.get().unwrap();
let user = format!(
"{}{}",
fake!(Internet.user_name),
fake!(Number.between(90, 9999))
);
let pw = format!(
"{}{}",
fake!(Name.name),
fake!(Number.between(10000, 99999))
);
let em = format!("{}@gmail.com", user);
let first = format!("{}", fake!(Name.first_name));
let last = format!("{}", fake!(Name.last_name));
create_user(&conn, &user, &pw, &em, &first, &last);
bar.inc(1);
user
}).collect();
Ok(usernames)
}
fn populate_categories(
pool: &Pool,
title: &str,
bar: &ProgressBar,
) -> Result<String, Box<dyn Error>> {
bar.set_message(&format!("Seeding 1 Test Category"));
let pool = pool.clone();
let conn = pool.get().unwrap();
create_category(&conn, title);
Ok(title.to_string())
}
fn populate_surveys(
pool: &Pool,
authors: &Vec<String>,
row_count: u32,
bar: &ProgressBar,
) -> Result<Vec<i32>, Box<dyn Error>> {
bar.set_message(&format!("Seeding {} surveys", row_count));
let survey_ids: Vec<i32> = authors.par_iter().map(|auth| {
let pool = pool.clone();
let conn = pool.get().unwrap();
let survey_title = format!("{}", fake!(Lorem.sentence(4, 8)));
let cat = "TestCategory";
let survey = create_survey(&conn, &auth, &survey_title, cat);
bar.inc(1);
survey.id
}).collect();
Ok(survey_ids)
}
fn populate_questions(
pool: &Pool,
survey_ids: &Vec<i32>,
row_count: u32,
bar: &ProgressBar,
) -> Result<Vec<i32>, Box<dyn Error>> {
bar.set_message(&format!("Seeding {} questions", row_count));
let question_ids: Vec<i32> = survey_ids.par_iter().map(|s_id| {
let pool = pool.clone();
let conn = pool.get().unwrap();
let q_title = format!("{}", fake!(Lorem.sentence(3, 7)));
let q_type = "multiple".to_string();
let question = create_question(&conn, *s_id, &q_type, &q_title);
bar.inc(1);
question.id
}).collect();
Ok(question_ids)
}
fn populate_choices(
pool: &Pool,
question_ids: &Vec<i32>,
row_count: u32,
bar: &ProgressBar,
) -> Result<Vec<i32>, Box<dyn Error>> {
bar.set_message(&format!("Seeding {} choices", (row_count * 4)));
let choice_ids: Vec<i32> = question_ids.par_iter().map(|q_id| {
(0..4).into_par_iter().map(|_| {
let pool = pool.clone();
let conn = pool.get().unwrap();
let c_title = format!("{}", fake!(Lorem.sentence(1, 4)));
let c_type = "text".to_string();
let choice = create_choice(&conn, *q_id, &c_type, &c_title);
bar.inc(1);
choice.id
}).collect()
}).collect::<Vec<Vec<i32>>>().concat();
Ok(choice_ids)
}
fn populate_votes(
pool: &Pool,
authors: &Vec<String>,
choice_ids: &Vec<i32>,
bar: &ProgressBar,
) -> Result<(), Box<dyn Error>> {
bar.set_message(&format!("{} users are voting", (authors.len())));
let mut rng = thread_rng();
let mut choice_idxs: Vec<usize> = (0..choice_ids.len()).collect();
let choice_slice: &mut [usize] = &mut choice_idxs;
let mut author_idxs: Vec<usize> = (0..authors.len()).collect();
let author_slice: &mut [usize] = &mut author_idxs;
choice_slice.shuffle(&mut rng);
author_slice.shuffle(&mut rng);
author_slice.par_iter().enumerate().for_each(|(i, rand_i)| {
let name = &authors[*rand_i];
if i < authors.len() - 1 {
(1..5 as usize).into_par_iter().for_each(|j| {
let pool = pool.clone();
let conn = pool.get().unwrap();
let c_id = choice_ids[choice_slice[(i + 1) * j]];
create_vote(&conn, c_id, name, 1);
bar.inc(1);
});
}
});
Ok(())
}
fn establish_connection() -> PgConnection {
dotenv().ok();
let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
PgConnection::establish(&database_url).expect(&format!("Error connecting to {}", database_url))
}
fn generate_pool() -> Pool {
dotenv().ok();
let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
pg_pool::init(&database_url)
}
fn create_user<'a>(
conn: &PgConnection,
user: &'a str,
pw: &'a str,
em: &'a str,
first: &'a str,
last: &'a str,
) -> User {
use self::schema::users;
let new_user = NewUser {
username: user,
password: pw,
email: em,
first_name: first,
last_name: last,
is_admin: false,
};
diesel::insert_into(users::table)
.values(&new_user)
.get_result(conn)
.expect("Error saving new user")
}
fn create_survey<'a>(conn: &PgConnection, auth: &'a str, survey_title: &'a str, cat: &'a str) -> Survey {
use self::schema::surveys;
let new_survey = NewSurvey {
author: auth,
title: survey_title,
published: true,
category: cat,
};
diesel::insert_into(surveys::table)
.values(&new_survey)
.get_result(conn)
.expect("Error saving new survey")
}
fn create_question<'a>(
conn: &PgConnection,
s_id: i32,
q_type: &'a str,
q_title: &'a str,
) -> Question {
use self::schema::questions;
let new_question = NewQuestion {
survey_id: s_id,
question_type: q_type,
title: q_title,
};
diesel::insert_into(questions::table)
.values(&new_question)
.get_result(conn)
.expect("Error saving new question")
}
fn create_choice<'a>(conn: &PgConnection, q_id: i32, c_type: &'a str, c_title: &'a str) -> Choice {
use self::schema::choices;
let new_choice = NewChoice {
question_id: q_id,
content_type: c_type,
title: c_title,
};
diesel::insert_into(choices::table)
.values(&new_choice)
.get_result(conn)
.expect("Error saving new choice")
}
fn create_vote<'a>(conn: &PgConnection, c_id: i32, name: &'a str, points: i32) -> Vote {
use self::schema::votes;
let new_vote = NewVote {
choice_id: c_id,
username: name,
score: points,
};
diesel::insert_into(votes::table)
.values(&new_vote)
.get_result(conn)
.expect("Error saving new vote")
}
fn create_category<'a>(conn: &PgConnection, title: &'a str) -> Category {
use self::schema::categories;
let new_category = NewCategory {
title,
};
diesel::insert_into(categories::table)
.values(&new_category)
.get_result(conn)
.expect("Error saving new vote")
}