hyaenidae/profiles/src/store/profile.rs

use super::{
    OwnerSource, Profile, ProfileChanges, ProfileImageChanges, StoreError, TermSearch, Undo,
};
use chrono::{DateTime, Utc};
use sled::{Db, Transactional, Tree};
use std::io::Cursor;
use uuid::Uuid;

#[derive(Clone, Debug)]
pub struct Store {
    profile_tree: Tree,
    handle_tree: Tree,
    owner_created_tree: Tree,
    created_tree: Tree,
    count_tree: Tree,
    handle_index: TermSearch,
}

#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
enum StoredOwnerSource {
    Local(Uuid),
    Remote(Uuid),
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
struct StoredProfile {
    id: Uuid,
    owner_source: StoredOwnerSource,
    handle: String,
    domain: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    display_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    display_name_source: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    description: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    description_source: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    icon: Option<Uuid>,
    #[serde(skip_serializing_if = "Option::is_none")]
    banner: Option<Uuid>,
    published: DateTime<Utc>,
    login_required: bool,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    #[serde(skip_serializing_if = "Option::is_none")]
    suspended_at: Option<DateTime<Utc>>,
}

impl Store {
    pub(super) fn build(db: &Db) -> Result<Self, sled::Error> {
        Ok(Store {
            profile_tree: db.open_tree("profiles/profiles")?,
            handle_tree: db.open_tree("profiles/profiles/handles")?,
            owner_created_tree: db.open_tree("/profiles/profiles/owner/created")?,
            created_tree: db.open_tree("/profiles/profiles/created")?,
            count_tree: db.open_tree("/profiles/profiles/count")?,
            handle_index: TermSearch::build("handle", 20, db)?,
        })
    }

    pub(crate) fn create(
        &self,
        source: OwnerSource,
        handle: &str,
        domain: &str,
        published: DateTime<Utc>,
    ) -> Result<Profile, StoreError> {
        let handle = handle.to_lowercase();

        let mut id;
        let mut stored_profile;

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

        let stored_source = match &source {
            OwnerSource::Local(id) => StoredOwnerSource::Local(*id),
            OwnerSource::Remote(id) => StoredOwnerSource::Remote(*id),
        };

        while {
            stored_profile_vec.clear();
            let writer = Cursor::new(&mut stored_profile_vec);
            id = Uuid::new_v4();
            stored_profile = StoredProfile {
                id,
                owner_source: stored_source.clone(),
                handle: handle.to_owned(),
                domain: domain.to_owned(),
                display_name: None,
                display_name_source: None,
                description: None,
                description_source: None,
                icon: None,
                banner: None,
                published,
                login_required: true,
                created_at: now,
                updated_at: now,
                suspended_at: None,
            };
            serde_json::to_writer(writer, &stored_profile)?;
            self.profile_tree
                .compare_and_swap(
                    id_profile_key(id),
                    None as Option<&[u8]>,
                    Some(stored_profile_vec.as_slice()),
                )?
                .is_err()
        } {}

        let source = source.clone();

        // ensure handle uniqueness
        match self.handle_tree.compare_and_swap(
            handle_id_key(&handle, domain),
            None as Option<&[u8]>,
            Some(id.as_bytes()),
        ) {
            Ok(Ok(_)) => (),
            Ok(Err(_)) => {
                self.profile_tree.remove(id_profile_key(id))?;
                return Err(StoreError::DoubleStore);
            }
            Err(e) => {
                self.profile_tree.remove(id_profile_key(id))?;
                return Err(e.into());
            }
        }

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

            created_tree.insert(created_profile_key(now, id).as_bytes(), id.as_bytes())?;
            owner_created_tree.insert(
                owner_created_profile_key(&source, now, id).as_bytes(),
                id.as_bytes(),
            )?;

            super::count(count_tree, &owner_profile_count_key(&source), |c| {
                c.saturating_add(1)
            })?;

            Ok(())
        });

        if let Err(e) = self.handle_index.insert(&handle) {
            log::error!("Failed to add {} to search index: {}", handle, e);
        }

        if let Err(e) = res {
            [&self.profile_tree, &self.handle_tree].transaction(move |trees| {
                let profile_tree = &trees[0];
                let handle_tree = &trees[1];

                profile_tree.remove(id_profile_key(id).as_bytes())?;
                handle_tree.remove(handle_id_key(&handle, domain).as_bytes())?;

                Ok(())
            })?;
            return Err(e.into());
        }

        Ok(stored_profile.into())
    }

    pub fn search<'a>(
        &'a self,
        handle_term: &'a str,
    ) -> impl DoubleEndedIterator<Item = Uuid> + 'a {
        self.handle_index
            .search(handle_term)
            .flat_map(move |handle| self.handle_iter(handle))
    }

    fn handle_iter<'a>(&'a self, handle: String) -> impl DoubleEndedIterator<Item = Uuid> + 'a {
        self.handle_tree
            .scan_prefix(handle_id_prefix(&handle))
            .values()
            .filter_map(|res| res.ok())
            .filter_map(|ivec| Uuid::from_slice(&ivec).ok())
    }

    pub fn count(&self, source: OwnerSource) -> Result<u64, StoreError> {
        match self.count_tree.get(owner_profile_count_key(&source))? {
            Some(ivec) => Ok(String::from_utf8_lossy(&ivec).parse::<u64>().unwrap_or(0)),
            None => Ok(0),
        }
    }

    pub(crate) fn update(&self, profile_changes: &ProfileChanges) -> Result<Profile, StoreError> {
        let stored_profile_ivec = match self
            .profile_tree
            .get(id_profile_key(profile_changes.id).as_bytes())?
        {
            Some(ivec) => ivec,
            None => return Err(StoreError::Missing),
        };

        let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?;

        if let Some(login_required) = profile_changes.login_required {
            stored_profile.login_required = login_required;
        }
        if let Some(display_name) = &profile_changes.display_name {
            stored_profile.display_name = Some(display_name.clone());
        }
        if let Some(display_name_source) = &profile_changes.display_name_source {
            stored_profile.display_name_source = Some(display_name_source.clone());
        }
        if let Some(description) = &profile_changes.description {
            stored_profile.description = Some(description.clone());
        }
        if let Some(description_source) = &profile_changes.description_source {
            stored_profile.description_source = Some(description_source.clone());
        }
        stored_profile.updated_at = Utc::now().into();

        let stored_profile_vec = serde_json::to_vec(&stored_profile)?;

        if self
            .profile_tree
            .compare_and_swap(
                id_profile_key(profile_changes.id).as_bytes(),
                Some(&stored_profile_ivec),
                Some(stored_profile_vec),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

        Ok(stored_profile.into())
    }

    pub(crate) fn update_images(
        &self,
        profile_image_changes: &ProfileImageChanges,
    ) -> Result<Profile, StoreError> {
        let stored_profile_ivec = match self
            .profile_tree
            .get(id_profile_key(profile_image_changes.id).as_bytes())?
        {
            Some(ivec) => ivec,
            None => return Err(StoreError::Missing),
        };

        let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?;

        if let Some(icon) = profile_image_changes.icon {
            stored_profile.icon = Some(icon);
        }
        if let Some(banner) = profile_image_changes.banner {
            stored_profile.banner = Some(banner);
        }
        stored_profile.updated_at = Utc::now().into();

        let stored_profile_vec = serde_json::to_vec(&stored_profile)?;

        if self
            .profile_tree
            .compare_and_swap(
                id_profile_key(profile_image_changes.id).as_bytes(),
                Some(&stored_profile_ivec),
                Some(stored_profile_vec),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

        Ok(stored_profile.into())
    }

    pub fn for_local(&self, id: Uuid) -> impl DoubleEndedIterator<Item = Uuid> {
        self.for_owner(OwnerSource::Local(id))
    }

    fn for_owner(&self, owner_source: OwnerSource) -> impl DoubleEndedIterator<Item = Uuid> {
        self.owner_created_tree
            .scan_prefix(owner_created_profile_prefix(&owner_source))
            .values()
            .filter_map(|res| res.ok())
            .filter_map(|ivec| Uuid::from_slice(&ivec).ok())
    }

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

    pub fn all(&self) -> impl DoubleEndedIterator<Item = Uuid> {
        self.created_tree
            .iter()
            .values()
            .filter_map(|res| res.ok())
            .filter_map(|ivec| Uuid::from_slice(&ivec).ok())
            .rev()
    }

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

        self.profile_tree
            .get(id_profile_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_profile_ivec| {
                let stored_profile: StoredProfile =
                    serde_json::from_slice(&stored_profile_ivec).ok()?;

                Some((
                    stored_profile.owner_source.into(),
                    stored_profile.created_at,
                ))
            })
            .into_iter()
            .flat_map(move |(source, created_at)| {
                let range_start = owner_created_profile_range_start(&source, 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.profile_tree
            .get(id_profile_key(id))
            .ok()
            .and_then(|opt| opt)
            .and_then(|stored_profile_ivec| {
                let stored_profile: StoredProfile =
                    serde_json::from_slice(&stored_profile_ivec).ok()?;

                Some((
                    stored_profile.owner_source.into(),
                    stored_profile.created_at,
                ))
            })
            .into_iter()
            .flat_map(move |(source, created_at)| {
                let range_end = owner_created_profile_range_start(&source, created_at);
                let range_end = range_end.as_bytes().to_vec();
                this.date_range(..range_end)
            })
            .rev()
    }

    pub fn is_local(&self, id: Uuid) -> Result<Option<bool>, StoreError> {
        if let Some(profile) = self.by_id(id)? {
            return Ok(Some(profile.owner_source.is_local()));
        }

        Ok(None)
    }

    pub fn by_id(&self, id: Uuid) -> Result<Option<Profile>, StoreError> {
        let stored_profile_ivec = match self.profile_tree.get(id_profile_key(id).as_bytes())? {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        let stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?;

        Ok(Some(stored_profile.into()))
    }

    pub fn by_handle(&self, handle: &str, domain: &str) -> Result<Option<Uuid>, StoreError> {
        let id_ivec = match self.handle_tree.get(handle_id_key(handle, domain))? {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        Ok(Uuid::from_slice(&id_ivec).ok())
    }

    pub(crate) fn suspend(&self, profile_id: Uuid) -> Result<Option<Undo<Profile>>, StoreError> {
        let stored_profile_ivec = match self
            .profile_tree
            .get(id_profile_key(profile_id).as_bytes())?
        {
            Some(ivec) => ivec,
            None => return Ok(None),
        };

        let now = Utc::now();
        let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?;
        stored_profile.banner = None;
        stored_profile.icon = None;
        stored_profile.description = None;
        stored_profile.display_name = None;
        stored_profile.suspended_at = Some(now);
        stored_profile.updated_at = now;
        let stored_profile_vec = serde_json::to_vec(&stored_profile)?;

        if self
            .profile_tree
            .compare_and_swap(
                id_profile_key(profile_id).as_bytes(),
                Some(&stored_profile_ivec),
                Some(stored_profile_vec),
            )?
            .is_err()
        {
            return Err(StoreError::DoubleStore);
        }

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

// Used to map id -> Profile
fn id_profile_key(id: Uuid) -> String {
    format!("/profile/{}/data", id)
}

// Used to map handle -> id
fn handle_id_key(handle: &str, domain: &str) -> String {
    format!("/handle/{}:{}/profile", handle, domain)
}

fn handle_id_prefix(handle: &str) -> String {
    format!("/handle/{}:", handle)
}

fn owner_created_profile_prefix(owner: &OwnerSource) -> String {
    format!("/owner/{}/created", owner)
}

// Used to fetch profiles for a given owner in a user-recognizalbe order
fn owner_created_profile_key(owner: &OwnerSource, created_at: DateTime<Utc>, id: Uuid) -> String {
    format!(
        "/owner/{}/created/{}/profile/{}",
        owner,
        created_at.to_rfc3339(),
        id
    )
}

fn owner_created_profile_range_start(owner: &OwnerSource, created_at: DateTime<Utc>) -> String {
    format!("/owner/{}/created/{}", owner, created_at.to_rfc3339(),)
}

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

fn owner_profile_count_key(owner: &OwnerSource) -> String {
    format!("/owner/{}/count", owner)
}

impl From<StoredOwnerSource> for OwnerSource {
    fn from(os: StoredOwnerSource) -> Self {
        match os {
            StoredOwnerSource::Local(id) => OwnerSource::Local(id),
            StoredOwnerSource::Remote(id) => OwnerSource::Remote(id),
        }
    }
}

impl From<StoredProfile> for Profile {
    fn from(sp: StoredProfile) -> Self {
        Profile {
            id: sp.id,
            owner_source: sp.owner_source.into(),
            handle: sp.handle.to_owned(),
            domain: sp.domain.to_owned(),
            display_name: sp.display_name,
            display_name_source: sp.display_name_source,
            description: sp.description,
            description_source: sp.description_source,
            icon: sp.icon,
            banner: sp.banner,
            published: sp.published,
            login_required: sp.login_required,
            suspended: sp.suspended_at.is_some(),
        }
    }
}