birdseed 0.2.0

// Licensed under the the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(proc_macro_derive_resolution_fallback)]

//! This crate's entire purpose is to seed the
//! [libellis](https://github.com/libellis/libellis-backend) database with fake data (generated by
//! [faker](https://github.com/tikotzky/faker-rs) library).
//!
//! ## Setup
//!
//! You MUST set a `PSQL_URL` environment variable to your libellis postgres database.
//!
//! ### Example
//! ```terminal
//! $ export PSQL_URL=postgres://username:password@localhost/
//! ```
//!
//! Note the ending forward slash which is necessary. You can optionally set this in a `.env` file
//! in the root folder of this project if you are running the project from a local folder.
//!
//! ## Installation
//!
//! You may install this in one of two ways. If you have `cargo` installed then it's very easy. If
//! not, you can install rust and cargo by following this very simple [cargo setup](https://doc.rust-lang.org/cargo/getting-started/installation.html) process.
//! Once you have cargo installed you can install this terminal application by running:
//!
//! ```terminal
//! $ cargo install birdseed
//! ```
//!
//! Optionally you may instead clone this repo and in the root directory build the release version
//! of this crate:
//!
//! ```terminal
//! $ git clone https://github.com/libellis/birdseed.git
//! $ cd birdseed
//! $ cargo build --release
//! ```
//!
//! ## Features
//!
//! ### `setup`
//!
//! You can setup the main libellis and libellis_test databases with this subcommand.  It will
//! attempt to drop both libellis and libellis_test before creating them so be careful! Only use
//! this if you don't need the data in your libellis database and want to start over, or are
//! creating your libellis databases for the first time.
//!
//! ```terminal
//! $ birdseed setup
//! ```
//!
//! ### `rebuild`
//!
//! You can rebuild all tables according to embedded diesel migrations. This drops each table (but
//! does not drop the database itself) and then rebuilds all tables. Note that you must already
//! have `libellis` and `libellis_test` databases for this to work. If you do not have those
//! databases run `birdseed setup` instead.
//!
//! ```terminal
//! $ birdseed rebuild
//! ```
//!
//! `rebuild` by default will rebuild all tables in both your main and test databases. If you would
//! like to specify to only rebuild one database, pass in 'main' or 'test' to the -database
//! argument flag:
//!
//! ### `fences`
//!
//! You can load in fence data from a geojson file with the fences subcommand:
//!
//! ```terminal
//! $ birdseed fences
//! ```
//!
//! By default it looks for a file called `fences.json` in the data folder from the root of this
//! crate. This folder only exists if you cloned the repo.  To specify a filepath yourself pass the
//! -f or -file flag after the fences subcommand:
//!
//! ```terminal
//! $ birdseed fences -f BerkeleyNeighborhoods.json
//! ```
//!
//! Note: This only works if you have a fences table - which would have been setup for you from the
//! most recent migrations when running `birdseed setup` or `birdseed migrate`.
//!
//! ### `feed`
//!
//! You can seed all databases with the `feed` subcommand:
//!
//! ```terminal
//! $ birdseed feed
//! ```
//!
//! We can specify a row count (overriding the default of 1000 rows):
//!
//! ```terminal
//! $ birdseed feed -r 10000
//! ```
//!
//! In this exampe we override the default of 1,000 rows and instead seed 10,000 rows.
//!
//! Note: What the row count really means is that we will seed row count amount of users, surveys
//! and questions, but row count * 4 amount of choices and votes.
//!
//! ```terminal
//! $ birdseed rebuild -database main
//! ```
//!
//! You can also use `-d` for shorthand:
//!
//! ```terminal
//! $ birdseed rebuild -d test
//! ```
//!
//! ### `migrate`
//!
//! To run migrations, use the migrate subcommand (this will update your database schema to the
//! most recent schema).
//!
//! ```terminal
//! $ birdseed migrate
//! ```
//!
//! By default this runs migrations on all databases. To run migrations on only main run:
//!
//! ```terminal
//! $ birdseed migrate -d main
//! ```
//!
//! To run migrations only on the test database run:
//!
//! ```terminal
//! $ birdseed migrate -d test
//! ```
//!
//! ### `clear`
//!
//! You can clear all tables with the `clear` subcommand:
//!
//! ```terminal
//! $ birdseed clear
//! ```
//!
//! ### `icecream`
//!
//! For fun and profit you can seed the database with an row count amount of users, a single poll
//! about icecream, and then populate that poll with fake votes from your newly faked user pool,
//! and have all of their votes counted from legitimate randomized locations within the city of San
//! Francisco.
//!
//! ```terminal
//! $ birdseed icecream
//! ```
//!
//! By default the row count is 1000, and can be overriden in the same way as when using the `feed`
//! subcommand:
//!
//! ```terminal
//! $ birdseed icecream -r 10000
//! ```
extern crate structopt;

#[macro_use]
extern crate diesel;
extern crate dotenv;

#[macro_use]
extern crate serde;

#[macro_use]
extern crate serde_json;

#[macro_use]
extern crate serde_derive;

#[macro_use]
extern crate diesel_migrations;

#[macro_use]
extern crate fake;

extern crate indicatif;
extern crate rand;

use std::fs::File;
use std::io::prelude::*;

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 indicatif::{ProgressBar, ProgressStyle};

/// This module provides crud for all our models, as well as populating fake data into each
/// table. It also contains the module for creating our pg pool.
pub mod db;

/// This module provides all of our insertable and queryable diesel models.
pub mod models;

mod schema;
/// Provides extra SQL functions not provided by default in diesel.
pub mod sql_functions;

use db::*;
use pg_pool::Pool;

embed_migrations!("./migrations");

#[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!"
)]
/// You can use birdseed to seed a libellis db with junk data!
pub enum Birdseed {
    #[structopt(name = "feed")]
    /// Seeds random data into all tables
    Feed {
        /// How many rows to inject
        #[structopt(short = "r", long = "rows", default_value = "1000")]
        row_count: u32,
    },

    #[structopt(name = "icecream")]
    /// Seeds random icrecream data into all tables
    Icecream {
        /// How many rows to inject
        #[structopt(short = "r", long = "rows", default_value = "1000")]
        row_count: u32,
    },

    #[structopt(name = "fences")]
    /// Loads in fences from a geojson file
    Fences {
        /// The file name to read from - default if not supplied is
        #[structopt(short = "f", long = "file", default_value = "data/fences.json")]
        filepath: String,
    },

    #[structopt(name = "setup")]
    /// Builds both libellis main and test databases and runs migrations
    Setup,

    #[structopt(name = "migrate")]
    /// Builds both libellis main and test databases and runs migrations
    Migrate {
        /// Which database to run migrations on, main, test, or all
        #[structopt(short = "d", long = "database", default_value = "all")]
        db: String,
    },

    #[structopt(name = "rebuild")]
    /// Rebuilds all tables per most recent schema (embedded in binary)
    Rebuild {
        /// Which database to run rebuild on, main, test, or all
        #[structopt(short = "d", long = "database", default_value = "all")]
        db: String,
    },

    #[structopt(name = "clear")]
    /// Clears all tables in libellis database
    Clear,
}

/// `run` will take in a Birdseed enum config (parsed in `main`) and either clear all tables or
/// populate all tables with number of rows specified in -r (1000 by default)
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::Icecream { row_count } => populate_icecream(row_count),
        Birdseed::Fences { filepath } => load_fences(filepath),
        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(())
}

// Create PSQL database specified.
fn setup_database(database: &str) {
    Command::new("createdb")
        .arg(database)
        .output()
        .expect("failed to create database");
}

// Drop PSQL database specified.
fn drop_database(database: &str) {
    Command::new("dropdb")
        .arg(database)
        .output()
        .expect("failed to drop database");
}

// Load all the fences stored in a geojson file by supplying a filepath.
fn load_fences(filepath: String) -> Result<(), Box<dyn Error>> {
    use std::io::BufReader;

    dotenv().ok();
    let base_url = env::var("PSQL_URL")?;
    env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));

    let pool = pg_pool::generate_pool();

    // read json file into a string
    let file = File::open(filepath)?;
    let mut buf_reader = BufReader::new(file);
    let mut contents = String::new();
    buf_reader.read_to_string(&mut contents)?;

    // load geo data into fences table
    fences::load_geo_data(&pool, &contents)?;

    Ok(())
}

// Kicks off populating all tables in main database and updating user
// with visual progress bar along the way
fn populate_all(row_count: u32) -> Result<(), Box<dyn Error>> {
    // get the base url and append it with the db name
    dotenv().ok();
    let base_url = env::var("PSQL_URL")?;
    env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));

    let pool = pg_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 = users::populate(&pool, row_count, &bar)?;

    let categories = categories::populate(&pool, row_count, &bar)?;

    let survey_ids = surveys::populate(&pool, &usernames, &categories, row_count, &bar)?;
    let question_ids = questions::populate(&pool, &survey_ids, row_count, &bar)?;
    let choice_ids = choices::populate(&pool, &question_ids, row_count, &bar)?;
    votes::populate(&pool, &usernames, &choice_ids, &bar)?;
    bar.finish();
    println!("\r\n             🐦 Birdseed has Finished Seeding! 🐦\r\n",);
    Ok(())
}

// Kicks off populating all tables in main database and updating user
// with visual progress bar along the way. This was hacked together
// quickly at a hackathon
//
// TODO: Break out into smaller functions.
fn populate_icecream(row_count: u32) -> Result<(), Box<dyn Error>> {
    // get the base url and append it with the db name
    dotenv().ok();
    let base_url = env::var("PSQL_URL")?;
    env::set_var("DATABASE_URL", &format!("{}{}", base_url, "libellis"));

    let pool = pg_pool::generate_pool();
    println!("\r\n                  🐦 Seeding All Tables 🐦\r\n",);

    let bar = ProgressBar::new((row_count * 9) as u64);
    bar.set_style(
        ProgressStyle::default_bar()
            .template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
            .progress_chars("##-"),
    );

    let usernames = users::populate(&pool, row_count, &bar)?;

    {
        let pool = pool.clone();
        let conn = pool.get().unwrap();
        categories::create(&conn, "food");
    }

    let survey_id;

    // scoped so the pool connection gets dropped automatically
    // inject the icecream survey
    {
        let pool = pool.clone();
        let conn = pool.get().unwrap();

        let survey_title = "What is your favorite icecream?";

        let cat = "food";

        let survey = surveys::create(&conn, &usernames[0], &survey_title, cat);

        survey_id = survey.id;
    }

    let question_id;
    // inject the one question for the icecream survey
    {
        let pool = pool.clone();
        let conn = pool.get().unwrap();

        let q_title = "What is your favorite icecream?";
        let q_type = "multiple".to_string();

        let question = questions::create(&conn, survey_id, &q_type, &q_title);

        question_id = question.id;
    }

    let choices = vec!["Strawberry", "Vanilla", "Chocolate"];
    let choice_ids;

    // inject our icecream flavors for fake users to fake vote on
    {
        let pool = pool.clone();
        let conn = pool.get().unwrap();

        choice_ids = choices
            .into_iter()
            .map(|choice| {
                let c_type = "text".to_string();
                let choice_struct = choices::create(&conn, question_id, &c_type, choice);
                choice_struct.id
            })
            .collect();
    }

    // populate fake votes with our newly created fake users
    votes::populate_icecream(&pool, &usernames, &choice_ids, &bar)?;
    bar.finish();
    println!("\r\n             🐦 Birdseed has Finished Seeding! 🐦\r\n",);
    Ok(())
}

// Runs migrations based on the database passed in (main, test or all)
fn migrate(database: &str) -> Result<(), Box<dyn Error>> {
    // get the base url and append it with the db name
    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(())
}

// Runs migrations on the main database.
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(())
}

// Runs migrations on the test database.
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(())
}

// Clears all tables in appropriate order and increments a progress bar with
// custom start and completion messages
fn clear_all() -> Result<(), Box<dyn Error>> {
    use self::schema::*;

    // get the base url and append it with the db name
    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(7);
    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(fences::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(())
}

// Drops all tables from libellis database
// Note that it's very intentional that we're not dealing with the result ->
// We don't care if a table doesn't exist, we simply want to attempt to drop
// ALL regardless so we are ready for a rebuild.  This takes care of situations
// where a user has manually deleted some but not all of their tables. Open a PR
// request if you have a better solution in mind.
#[allow(unused_must_use)]
fn drop_all(conn: &PgConnection) {
    let drop_statements = vec![
        "DROP VIEW users_votes",
        "DROP TABLE votes",
        "DROP TABLE fences cascade",
        "DROP EXTENSION postgis",
        "DROP TABLE choices",
        "DROP TABLE questions",
        "DROP TABLE surveys",
        "DROP TABLE categories",
        "DROP TABLE users",
        "DROP TABLE __diesel_schema_migrations",
    ];

    let bar = ProgressBar::new(drop_statements.len() as u64);
    bar.set_style(
        ProgressStyle::default_bar()
            .template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
            .progress_chars("##-"),
    );

    drop_statements.iter().for_each(|statement| {
        conn.execute(statement);
        bar.inc(1);
    });

    bar.finish();
}

// Rebuilds all tables per most recent embedded diesel migrations
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(())
}

// Drops all tables in the main database and rebuilds them by running embedded migrations
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(())
}

// Drops all tables in the test database and rebuilds them by running embedded migrations
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(())
}

// Establishes a connection to the libellis postgres database on your machine, as specified by your
// DATABASE_URL environment variable. Returns a single PgConnection
//
// Likely don't need this farther down the road once all functions have been re-factored to use
// a pool.
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))
}