haematite-cli 0.21.8

//! Mirrored edge keys read and changed at one captured branch version.

use std::collections::BTreeSet;
use std::num::NonZeroUsize;

use haematite::db::ReadOnlyDatabase;
use serde_json::Value;

use super::mutations::PreparedWrite;
use super::reads::{Page, PageLimits, VisibleRowPage, bounded_key_page, visible_row_page};
use super::{Captured, Failure, document_key, identifier, reads};
use crate::serve::access::Access;
use crate::serve::protocol::{Edge, EdgeDelete, EdgeImport, Node};

#[derive(Clone, Copy)]
enum Direction {
    Out,
    In,
}

impl Direction {
    const fn tag(self) -> &'static str {
        match self {
            Self::Out => "e",
            Self::In => "i",
        }
    }
}

fn validate_node(node: &Node) -> Result<(), Failure> {
    identifier(&node.collection)?;
    identifier(&node.document_id)
}

pub(super) fn keys(from: &Node, relation: &str, to: &Node) -> Result<(Vec<u8>, Vec<u8>), Failure> {
    validate_node(from)?;
    identifier(relation)?;
    validate_node(to)?;
    Ok((
        format!(
            "e\0{}\0{}\0{relation}\0{}\0{}",
            from.collection, from.document_id, to.collection, to.document_id
        )
        .into_bytes(),
        format!(
            "i\0{}\0{}\0{relation}\0{}\0{}",
            to.collection, to.document_id, from.collection, from.document_id
        )
        .into_bytes(),
    ))
}

pub(super) fn put(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    edge: &Edge,
) -> Result<Vec<PreparedWrite>, Failure> {
    let (out, incoming) = keys(&edge.from, &edge.relation, &edge.to)?;
    if !edge.properties.is_object() {
        return Err(Failure::new(
            "edge_properties",
            "edge properties must be a JSON object",
        ));
    }
    for (name, node) in [("from", &edge.from), ("to", &edge.to)] {
        if !reads::document_exists_for_constraint(reader, at, &node.collection, &node.document_id)?
        {
            return Err(Failure {
                path: Some(format!("/{name}")),
                ..Failure::new(
                    "document_missing",
                    format!(
                        "edge {name} document {}/{} does not exist at the captured version",
                        node.collection, node.document_id,
                    ),
                )
            });
        }
    }
    super::relations::validate_edge(reader, at, edge)?;
    let bytes = serde_json::to_vec(&edge.properties).map_err(Failure::storage)?;
    Ok(vec![(out, Some(bytes.clone())), (incoming, Some(bytes))])
}

/// `edges.put_many` (REL0): each edge checked as `edge.put` checks it
/// against the stored edges, and against the edges before it in the batch
/// for its relation's cardinality, in one pass; the first refusal names its
/// row and nothing is written.
pub(super) fn put_many(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    import: &EdgeImport,
) -> Result<Vec<PreparedWrite>, Failure> {
    if import.edges.is_empty() {
        return Err(Failure {
            code: "import_empty".to_owned(),
            message: "an edge import must contain at least one edge".to_owned(),
            path: Some("/edges".to_owned()),
            cause: None,
        });
    }
    let mut writes = Vec::with_capacity(import.edges.len().saturating_mul(2));
    let mut batch = super::relations::Batch::default();
    let mut identities = BTreeSet::new();
    for (row, edge) in import.edges.iter().enumerate() {
        let in_row = |failure: Failure| Failure {
            path: Some(format!(
                "/edges/{row}{}",
                failure.path.as_deref().unwrap_or("")
            )),
            ..failure
        };
        let (out, _) = keys(&edge.from, &edge.relation, &edge.to).map_err(in_row)?;
        if !identities.insert(out) {
            return Err(in_row(Failure::new(
                "import_duplicate",
                "repeated relation and endpoints in edge import",
            )));
        }
        if super::relations::get(reader, at, &edge.relation)
            .map_err(in_row)?
            .is_none()
        {
            return Err(in_row(Failure {
                path: Some("/relation".to_owned()),
                ..Failure::new("relation_missing", &edge.relation)
            }));
        }
        writes.extend(put(reader, at, edge).map_err(in_row)?);
        batch.admit(reader, at, edge).map_err(in_row)?;
    }
    Ok(writes)
}

pub(super) fn delete(edge: &EdgeDelete) -> Result<Vec<PreparedWrite>, Failure> {
    let (out, incoming) = keys(&edge.from, &edge.relation, &edge.to)?;
    Ok(vec![(out, None), (incoming, None)])
}

fn node_prefix(
    direction: Direction,
    node: &Node,
    relation: Option<&str>,
) -> Result<Vec<u8>, Failure> {
    validate_node(node)?;
    let mut prefix = format!(
        "{}\0{}\0{}\0",
        direction.tag(),
        node.collection,
        node.document_id
    )
    .into_bytes();
    if let Some(relation) = relation {
        identifier(relation)?;
        prefix.extend_from_slice(relation.as_bytes());
        prefix.push(0);
    }
    Ok(prefix)
}

pub(super) fn decode_key(key: &[u8]) -> Result<EdgeDelete, Failure> {
    if key.len() > 6 * 513 {
        return Err(Failure::new(
            "edge_key",
            "edge key exceeds its component bounds",
        ));
    }
    let text = std::str::from_utf8(key).map_err(Failure::storage)?;
    let parts: Vec<_> = text.split('\0').collect();
    let [
        tag,
        first_collection,
        first_id,
        relation,
        second_collection,
        second_id,
    ] = parts.as_slice()
    else {
        return Err(Failure::new(
            "edge_key",
            "edge key must have five identifier parts",
        ));
    };
    for part in [
        first_collection,
        first_id,
        relation,
        second_collection,
        second_id,
    ] {
        identifier(part)?;
    }
    let first = Node {
        collection: (*first_collection).to_owned(),
        document_id: (*first_id).to_owned(),
    };
    let second = Node {
        collection: (*second_collection).to_owned(),
        document_id: (*second_id).to_owned(),
    };
    let (from, to) = match *tag {
        "e" => (first, second),
        "i" => (second, first),
        _ => return Err(Failure::new("edge_key", "unknown edge key direction")),
    };
    Ok(EdgeDelete {
        relation: (*relation).to_owned(),
        from,
        to,
    })
}

pub(super) fn page_start(prefix: &[u8], after: Option<&str>) -> Result<Vec<u8>, Failure> {
    let mut from = prefix.to_vec();
    if let Some(after) = after {
        if after.len() > 5 * 513 {
            return Err(Failure::new(
                "invalid_continuation",
                "edge continuation is too large",
            ));
        }
        from.extend_from_slice(after.as_bytes());
        decode_key(&from)?;
        from.push(0);
    }
    Ok(from)
}

pub(super) fn decode_edge(prefix: &[u8], suffix: &str, bytes: &[u8]) -> Result<Edge, Failure> {
    let mut key = prefix.to_vec();
    key.extend_from_slice(suffix.as_bytes());
    let identity = decode_key(&key)?;
    let properties: Value = serde_json::from_slice(bytes).map_err(Failure::storage)?;
    if !properties.is_object() {
        return Err(Failure::new(
            "edge_properties",
            "stored edge properties are not a JSON object",
        ));
    }
    Ok(Edge {
        relation: identity.relation,
        from: identity.from,
        to: identity.to,
        properties,
    })
}

fn effective_bytes(limit: Option<NonZeroUsize>) -> Result<NonZeroUsize, Failure> {
    let bytes = limit.map_or(usize::MAX, NonZeroUsize::get);
    NonZeroUsize::new(bytes)
        .ok_or_else(|| Failure::new("response_limit", "edge byte limit is zero"))
}

fn page(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    prefix: &[u8],
    after: Option<&str>,
    limit: NonZeroUsize,
    byte_limit: Option<NonZeroUsize>,
) -> Result<Page<Edge>, Failure> {
    let (entries, next_after) = bounded_key_page(
        reader,
        at,
        prefix,
        &page_start(prefix, after)?,
        PageLimits {
            rows: limit,
            bytes: effective_bytes(byte_limit)?,
        },
        |suffix, bytes| decode_edge(prefix, &suffix, &bytes),
        |edge| {
            serde_json::to_vec(edge)
                .map(|bytes| bytes.len())
                .map_err(Failure::storage)
        },
    )?;
    Ok(Page {
        branch: at.branch.clone(),
        version: at.version,
        entries,
        next_after,
    })
}

pub(in crate::serve) fn out(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    node: &Node,
    relation: Option<&str>,
    after: Option<&str>,
    limit: NonZeroUsize,
    byte_limit: Option<NonZeroUsize>,
) -> Result<Page<Edge>, Failure> {
    page(
        reader,
        at,
        &node_prefix(Direction::Out, node, relation)?,
        after,
        limit,
        byte_limit,
    )
}

pub(in crate::serve) fn incoming(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    node: &Node,
    relation: Option<&str>,
    after: Option<&str>,
    limit: NonZeroUsize,
    byte_limit: Option<NonZeroUsize>,
) -> Result<Page<Edge>, Failure> {
    page(
        reader,
        at,
        &node_prefix(Direction::In, node, relation)?,
        after,
        limit,
        byte_limit,
    )
}

pub(in crate::serve) fn export(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    after: Option<&str>,
    limit: NonZeroUsize,
    byte_limit: Option<NonZeroUsize>,
) -> Result<Page<Edge>, Failure> {
    page(reader, at, b"e\0", after, limit, byte_limit)
}

pub(in crate::serve) struct VisibleEdges<'a> {
    pub access: &'a Access,
    pub node: &'a Node,
    pub relation: Option<&'a str>,
    pub after: Option<&'a str>,
    pub limits: PageLimits,
}

pub(in crate::serve) fn visible_out(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    request: &VisibleEdges<'_>,
) -> Result<Page<Edge>, Failure> {
    visible_at_node(reader, at, Direction::Out, request)
}

pub(in crate::serve) fn visible_incoming(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    request: &VisibleEdges<'_>,
) -> Result<Page<Edge>, Failure> {
    visible_at_node(reader, at, Direction::In, request)
}

fn visible_at_node(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    direction: Direction,
    request: &VisibleEdges<'_>,
) -> Result<Page<Edge>, Failure> {
    let prefix = node_prefix(direction, request.node, request.relation)?;
    let from = page_start(&prefix, request.after)?;
    // The same empty page hides denial and absence without scanning incident keys.
    if !request
        .access
        .can_read_document(&request.node.collection, &request.node.document_id)
        || reads::get_bytes(
            reader,
            at,
            &document_key(&request.node.collection, &request.node.document_id),
        )?
        .is_none()
    {
        return Ok(Page {
            branch: at.branch.clone(),
            version: at.version,
            entries: Vec::new(),
            next_after: None,
        });
    }
    visible_page(
        reader,
        at,
        request.access,
        VisibleRowPage {
            prefix: &prefix,
            from: &from,
            limits: request.limits,
        },
    )
}

pub(in crate::serve) fn visible_export(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    access: &Access,
    after: Option<&str>,
    limits: PageLimits,
) -> Result<Page<Edge>, Failure> {
    visible_page(
        reader,
        at,
        access,
        VisibleRowPage {
            prefix: b"e\0",
            from: &page_start(b"e\0", after)?,
            limits,
        },
    )
}

fn visible_page(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    access: &Access,
    request: VisibleRowPage<'_>,
) -> Result<Page<Edge>, Failure> {
    let prefix = request.prefix;
    let (entries, next_after) = visible_row_page(
        reader,
        at,
        request,
        |suffix| {
            let mut key = prefix.to_vec();
            key.extend_from_slice(suffix.as_bytes());
            let identity = decode_key(&key)?;
            Ok(
                access.can_read_document(&identity.from.collection, &identity.from.document_id)
                    && access.can_read_document(&identity.to.collection, &identity.to.document_id),
            )
        },
        |suffix, bytes| decode_edge(prefix, &suffix, &bytes),
        |edge| {
            serde_json::to_vec(edge)
                .map(|bytes| bytes.len())
                .map_err(Failure::storage)
        },
    )?;
    Ok(Page {
        branch: at.branch.clone(),
        version: at.version,
        entries,
        next_after,
    })
}

pub(super) fn removals_for_document(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    document_id: &str,
) -> Result<Vec<PreparedWrite>, Failure> {
    let node = Node {
        collection: collection.to_owned(),
        document_id: document_id.to_owned(),
    };
    let limit = NonZeroUsize::new(128)
        .ok_or_else(|| Failure::new("page_limit", "incident edge page limit is zero"))?;
    let mut removed = BTreeSet::new();
    for direction in [Direction::Out, Direction::In] {
        let prefix = node_prefix(direction, &node, None)?;
        let mut after = None;
        loop {
            let (identities, next) = bounded_key_page(
                reader,
                at,
                &prefix,
                &page_start(&prefix, after.as_deref())?,
                PageLimits {
                    rows: limit,
                    bytes: NonZeroUsize::MAX,
                },
                |suffix, _bytes| {
                    let mut key = prefix.clone();
                    key.extend_from_slice(suffix.as_bytes());
                    decode_key(&key)
                },
                |identity| {
                    serde_json::to_vec(identity)
                        .map(|bytes| bytes.len())
                        .map_err(Failure::storage)
                },
            )?;
            for identity in identities {
                for (key, _) in delete(&identity)? {
                    removed.insert(key);
                }
            }
            let Some(next) = next else { break };
            after = Some(next);
        }
    }
    Ok(removed.into_iter().map(|key| (key, None)).collect())
}

#[cfg(test)]
mod tests {
    use std::error::Error;

    use haematite::db::DatabaseConfig;
    use haematite::store::NodeCacheBudget;
    use serde_json::{json, value::RawValue};

    use super::removals_for_document;
    use crate::serve::protocol::{
        Change, CollectionWrite, DocumentWrite, Edge, Mutation, Node, Selection,
    };
    use crate::serve::storage::{Failure, Storage};

    fn apply(storage: &mut Storage, id: &str, change: Change) -> Result<(), Failure> {
        let expected = storage
            .capture(&Selection {
                branch: "main".to_owned(),
                version: None,
            })?
            .version;
        storage.mutate(&Mutation {
            request_id: id.to_owned(),
            branch: "main".to_owned(),
            expected,
            change,
        })?;
        Ok(())
    }

    #[test]
    fn incident_removals_include_both_mirrors_and_deduplicate_a_self_edge()
    -> Result<(), Box<dyn Error>> {
        let directory = tempfile::tempdir()?;
        let mut storage = Storage::create(
            DatabaseConfig {
                data_dir: directory.path().join("database"),
                shard_count: 2,
                distributed: None,
                executor_threads: Some(2),
                node_cache_budget: Some(NodeCacheBudget::Unlimited),
                retention: haematite::Retention::KeepAll,
            },
            NodeCacheBudget::Unlimited,
            fixture_long_wait_interval(),
        )?;
        apply(
            &mut storage,
            "collection",
            Change::Collection(CollectionWrite {
                collection: "records".to_owned(),
                schema: None,
            }),
        )?;
        for id in ["a", "b", "c"] {
            apply(
                &mut storage,
                id,
                Change::Put(DocumentWrite {
                    collection: "records".to_owned(),
                    document_id: id.to_owned(),
                    document: RawValue::from_string("{}".to_owned())?,
                }),
            )?;
        }
        for (index, (from, to)) in [("a", "b"), ("b", "a"), ("a", "a"), ("c", "b")]
            .into_iter()
            .enumerate()
        {
            apply(
                &mut storage,
                &format!("edge-{index}"),
                Change::EdgePut(Edge {
                    relation: "links".to_owned(),
                    from: Node {
                        collection: "records".to_owned(),
                        document_id: from.to_owned(),
                    },
                    to: Node {
                        collection: "records".to_owned(),
                        document_id: to.to_owned(),
                    },
                    properties: json!({}),
                }),
            )?;
        }
        let at = storage.capture(&Selection {
            branch: "main".to_owned(),
            version: None,
        })?;
        let removals = removals_for_document(&storage.reader()?, &at, "records", "a")?;
        let expected: Vec<_> = [
            "e\0records\0a\0links\0records\0a",
            "e\0records\0a\0links\0records\0b",
            "e\0records\0b\0links\0records\0a",
            "i\0records\0a\0links\0records\0a",
            "i\0records\0a\0links\0records\0b",
            "i\0records\0b\0links\0records\0a",
        ]
        .into_iter()
        .map(|key| (key.as_bytes().to_vec(), None))
        .collect();
        assert_eq!(removals, expected);
        assert_eq!(
            storage
                .capture(&Selection {
                    branch: "main".to_owned(),
                    version: None
                })?
                .version,
            at.version
        );
        assert!(removals_for_document(&storage.reader()?, &at, "records", "absent")?.is_empty());
        Ok(())
    }

    const fn fixture_long_wait_interval() -> std::time::Duration {
        std::time::Duration::from_secs(1)
    }
}

#[cfg(test)]
mod nonzero_tests {
    use super::{Edge, Node, NonZeroUsize, effective_bytes, export};

    #[test]
    fn nonzero_edge_byte_limit_readback() -> Result<(), Box<dyn std::error::Error>> {
        for (given, expected) in [
            (None, usize::MAX),
            (Some(1), 1),
            (Some(37), 37),
            (Some(usize::MAX), usize::MAX),
        ] {
            let given = given
                .map(|value| NonZeroUsize::new(value).ok_or("zero byte limit"))
                .transpose()?;
            assert_eq!(
                effective_bytes(std::hint::black_box(given))?.get(),
                expected
            );
        }
        Ok(())
    }

    #[test]
    fn nonzero_edge_page_stops_at_byte_limit() -> Result<(), Box<dyn std::error::Error>> {
        use crate::serve::protocol::{Change, CollectionWrite, DocumentWrite, Mutation, Selection};
        use crate::serve::storage::Storage;
        use haematite::db::DatabaseConfig;
        use haematite::store::NodeCacheBudget;
        let directory = tempfile::tempdir()?;
        let mut storage = Storage::create(
            DatabaseConfig {
                data_dir: directory.path().join("database"),
                shard_count: 1,
                distributed: None,
                executor_threads: Some(1),
                node_cache_budget: Some(NodeCacheBudget::Unlimited),
                retention: haematite::Retention::KeepAll,
            },
            NodeCacheBudget::Unlimited,
            fixture_long_wait_interval(),
        )?;
        let selection = Selection {
            branch: "main".to_owned(),
            version: None,
        };
        let mut changes = vec![Change::Collection(CollectionWrite {
            collection: "records".to_owned(),
            schema: None,
        })];
        for id in ["a", "b", "c"] {
            changes.push(Change::Put(DocumentWrite {
                collection: "records".to_owned(),
                document_id: id.to_owned(),
                document: serde_json::value::RawValue::from_string("{}".to_owned())?,
            }));
        }
        for id in ["b", "c"] {
            changes.push(Change::EdgePut(Edge {
                from: Node {
                    collection: "records".to_owned(),
                    document_id: "a".to_owned(),
                },
                relation: "links".to_owned(),
                to: Node {
                    collection: "records".to_owned(),
                    document_id: id.to_owned(),
                },
                properties: serde_json::json!({}),
            }));
        }
        for (index, change) in changes.into_iter().enumerate() {
            let expected = storage.capture(&selection)?.version;
            storage.mutate(&Mutation {
                request_id: format!("byte/{index}"),
                branch: "main".to_owned(),
                expected,
                change,
            })?;
        }
        let reader = storage.reader()?;
        let at = storage.capture(&selection)?;
        let all = export(&reader, &at, None, NonZeroUsize::MAX, None)?;
        assert_eq!(all.entries.len(), 2);
        for bytes in [None, Some(NonZeroUsize::MAX)] {
            let read = export(&reader, &at, None, NonZeroUsize::MAX, bytes)?;
            assert_eq!(
                serde_json::to_value(read.entries)?,
                serde_json::to_value(&all.entries)?
            );
        }
        let first = serde_json::to_vec(&all.entries[0])?.len();
        let bytes = NonZeroUsize::new(first).ok_or("empty serialized edge")?;
        let read = export(&reader, &at, None, NonZeroUsize::MAX, Some(bytes))?;
        assert_eq!(read.entries.len(), 1);
        assert_eq!(
            serde_json::to_value(&read.entries[0])?,
            serde_json::to_value(&all.entries[0])?
        );
        assert!(read.next_after.is_some());
        let refused = export(
            &reader,
            &at,
            None,
            NonZeroUsize::MAX,
            Some(NonZeroUsize::MIN),
        )
        .err()
        .ok_or("minimum byte budget accepted")?;
        assert_eq!(refused.code, "response_limit");
        Ok(())
    }

    const fn fixture_long_wait_interval() -> std::time::Duration {
        std::time::Duration::from_secs(1)
    }
}

#[cfg(test)]
#[path = "edges_many_tests.rs"]
mod many_tests;