hyaenidae/profiles/src/store/comment.rs

use super::{Comment, CommentChanges, ReplyComment, StoreError, SubmissionComment, Undo};
use chrono::{DateTime, Utc};
use sled::{Db, Transactional, Tree};
use std::io::Cursor;
use uuid::Uuid;

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

#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
struct StoredComment<'a> {
    id: Uuid,
    submission_id: Uuid,
    profile_id: Uuid,
    comment_id: Option<Uuid>,
    body: &'a str,
    published: DateTime<Utc>,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
}

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

    pub fn create(
        &self,
        submission_id: Uuid,
        profile_id: Uuid,
        comment_id: Option<Uuid>,
        body: &str,
        published: DateTime<Utc>,
    ) -> Result<Comment, StoreError> {
        let mut id;
        let mut stored_comment;

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

        while {
            stored_comment_vec.clear();
            let cursor = Cursor::new(&mut stored_comment_vec);
            id = Uuid::new_v4();
            stored_comment = StoredComment {
                id,
                submission_id,
                profile_id,
                comment_id,
                body,
                published,
                created_at: now,
                updated_at: now,
            };
            serde_json::to_writer(cursor, &stored_comment)?;
            self.comment_tree
                .compare_and_swap(
                    id_comment_key(id),
                    None as Option<&[u8]>,
                    Some(stored_comment_vec.as_slice()),
                )?
                .is_err()
        } {}

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

            submission_tree.insert(
                submission_id_key(submission_id, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;
            submission_created_tree.insert(
                submission_id_created_comment_key(submission_id, now, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;
            created_tree.insert(
                created_comment_key(now, id).as_bytes(),
                id.to_string().as_bytes(),
            )?;

            super::count(
                count_tree,
                &submission_id_comment_count_key(submission_id),
                |c| c.saturating_add(1),
            )?;

            Ok(())
        });

        if let Err(e) = res {
            self.comment_tree.remove(id_comment_key(id))?;
            return Err(e.into());
        }

        Ok(stored_comment.into())
    }

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

    pub fn update(&self, changes: &CommentChanges) -> Result<Comment, StoreError> {
        let stored_comment_ivec = match self.comment_tree.get(id_comment_key(changes.id))? {
            Some(ivec) => ivec,
            None => return Err(StoreError::Missing),
        };

        let mut stored_comment: StoredComment = serde_json::from_slice(&stored_comment_ivec)?;

        if let Some(body) = changes.body.as_ref() {
            stored_comment.body = &body;
        }
        stored_comment.updated_at = Utc::now().into();

        let stored_comment_vec = serde_json::to_vec(&stored_comment)?;

        if self
            .comment_tree
            .compare_and_swap(
                id_comment_key(changes.id),
                Some(&stored_comment_ivec),
                Some(stored_comment_vec.as_slice()),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

        Ok(stored_comment.into())
    }

    pub fn for_submission(&self, submission_id: Uuid) -> impl DoubleEndedIterator<Item = Uuid> {
        self.submission_tree
            .scan_prefix(submission_id_prefix(submission_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.submission_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.comment_tree
            .get(id_comment_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_comment_ivec| {
                let stored_comment: StoredComment =
                    serde_json::from_slice(&stored_comment_ivec).ok()?;

                Some((stored_comment.submission_id, stored_comment.created_at))
            })
            .into_iter()
            .flat_map(move |(submission_id, created_at)| {
                let range_start =
                    submission_id_created_comment_range_start(submission_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.comment_tree
            .get(id_comment_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_comment_ivec| {
                let stored_comment: StoredComment =
                    serde_json::from_slice(&stored_comment_ivec).ok()?;

                Some((stored_comment.submission_id, stored_comment.created_at))
            })
            .into_iter()
            .flat_map(move |(submission_id, created_at)| {
                let range_end =
                    submission_id_created_comment_range_start(submission_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<Comment>, StoreError> {
        let stored_comment_ivec = match self.comment_tree.get(id_comment_key(id))? {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        let stored_comment: StoredComment = serde_json::from_slice(&stored_comment_ivec)?;

        Ok(Some(stored_comment.into()))
    }

    pub fn delete(&self, comment_id: Uuid) -> Result<Option<Undo<Comment>>, StoreError> {
        let stored_comment_ivec = match self.comment_tree.get(id_comment_key(comment_id))? {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        let stored_comment: StoredComment = serde_json::from_slice(&stored_comment_ivec)?;
        let created_at = stored_comment.created_at;
        let submission_id = stored_comment.submission_id;

        let comment = stored_comment.into();

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

            created_tree.remove(created_comment_key(created_at, comment_id).as_bytes())?;
            submission_created_tree.remove(
                submission_id_created_comment_key(submission_id, created_at, comment_id).as_bytes(),
            )?;
            submission_tree.remove(submission_id_key(submission_id, comment_id).as_bytes())?;
            comment_tree.remove(id_comment_key(comment_id).as_bytes())?;

            super::count(
                count_tree,
                &submission_id_comment_count_key(submission_id),
                |c| c.saturating_sub(1),
            )?;

            Ok(())
        })?;

        Ok(Some(Undo(comment)))
    }
}

// Maps id -> Comment
fn id_comment_key(id: Uuid) -> String {
    format!("/comment/{}/data", id)
}

// Maps submission_id -> id
fn submission_id_key(submission_id: Uuid, id: Uuid) -> String {
    format!("/submission/{}/comment/{}", submission_id, id)
}

fn submission_id_prefix(submission_id: Uuid) -> String {
    format!("/submission/{}/comment", submission_id)
}

// Maps submission_id + created_at -> id
fn submission_id_created_comment_key(
    submission_id: Uuid,
    created_at: DateTime<Utc>,
    id: Uuid,
) -> String {
    format!(
        "/submission/{}/created/{}/comment/{}",
        submission_id,
        created_at.to_rfc3339(),
        id,
    )
}

fn submission_id_created_comment_range_start(
    submission_id: Uuid,
    created_at: DateTime<Utc>,
) -> String {
    format!(
        "/submission/{}/created/{}",
        submission_id,
        created_at.to_rfc3339()
    )
}

// Maps created_at -> id
fn created_comment_key(created_at: DateTime<Utc>, id: Uuid) -> String {
    format!("/created/{}/comment/{}", created_at.to_rfc3339(), id)
}

// Maps submission -> comment count
fn submission_id_comment_count_key(submission_id: Uuid) -> String {
    format!("/submission/{}/count", submission_id)
}

impl<'a> From<StoredComment<'a>> for Comment {
    fn from(sc: StoredComment<'a>) -> Self {
        if let Some(comment_id) = sc.comment_id {
            Comment::Reply(ReplyComment {
                id: sc.id,
                submission_id: sc.submission_id,
                profile_id: sc.profile_id,
                comment_id,
                body: sc.body.to_owned(),
                published: sc.published,
            })
        } else {
            Comment::Submission(SubmissionComment {
                id: sc.id,
                submission_id: sc.submission_id,
                profile_id: sc.profile_id,
                body: sc.body.to_owned(),
                published: sc.published,
            })
        }
    }
}