blogdb/src/assets.rs

use std::{fmt::Debug, path::PathBuf, str::FromStr};

use anyhow::anyhow;
use infer::Type;
use mime_guess::Mime;
use rand::random;
use sqlx::prelude::FromRow;

use crate::BlogDb;

#[derive(Clone, FromRow)]
pub struct Asset {
    pub slug: String,
    pub mime: String,
    pub data: Vec<u8>,
}

impl Debug for Asset {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Asset(slug: {}, mime: {})", self.slug, self.mime)?;
        Ok(())
    }
}

const SUPPORTED_MIME_TYPES: &[&str] = &["image/jpeg", "font/woff", "font/woff2"];

impl BlogDb {
    /// Add an asset with the given name and data.
    ///
    /// If the asset shares a name with another asset in the parent path, the name will be prefixed with a unique id. Use the returned [`Asset`] name to access this asset in the future.
    pub async fn add_asset<S, D>(&self, slug: S, data: D) -> anyhow::Result<Asset>
    where
        S: AsRef<str>,
        D: Into<Vec<u8>>,
    {
        let data = data.into();

        let slug = PathBuf::from(slug.as_ref());
        let asset_ext: &str;
        let mime = match slug.extension() {
            Some(_) => {
                asset_ext = "";
                match mime_guess::from_path(&slug).first() {
                    Some(mime) => &mime.to_string(),
                    None => match infer::get(&data) {
                        Some(mime) => mime.mime_type(),
                        None => {
                            return Err(anyhow::Error::msg(
                                "Couldn't get the mime type of the provided asset",
                            ))
                        }
                    },
                }
            }
            None => match infer::get(&data) {
                Some(mime) => {
                    asset_ext = mime.extension();
                    mime.mime_type()
                }
                None => {
                    return Err(anyhow::Error::msg(
                        "Couldn't get the mime type of the provided asset",
                    ))
                }
            },
        };

        let mut asset_name = slug
            .parent()
            .map(|parent| parent.to_str().unwrap())
            .unwrap_or("")
            .to_string();
        asset_name.push_str(&random::<i64>().to_string());
        if asset_ext.is_empty() {
            asset_name.push_str(&format!("-{}", slug.file_name().unwrap().to_str().unwrap()));
        } else {
            asset_name.push_str(&format!(
                "-{}.{asset_ext}",
                slug.file_stem().unwrap().to_str().unwrap()
            ));
        }

        let asset: Asset =
            sqlx::query_as("INSERT INTO assets (slug, mime, data) VALUES (?,?,?) RETURNING *")
                .bind(&asset_name)
                .bind(&mime)
                .bind(&data)
                .fetch_one(&self.db)
                .await?;
        Ok(asset)
    }

    async fn asset_path_collides<S>(&self, slug: S) -> bool
    where
        S: AsRef<str>,
    {
        self.get_asset(slug).await.is_ok()
    }

    async fn mime_is_supported(&self, r#type: &Type) -> bool {
        SUPPORTED_MIME_TYPES.contains(&r#type.mime_type())
    }

    /// Adds a collection of asset tuples to the db
    pub async fn add_assets<A, S, D>(&self, assets: A) -> anyhow::Result<()>
    where
        A: IntoIterator<Item = (S, D)>,
        S: AsRef<str>,
        D: Into<Vec<u8>>,
    {
        let assets = assets.into_iter();
        for asset in assets.into_iter() {
            let slug = asset.0;
            let data = asset.1.into();

            let path = PathBuf::from(slug.as_ref());
            let mime = match path.extension() {
                Some(ext) => match ext.to_str() {
                    Some(str) => match mime_guess::from_ext(str).first() {
                        Some(mime) => mime.to_string(),
                        None => match infer::get(&data) {
                            Some(ext) => ext.mime_type().to_string(),
                            None => "text/plain".to_string(),
                        },
                    },
                    None => match infer::get(&data) {
                        Some(ext) => ext.mime_type().to_string(),
                        None => "text/plain".to_string(),
                    },
                },
                None => match infer::get(&data) {
                    Some(ext) => ext.mime_type().to_string(),
                    None => "text/plain".to_string(),
                },
            };

            let result =
                sqlx::query("REPLACE INTO assets (slug, mime, data) VALUES (?,?,?) RETURNING *")
                    .bind(slug.as_ref())
                    .bind(mime)
                    .bind(data)
                    .execute(&self.db)
                    .await;
        }
        Ok(())
    }

    /// Get asset by slug
    pub async fn get_asset<S>(&self, slug: S) -> anyhow::Result<Asset>
    where
        S: AsRef<str>,
    {
        let asset = sqlx::query_as("SELECT * FROM assets WHERE slug=?")
            .bind(slug.as_ref())
            .fetch_one(&self.db)
            .await;
        // let asset: Asset = sqlx::query_as("SELECT * FROM assets WHERE slug=?")
        //     .bind(slug.as_ref())
        //     .fetch_one(&self.db)
        //     .await?;
        Ok(asset?)
    }

    /// Get all assets
    pub async fn get_assets(&self) -> anyhow::Result<Vec<Asset>> {
        let assets: Vec<Asset> = sqlx::query_as("SELECT * FROM assets")
            .fetch_all(&self.db)
            .await?;
        Ok(assets)
    }

    /// Delete asset by name
    pub async fn delete_asset<S>(&self, name: S) -> anyhow::Result<()>
    where
        S: AsRef<str>,
    {
        sqlx::query("DELETE FROM assets WHERE name=?")
            .bind(name.as_ref())
            .execute(&self.db)
            .await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::util::tests::*;
    use sqlx::SqlitePool;

    #[sqlx::test]
    async fn assets(pool: SqlitePool) -> anyhow::Result<()> {
        let (db, session_id) = get_init_db(pool).await?;
        let slug = "picture";
        // PNG magic numbers
        let data = &[0x89, 0x50, 0x4e, 0x47];
        let asset = db.add_asset(slug, data).await;
        assert!(asset.is_ok());

        Ok(())
    }
}