blob-store-app/src/shard/mod.rs

mod fn_get;
mod fn_migrate;
mod fn_store;
pub mod shard_error;

pub use fn_get::GetResult;
pub use fn_store::StoreResult;

use crate::{sha256::Sha256, shard::shard_error::ShardError};
use axum::body::Bytes;
use rusqlite::{params, types::FromSql, OptionalExtension};
use std::error::Error;
use tokio_rusqlite::Connection;
use tracing::{debug, error};

pub type UtcDateTime = chrono::DateTime<chrono::Utc>;

#[derive(Clone)]
pub struct Shard {
    id: usize,
    conn: Connection,
}

impl Shard {
    pub async fn open(id: usize, conn: Connection) -> Result<Self, Box<dyn Error>> {
        let shard = Self { id, conn };
        shard.migrate().await?;
        Ok(shard)
    }

    pub async fn close(self) -> Result<(), Box<dyn Error>> {
        self.conn.close().await.map_err(|e| e.into())
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub async fn db_size_bytes(&self) -> Result<usize, Box<dyn Error>> {
        self.query_single_row(
            "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()",
        )
        .await
    }

    async fn query_single_row<T: FromSql + Send + 'static>(
        &self,
        query: &'static str,
    ) -> Result<T, Box<dyn Error>> {
        self.conn
            .call(move |conn| {
                let value: T = conn.query_row(query, [], |row| row.get(0))?;
                Ok(value)
            })
            .await
            .map_err(|e| e.into())
    }

    pub async fn num_entries(&self) -> Result<usize, Box<dyn Error>> {
        get_num_entries(&self.conn).await.map_err(|e| e.into())
    }
}

async fn get_num_entries(conn: &Connection) -> Result<usize, tokio_rusqlite::Error> {
    conn.call(|conn| {
        let count: usize = conn.query_row("SELECT COUNT(*) FROM entries", [], |row| row.get(0))?;
        Ok(count)
    })
    .await
}

#[cfg(test)]
pub mod test {
    use super::StoreResult;
    use crate::sha256::Sha256;

    pub async fn make_shard() -> super::Shard {
        let conn = tokio_rusqlite::Connection::open_in_memory().await.unwrap();
        super::Shard::open(0, conn).await.unwrap()
    }

    #[tokio::test]
    async fn test_num_entries() {
        let shard = make_shard().await;
        let num_entries = shard.num_entries().await.unwrap();
        assert_eq!(num_entries, 0);
    }

    #[tokio::test]
    async fn test_db_size_bytes() {
        let shard = make_shard().await;
        let db_size = shard.db_size_bytes().await.unwrap();
        assert!(db_size > 0);
    }

    #[tokio::test]
    async fn test_not_found_get() {
        let shard = make_shard().await;
        let sha256 = Sha256::from_bytes("hello, world!".as_bytes());
        let get_result = shard.get(sha256).await.unwrap();
        assert!(get_result.is_none());
    }

    #[tokio::test]
    async fn test_store_and_get() {
        let shard = make_shard().await;
        let data = "hello, world!".as_bytes();
        let sha256 = Sha256::from_bytes(data);
        let store_result = shard
            .store(sha256, "text/plain".to_string(), data.into())
            .await
            .unwrap();
        assert_eq!(
            store_result,
            StoreResult::Created {
                data_size: data.len(),
                stored_size: data.len()
            }
        );
        assert_eq!(shard.num_entries().await.unwrap(), 1);

        let get_result = shard.get(sha256).await.unwrap().unwrap();
        assert_eq!(get_result.content_type, "text/plain");
        assert_eq!(get_result.data, data);
        assert_eq!(get_result.stored_size, data.len());
    }

    #[tokio::test]
    async fn test_store_duplicate() {
        let shard = make_shard().await;
        let data = "hello, world!".as_bytes();
        let sha256 = Sha256::from_bytes(data);

        let store_result = shard
            .store(sha256, "text/plain".to_string(), data.into())
            .await
            .unwrap();
        assert_eq!(
            store_result,
            StoreResult::Created {
                data_size: data.len(),
                stored_size: data.len()
            }
        );

        let store_result = shard
            .store(sha256, "text/plain".to_string(), data.into())
            .await
            .unwrap();
        assert_eq!(
            store_result,
            StoreResult::Exists {
                data_size: data.len(),
                stored_size: data.len()
            }
        );
        assert_eq!(shard.num_entries().await.unwrap(), 1);
    }
}