what2watch/src/import_utils.rs

171 lines
5.2 KiB
Rust
Raw Normal View History

2023-07-28 23:15:27 +00:00
use std::time::{Duration, Instant};
use julid::Julid;
use sqlx::{query_scalar, SqlitePool};
2023-07-06 15:48:42 +00:00
2023-07-28 23:15:27 +00:00
use crate::{util::year_to_epoch, ShowKind, User, Watch, WatchQuest};
2023-07-09 04:21:12 +00:00
const USER_EXISTS_QUERY: &str = "select count(*) from users where id = $1";
//-************************************************************************
// the omega user is the system ID, but has no actual power in the app
//-************************************************************************
2023-07-28 23:15:27 +00:00
const OMEGA_ID: Julid = Julid::omega();
const BULK_INSERT: usize = 2_000;
#[derive(Debug, sqlx::FromRow, Clone)]
2024-01-16 00:00:24 +00:00
pub struct ImportImdbMovie {
2023-07-03 22:20:19 +00:00
pub title: String,
pub year: Option<String>,
2023-10-22 20:35:41 +00:00
pub runtime: Option<String>,
2024-01-16 00:00:24 +00:00
pub aid: String,
2023-07-03 22:20:19 +00:00
}
2024-01-16 00:00:24 +00:00
impl From<ImportImdbMovie> for Watch {
fn from(value: ImportImdbMovie) -> Self {
2023-07-03 22:20:19 +00:00
Watch {
2023-07-29 19:27:12 +00:00
id: OMEGA_ID, // this is ignored by the inserter
2023-07-03 22:20:19 +00:00
title: value.title,
release_date: year_to_epoch(value.year.as_deref()),
2023-10-22 20:35:41 +00:00
length: value.runtime.and_then(|v| v.parse::<i64>().ok()),
2023-07-03 22:20:19 +00:00
kind: ShowKind::Movie,
2024-01-16 00:00:24 +00:00
metadata_url: Some(format!("https://imdb.com/title/{}/", &value.aid)),
2023-07-28 23:15:27 +00:00
added_by: OMEGA_ID,
2023-07-03 22:20:19 +00:00
}
}
}
2024-01-16 00:00:24 +00:00
impl From<&ImportImdbMovie> for Watch {
fn from(value: &ImportImdbMovie) -> Self {
2023-07-04 18:30:42 +00:00
Watch {
2023-07-29 19:27:12 +00:00
id: OMEGA_ID,
2023-07-04 18:30:42 +00:00
title: value.title.to_string(),
release_date: year_to_epoch(value.year.as_deref()),
2023-10-22 20:35:41 +00:00
length: value.runtime.as_ref().and_then(|v| v.parse::<i64>().ok()),
2023-07-04 18:30:42 +00:00
kind: ShowKind::Movie,
2024-01-16 00:00:24 +00:00
metadata_url: Some(format!("https://imdb.com/title/{}/", value.aid)),
2023-07-28 23:15:27 +00:00
added_by: OMEGA_ID,
2023-07-04 18:30:42 +00:00
}
}
}
//-************************************************************************
// utility functions for building CLI tools, currently just for benchmarking
//-************************************************************************
pub async fn add_watch_quests(pool: &SqlitePool, quests: &[WatchQuest]) -> Result<(), ()> {
2023-07-09 05:00:26 +00:00
let mut builder = sqlx::QueryBuilder::new("insert into watch_quests (user, watch) ");
builder.push_values(quests, |mut b, quest| {
let user = quest.user;
let watch = quest.watch;
//eprintln!("{user}, {watch}");
b.push_bind(user).push_bind(watch);
});
let q = builder.build();
q.execute(pool).await.map_err(|e| {
dbg!(e);
})?;
Ok(())
}
2023-07-28 23:15:27 +00:00
pub async fn add_users(db_pool: &SqlitePool, users: &[User]) -> Result<Duration, ()> {
let mut builder =
2023-07-28 23:15:27 +00:00
sqlx::QueryBuilder::new("insert into users (username, displayname, email, pwhash)");
2023-07-28 23:15:27 +00:00
let start = Instant::now();
builder.push_values(users.iter(), |mut b, user| {
2023-07-28 23:15:27 +00:00
b.push_bind(&user.username)
.push_bind(&user.displayname)
.push_bind(&user.email)
.push_bind(&user.pwhash);
});
let q = builder.build();
q.execute(db_pool).await.map_err(|_| ())?;
2023-07-28 23:15:27 +00:00
let end = Instant::now();
let dur = end - start;
Ok(dur)
}
2024-01-16 00:00:24 +00:00
pub async fn add_imdb_movies(
2023-07-09 05:00:26 +00:00
w2w_db: &SqlitePool,
movie_db: &SqlitePool,
2023-07-28 23:15:27 +00:00
num: u32,
) -> Result<Duration, ()> {
2024-01-16 00:00:24 +00:00
const IMDB_MOVIE_QUERY: &str =
"select * from movie_titles where porn = 0 order by year, title asc limit ?";
2023-07-28 23:15:27 +00:00
let omega = ensure_omega(w2w_db).await;
2024-01-16 00:00:24 +00:00
let movies: Vec<ImportImdbMovie> = sqlx::query_as(IMDB_MOVIE_QUERY)
2023-07-28 23:15:27 +00:00
.bind(num)
2023-07-06 15:48:42 +00:00
.fetch_all(movie_db)
.await
.unwrap();
2023-07-28 23:15:27 +00:00
let start = Instant::now();
for movies in movies.as_slice().chunks(BULK_INSERT) {
let mut builder = sqlx::QueryBuilder::new(
2024-01-16 00:00:24 +00:00
"insert into watches (kind, title, length, release_date, added_by, metadata_url) ",
);
builder.push_values(movies, |mut b, movie| {
2024-01-16 00:00:24 +00:00
let movie: Watch = movie.into();
2023-07-28 23:15:27 +00:00
b.push_bind(ShowKind::Movie)
2024-01-16 00:00:24 +00:00
.push_bind(movie.title)
.push_bind(movie.length)
.push_bind(movie.release_date)
.push_bind(omega)
.push_bind(movie.metadata_url);
});
let q = builder.build();
2023-07-09 05:00:26 +00:00
q.execute(w2w_db).await.map_err(|_| ())?;
}
2023-07-28 23:15:27 +00:00
let end = Instant::now();
let dur = end - start;
2023-07-28 23:15:27 +00:00
Ok(dur)
}
2023-07-28 23:15:27 +00:00
pub async fn ensure_omega(db_pool: &SqlitePool) -> Julid {
if !check_omega_exists(db_pool).await {
2024-01-07 00:54:11 +00:00
User::omega().try_insert(db_pool).await.unwrap();
}
2023-07-28 23:15:27 +00:00
OMEGA_ID
}
async fn check_omega_exists(db_pool: &SqlitePool) -> bool {
let count = query_scalar(USER_EXISTS_QUERY)
2023-07-29 19:27:12 +00:00
.bind(OMEGA_ID)
.fetch_one(db_pool)
.await
.unwrap_or(0);
count > 0
}
//-************************************************************************
//tests
//-************************************************************************
#[cfg(test)]
mod test {
use tokio::runtime::Runtime;
use super::*;
#[test]
fn ensure_omega_user() {
let p = crate::db::get_db_pool();
let rt = Runtime::new().unwrap();
rt.block_on(async {
2024-01-07 00:54:11 +00:00
dbg!("checking omega");
assert!(!check_omega_exists(&p).await);
2024-01-07 00:54:11 +00:00
dbg!("no omega");
ensure_omega(&p).await;
});
2024-01-07 00:54:11 +00:00
dbg!("maybe omega");
2023-07-28 23:15:27 +00:00
assert!(rt.block_on(check_omega_exists(&p)));
}
}