src/db.rs

Rama main · réplica de lectura del repositorio original.

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<T> {
    fn convert(&self) -> T;
}

impl MediaTypeConverter<i32> 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<Pool<Sqlite>> {
    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<HashMap<String, uuid::Uuid>> {
    tracing::debug!("Collect all DB entries UUIDs");
    let res = sqlx::query_as::<Sqlite, MediaHashRow>(
        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<u64> {
    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<i32>,
    pub favorite: Option<i32>,
    pub archived: Option<i32>,
    pub missing: Option<i32>,
}

pub async fn media_count(pool: &SqlitePool) -> Result<MediaCount> {
    let res = sqlx::query_as::<Sqlite, MediaCount>(
        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<bool> {
    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<bool> {
    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<Utc>,
    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<bool>,
    pub taken_after: Option<DateTime<Utc>>,
    pub taken_before: Option<DateTime<Utc>>,
    pub collection: Option<String>,
}

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<Filter>,
    order: impl Into<Order>,
    limit: usize,
) -> Result<Vec<MediaRow>> {
    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::<Sqlite, MediaRow>(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<Filter>,
    order: impl Into<Order>,
    uuid: &uuid::Uuid,
) -> Result<(Option<MediaRow>, Option<MediaRow>, Option<MediaRow>)> {
    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::<Sqlite, MediaRow>(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<Option<String>> {
    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<Option<uuid::Uuid>> {
    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<Vec<String>> {
    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)
}