GetCollections::into_stream

Author: rylev

sdk/data_cosmos/examples/collection.rs

@@ -1,4 +1,3 @@
-use azure_core::prelude::*;
 use azure_data_cosmos::prelude::*;
 use futures::stream::StreamExt;
 use std::error::Error;
@@ -60,12 +59,10 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
 
     for db in databases.databases {
         let database_client = client.clone().into_database_client(db.id.clone());
-        let collections = Box::pin(
-            database_client.list_collections(Context::new(), ListCollectionsOptions::new()),
-        )
-        .next()
-        .await
-        .unwrap()?;
+        let collections = Box::pin(database_client.list_collections().into_stream())
+            .next()
+            .await
+            .unwrap()?;
         println!(
             "database {} has {} collection(s)",
             db.id,

sdk/data_cosmos/examples/create_delete_database.rs

@@ -1,4 +1,3 @@
-use azure_core::Context;
 use azure_data_cosmos::prelude::*;
 use futures::stream::StreamExt;
 use std::error::Error;
@@ -61,7 +60,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         let get_collection_response = db_collection.get_collection().into_future().await?;
         println!("get_collection_response == {:#?}", get_collection_response);
 
-        let stream = db_client.list_collections(Context::new(), ListCollectionsOptions::new());
+        let stream = db_client.list_collections().into_stream();
         let mut stream = Box::pin(stream);
         while let Some(res) = stream.next().await {
             let res = res?;

sdk/data_cosmos/examples/database_00.rs

@@ -25,11 +25,10 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         println!("database == {:?}", db);
         let database = client.clone().into_database_client(db.name().to_owned());
 
-        let collections =
-            Box::pin(database.list_collections(Context::new(), ListCollectionsOptions::new()))
-                .next()
-                .await
-                .unwrap()?;
+        let collections = Box::pin(database.list_collections().into_stream())
+            .next()
+            .await
+            .unwrap()?;
         for collection in collections.collections {
             println!("collection == {:?}", collection);
             let collection_client = database.clone().into_collection_client(collection.id);

sdk/data_cosmos/examples/database_01.rs

@@ -1,4 +1,3 @@
-use azure_core::Context;
 use azure_data_cosmos::prelude::*;
 use futures::stream::StreamExt;
 use std::error::Error;
@@ -18,11 +17,10 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
     let database_client = client.into_database_client("pollo");
     println!("database_name == {}", database_client.database_name());
 
-    let collections =
-        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .unwrap()?;
+    let collections = Box::pin(database_client.list_collections().into_stream())
+        .next()
+        .await
+        .unwrap()?;
     println!("collections == {:#?}", collections);
 
     let collection_client = database_client.into_collection_client("cnt");

sdk/data_cosmos/examples/document_00.rs

@@ -85,7 +85,8 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
             client
                 .clone()
                 .into_database_client(database.id.clone())
-                .list_collections(Context::new(), ListCollectionsOptions::new()),
+                .list_collections()
+                .into_stream(),
         )
         .next()
         .await

sdk/data_cosmos/src/clients/database_client.rs

@@ -68,61 +68,8 @@ impl DatabaseClient {
     }
 
     /// List collections in the database
-    pub fn list_collections(
-        &self,
-        ctx: Context,
-        options: ListCollectionsOptions,
-    ) -> impl Stream<Item = crate::Result<ListCollectionsResponse>> + '_ {
-        unfold(State::Init, move |state: State| {
-            let this = self.clone();
-            let ctx = ctx.clone();
-            let options = options.clone();
-            async move {
-                let response = match state {
-                    State::Init => {
-                        let mut request = this.cosmos_client().prepare_request_pipeline(
-                            &format!("dbs/{}/colls", this.database_name()),
-                            http::Method::GET,
-                        );
-
-                        r#try!(options.decorate_request(&mut request));
-                        let response = r#try!(
-                            this.pipeline()
-                                .send(ctx.clone().insert(ResourceType::Collections), &mut request)
-                                .await
-                        );
-                        ListCollectionsResponse::try_from(response).await
-                    }
-                    State::Continuation(continuation_token) => {
-                        let continuation = Continuation::new(continuation_token.as_str());
-                        let mut request = this.cosmos_client().prepare_request_pipeline(
-                            &format!("dbs/{}/colls", self.database_name()),
-                            http::Method::GET,
-                        );
-
-                        r#try!(options.decorate_request(&mut request));
-                        r#try!(continuation.add_as_header2(&mut request));
-                        let response = r#try!(
-                            this.pipeline()
-                                .send(ctx.clone().insert(ResourceType::Collections), &mut request)
-                                .await
-                        );
-                        ListCollectionsResponse::try_from(response).await
-                    }
-                    State::Done => return None,
-                };
-
-                let response = r#try!(response);
-
-                let next_state = response
-                    .continuation_token
-                    .clone()
-                    .map(State::Continuation)
-                    .unwrap_or(State::Done);
-
-                Some((Ok(response), next_state))
-            }
-        })
+    pub fn list_collections(&self) -> ListCollectionsBuilder {
+        ListCollectionsBuilder::new(self.clone())
     }
 
     /// Create a collection

sdk/data_cosmos/src/operations/list_collections.rs

@@ -5,34 +5,122 @@ use crate::ResourceQuota;
 use azure_core::collect_pinned_stream;
 use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::prelude::*;
-use azure_core::Request as HttpRequest;
 use azure_core::Response as HttpResponse;
 use chrono::{DateTime, Utc};
+use futures::stream::unfold;
+use futures::Stream;
 
+/// Macro for short cutting a stream on error
+macro_rules! r#try {
+    ($expr:expr $(,)?) => {
+        match $expr {
+            Result::Ok(val) => val,
+            Result::Err(err) => {
+                return Some((Err(err.into()), State::Done));
+            }
+        }
+    };
+}
+
+/// Stream state
+#[derive(Debug, Clone, PartialEq)]
+enum State {
+    Init,
+    Continuation(String),
+    Done,
+}
 #[derive(Debug, Clone)]
-pub struct ListCollectionsOptions {
+pub struct ListCollectionsBuilder {
+    client: DatabaseClient,
     consistency_level: Option<ConsistencyLevel>,
     max_item_count: MaxItemCount,
+    context: Context,
 }
 
-impl ListCollectionsOptions {
-    pub fn new() -> Self {
+impl ListCollectionsBuilder {
+    pub(crate) fn new(client: DatabaseClient) -> Self {
         Self {
+            client,
             max_item_count: MaxItemCount::new(-1),
             consistency_level: None,
+            context: Context::new(),
         }
     }
 
     setters! {
         consistency_level: ConsistencyLevel => Some(consistency_level),
         max_item_count: i32 => MaxItemCount::new(max_item_count),
+        context: Context => context,
     }
 
-    pub fn decorate_request(&self, request: &mut HttpRequest) -> crate::Result<()> {
-        azure_core::headers::add_optional_header2(&self.consistency_level, request)?;
-        azure_core::headers::add_mandatory_header2(&self.max_item_count, request)?;
+    pub fn into_stream(
+        self,
+    ) -> impl Stream<Item = crate::Result<ListCollectionsResponse>> + 'static {
+        unfold(State::Init, move |state: State| {
+            let this = self.clone();
+            let ctx = self.context.clone();
+            async move {
+                let response = match state {
+                    State::Init => {
+                        let mut request = this.client.cosmos_client().prepare_request_pipeline(
+                            &format!("dbs/{}/colls", this.client.database_name()),
+                            http::Method::GET,
+                        );
+
+                        r#try!(azure_core::headers::add_optional_header2(
+                            &this.consistency_level,
+                            &mut request,
+                        ));
+                        r#try!(azure_core::headers::add_mandatory_header2(
+                            &this.max_item_count,
+                            &mut request,
+                        ));
+                        let response = r#try!(
+                            this.client
+                                .pipeline()
+                                .send(ctx.clone().insert(ResourceType::Collections), &mut request)
+                                .await
+                        );
+                        ListCollectionsResponse::try_from(response).await
+                    }
+                    State::Continuation(continuation_token) => {
+                        let continuation = Continuation::new(continuation_token.as_str());
+                        let mut request = this.client.cosmos_client().prepare_request_pipeline(
+                            &format!("dbs/{}/colls", this.client.database_name()),
+                            http::Method::GET,
+                        );
+
+                        r#try!(azure_core::headers::add_optional_header2(
+                            &this.consistency_level,
+                            &mut request,
+                        ));
+                        r#try!(azure_core::headers::add_mandatory_header2(
+                            &this.max_item_count,
+                            &mut request,
+                        ));
+                        r#try!(continuation.add_as_header2(&mut request));
+                        let response = r#try!(
+                            this.client
+                                .pipeline()
+                                .send(ctx.clone().insert(ResourceType::Collections), &mut request)
+                                .await
+                        );
+                        ListCollectionsResponse::try_from(response).await
+                    }
+                    State::Done => return None,
+                };
 
-        Ok(())
+                let response = r#try!(response);
+
+                let next_state = response
+                    .continuation_token
+                    .clone()
+                    .map(State::Continuation)
+                    .unwrap_or(State::Done);
+
+                Some((Ok(response), next_state))
+            }
+        })
     }
 }
 
@@ -90,3 +178,16 @@ impl ListCollectionsResponse {
         })
     }
 }
+
+/// The future returned by calling `into_future` on the builder.
+pub type ListCollections =
+    futures::future::BoxFuture<'static, crate::Result<ListCollectionsResponse>>;
+
+#[cfg(feature = "into_future")]
+impl std::future::IntoFuture for ListCollectionsBuilder {
+    type Future = ListCollections;
+    type Output = <ListCollections as std::future::Future>::Output;
+    fn into_future(self) -> Self::Future {
+        Self::into_future(self)
+    }
+}

sdk/data_cosmos/tests/cosmos_collection.rs

@@ -27,12 +27,11 @@ async fn create_and_delete_collection() {
         .into_future()
         .await
         .unwrap();
-    let collections =
-        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .unwrap()
-            .unwrap();
+    let collections = Box::pin(database_client.list_collections().into_stream())
+        .next()
+        .await
+        .unwrap()
+        .unwrap();
     assert!(collections.collections.len() == 1);
 
     // try to get the previously created collection
@@ -60,12 +59,11 @@ async fn create_and_delete_collection() {
         .into_future()
         .await
         .unwrap();
-    let collections =
-        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .unwrap()
-            .unwrap();
+    let collections = Box::pin(database_client.list_collections().into_stream())
+        .next()
+        .await
+        .unwrap()
+        .unwrap();
     assert!(collections.collections.len() == 0);
 
     database_client
@@ -105,12 +103,11 @@ async fn replace_collection() {
         .await
         .unwrap();
 
-    let collections =
-        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .unwrap()
-            .unwrap();
+    let collections = Box::pin(database_client.list_collections().into_stream())
+        .next()
+        .await
+        .unwrap()
+        .unwrap();
     assert_eq!(collections.collections.len(), 1);
     assert_eq!(
         collection.collection.indexing_policy,
@@ -152,12 +149,11 @@ async fn replace_collection() {
         .await
         .unwrap();
 
-    let collections =
-        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .unwrap()
-            .unwrap();
+    let collections = Box::pin(database_client.list_collections().into_stream())
+        .next()
+        .await
+        .unwrap()
+        .unwrap();
     assert_eq!(collections.collections.len(), 1);
     let eps: Vec<&ExcludedPath> = collections.collections[0]
         .indexing_policy

sdk/data_cosmos/tests/create_database_and_collection.rs

@@ -15,7 +15,6 @@ async fn create_database_and_collection() -> Result<(), BoxedError> {
 
     let client = setup::initialize("create_database_and_collection")?;
     let database_name = "test-create-database-and-collection";
-    let context = Context::new();
 
     // create database!
     log::info!("Creating a database with name '{}'...", database_name);
@@ -39,11 +38,10 @@ async fn create_database_and_collection() -> Result<(), BoxedError> {
 
     // list collections!
     log::info!("Listing all collections...");
-    let collections =
-        Box::pin(db_client.list_collections(context.clone(), ListCollectionsOptions::new()))
-            .next()
-            .await
-            .expect("No collection page")?;
+    let collections = Box::pin(db_client.list_collections().into_stream())
+        .next()
+        .await
+        .expect("No collection page")?;
     assert_eq!(collections.count, 1);
     log::info!("Successfully listed collections");
     log::debug!("The list_collection response: {:#?}", collections);