hyaenidae/profiles/src/store/submission.rs

use super::{StoreError, Submission, SubmissionChanges, SubmissionFileChanges, Undo, Visibility};
use chrono::{DateTime, Utc};
use sled::{Db, Transactional, Tree};
use std::io::Cursor;
use uuid::Uuid;

#[derive(Clone, Debug)]
pub struct Store {
    submission_tree: Tree,
    profile_tree: Tree,
    profile_created_tree: Tree,
    created_tree: Tree,
    count_tree: Tree,
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
struct StoredSubmission<'a> {
    id: Uuid,
    profile_id: Uuid,
    title: &'a str,
    description: Option<&'a str>,
    files: Vec<Uuid>,
    published: Option<DateTime<Utc>>,
    visibility: Visibility,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
}

impl Store {
    pub fn build(db: &Db) -> Result<Self, sled::Error> {
        Ok(Store {
            submission_tree: db.open_tree("profiles/submissions")?,
            profile_tree: db.open_tree("profiles/submissions/profile")?,
            profile_created_tree: db.open_tree("/profiles/submissions/profile/created")?,
            created_tree: db.open_tree("/profiles/submissions/created")?,
            count_tree: db.open_tree("/profiles/submissions/count")?,
        })
    }

    pub fn create(
        &self,
        profile_id: Uuid,
        title: &str,
        visibility: Visibility,
    ) -> Result<Submission, StoreError> {
        let mut id;
        let mut stored_submission;

        let now = Utc::now().into();
        let mut stored_submission_vec = vec![];

        while {
            stored_submission_vec.clear();
            let writer = Cursor::new(&mut stored_submission_vec);
            id = Uuid::new_v4();
            stored_submission = StoredSubmission {
                id,
                profile_id,
                title,
                description: None,
                files: vec![],
                published: None,
                visibility,
                created_at: now,
                updated_at: now,
            };
            serde_json::to_writer(writer, &stored_submission)?;
            self.submission_tree
                .compare_and_swap(
                    id_submission_key(id).as_bytes(),
                    None as Option<&[u8]>,
                    Some(stored_submission_vec.as_slice()),
                )?
                .is_err()
        } {}

        let res = [
            &self.profile_tree,
            &self.created_tree,
            &self.profile_created_tree,
            &self.count_tree,
        ]
        .transaction(move |trees| {
            let profile_tree = &trees[0];
            let created_tree = &trees[1];
            let profile_created_tree = &trees[2];
            let count_tree = &trees[3];

            profile_tree.insert(
                profile_id_key(profile_id, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;
            created_tree.insert(
                created_submission_key(now, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;
            profile_created_tree.insert(
                profile_id_created_submission_key(profile_id, now, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;

            super::count(
                count_tree,
                &profile_id_submission_count_key(profile_id),
                |c| c.saturating_add(1),
            )?;

            Ok(())
        });

        if let Err(e) = res {
            self.submission_tree
                .remove(id_submission_key(id).as_bytes())?;
            return Err(e.into());
        }

        Ok(stored_submission.into())
    }

    pub fn count(&self, profile_id: Uuid) -> Result<u64, StoreError> {
        match self
            .count_tree
            .get(profile_id_submission_count_key(profile_id))?
        {
            Some(ivec) => Ok(String::from_utf8_lossy(&ivec)
                .parse::<u64>()
                .expect("Count is valid")),
            None => Ok(0),
        }
    }

    pub fn update(&self, changes: &SubmissionChanges) -> Result<Submission, StoreError> {
        let stored_submission_ivec =
            match self.submission_tree.get(id_submission_key(changes.id))? {
                Some(ivec) => ivec,
                None => return Err(StoreError::Missing),
            };

        let mut stored_submission: StoredSubmission =
            serde_json::from_slice(&stored_submission_ivec)?;

        if let Some(title) = changes.title.as_ref() {
            stored_submission.title = &title;
        }
        if let Some(description) = changes.description.as_ref() {
            stored_submission.description = Some(&description);
        }
        stored_submission.published = changes.published;
        stored_submission.updated_at = Utc::now();

        let stored_submission_vec = serde_json::to_vec(&stored_submission)?;

        if self
            .submission_tree
            .compare_and_swap(
                id_submission_key(changes.id),
                Some(&stored_submission_ivec),
                Some(stored_submission_vec.as_slice()),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

        Ok(stored_submission.into())
    }

    pub fn update_files(&self, changes: &SubmissionFileChanges) -> Result<Submission, StoreError> {
        let stored_submission_ivec =
            match self.submission_tree.get(id_submission_key(changes.id))? {
                Some(ivec) => ivec,
                None => return Err(StoreError::Missing),
            };

        let mut stored_submission: StoredSubmission =
            serde_json::from_slice(&stored_submission_ivec)?;

        stored_submission.files = changes.files.clone();
        stored_submission.updated_at = Utc::now();

        let stored_submission_vec = serde_json::to_vec(&stored_submission)?;

        if self
            .submission_tree
            .compare_and_swap(
                id_submission_key(changes.id),
                Some(&stored_submission_ivec),
                Some(stored_submission_vec.as_slice()),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

        Ok(stored_submission.into())
    }

    pub fn for_profile(&self, profile_id: Uuid) -> impl DoubleEndedIterator<Item = Uuid> {
        self.profile_tree
            .scan_prefix(profile_id_prefix(profile_id))
            .values()
            .filter_map(|res| res.ok())
            .filter_map(|ivec| {
                let id_str = String::from_utf8_lossy(&ivec);
                id_str.parse::<Uuid>().ok()
            })
    }

    fn date_range<K>(
        &self,
        range: impl std::ops::RangeBounds<K>,
    ) -> impl DoubleEndedIterator<Item = Uuid>
    where
        K: AsRef<[u8]>,
    {
        self.profile_created_tree
            .range(range)
            .values()
            .filter_map(|res| res.ok())
            .filter_map(move |ivec| {
                let id_str = String::from_utf8_lossy(&ivec);
                id_str.parse::<Uuid>().ok()
            })
    }

    pub fn newer_than(&self, id: Uuid) -> impl DoubleEndedIterator<Item = Uuid> {
        let this = self.clone();

        self.submission_tree
            .get(id_submission_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_submission_ivec| {
                let stored_submission: StoredSubmission =
                    serde_json::from_slice(&stored_submission_ivec).ok()?;

                Some((stored_submission.profile_id, stored_submission.created_at))
            })
            .into_iter()
            .flat_map(move |(profile_id, created_at)| {
                let range_start = profile_id_created_submission_range_start(profile_id, created_at);
                let range_start = range_start.as_bytes().to_vec();
                this.date_range(range_start..)
            })
    }

    pub fn older_than(&self, id: Uuid) -> impl DoubleEndedIterator<Item = Uuid> {
        let this = self.clone();

        self.submission_tree
            .get(id_submission_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_submission_ivec| {
                let stored_submission: StoredSubmission =
                    serde_json::from_slice(&stored_submission_ivec).ok()?;

                Some((stored_submission.profile_id, stored_submission.created_at))
            })
            .into_iter()
            .flat_map(move |(profile_id, created_at)| {
                let range_end = profile_id_created_submission_range_start(profile_id, created_at);
                let range_end = range_end.as_bytes().to_vec();
                this.date_range(..range_end)
            })
    }

    pub fn by_id(&self, id: Uuid) -> Result<Option<Submission>, StoreError> {
        let stored_submission_ivec = match self.submission_tree.get(id_submission_key(id))? {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        let stored_submission: StoredSubmission = serde_json::from_slice(&stored_submission_ivec)?;

        Ok(Some(stored_submission.into()))
    }

    pub fn delete(&self, submission_id: Uuid) -> Result<Option<Undo<Submission>>, StoreError> {
        let stored_submission_ivec =
            match self.submission_tree.get(id_submission_key(submission_id))? {
                Some(ivec) => ivec,
                None => return Ok(None),
            };

        let stored_submission: StoredSubmission = serde_json::from_slice(&stored_submission_ivec)?;

        let id = submission_id;
        let profile_id = stored_submission.profile_id;
        let created_at = stored_submission.created_at;

        [
            &self.submission_tree,
            &self.profile_tree,
            &self.profile_created_tree,
            &self.created_tree,
            &self.count_tree,
        ]
        .transaction(move |trees| {
            let submission_tree = &trees[0];
            let profile_tree = &trees[1];
            let profile_created_tree = &trees[2];
            let created_tree = &trees[3];
            let count_tree = &trees[4];

            submission_tree.remove(id_submission_key(id).as_bytes())?;
            profile_tree.remove(profile_id_key(profile_id, id).as_bytes())?;
            profile_created_tree
                .remove(profile_id_created_submission_key(profile_id, created_at, id).as_bytes())?;
            created_tree.remove(created_submission_key(created_at, id).as_bytes())?;

            super::count(
                count_tree,
                &profile_id_submission_count_key(profile_id),
                |c| c.saturating_sub(1),
            )?;

            Ok(())
        })?;

        Ok(Some(Undo(stored_submission.into())))
    }
}

// Used to map id -> Submission
fn id_submission_key(id: Uuid) -> String {
    format!("/submission/{}/data", id)
}

// Used to map profile_id -> id
fn profile_id_key(profile_id: Uuid, id: Uuid) -> String {
    format!("/profile/{}/submission/{}", profile_id, id)
}

fn profile_id_prefix(profile_id: Uuid) -> String {
    format!("/profile/{}/submission", profile_id)
}

// Used to fetch submissions for a given profile in a user-recognizalbe order
fn profile_id_created_submission_key(
    profile_id: Uuid,
    created_at: DateTime<Utc>,
    id: Uuid,
) -> String {
    format!(
        "/profile/{}/created/{}/submission/{}",
        profile_id,
        created_at.to_rfc3339(),
        id
    )
}

fn profile_id_created_submission_range_start(
    profile_id: Uuid,
    created_at: DateTime<Utc>,
) -> String {
    format!(
        "/profile/{}/created/{}",
        profile_id,
        created_at.to_rfc3339()
    )
}

// Used to map created_at -> id
fn created_submission_key(created_at: DateTime<Utc>, id: Uuid) -> String {
    format!("/created/{}/submission/{}", created_at.to_rfc3339(), id)
}

fn profile_id_submission_count_key(profile_id: Uuid) -> String {
    format!("/profile/{}/count", profile_id)
}

impl<'a> From<StoredSubmission<'a>> for Submission {
    fn from(ss: StoredSubmission<'a>) -> Self {
        Submission {
            id: ss.id,
            profile_id: ss.profile_id,
            title: ss.title.to_owned(),
            description: ss.description.map(|d| d.to_owned()),
            files: ss.files.clone(),
            published: ss.published,
            visibility: ss.visibility,
        }
    }
}