Finish Cosmos database client (#455)

* Do delete database operation

* Move list_collections to pipeline

* Add test for listing collection

* Format recorded request/responses

* Delete database in the test

Author: rylev

sdk/cosmos/examples/database_01.rs

@@ -1,5 +1,6 @@
 use azure_core::Context;
 use azure_cosmos::prelude::*;
+use futures::stream::StreamExt;
 use std::error::Error;
 
 #[tokio::main]
@@ -17,7 +18,11 @@ 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 = database_client.list_collections().execute().await?;
+    let collections =
+        Box::pin(database_client.list_collections(Context::new(), ListCollectionsOptions::new()))
+            .next()
+            .await
+            .unwrap()?;
     println!("collections == {:#?}", collections);
 
     let collection_client = database_client.into_collection_client("cnt");

sdk/cosmos/src/clients/cosmos_client.rs

@@ -307,7 +307,16 @@ impl CosmosClient {
                 time,
             )
         };
-        self.prepare_request_with_signature(uri_path, http_method, time, &auth)
+        trace!("prepare_request::auth == {:?}", auth);
+        let uri = format!("{}/{}", self.cloud_location.url(), uri_path);
+        debug!("building request. uri: {}", uri);
+
+        RequestBuilder::new()
+            .method(http_method)
+            .uri(uri)
+            .header(HEADER_DATE, time.to_string())
+            .header(HEADER_VERSION, HeaderValue::from_static(AZURE_VERSION))
+            .header(header::AUTHORIZATION, auth)
     }
 
     /// Prepares' an `azure_core::Request`. This function will
@@ -332,28 +341,6 @@ impl CosmosClient {
             .into()
     }
 
-    fn prepare_request_with_signature(
-        &self,
-        uri_path: &str,
-        http_method: http::Method,
-        time_nonce: TimeNonce,
-        signature: &str,
-    ) -> RequestBuilder {
-        trace!("prepare_request::auth == {:?}", signature);
-        let uri = format!("{}/{}", self.cloud_location.url(), uri_path);
-        debug!(
-            "cosmos::client::prepare_request_with_resource_signature::uri == {:?}",
-            uri
-        );
-
-        RequestBuilder::new()
-            .method(http_method)
-            .uri(uri)
-            .header(HEADER_DATE, time_nonce.to_string())
-            .header(HEADER_VERSION, HeaderValue::from_static(AZURE_VERSION))
-            .header(header::AUTHORIZATION, signature)
-    }
-
     pub(crate) fn pipeline(&self) -> &Pipeline<CosmosContext> {
         &self.pipeline
     }

sdk/cosmos/src/clients/database_client.rs

@@ -2,13 +2,33 @@ use super::*;
 use crate::authorization_policy::CosmosContext;
 use crate::operations::*;
 use crate::resources::ResourceType;
-use crate::{requests, ReadonlyString};
+use crate::ReadonlyString;
 use azure_core::pipeline::Pipeline;
 use azure_core::prelude::Continuation;
-use azure_core::{AddAsHeader, Context, HttpClient, PipelineContext};
+use azure_core::{AddAsHeader, Context, PipelineContext};
 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,
+}
+
 /// A client for Cosmos database resources.
 #[derive(Debug, Clone)]
 pub struct DatabaseClient {
@@ -59,14 +79,90 @@ impl DatabaseClient {
         Ok(GetDatabaseResponse::try_from(response).await?)
     }
 
-    /// List collections in the database
-    pub fn list_collections(&self) -> requests::ListCollectionsBuilder<'_> {
-        requests::ListCollectionsBuilder::new(self)
+    /// Delete the database
+    pub async fn delete_database(
+        &self,
+        ctx: Context,
+        options: DeleteDatabaseOptions,
+    ) -> crate::Result<DeleteDatabaseResponse> {
+        let mut request = self
+            .cosmos_client()
+            .prepare_request_pipeline(&format!("dbs/{}", self.database_name()), http::Method::GET);
+        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Databases.into());
+
+        options.decorate_request(&mut request)?;
+        let response = self
+            .pipeline()
+            .send(&mut pipeline_context, &mut request)
+            .await?
+            .validate(http::StatusCode::OK)
+            .await?;
+
+        Ok(DeleteDatabaseResponse::try_from(response).await?)
     }
 
-    /// Delete the database
-    pub fn delete_database(&self) -> requests::DeleteDatabaseBuilder<'_> {
-        requests::DeleteDatabaseBuilder::new(self)
+    /// 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,
+                        );
+                        let mut pipeline_context =
+                            PipelineContext::new(ctx.clone(), ResourceType::Collections.into());
+
+                        r#try!(options.decorate_request(&mut request));
+                        let response = r#try!(
+                            this.pipeline()
+                                .send(&mut pipeline_context, &mut request)
+                                .await
+                        );
+                        let response = r#try!(response.validate(http::StatusCode::OK).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,
+                        );
+                        let mut pipeline_context =
+                            PipelineContext::new(ctx.clone(), ResourceType::Collections.into());
+
+                        r#try!(options.decorate_request(&mut request));
+                        r#try!(continuation.add_as_header2(&mut request));
+                        let response = r#try!(
+                            this.pipeline()
+                                .send(&mut pipeline_context, &mut request)
+                                .await
+                        );
+                        let response = r#try!(response.validate(http::StatusCode::OK).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))
+            }
+        })
     }
 
     /// Create a collection
@@ -99,24 +195,6 @@ impl DatabaseClient {
         ctx: Context,
         options: ListUsersOptions,
     ) -> impl Stream<Item = crate::Result<ListUsersResponse>> + '_ {
-        macro_rules! r#try {
-            ($expr:expr $(,)?) => {
-                match $expr {
-                    Result::Ok(val) => val,
-                    Result::Err(err) => {
-                        return Some((Err(err.into()), State::Done));
-                    }
-                }
-            };
-        }
-
-        #[derive(Debug, Clone, PartialEq)]
-        enum State {
-            Init,
-            Continuation(String),
-            Done,
-        }
-
         unfold(State::Init, move |state: State| {
             let this = self.clone();
             let ctx = ctx.clone();
@@ -188,21 +266,6 @@ impl DatabaseClient {
         UserClient::new(self, user_name)
     }
 
-    pub(crate) fn prepare_request_with_database_name(
-        &self,
-        method: http::Method,
-    ) -> http::request::Builder {
-        self.cosmos_client().prepare_request(
-            &format!("dbs/{}", self.database_name()),
-            method,
-            ResourceType::Databases,
-        )
-    }
-
-    pub(crate) fn http_client(&self) -> &dyn HttpClient {
-        self.cosmos_client().http_client()
-    }
-
     fn pipeline(&self) -> &Pipeline<CosmosContext> {
         self.cosmos_client.pipeline()
     }

sdk/cosmos/src/responses/delete_database_response.rs renamed to sdk/cosmos/src/operations/delete_database.rs

@@ -1,7 +1,31 @@
 use crate::headers::from_headers::*;
+use crate::prelude::*;
 use crate::ResourceQuota;
 use azure_core::headers::session_token_from_headers;
-use http::response::Response;
+use azure_core::Request as HttpRequest;
+use azure_core::Response as HttpResponse;
+
+#[derive(Debug, Clone)]
+pub struct DeleteDatabaseOptions {
+    consistency_level: Option<ConsistencyLevel>,
+}
+
+impl DeleteDatabaseOptions {
+    pub fn new() -> Self {
+        Self {
+            consistency_level: None,
+        }
+    }
+
+    setters! {
+        consistency_level: ConsistencyLevel => Some(consistency_level),
+    }
+
+    pub fn decorate_request(&self, request: &mut HttpRequest) -> crate::Result<()> {
+        azure_core::headers::add_optional_header2(&self.consistency_level, request)?;
+        Ok(())
+    }
+}
 
 #[derive(Debug, Clone)]
 pub struct DeleteDatabaseResponse {
@@ -12,10 +36,8 @@ pub struct DeleteDatabaseResponse {
     pub resource_usage: Vec<ResourceQuota>,
 }
 
-impl std::convert::TryFrom<Response<bytes::Bytes>> for DeleteDatabaseResponse {
-    type Error = crate::Error;
-
-    fn try_from(response: Response<bytes::Bytes>) -> Result<Self, Self::Error> {
+impl DeleteDatabaseResponse {
+    pub async fn try_from(response: HttpResponse) -> crate::Result<Self> {
         let headers = response.headers();
 
         let charge = request_charge_from_headers(headers)?;

sdk/cosmos/src/operations/list_collections.rs

@@ -0,0 +1,92 @@
+use crate::headers::from_headers::*;
+use crate::prelude::*;
+use crate::resources::Collection;
+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};
+
+#[derive(Debug, Clone)]
+pub struct ListCollectionsOptions {
+    consistency_level: Option<ConsistencyLevel>,
+    max_item_count: MaxItemCount,
+}
+
+impl ListCollectionsOptions {
+    pub fn new() -> Self {
+        Self {
+            max_item_count: MaxItemCount::new(-1),
+            consistency_level: None,
+        }
+    }
+
+    setters! {
+        consistency_level: ConsistencyLevel => Some(consistency_level),
+        max_item_count: i32 => MaxItemCount::new(max_item_count),
+    }
+
+    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)?;
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct ListCollectionsResponse {
+    pub rid: String,
+    pub collections: Vec<Collection>,
+    pub count: u32,
+    pub last_state_change: DateTime<Utc>,
+    pub resource_quota: Vec<ResourceQuota>,
+    pub resource_usage: Vec<ResourceQuota>,
+    pub schema_version: String,
+    pub alt_content_path: String,
+    pub content_path: String,
+    pub charge: f64,
+    pub service_version: String,
+    pub activity_id: uuid::Uuid,
+    pub session_token: String,
+    pub gateway_version: String,
+    pub continuation_token: Option<String>,
+}
+
+impl ListCollectionsResponse {
+    pub async fn try_from(response: HttpResponse) -> crate::Result<Self> {
+        let (_status_code, headers, pinned_stream) = response.deconstruct();
+        let body = collect_pinned_stream(pinned_stream).await?;
+
+        #[derive(Deserialize, Debug)]
+        pub struct Response {
+            _rid: String,
+            #[serde(rename = "DocumentCollections")]
+            pub collections: Vec<Collection>,
+            #[serde(rename = "_count")]
+            pub count: u32,
+        }
+
+        let response: Response = serde_json::from_slice(&*body)?;
+
+        Ok(Self {
+            rid: response._rid,
+            collections: response.collections,
+            count: response.count,
+            last_state_change: last_state_change_from_headers(&headers)?,
+            resource_quota: resource_quota_from_headers(&headers)?,
+            resource_usage: resource_usage_from_headers(&headers)?,
+            schema_version: schema_version_from_headers(&headers)?.to_owned(),
+            alt_content_path: alt_content_path_from_headers(&headers)?.to_owned(),
+            content_path: content_path_from_headers(&headers)?.to_owned(),
+            charge: request_charge_from_headers(&headers)?,
+            service_version: service_version_from_headers(&headers)?.to_owned(),
+            activity_id: activity_id_from_headers(&headers)?,
+            session_token: session_token_from_headers(&headers)?,
+            gateway_version: gateway_version_from_headers(&headers)?.to_owned(),
+            continuation_token: continuation_token_from_headers_optional(&headers)?,
+        })
+    }
+}