Merge pull request #221 from vbmade2000/list_col_header_193

crates/bonsaidb-client/src/client/remote_database.rs

@@ -4,7 +4,7 @@ use async_trait::async_trait;
 use bonsaidb_core::{
     connection::{AccessPolicy, Connection, QueryKey, Range, Sort},
     custom_api::CustomApi,
-    document::{AnyDocumentId, OwnedDocument},
+    document::{AnyDocumentId, Header, OwnedDocument},
     key::Key,
     networking::{DatabaseRequest, DatabaseResponse, Request, Response},
     schema::{
@@ -156,6 +156,40 @@ impl<A: CustomApi> Connection for RemoteDatabase<A> {
         }
     }
 
+    async fn list_headers<C, R, PrimaryKey>(
+        &self,
+        ids: R,
+        order: Sort,
+        limit: Option<u32>,
+    ) -> Result<Vec<Header>, bonsaidb_core::Error>
+    where
+        C: Collection,
+        R: Into<Range<PrimaryKey>> + Send,
+        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
+    {
+        match self
+            .client
+            .send_request(Request::Database {
+                database: self.name.to_string(),
+                request: DatabaseRequest::ListHeaders {
+                    collection: C::collection_name(),
+                    ids: ids.into().map_result(|id| id.into().to_document_id())?,
+                    order,
+                    limit,
+                },
+            })
+            .await?
+        {
+            Response::Database(DatabaseResponse::DocumentHeaders(document_headers)) => {
+                Ok(document_headers)
+            }
+            Response::Error(err) => Err(err),
+            other => Err(bonsaidb_core::Error::Networking(
+                bonsaidb_core::networking::Error::UnexpectedResponse(format!("{:?}", other)),
+            )),
+        }
+    }
+
     async fn count<C, R, PrimaryKey>(&self, ids: R) -> Result<u64, bonsaidb_core::Error>
     where
         C: Collection,

crates/bonsaidb-core/src/connection.rs

@@ -11,7 +11,8 @@ use zeroize::Zeroize;
 use crate::schema::Nameable;
 use crate::{
     document::{
-        AnyDocumentId, CollectionDocument, CollectionHeader, Document, HasHeader, OwnedDocument,
+        AnyDocumentId, CollectionDocument, CollectionHeader, Document, HasHeader, Header,
+        OwnedDocument,
     },
     key::{IntoPrefixRange, Key},
     permissions::Permissions,
@@ -189,6 +190,27 @@ pub trait Connection: Send + Sync {
         R: Into<Range<PrimaryKey>> + Send,
         PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send;
 
+    /// Retrieves all documents within the range of `ids`. To retrieve all
+    /// documents, pass in `..` for `ids`.
+    ///
+    /// This is the lower-level API. For better ergonomics, consider using one
+    /// of:
+    ///
+    /// - [`SerializedCollection::all()`]
+    /// - [`self.collection::<Collection>().all()`](Collection::all)
+    /// - [`SerializedCollection::list_headers()`]
+    /// - [`self.collection::<Collection>().list_headers()`](Collection::list_headers)
+    async fn list_headers<C, R, PrimaryKey>(
+        &self,
+        ids: R,
+        order: Sort,
+        limit: Option<u32>,
+    ) -> Result<Vec<Header>, Error>
+    where
+        C: schema::Collection,
+        R: Into<Range<PrimaryKey>> + Send,
+        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send;
+
     /// Counts the number of documents within the range of `ids`.
     ///
     /// This is the lower-level API. For better ergonomics, consider using
@@ -904,6 +926,44 @@ where
             _ => unreachable!("Attempted to use after retrieving the result"),
         }
     }
+
+    /// Returns the list of headers for documents contained within the range.
+    ///
+    /// Order and limit are ignored if they were set.
+    ///
+    /// ```rust
+    /// # bonsaidb_core::__doctest_prelude!();
+    /// # fn test_fn<C: Connection>(db: &C) -> Result<(), Error> {
+    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
+    /// println!(
+    ///     "Number of documents with id 42 or larger: {:?}",
+    ///     db.collection::<MyCollection>().list(42..).headers().await?
+    /// );
+    /// println!(
+    ///     "Number of documents in MyCollection: {:?}",
+    ///     db.collection::<MyCollection>().all().headers().await?
+    /// );
+    /// # Ok(())
+    /// # })
+    /// # }
+    /// ```
+    pub async fn headers(self) -> Result<Vec<Header>, Error> {
+        match self.state {
+            ListState::Pending(Some(ListBuilder {
+                collection,
+                range,
+                sort,
+                limit,
+                ..
+            })) => {
+                collection
+                    .connection
+                    .list_headers::<Cl, _, _>(range, sort, limit)
+                    .await
+            }
+            _ => unreachable!("Attempted to use after retrieving the result"),
+        }
+    }
 }
 
 impl<'a, Cn, Cl> Future for List<'a, Cn, Cl>

crates/bonsaidb-core/src/document/id.rs

@@ -27,7 +27,7 @@ impl Deref for DocumentId {
 
 impl Ord for DocumentId {
     fn cmp(&self, other: &Self) -> Ordering {
-        (&**self).cmp(&**other)
+        (**self).cmp(&**other)
     }
 }
 
@@ -91,7 +91,7 @@ impl Display for DocumentId {
 
 impl Hash for DocumentId {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        (&**self).hash(state);
+        (**self).hash(state);
     }
 }
 

crates/bonsaidb-core/src/networking.rs

@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     connection::{AccessPolicy, Authenticated, Database, QueryKey, Range, Sort},
-    document::{DocumentId, OwnedDocument},
+    document::{DocumentId, Header, OwnedDocument},
     keyvalue::{KeyOperation, Output},
     schema::{
         self,
@@ -168,6 +168,18 @@ pub enum DatabaseRequest {
         /// The maximum number of results to return.
         limit: Option<u32>,
     },
+    /// Retrieve headers of multiple documents.
+    #[cfg_attr(feature = "actionable-traits", actionable(protection = "simple"))]
+    ListHeaders {
+        /// The collection of the documents.
+        collection: CollectionName,
+        /// The range of ids to list.
+        ids: Range<DocumentId>,
+        /// The order for the query into the collection.
+        order: Sort,
+        /// The maximum number of results to return.
+        limit: Option<u32>,
+    },
     /// Counts the number of documents in the specified range.
     #[cfg_attr(feature = "actionable-traits", actionable(protection = "simple"))]
     Count {
@@ -341,6 +353,8 @@ pub enum ServerResponse {
 pub enum DatabaseResponse {
     /// One or more documents.
     Documents(Vec<OwnedDocument>),
+    /// One or more document headers.
+    DocumentHeaders(Vec<Header>),
     /// A result count.
     Count(u64),
     /// Results of [`DatabaseRequest::ApplyTransaction`].

crates/bonsaidb-core/src/permissions/bonsai.rs

@@ -170,6 +170,11 @@ pub enum DocumentAction {
     /// [`collection_resource_name()`] for the format of collection resource
     /// names.
     List,
+    /// Allows listing documents through
+    /// [`Connection::list_headers()`](crate::connection::Connection::list_headers). See
+    /// [`collection_resource_name()`] for the format of collection resource
+    /// names.
+    ListHeaders,
     /// Allows counting documents through
     /// [`Connection::count()`](crate::connection::Connection::count). See
     /// [`collection_resource_name()`] for the format of collection resource

crates/bonsaidb-core/src/schema/collection.rs

@@ -9,7 +9,7 @@ use transmog_pot::Pot;
 use crate::{
     connection::{self, Connection, Range},
     document::{
-        AnyDocumentId, BorrowedDocument, CollectionDocument, Document, DocumentId, KeyId,
+        AnyDocumentId, BorrowedDocument, CollectionDocument, Document, DocumentId, Header, KeyId,
         OwnedDocument, OwnedDocuments,
     },
     key::{IntoPrefixRange, Key},
@@ -1279,6 +1279,30 @@ where
     pub async fn count(self) -> Result<u64, Error> {
         self.0.count().await
     }
+
+    /// Returns the number of list of document headers contained within the range.
+    ///
+    /// Order and limit are ignored if they were set.
+    ///
+    /// ```rust
+    /// # bonsaidb_core::__doctest_prelude!();
+    /// # fn test_fn<C: Connection>(db: &C) -> Result<(), Error> {
+    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
+    /// println!(
+    ///     "Number of documents with id 42 or larger: {:?}",
+    ///     MyCollection::list(42.., db).headers().await?
+    /// );
+    /// println!(
+    ///     "Number of documents in MyCollection: {:?}",
+    ///     MyCollection::all(db).headers().await?
+    /// );
+    /// # Ok(())
+    /// # })
+    /// # }
+    /// ```
+    pub async fn headers(self) -> Result<Vec<Header>, Error> {
+        self.0.headers().await
+    }
 }
 
 impl<'a, Cn, Cl> Future for List<'a, Cn, Cl>

crates/bonsaidb-core/src/test_util.rs

@@ -881,6 +881,14 @@ pub async fn list_tests<C: Connection>(db: &C) -> anyhow::Result<()> {
     assert_eq!(both_docs[0].contents.value, doc1_value.value);
     assert_eq!(both_docs[1].contents.value, doc2_value.value);
 
+    let both_headers = db
+        .collection::<Basic>()
+        .list(doc1.id..=doc2.id)
+        .headers()
+        .await?;
+
+    assert_eq!(both_headers.len(), 2);
+
     let one_doc = Basic::list(doc1.id..doc2.id, db).await?;
     assert_eq!(one_doc.len(), 1);
 

crates/bonsaidb-local/src/database.rs

@@ -326,6 +326,18 @@ impl Database {
         self.list(ids, order, limit, collection).await
     }
 
+    #[cfg(feature = "internal-apis")]
+    #[doc(hidden)]
+    pub async fn list_headers_from_collection(
+        &self,
+        ids: Range<DocumentId>,
+        order: Sort,
+        limit: Option<u32>,
+        collection: &CollectionName,
+    ) -> Result<Vec<Header>, bonsaidb_core::Error> {
+        self.list_headers(ids, order, limit, collection).await
+    }
+
     #[cfg(feature = "internal-apis")]
     #[doc(hidden)]
     pub async fn count_from_collection(
@@ -594,6 +606,67 @@ impl Database {
         .unwrap()
     }
 
+    pub(crate) async fn list_headers(
+        &self,
+        ids: Range<DocumentId>,
+        sort: Sort,
+        limit: Option<u32>,
+        collection: &CollectionName,
+    ) -> Result<Vec<Header>, bonsaidb_core::Error> {
+        let task_self = self.clone();
+        let collection = collection.clone();
+        tokio::task::spawn_blocking(move || {
+            let tree = task_self
+                .data
+                .context
+                .roots
+                .tree(task_self.collection_tree::<Versioned, _>(
+                    &collection,
+                    document_tree_name(&collection),
+                )?)
+                .map_err(Error::from)?;
+            let mut found_headers = Vec::new();
+            let mut keys_read = 0;
+            let ids = DocumentIdRange(ids);
+            tree.scan(
+                &ids.borrow_as_bytes(),
+                match sort {
+                    Sort::Ascending => true,
+                    Sort::Descending => false,
+                },
+                |_, _, _| ScanEvaluation::ReadData,
+                |_, _| {
+                    if let Some(limit) = limit {
+                        if keys_read >= limit {
+                            return ScanEvaluation::Stop;
+                        }
+
+                        keys_read += 1;
+                    }
+                    ScanEvaluation::ReadData
+                },
+                |_, _, doc| {
+                    found_headers.push(
+                        deserialize_document(&doc)
+                            .map(BorrowedDocument::into_owned)
+                            .map(|doc| doc.header)
+                            .map_err(AbortError::Other)?,
+                    );
+                    Ok(())
+                },
+            )
+            .map_err(|err| match err {
+                AbortError::Other(err) => err,
+                AbortError::Nebari(err) => crate::Error::from(err),
+            })
+            .unwrap();
+
+            Ok(found_headers)
+        })
+        .await
+        .unwrap()
+    }
+
     pub(crate) async fn count(
         &self,
         ids: Range<DocumentId>,
@@ -1384,6 +1457,27 @@ impl Connection for Database {
         .await
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(ids, order, limit)))]
+    async fn list_headers<C, R, PrimaryKey>(
+        &self,
+        ids: R,
+        order: Sort,
+        limit: Option<u32>,
+    ) -> Result<Vec<Header>, bonsaidb_core::Error>
+    where
+        C: schema::Collection,
+        R: Into<Range<PrimaryKey>> + Send,
+        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
+    {
+        self.list_headers(
+            ids.into().map_result(|id| id.into().to_document_id())?,
+            order,
+            limit,
+            &C::collection_name(),
+        )
+        .await
+    }
+
     #[cfg_attr(feature = "tracing", tracing::instrument(skip(ids)))]
     async fn count<C, R, PrimaryKey>(&self, ids: R) -> Result<u64, bonsaidb_core::Error>
     where