pique/_experiments/2024-03-02-database-benchmark/src/bin/main.rs
Nicole Tietz-Sokolskaya 77d4ebb371 Add ADR for our choice of SQLite as our primary database,
backed up by experiments demonstrating that SQLite will meet all of our
requirements.

This also introduces ADRs in the repo, and adds a README in preparation
making the repository public.
2024-03-16 11:12:46 -04:00

84 lines
2.1 KiB
Rust

use bench::data::random_entities;
use env_logger::{Builder, Env};
use log::info;
use migration::{Migrator, MigratorTrait};
use rand::prelude::*;
use sea_orm::prelude::*;
use sea_orm::sea_query::{Func, SimpleExpr};
use sea_orm::ConnectOptions;
use sea_orm::{Database, QuerySelect};
use entity::prelude::*;
async fn run() -> Result<(), anyhow::Error> {
dotenvy::dotenv()?;
let db_url = std::env::var("DATABASE_URL")?;
info!("starting db");
let opts = ConnectOptions::new(db_url);
let db = Database::connect(opts).await?;
Migrator::refresh(&db).await?;
let db = &db;
info!("connected to db");
info!("starting data load");
let pages = random_entities(1000, 1_000_000);
info!("finished data load");
info!("starting db insert");
for chunk in pages.chunks(5000) {
let _ = Page::insert_many(chunk.to_vec()).exec(db).await?;
}
info!("finished db insert");
let length_expr: SimpleExpr = Func::char_length(Expr::col((
entity::page::Entity,
entity::page::Column::Text,
)))
.into();
info!("fetching big row count");
let mut large_row_ids: Vec<i32> = entity::page::Entity::find()
.filter(length_expr.binary(migration::BinOper::GreaterThan, Expr::val(8 * 1024)))
.column(entity::page::Column::Id)
.into_tuple()
.all(db)
.await?;
info!("counted {} big rows", large_row_ids.len());
let num_rows = Page::find().count(db).await?;
info!("inserted {} rows", num_rows);
let mut rng = thread_rng();
large_row_ids.shuffle(&mut rng);
info!("starting");
let mut bytes_read = 0;
for id in large_row_ids.iter().take(1000) {
let row = Page::find_by_id(*id).one(db).await?.unwrap();
bytes_read += row.text.len();
}
println!("read {} bytes", bytes_read);
info!("done");
Ok(())
}
#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
init_logger();
run().await?;
Ok(())
}
fn init_logger() {
let env = Env::default().filter_or("BENCH_LOG_LEVEL", "info,sqlx=error");
Builder::from_env(env).format_timestamp_millis().init();
}