use std::{collections::HashMap, fmt::Display}; use crate::media::{ProcessedMedia, ProcessedMediaType}; use chrono::{DateTime, Utc}; use color_eyre::{Result, eyre::Context}; use sqlx::{ Pool, Sqlite, SqlitePool, migrate::MigrateDatabase, query::QueryAs, sqlite::SqliteArguments, }; pub trait MediaTypeConverter { fn convert(&self) -> T; } impl MediaTypeConverter for ProcessedMediaType { fn convert(&self) -> i32 { match self { Self::Image => 0, Self::Video => 1, } } } #[allow(clippy::module_name_repetitions, clippy::cognitive_complexity)] pub async fn setup_db(db_file: &str) -> Result> { tracing::debug!("Setup db"); tracing::debug!("Ensure db exists"); if !Sqlite::database_exists(db_file).await.unwrap_or(false) { Sqlite::create_database(db_file).await?; } tracing::debug!("Create pool and run migrations"); let pool = SqlitePool::connect(db_file).await?; sqlx::migrate!("./migrations").run(&pool).await?; tracing::debug!("DB setup: {:?}", db_file); Ok(pool) } pub async fn media_insert(pool: &SqlitePool, processed_media: ProcessedMedia) -> Result<()> { let media_type = processed_media.media_type.convert(); match &processed_media.data { Some(data) => { sqlx::query!( r#" INSERT OR REPLACE INTO media (id, path, taken_at, type, walked) VALUES (?1, ?2, ?3, ?4, 1) "#, processed_media.uuid, processed_media.path, data.taken_at, media_type, ) } None => { sqlx::query!( r#" UPDATE media SET walked = 1, path = ?2 WHERE ID = ?1 "#, processed_media.uuid, processed_media.path, ) } } .execute(pool) .await?; Ok(()) } #[derive(sqlx::FromRow)] pub struct MediaHashRow { pub id: uuid::Uuid, pub path: String, } pub async fn media_hashmap(pool: &SqlitePool) -> Result> { tracing::debug!("Collect all DB entries UUIDs"); let res = sqlx::query_as::( r" SELECT id, path FROM media ", ) .fetch_all(pool) .await?; Ok(res.into_iter().map(|e| (e.path, e.id)).collect()) } pub async fn media_mark_unwalked(pool: &SqlitePool) -> Result<()> { sqlx::query!( r#" UPDATE media SET walked = 0 "# ) .execute(pool) .await?; Ok(()) } pub async fn media_mark_missing(pool: &SqlitePool) -> Result { let res = sqlx::query!( r#" UPDATE media SET missing = 1 WHERE walked = 0; UPDATE media SET missing = 0 WHERE walked = 1; "# ) .execute(pool) .await?; Ok(res.rows_affected()) } #[derive(sqlx::FromRow, Debug)] pub struct MediaCount { pub count: i32, pub walked: Option, pub favorite: Option, pub archived: Option, pub missing: Option, } pub async fn media_count(pool: &SqlitePool) -> Result { let res = sqlx::query_as::( r" SELECT COUNT(*) as count, SUM(walked) as walked, SUM(favorite) as favorite, SUM(archived) as archived, SUM(missing) as missing FROM media ", ) .fetch_one(pool) .await?; Ok(res) } pub async fn media_archive(pool: &SqlitePool, uuid: &uuid::Uuid, archive: bool) -> Result { let res = sqlx::query!( r#" UPDATE media SET archived = ?2 WHERE id = ?1 "#, uuid, archive, ) .execute(pool) .await?; Ok(res.rows_affected() > 0) } pub async fn media_favorite(pool: &SqlitePool, uuid: &uuid::Uuid, archive: bool) -> Result { let res = sqlx::query!( r#" UPDATE media SET favorite = ?2 WHERE id = ?1 "#, uuid, archive, ) .execute(pool) .await?; Ok(res.rows_affected() > 0) } #[derive(sqlx::FromRow, Debug)] pub struct MediaRow { pub id: uuid::Uuid, pub path: String, pub taken_at: DateTime, pub media_type: i32, pub archived: bool, pub favorite: bool, } pub enum Order { Asc, Desc, } impl Display for Order { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Asc => f.write_str("ORDER BY taken_at ASC")?, Self::Desc => f.write_str("ORDER BY taken_at DESC")?, } Ok(()) } } pub struct Filter { pub archived: bool, pub favorite: Option, pub taken_after: Option>, pub taken_before: Option>, pub collection: Option, } impl Display for Filter { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str("WHERE NOT missing AND archived = ?")?; if self.favorite.is_some() { f.write_str(" AND favorite = ?")?; } if self.taken_after.is_some() { f.write_str(" AND taken_at > ?")?; } if self.taken_before.is_some() { f.write_str(" AND taken_at < ?")?; } if self.collection.is_some() { f.write_str(" AND path LIKE '%' || ? || '%'")?; } Ok(()) } } impl Filter { fn bind<'a>( &'a self, query: QueryAs<'a, Sqlite, MediaRow, SqliteArguments>, ) -> QueryAs<'a, Sqlite, MediaRow, SqliteArguments> { let mut query = query.bind(self.archived); if let Some(favorite) = self.favorite { query = query.bind(favorite); } if let Some(taken_after) = self.taken_after { query = query.bind(taken_after); } if let Some(taken_before) = self.taken_before { query = query.bind(taken_before); } if let Some(collection) = &self.collection { query = query.bind(collection); } query } } #[allow(clippy::future_not_send)] pub async fn media_list( pool: &SqlitePool, filter: impl Into, order: impl Into, limit: usize, ) -> Result> { let filter = filter.into(); let order = order.into(); let query = format!( r" -- sql SELECT id, path, taken_at, type as media_type, archived, favorite FROM media {filter} {order} LIMIT ? " ); let mut query = sqlx::query_as::(sqlx::AssertSqlSafe(query)); query = filter.bind(query); query = query.bind(i32::try_from(limit).expect("Failed to convert limit")); query.fetch_all(pool).await.wrap_err("Failed to fetch rows") } #[allow(clippy::future_not_send)] pub async fn media_get( pool: &SqlitePool, filter: impl Into, order: impl Into, uuid: &uuid::Uuid, ) -> Result<(Option, Option, Option)> { let filter = filter.into(); let order = order.into(); let query = format!( r" -- sql WITH NR_MED AS ( SELECT *, ROW_NUMBER() OVER ({order}) AS RN FROM media {filter} ) SELECT id, path, taken_at, type as media_type, archived, favorite FROM NR_MED WHERE RN IN ( SELECT RN+i FROM NR_MED CROSS JOIN (SELECT -1 AS i UNION ALL SELECT 0 UNION ALL SELECT 1) n WHERE id = ? ) " ); let mut query = sqlx::query_as::(sqlx::AssertSqlSafe(query)); query = filter.bind(query); query = query.bind(uuid); let res = query .fetch_all(pool) .await .wrap_err("Failed to fetch rows")?; match res.len() { 1 => { let mut res = res.into_iter(); Ok((None, res.next(), None)) } 2 => { let mut res = res.into_iter(); let a = res.next().expect("First media should be there"); let b = res.next().expect("Second media should be there"); if a.id == *uuid { Ok((None, Some(a), Some(b))) } else { Ok((Some(a), Some(b), None)) } } 3 => { let mut res = res.into_iter(); Ok((res.next(), res.next(), res.next())) } _ => Err(color_eyre::eyre::eyre!("Media query returned bad result")), } } pub async fn media_get_path(pool: &SqlitePool, uuid: &uuid::Uuid) -> Result> { let res = sqlx::query_scalar::<_, String>( r" SELECT path FROM media WHERE id = ? ", ) .bind(uuid) .fetch_optional(pool) .await?; Ok(res) } pub async fn media_get_uuid_by_path(pool: &SqlitePool, path: &str) -> Result> { let res = sqlx::query_scalar::<_, uuid::Uuid>( r" SELECT id FROM media WHERE path = ? ", ) .bind(path) .fetch_optional(pool) .await?; Ok(res) } pub async fn collections_search( pool: &SqlitePool, media_dir: &str, search_term: &str, ) -> Result> { let mut search_pattern = String::from("%"); for c in search_term.chars().filter(|c| !c.is_whitespace()) { search_pattern.push(c); search_pattern.push('%'); } let offset = i32::try_from(media_dir.len()).expect("Failed to convert media_dir length") + if media_dir.ends_with('/') { 1 } else { 2 }; let mut res = sqlx::query_scalar::<_, String>( r" WITH RECURSIVE folders (folder, rest) AS ( SELECT '', SUBSTR(path, ?1) FROM media UNION ALL SELECT folder || SUBSTR(rest, 1, INSTR(rest, '/')), SUBSTR(rest, INSTR(rest, '/') + 1) FROM folders WHERE INSTR(rest, '/') > 0 ) SELECT DISTINCT folder FROM folders WHERE folder != '' AND LOWER(folder) LIKE LOWER(?2) ORDER BY LENGTH(folder), folder COLLATE NOCASE ", ) .bind(offset) .bind(search_pattern) .fetch_all(pool) .await?; if search_term == "/" { res.insert(0, "/".to_string()); } Ok(res) }