diff --git a/src/bin/import_imdb.rs b/src/bin/import_imdb.rs index 02925bc..696f2a3 100644 --- a/src/bin/import_imdb.rs +++ b/src/bin/import_imdb.rs @@ -2,7 +2,7 @@ use std::{ffi::OsString, time::Duration}; use clap::Parser; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; -use what2watch::{get_db_pool, import_utils::add_omega_watches}; +use what2watch::{get_db_pool, import_utils::add_imdb_movies}; #[derive(Debug, Parser)] struct Cli { @@ -11,7 +11,7 @@ struct Cli { #[clap( long, short, - help = "Number of movies to add in the benchmark.", + help = "Number of movies to add.", default_value_t = 10_000 )] pub number: u32, @@ -46,7 +46,7 @@ fn main() { .unwrap(); rt.block_on(async { - let dur = add_omega_watches(&w2w_db, &movie_db, num).await.unwrap(); + let dur = add_imdb_movies(&w2w_db, &movie_db, num).await.unwrap(); let rows: i32 = sqlx::query_scalar("select count(*) from watches") .fetch_one(&w2w_db) .await diff --git a/src/import_utils.rs b/src/import_utils.rs index 2acd200..3cf3ec5 100644 --- a/src/import_utils.rs +++ b/src/import_utils.rs @@ -7,8 +7,6 @@ use crate::{util::year_to_epoch, ShowKind, User, Watch, WatchQuest}; const USER_EXISTS_QUERY: &str = "select count(*) from users where id = $1"; -const MOVIE_QUERY: &str = "select * from movie_titles order by year, title asc limit ?"; - //-************************************************************************ // the omega user is the system ID, but has no actual power in the app //-************************************************************************ @@ -16,35 +14,36 @@ const OMEGA_ID: Julid = Julid::omega(); const BULK_INSERT: usize = 2_000; #[derive(Debug, sqlx::FromRow, Clone)] -pub struct ImportMovieOmega { +pub struct ImportImdbMovie { pub title: String, pub year: Option, pub runtime: Option, + pub aid: String, } -impl From for Watch { - fn from(value: ImportMovieOmega) -> Self { +impl From for Watch { + fn from(value: ImportImdbMovie) -> Self { Watch { id: OMEGA_ID, // this is ignored by the inserter title: value.title, release_date: year_to_epoch(value.year.as_deref()), length: value.runtime.and_then(|v| v.parse::().ok()), kind: ShowKind::Movie, - metadata_url: None, + metadata_url: Some(format!("https://imdb.com/title/{}/", &value.aid)), added_by: OMEGA_ID, } } } -impl From<&ImportMovieOmega> for Watch { - fn from(value: &ImportMovieOmega) -> Self { +impl From<&ImportImdbMovie> for Watch { + fn from(value: &ImportImdbMovie) -> Self { Watch { id: OMEGA_ID, title: value.title.to_string(), release_date: year_to_epoch(value.year.as_deref()), length: value.runtime.as_ref().and_then(|v| v.parse::().ok()), kind: ShowKind::Movie, - metadata_url: None, + metadata_url: Some(format!("https://imdb.com/title/{}/", value.aid)), added_by: OMEGA_ID, } } @@ -89,14 +88,17 @@ pub async fn add_users(db_pool: &SqlitePool, users: &[User]) -> Result Result { + const IMDB_MOVIE_QUERY: &str = + "select * from movie_titles where porn = 0 order by year, title asc limit ?"; + let omega = ensure_omega(w2w_db).await; - let movies: Vec = sqlx::query_as(MOVIE_QUERY) + let movies: Vec = sqlx::query_as(IMDB_MOVIE_QUERY) .bind(num) .fetch_all(movie_db) .await @@ -105,16 +107,17 @@ pub async fn add_omega_watches( let start = Instant::now(); for movies in movies.as_slice().chunks(BULK_INSERT) { let mut builder = sqlx::QueryBuilder::new( - "insert into watches (kind, title, length, release_date, added_by) ", + "insert into watches (kind, title, length, release_date, added_by, metadata_url) ", ); builder.push_values(movies, |mut b, movie| { - let title = &movie.title; + let movie: Watch = movie.into(); b.push_bind(ShowKind::Movie) - .push_bind(title) - .push_bind(movie.runtime.as_ref().and_then(|l| l.parse::().ok())) - .push_bind(year_to_epoch(movie.year.as_deref())) - .push_bind(omega); + .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(); q.execute(w2w_db).await.map_err(|_| ())?;