diff --git a/crates/bonsaidb-client/src/client/remote_database.rs b/crates/bonsaidb-client/src/client/remote_database.rs index 7f5dc55747..13ec3f1e6e 100644 --- a/crates/bonsaidb-client/src/client/remote_database.rs +++ b/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 Connection for RemoteDatabase { } } + async fn list_headers( + &self, + ids: R, + order: Sort, + limit: Option, + ) -> Result, bonsaidb_core::Error> + where + C: Collection, + R: Into> + Send, + PrimaryKey: Into> + 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(&self, ids: R) -> Result where C: Collection, diff --git a/crates/bonsaidb-core/src/connection.rs b/crates/bonsaidb-core/src/connection.rs index 115c753f45..808c71c023 100644 --- a/crates/bonsaidb-core/src/connection.rs +++ b/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> + Send, PrimaryKey: Into> + 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::().all()`](Collection::all) + /// - [`SerializedCollection::list_headers()`] + /// - [`self.collection::().list_headers()`](Collection::list_headers) + async fn list_headers( + &self, + ids: R, + order: Sort, + limit: Option, + ) -> Result, Error> + where + C: schema::Collection, + R: Into> + Send, + PrimaryKey: Into> + 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(db: &C) -> Result<(), Error> { + /// # tokio::runtime::Runtime::new().unwrap().block_on(async { + /// println!( + /// "Number of documents with id 42 or larger: {:?}", + /// db.collection::().list(42..).headers().await? + /// ); + /// println!( + /// "Number of documents in MyCollection: {:?}", + /// db.collection::().all().headers().await? + /// ); + /// # Ok(()) + /// # }) + /// # } + /// ``` + pub async fn headers(self) -> Result, Error> { + match self.state { + ListState::Pending(Some(ListBuilder { + collection, + range, + sort, + limit, + .. + })) => { + collection + .connection + .list_headers::(range, sort, limit) + .await + } + _ => unreachable!("Attempted to use after retrieving the result"), + } + } } impl<'a, Cn, Cl> Future for List<'a, Cn, Cl> diff --git a/crates/bonsaidb-core/src/document/id.rs b/crates/bonsaidb-core/src/document/id.rs index 284ea11541..080e6fca7b 100644 --- a/crates/bonsaidb-core/src/document/id.rs +++ b/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(&self, state: &mut H) { - (&**self).hash(state); + (**self).hash(state); } } diff --git a/crates/bonsaidb-core/src/networking.rs b/crates/bonsaidb-core/src/networking.rs index c3269322c5..69735b51dc 100644 --- a/crates/bonsaidb-core/src/networking.rs +++ b/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, }, + /// 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, + /// The order for the query into the collection. + order: Sort, + /// The maximum number of results to return. + limit: Option, + }, /// 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), + /// One or more document headers. + DocumentHeaders(Vec
), /// A result count. Count(u64), /// Results of [`DatabaseRequest::ApplyTransaction`]. diff --git a/crates/bonsaidb-core/src/permissions/bonsai.rs b/crates/bonsaidb-core/src/permissions/bonsai.rs index 24fa0fd4e5..3c4cdd96e6 100644 --- a/crates/bonsaidb-core/src/permissions/bonsai.rs +++ b/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 diff --git a/crates/bonsaidb-core/src/schema/collection.rs b/crates/bonsaidb-core/src/schema/collection.rs index f251ddf7df..ffc511036e 100644 --- a/crates/bonsaidb-core/src/schema/collection.rs +++ b/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 { 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(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, Error> { + self.0.headers().await + } } impl<'a, Cn, Cl> Future for List<'a, Cn, Cl> diff --git a/crates/bonsaidb-core/src/test_util.rs b/crates/bonsaidb-core/src/test_util.rs index ef906200e6..2f21dc7864 100644 --- a/crates/bonsaidb-core/src/test_util.rs +++ b/crates/bonsaidb-core/src/test_util.rs @@ -881,6 +881,14 @@ pub async fn list_tests(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::() + .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); diff --git a/crates/bonsaidb-local/src/database.rs b/crates/bonsaidb-local/src/database.rs index dd858a5dbd..7f4e0e7c64 100644 --- a/crates/bonsaidb-local/src/database.rs +++ b/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, + order: Sort, + limit: Option, + collection: &CollectionName, + ) -> Result, 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, + sort: Sort, + limit: Option, + collection: &CollectionName, + ) -> Result, 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::( + &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, @@ -1384,6 +1457,27 @@ impl Connection for Database { .await } + #[cfg_attr(feature = "tracing", tracing::instrument(skip(ids, order, limit)))] + async fn list_headers( + &self, + ids: R, + order: Sort, + limit: Option, + ) -> Result, bonsaidb_core::Error> + where + C: schema::Collection, + R: Into> + Send, + PrimaryKey: Into> + 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(&self, ids: R) -> Result where diff --git a/crates/bonsaidb-server/src/server.rs b/crates/bonsaidb-server/src/server.rs index 80a4efa0c1..188c3998e4 100644 --- a/crates/bonsaidb-server/src/server.rs +++ b/crates/bonsaidb-server/src/server.rs @@ -1547,6 +1547,42 @@ impl<'s, B: Backend> bonsaidb_core::networking::ListHandler for DatabaseDispatch } } +#[async_trait] +impl<'s, B: Backend> bonsaidb_core::networking::ListHeadersHandler for DatabaseDispatcher<'s, B> { + type Action = BonsaiAction; + + async fn resource_name<'a>( + &'a self, + collection: &'a CollectionName, + _ids: &'a Range, + _order: &'a Sort, + _limit: &'a Option, + ) -> Result, Error> { + Ok(collection_resource_name(&self.name, collection)) + } + + fn action() -> Self::Action { + BonsaiAction::Database(DatabaseAction::Document(DocumentAction::ListHeaders)) + } + + async fn handle_protected( + &self, + _permissions: &Permissions, + collection: CollectionName, + ids: Range, + order: Sort, + limit: Option, + ) -> Result>, Error> { + let documents = self + .database + .list_headers_from_collection(ids, order, limit, &collection) + .await?; + Ok(Response::Database(DatabaseResponse::DocumentHeaders( + documents, + ))) + } +} + #[async_trait] impl<'s, B: Backend> bonsaidb_core::networking::CountHandler for DatabaseDispatcher<'s, B> { type Action = BonsaiAction; diff --git a/crates/bonsaidb-server/src/server/database.rs b/crates/bonsaidb-server/src/server/database.rs index 21dc7c2294..8f8c60753d 100644 --- a/crates/bonsaidb-server/src/server/database.rs +++ b/crates/bonsaidb-server/src/server/database.rs @@ -4,7 +4,7 @@ use async_trait::async_trait; use bonsaidb_core::{ circulate::Message, connection::{AccessPolicy, QueryKey, Range, Sort}, - document::{AnyDocumentId, OwnedDocument}, + document::{AnyDocumentId, Header, OwnedDocument}, keyvalue::KeyValue, pubsub::{PubSub, Subscriber}, schema::{self, view::map::MappedDocuments, Map, MappedValue, SerializedView}, @@ -135,6 +135,22 @@ impl bonsaidb_core::connection::Connection for ServerDatabase { self.db.list::(ids, order, limit).await } + async fn list_headers( + &self, + ids: R, + order: Sort, + limit: Option, + ) -> Result, bonsaidb_core::Error> + where + C: schema::Collection, + R: Into> + Send, + PrimaryKey: Into> + Send, + { + self.db + .list_headers::(ids, order, limit) + .await + } + async fn count(&self, ids: R) -> Result where C: schema::Collection, diff --git a/crates/bonsaidb/src/any_connection.rs b/crates/bonsaidb/src/any_connection.rs index 1d7164adec..4a03791117 100644 --- a/crates/bonsaidb/src/any_connection.rs +++ b/crates/bonsaidb/src/any_connection.rs @@ -4,7 +4,7 @@ use bonsaidb_core::connection::{Authenticated, Authentication}; use bonsaidb_core::{ async_trait::async_trait, connection::{self, AccessPolicy, Connection, QueryKey, Range, Sort, StorageConnection}, - document::{AnyDocumentId, OwnedDocument}, + document::{AnyDocumentId, Header, OwnedDocument}, schema::{ view::map::MappedDocuments, Collection, Map, MappedValue, Nameable, Schema, SchemaName, SerializedView, @@ -259,6 +259,23 @@ impl Connection for AnyDatabase { } } + async fn list_headers( + &self, + ids: R, + order: Sort, + limit: Option, + ) -> Result, bonsaidb_core::Error> + where + C: Collection, + R: Into> + Send, + PrimaryKey: Into> + Send, + { + match self { + Self::Local(server) => server.list_headers::(ids, order, limit).await, + Self::Networked(client) => client.list_headers::(ids, order, limit).await, + } + } + async fn count(&self, ids: R) -> Result where C: Collection,