s3m-core 0.18.0

use crate::{
    s3::{S3, actions, checksum::Checksum},
    stream::{part::Part, state},
};
use anyhow::Result;
use rkyv::{from_bytes, rancor::Error as RkyvError, to_bytes};
use std::{
    collections::BTreeMap,
    convert::Into,
    path::{Path, PathBuf},
};

pub const DB_PARTS: &str = "parts";
pub const DB_UPLOADED: &str = "uploaded parts";

#[derive(Debug, Clone)]
pub struct Db {
    storage: sled::Db,
    key: String,
    state_dir: PathBuf,
}

impl Db {
    #[must_use]
    pub fn state_dir(path: &Path, checksum: &str) -> PathBuf {
        path.join("streams").join(checksum)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not create the db
    pub fn new(s3: &S3, key: &str, checksum: &str, mtime: u128, path: &Path) -> Result<Self> {
        let key = format!("{} {} {}", &s3.hash()[0..8], key, mtime);
        let state_dir = Self::state_dir(path, checksum);

        log::debug!("db key: [{}], path: {}", &key, path.display());

        let storage = sled::Config::new()
            .path(&state_dir)
            .use_compression(false)
            .mode(sled::Mode::LowSpace)
            .open()?;
        Ok(Self {
            storage,
            key,
            state_dir,
        })
    }

    /// # Errors
    ///
    /// Will return `Err` if can not open an existing db
    pub fn open_existing(path: &Path, key: String) -> Result<Self> {
        let storage = sled::Config::new()
            .path(path)
            .use_compression(false)
            .mode(sled::Mode::LowSpace)
            .open()?;
        Ok(Self {
            storage,
            key,
            state_dir: path.to_path_buf(),
        })
    }

    #[must_use]
    pub const fn db(&self) -> &sled::Db {
        &self.storage
    }

    #[must_use]
    pub fn state_key(&self) -> &str {
        &self.key
    }

    /// # Errors
    ///
    /// Will return `Err` if can not query the db
    pub fn check(&self) -> Result<Option<String>> {
        let etag = &self
            .storage
            .get(format!("etag {}", &self.key).as_bytes())?
            .map(|s| String::from_utf8(s.to_vec()).map(|s| format!("ETag: {s}")))
            .transpose()?;
        Ok(etag.clone())
    }

    /// # Errors
    ///
    /// Will return `Err` if can not query the db
    pub fn upload_id(&self) -> Result<Option<String>> {
        let uid = &self
            .storage
            .get(&self.key)?
            .map(|s| String::from_utf8(s.to_vec()))
            .transpose()?;
        Ok(uid.clone())
    }

    /// # Errors
    ///
    /// Will return `Err` if can not open the tree
    pub fn db_parts(&self) -> Result<sled::Tree> {
        Ok(self.storage.open_tree(DB_PARTS)?)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not open the tree
    pub fn db_uploaded(&self) -> Result<sled::Tree> {
        Ok(self.storage.open_tree(DB_UPLOADED)?)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not insert the `upload_id`
    pub fn save_upload_id(&self, uid: &str) -> Result<Option<sled::IVec>> {
        let previous = self.storage.insert(&self.key, uid)?;
        self.storage.flush()?;
        state::touch_metadata(&self.state_dir)?;
        Ok(previous)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not insert the itag
    pub fn save_etag(&self, etag: &str) -> Result<Option<sled::IVec>> {
        let previous = self
            .storage
            .insert(format!("etag {}", &self.key).as_bytes(), etag)?;
        self.storage.flush()?;
        state::touch_metadata(&self.state_dir)?;
        Ok(previous)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not `flush`
    pub fn flush(&self) -> Result<usize> {
        let flushed = self.storage.flush()?;
        state::touch_metadata(&self.state_dir)?;
        Ok(flushed)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not insert a `Part`
    pub fn create_part(
        &self,
        number: u16,
        seek: u64,
        chunk: u64,
        checksum: Option<Checksum>,
    ) -> Result<Option<sled::IVec>> {
        let part = Part::new(number, seek, chunk, checksum);
        let bytes = to_bytes::<RkyvError>(&part)?;
        let previous = self
            .db_parts()?
            .insert(number.to_be_bytes(), bytes.as_slice())?;
        state::touch_metadata(&self.state_dir)?;
        Ok(previous)
    }

    /// # Errors
    ///
    /// Will return `Err` if can not insert a `Part`
    pub fn get_part(&self, number: u16) -> Result<Option<Part>> {
        let part = &self
            .db_parts()?
            .get(number.to_be_bytes())?
            .map(|part| from_bytes::<Part, RkyvError>(&part))
            .transpose()?;
        Ok(part.clone())
    }

    /// # Errors
    ///
    /// Will return `Err` if can not create the `BTreeMap<u16, actions::Part>`
    pub fn uploaded_parts(&self) -> Result<BTreeMap<u16, actions::Part>> {
        self.db_uploaded()?
            .into_iter()
            .values()
            .flat_map(|part| {
                part.map(|part| {
                    from_bytes::<Part, RkyvError>(&part)
                        .map(|p| {
                            (
                                p.get_number(),
                                actions::Part {
                                    etag: p.get_etag().to_string(),
                                    number: p.get_number(),
                                    checksum: p.get_checksum(),
                                },
                            )
                        })
                        .map_err(Into::into)
                })
            })
            .collect::<Result<BTreeMap<u16, actions::Part>>>()
    }
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::indexing_slicing,
    clippy::unnecessary_wraps
)]
mod tests {
    use super::*;
    use crate::s3::{Region, S3, credentials::Credentials};
    use anyhow::Result;
    use secrecy::SecretString;
    use std::path::PathBuf;
    use tempfile::tempdir;

    #[test]
    fn test_db() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;
        assert_eq!(db.db().len(), 0);
        assert_eq!(db.db_parts()?.len(), 0);
        assert_eq!(db.db_uploaded()?.len(), 0);
        assert_eq!(db.check()?, None);
        assert_eq!(db.upload_id()?, None);
        assert_eq!(db.save_upload_id("uid")?, None);
        assert_eq!(db.upload_id()?, Some("uid".to_string()));
        assert_eq!(db.save_etag("etag")?, None);
        assert!(db.flush().is_ok());
        assert_eq!(db.flush()?, 0);
        assert_eq!(db.check()?, Some("ETag: etag".to_string()));
        assert_eq!(db.upload_id()?, Some("uid".to_string()));
        assert_eq!(db.db().len(), 2);
        assert_eq!(db.db_parts()?.len(), 0);
        assert_eq!(db.db_uploaded()?.len(), 0);
        assert_eq!(db.create_part(1, 0, 0, None)?, None);
        assert_eq!(db.create_part(2, 0, 0, None)?, None);
        assert_eq!(db.db().len(), 2);
        assert_eq!(db.db_parts()?.len(), 2);
        assert_eq!(db.db_uploaded()?.len(), 0);
        assert_eq!(db.get_part(1)?.unwrap().get_number(), 1);
        assert_eq!(db.get_part(2)?.unwrap().get_number(), 2);
        assert_eq!(db.uploaded_parts()?.len(), 0);
        assert!(db.flush().is_ok());
        assert_eq!(db.db_parts()?.len(), 2);
        assert_eq!(db.db_uploaded()?.len(), 0);
        assert_eq!(db.uploaded_parts()?.len(), 0);
        assert_eq!(db.get_part(1)?.unwrap().get_number(), 1);
        assert_eq!(db.get_part(2)?.unwrap().get_number(), 2);
        Ok(())
    }

    #[test]
    fn test_db_parts() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;
        assert_eq!(db.db_parts()?.len(), 0);
        assert_eq!(db.create_part(1, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 1);
        assert_eq!(db.create_part(2, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 2);
        assert_eq!(db.create_part(3, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 3);
        assert_eq!(db.create_part(4, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 4);
        assert_eq!(db.create_part(5, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 5);
        assert_eq!(db.create_part(6, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 6);
        assert_eq!(db.create_part(7, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 7);
        assert_eq!(db.create_part(8, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 8);
        assert_eq!(db.create_part(9, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 9);
        assert_eq!(db.create_part(10, 0, 0, None)?, None);
        assert_eq!(db.db_parts()?.len(), 10);
        assert_eq!(db.db_parts()?.iter().count(), 10);
        Ok(())
    }

    #[test]
    fn test_db_uploaded() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;
        assert_eq!(db.db_uploaded()?.len(), 0);
        assert_eq!(
            db.db_uploaded()?
                .insert(1_i32.to_be_bytes(), "etag".as_bytes())?,
            None
        );
        assert_eq!(db.db_uploaded()?.len(), 1);
        assert_eq!(
            db.db_uploaded()?
                .insert(2_i32.to_be_bytes(), "etag".as_bytes())?,
            None
        );
        assert_eq!(db.db_uploaded()?.len(), 2);
        Ok(())
    }

    #[test]
    fn test_db_part_rkyv_roundtrip() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;

        // Create a part with all fields populated
        let checksum =
            crate::s3::checksum::Checksum::new(crate::s3::checksum::ChecksumAlgorithm::Sha256);
        db.create_part(1, 1024, 5_242_880, Some(checksum))?;

        // Retrieve and verify
        let part = db.get_part(1)?.unwrap();
        assert_eq!(part.get_number(), 1);
        assert_eq!(part.get_seek(), 1024);
        assert_eq!(part.get_chunk(), 5_242_880);
        assert!(part.get_checksum().is_some());
        assert_eq!(
            part.get_checksum().unwrap().algorithm,
            crate::s3::checksum::ChecksumAlgorithm::Sha256
        );
        Ok(())
    }

    #[test]
    fn test_db_part_rkyv_multiple_parts() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;

        // Create multiple parts with different checksums
        for i in 1..=10 {
            let checksum =
                crate::s3::checksum::Checksum::new(crate::s3::checksum::ChecksumAlgorithm::Crc32c);
            db.create_part(i, u64::from(i) * 5_242_880, 5_242_880, Some(checksum))?;
        }

        // Verify all parts
        for i in 1..=10 {
            let part = db.get_part(i)?.unwrap();
            assert_eq!(part.get_number(), i);
            assert_eq!(part.get_seek(), u64::from(i) * 5_242_880);
            assert_eq!(part.get_chunk(), 5_242_880);
        }
        Ok(())
    }

    #[test]
    fn test_db_part_rkyv_without_checksum() -> Result<()> {
        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;

        // Create part without checksum
        db.create_part(1, 0, 1024, None)?;

        let part = db.get_part(1)?.unwrap();
        assert_eq!(part.get_number(), 1);
        assert!(part.get_checksum().is_none());
        Ok(())
    }

    #[test]
    fn test_db_uploaded_parts_rkyv() -> Result<()> {
        use rkyv::{rancor::Error as RkyvError, to_bytes};

        let dir = tempdir()?;
        let path = PathBuf::from(dir.path());
        let s3 = S3::new(
            &Credentials::new(
                "AKIAIOSFODNN7EXAMPLE",
                &SecretString::new("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".into()),
            ),
            &"us-west-1".parse::<Region>().unwrap(),
            Some("awsexamplebucket1".to_string()),
            false,
        );
        let db = Db::new(&s3, "key", "checksum", 0, &path)?;

        // Manually insert serialized parts into uploaded tree
        let mut checksum =
            crate::s3::checksum::Checksum::new(crate::s3::checksum::ChecksumAlgorithm::Crc32c);
        checksum.checksum = "yZRlqg==".to_string();

        let part = Part::new(1, 0, 5_242_880, Some(checksum))
            .set_etag("\"d41d8cd98f00b204e9800998ecf8427e\"".to_string());

        let bytes = to_bytes::<RkyvError>(&part)?;
        db.db_uploaded()?
            .insert(1_u16.to_be_bytes(), bytes.as_slice())?;

        // Verify uploaded_parts() can deserialize
        let uploaded = db.uploaded_parts()?;
        assert_eq!(uploaded.len(), 1);
        let action_part = uploaded.get(&1).unwrap();
        assert_eq!(action_part.number, 1);
        assert_eq!(action_part.etag, "\"d41d8cd98f00b204e9800998ecf8427e\"");
        assert!(action_part.checksum.is_some());
        Ok(())
    }
}