haematite-cli 0.21.3

#![forbid(unsafe_code)]
//! Deleting a document removes its edges, in both directions, and its
//! vectors in the same version (D7), and leaves every other record alone.

use std::error::Error;
use std::num::NonZeroU32;

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use haematite_cli::serve::protocol::{
    Change, CollectionWrite, DocumentDelete, DocumentRead, DocumentWrite, Edge, EdgePage,
    GraphPage, Mutation, Node, Operation, Selection, VectorRead, VectorWrite, Version,
};
use haematite_cli::serve::service::Service;
use haematite_cli::serve::storage::Storage;
use serde_json::{Value, json, value::RawValue};

type TestResult<T = ()> = Result<T, Box<dyn Error>>;

fn limit() -> TestResult<NonZeroU32> {
    Ok(NonZeroU32::new(16).ok_or("page limit is zero")?)
}

fn at(version: Option<Version>) -> Selection {
    Selection {
        branch: "main".to_owned(),
        version,
    }
}

fn node(document_id: &str) -> Node {
    Node {
        collection: "records".to_owned(),
        document_id: document_id.to_owned(),
    }
}

fn edge(from: &str, to: &str) -> Edge {
    Edge {
        relation: "links".to_owned(),
        from: node(from),
        to: node(to),
        properties: json!({}),
    }
}

fn apply(service: &Service, id: &str, change: Change) -> TestResult<Version> {
    let listed = service.dispatch(Operation::GraphExport(GraphPage {
        selection: at(None),
        after: None,
        limit: limit()?,
        relations: None,
        collections: None,
        members: None,
        records: None,
        unrelated: None,
    }))?;
    let expected: Version = serde_json::from_value(listed["selection"]["version"].clone())?;
    let receipt = service.dispatch(Operation::Mutate(Mutation {
        request_id: id.to_owned(),
        branch: "main".to_owned(),
        expected,
        change,
    }))?;
    Ok(serde_json::from_value(receipt["version"].clone())?)
}

fn edges(service: &Service, version: Option<Version>) -> TestResult<Vec<Edge>> {
    let answer = service.dispatch(Operation::GraphExport(GraphPage {
        selection: at(version),
        after: None,
        limit: limit()?,
        relations: None,
        collections: None,
        members: None,
        records: None,
        unrelated: None,
    }))?;
    assert_eq!(answer["next_after"], Value::Null);
    Ok(serde_json::from_value(answer["edges"].clone())?)
}

fn at_node(service: &Service, incoming: bool, document_id: &str) -> TestResult<Vec<Edge>> {
    let page = EdgePage {
        selection: at(None),
        node: node(document_id),
        relation: None,
        after: None,
        limit: limit()?,
    };
    let operation = if incoming {
        Operation::EdgesIn(page)
    } else {
        Operation::EdgesOut(page)
    };
    let answer = service.dispatch(operation)?;
    Ok(serde_json::from_value(answer["edges"].clone())?)
}

fn vector(service: &Service, version: Option<Version>, id: &str, field: &str) -> TestResult<Value> {
    let answer = service.dispatch(Operation::VectorGet(VectorRead {
        selection: at(version),
        collection: "records".to_owned(),
        document_id: id.to_owned(),
        field: field.to_owned(),
    }))?;
    Ok(answer["vector"].clone())
}

fn served() -> Result<(tempfile::TempDir, Service), Box<dyn std::error::Error>> {
    let directory = tempfile::tempdir()?;
    let 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(),
    )?;
    let service = Service::new(storage, limit()?)?;
    Ok((directory, service))
}

#[test]
fn deleting_a_document_removes_its_edges_both_ways_and_its_vectors_and_nothing_else() -> TestResult
{
    let (_directory, service) = served()?;
    apply(
        &service,
        "declare",
        Change::Collection(CollectionWrite {
            collection: "records".to_owned(),
            schema: None,
        }),
    )?;
    for id in ["a", "b", "c"] {
        apply(
            &service,
            &format!("document-{id}"),
            Change::Put(DocumentWrite {
                collection: "records".to_owned(),
                document_id: id.to_owned(),
                document: RawValue::from_string("{}".to_owned())?,
            }),
        )?;
    }
    // a -> b, b -> a, a -> a and c -> b: three touch a, one does not.
    for (index, (from, to)) in [("a", "b"), ("b", "a"), ("a", "a"), ("c", "b")]
        .into_iter()
        .enumerate()
    {
        apply(
            &service,
            &format!("edge-{index}"),
            Change::EdgePut(edge(from, to)),
        )?;
    }
    for (id, field, members) in [
        ("a", "embedding", vec![1.0_f32, 0.0]),
        ("a", "summary", vec![0.5_f32, 0.5, 0.5]),
        ("b", "embedding", vec![0.0_f32, 1.0]),
    ] {
        apply(
            &service,
            &format!("vector-{id}-{field}"),
            Change::VectorPut(VectorWrite {
                collection: "records".to_owned(),
                document_id: id.to_owned(),
                field: field.to_owned(),
                vector: members,
            }),
        )?;
    }
    assert_eq!(edges(&service, None)?.len(), 4);
    assert_eq!(vector(&service, None, "a", "embedding")?, json!([1.0, 0.0]));
    assert_eq!(
        vector(&service, None, "a", "summary")?,
        json!([0.5, 0.5, 0.5])
    );

    let before = apply(
        &service,
        "unrelated",
        Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: "c".to_owned(),
            document: RawValue::from_string(r#"{"touched":true}"#.to_owned())?,
        }),
    )?;
    let deleted = apply(
        &service,
        "delete-a",
        Change::Delete(DocumentDelete {
            collection: "records".to_owned(),
            document_id: "a".to_owned(),
        }),
    )?;
    assert_ne!(deleted, before);

    let gone = service.dispatch(Operation::Get(DocumentRead {
        selection: at(None),
        collection: "records".to_owned(),
        document_id: "a".to_owned(),
    }))?;
    assert_eq!(gone["document"], Value::Null);
    // Only c -> b is left, and it is read from both of its ends.
    assert_eq!(edges(&service, None)?, vec![edge("c", "b")]);
    assert!(at_node(&service, false, "a")?.is_empty());
    assert!(at_node(&service, true, "a")?.is_empty());
    assert!(at_node(&service, false, "b")?.is_empty());
    assert_eq!(at_node(&service, true, "b")?, vec![edge("c", "b")]);
    assert_eq!(at_node(&service, false, "c")?, vec![edge("c", "b")]);
    assert_eq!(vector(&service, None, "a", "embedding")?, Value::Null);
    assert_eq!(vector(&service, None, "a", "summary")?, Value::Null);
    assert_eq!(vector(&service, None, "b", "embedding")?, json!([0.0, 1.0]));

    // The version before the delete still answers everything it held.
    assert_eq!(edges(&service, Some(before))?.len(), 4);
    assert_eq!(
        vector(&service, Some(before), "a", "embedding")?,
        json!([1.0, 0.0])
    );
    Ok(())
}

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