Document operations (#476)

* Move replace_document to pipeline

* Move delete document operation to pipeline

* Address feedback

Author: rylev

sdk/cosmos/examples/document_entries_00.rs

@@ -147,10 +147,13 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
     let replace_document_response = client
         .clone()
         .into_document_client(id.clone(), &id)?
-        .replace_document()
-        .consistency_level(ConsistencyLevel::from(&response))
-        .if_match_condition(IfMatchCondition::Match(&doc.etag)) // use optimistic concurrency check
-        .execute(&doc.document)
+        .replace_document(
+            Context::new(),
+            &doc.document,
+            ReplaceDocumentOptions::new()
+                .consistency_level(ConsistencyLevel::from(&response))
+                .if_match_condition(IfMatchCondition::Match(&doc.etag)), // use optimistic concurrency check
+        )
         .await?;
 
     println!(
@@ -180,9 +183,10 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         client
             .clone()
             .into_document_client(id.clone(), &id)?
-            .delete_document()
-            .consistency_level(&response)
-            .execute()
+            .delete_document(
+                Context::new(),
+                DeleteDocumentOptions::new().consistency_level(&response),
+            )
             .await?;
     }
     println!("Cleaned up");

sdk/cosmos/examples/document_entries_01.rs

@@ -107,9 +107,11 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
 
     let replace_document_response = client
         .into_document_client(doc.id.clone(), &doc.id)?
-        .replace_document()
-        .consistency_level(&query_documents_response)
-        .execute(&doc)
+        .replace_document(
+            Context::new(),
+            &doc,
+            ReplaceDocumentOptions::new().consistency_level(&query_documents_response),
+        )
         .await?;
     println!(
         "replace_document_response == {:#?}",

sdk/cosmos/examples/readme.rs

@@ -150,10 +150,12 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         collection_client
             .clone()
             .into_document_client(document.result.id.clone(), &document.result.a_number)?
-            .delete_document()
-            .consistency_level(session_token.clone())
-            .if_match_condition(&document.document_attributes)
-            .execute()
+            .delete_document(
+                Context::new(),
+                DeleteDocumentOptions::new()
+                    .consistency_level(session_token.clone())
+                    .if_match_condition(&document.document_attributes),
+            )
             .await?;
     }
 

sdk/cosmos/examples/remove_all_documents.rs

@@ -1,3 +1,4 @@
+use azure_core::prelude::*;
 use azure_cosmos::prelude::*;
 use futures::stream::StreamExt;
 use serde_json::Value;
@@ -73,8 +74,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         client
             .clone()
             .into_document_client(id.clone(), &partition_key)?
-            .delete_document()
-            .execute()
+            .delete_document(Context::new(), DeleteDocumentOptions::new())
             .await?;
     }
 

sdk/cosmos/src/clients/document_client.rs

@@ -1,5 +1,5 @@
 use super::{AttachmentClient, CollectionClient, CosmosClient, DatabaseClient};
-use crate::prelude::{GetDocumentOptions, GetDocumentResponse};
+use crate::operations::*;
 use crate::resources::ResourceType;
 use crate::{requests, ReadonlyString};
 use azure_core::{Context, HttpClient, PipelineContext, Request};
@@ -16,10 +16,11 @@ pub struct DocumentClient {
 
 impl DocumentClient {
     /// This function creates a new instance of a DocumentClient. A document is identified by its
-    /// primary key and its partition key. Partition key is eagerly evaluated: the json
-    /// representation is generated as soon as you call the `new` function. This avoids doing the
-    /// serialization over and over, saving time. It also releases the borrow since the serialized
-    /// string is owned by the DocumentClient.
+    /// primary key and its partition key.
+    ///
+    /// Partition key is eagerly evaluated: the json representation is generated as soon as you
+    /// call the `new` function. This avoids doing the serialization over and over, saving time.
+    /// It also releases the borrow since the serialized string is owned by the `DocumentClient`.
     pub(crate) fn new<S: Into<String>, PK: Serialize>(
         collection_client: CollectionClient,
         document_name: S,
@@ -57,11 +58,6 @@ impl DocumentClient {
         &self.partition_key_serialized
     }
 
-    /// replace a document in a collection
-    pub fn replace_document<'a>(&'a self) -> requests::ReplaceDocumentBuilder<'a, '_> {
-        requests::ReplaceDocumentBuilder::new(self)
-    }
-
     /// Get a document
     pub async fn get_document<T>(
         &self,
@@ -72,7 +68,7 @@ impl DocumentClient {
         T: DeserializeOwned,
     {
         let mut request = self.prepare_request_pipeline_with_document_name(http::Method::GET);
-        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Databases.into());
+        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Documents.into());
 
         options.decorate_request(&mut request)?;
 
@@ -85,9 +81,45 @@ impl DocumentClient {
         GetDocumentResponse::try_from(response).await
     }
 
+    /// replace a document in a collection
+    pub async fn replace_document<T: Serialize>(
+        &self,
+        ctx: Context,
+        document: &T,
+        options: ReplaceDocumentOptions<'_>,
+    ) -> crate::Result<ReplaceDocumentResponse> {
+        let mut request = self.prepare_request_pipeline_with_document_name(http::Method::PUT);
+        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Documents.into());
+
+        options.decorate_request(&mut request, document, self.partition_key_serialized())?;
+
+        let response = self
+            .cosmos_client()
+            .pipeline()
+            .send(&mut pipeline_context, &mut request)
+            .await?;
+
+        ReplaceDocumentResponse::try_from(response).await
+    }
+
     /// Delete a document
-    pub fn delete_document(&self) -> requests::DeleteDocumentBuilder<'_> {
-        requests::DeleteDocumentBuilder::new(self)
+    pub async fn delete_document(
+        &self,
+        ctx: Context,
+        options: DeleteDocumentOptions<'_>,
+    ) -> crate::Result<DeleteDocumentResponse> {
+        let mut request = self.prepare_request_pipeline_with_document_name(http::Method::DELETE);
+        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Documents.into());
+
+        options.decorate_request(&mut request, self.partition_key_serialized())?;
+
+        let response = self
+            .cosmos_client()
+            .pipeline()
+            .send(&mut pipeline_context, &mut request)
+            .await?;
+
+        DeleteDocumentResponse::try_from(response).await
     }
 
     /// List all attachments for a document
@@ -103,22 +135,6 @@ impl DocumentClient {
         AttachmentClient::new(self, attachment_name)
     }
 
-    pub(crate) fn prepare_request_with_document_name(
-        &self,
-        method: http::Method,
-    ) -> http::request::Builder {
-        self.cosmos_client().prepare_request(
-            &format!(
-                "dbs/{}/colls/{}/docs/{}",
-                self.database_client().database_name(),
-                self.collection_client().collection_name(),
-                self.document_name()
-            ),
-            method,
-            ResourceType::Documents,
-        )
-    }
-
     fn prepare_request_pipeline_with_document_name(&self, method: http::Method) -> Request {
         self.cosmos_client().prepare_request_pipeline(
             &format!(

sdk/cosmos/src/operations/create_document.rs

@@ -19,7 +19,7 @@ pub struct CreateDocumentOptions<'a> {
     if_match_condition: Option<IfMatchCondition<'a>>,
     if_modified_since: Option<IfModifiedSince<'a>>,
     consistency_level: Option<ConsistencyLevel>,
-    allow_tentative_writes: TenativeWritesAllowance,
+    allow_tentative_writes: TentativeWritesAllowance,
     partition_key: Option<String>,
 }
 
@@ -31,7 +31,7 @@ impl<'a> CreateDocumentOptions<'a> {
             if_match_condition: None,
             if_modified_since: None,
             consistency_level: None,
-            allow_tentative_writes: TenativeWritesAllowance::Deny,
+            allow_tentative_writes: TentativeWritesAllowance::Deny,
             partition_key: None,
         }
     }
@@ -40,7 +40,7 @@ impl<'a> CreateDocumentOptions<'a> {
         consistency_level: ConsistencyLevel => Some(consistency_level),
         if_match_condition: IfMatchCondition<'a> => Some(if_match_condition),
         if_modified_since: &'a DateTime<Utc> => Some(IfModifiedSince::new(if_modified_since)),
-        allow_tentative_writes: TenativeWritesAllowance,
+        allow_tentative_writes: TentativeWritesAllowance,
         is_upsert: bool => if is_upsert { IsUpsert::Yes } else { IsUpsert::No },
         indexing_directive: IndexingDirective,
     }
@@ -62,10 +62,10 @@ impl<'a> CreateDocumentOptions<'a> {
         DOC: Serialize + CosmosEntity<'b>,
     {
         let serialized = serde_json::to_string(document)?;
-        let partition_key = self
-            .partition_key
-            .clone()
-            .unwrap_or_else(|| serialize_partition_key(&document.partition_key()).unwrap());
+        let partition_key = match &self.partition_key {
+            Some(s) => s.clone(),
+            None => serialize_partition_key(&document.partition_key())?,
+        };
 
         add_as_partition_key_header_serialized2(&partition_key, req);
         azure_core::headers::add_optional_header2(&self.if_match_condition, req)?;

sdk/cosmos/src/operations/delete_document.rs

@@ -0,0 +1,75 @@
+use crate::prelude::*;
+
+use azure_core::prelude::*;
+use azure_core::{Request as HttpRequest, Response as HttpResponse};
+use chrono::{DateTime, Utc};
+
+#[derive(Debug, Clone)]
+pub struct DeleteDocumentOptions<'a> {
+    if_match_condition: Option<IfMatchCondition<'a>>,
+    if_modified_since: Option<IfModifiedSince<'a>>,
+    consistency_level: Option<ConsistencyLevel>,
+    allow_tentative_writes: TentativeWritesAllowance,
+}
+
+impl<'a> DeleteDocumentOptions<'a> {
+    pub fn new() -> DeleteDocumentOptions<'a> {
+        Self {
+            if_match_condition: None,
+            if_modified_since: None,
+            consistency_level: None,
+            allow_tentative_writes: TentativeWritesAllowance::Deny,
+        }
+    }
+
+    setters! {
+        consistency_level: ConsistencyLevel => Some(consistency_level),
+        if_match_condition: IfMatchCondition<'a> => Some(if_match_condition),
+        allow_tentative_writes: TentativeWritesAllowance,
+        if_modified_since: &'a DateTime<Utc> => Some(IfModifiedSince::new(if_modified_since)),
+    }
+
+    pub fn decorate_request(
+        &self,
+        request: &mut HttpRequest,
+        serialized_partition_key: &str,
+    ) -> crate::Result<()> {
+        azure_core::headers::add_optional_header2(&self.if_match_condition, request)?;
+        azure_core::headers::add_optional_header2(&self.if_modified_since, request)?;
+        azure_core::headers::add_optional_header2(&self.consistency_level, request)?;
+        azure_core::headers::add_mandatory_header2(&self.allow_tentative_writes, request)?;
+
+        crate::cosmos_entity::add_as_partition_key_header_serialized2(
+            serialized_partition_key,
+            request,
+        );
+
+        Ok(())
+    }
+}
+
+use crate::headers::from_headers::*;
+use azure_core::headers::session_token_from_headers;
+
+#[derive(Debug, Clone)]
+pub struct DeleteDocumentResponse {
+    pub charge: f64,
+    pub activity_id: uuid::Uuid,
+    pub session_token: String,
+}
+
+impl DeleteDocumentResponse {
+    pub async fn try_from(response: HttpResponse) -> crate::Result<Self> {
+        let (_status_code, headers, _pinned_stream) = response.deconstruct();
+
+        let charge = request_charge_from_headers(&headers)?;
+        let activity_id = activity_id_from_headers(&headers)?;
+        let session_token = session_token_from_headers(&headers)?;
+
+        Ok(Self {
+            charge,
+            activity_id,
+            session_token,
+        })
+    }
+}

sdk/cosmos/src/operations/get_document.rs

@@ -33,7 +33,6 @@ impl<'a> GetDocumentOptions<'a> {
     }
 
     pub(crate) fn decorate_request(&self, request: &mut HttpRequest) -> crate::Result<()> {
-        // add trait headers
         azure_core::headers::add_optional_header2(&self.if_match_condition, request)?;
         azure_core::headers::add_optional_header2(&self.if_modified_since, request)?;
         azure_core::headers::add_optional_header2(&self.consistency_level, request)?;

sdk/cosmos/src/operations/mod.rs

@@ -9,6 +9,7 @@ mod create_permission;
 mod create_user;
 mod delete_collection;
 mod delete_database;
+mod delete_document;
 mod delete_permission;
 mod delete_user;
 mod get_collection;
@@ -20,6 +21,7 @@ mod list_collections;
 mod list_databases;
 mod list_users;
 mod replace_collection;
+mod replace_document;
 mod replace_permission;
 mod replace_user;
 
@@ -30,6 +32,7 @@ pub use create_permission::*;
 pub use create_user::*;
 pub use delete_collection::*;
 pub use delete_database::*;
+pub use delete_document::*;
 pub use delete_permission::*;
 pub use delete_user::*;
 pub use get_collection::*;
@@ -41,5 +44,6 @@ pub use list_collections::*;
 pub use list_databases::*;
 pub use list_users::*;
 pub use replace_collection::*;
+pub use replace_document::*;
 pub use replace_permission::*;
 pub use replace_user::*;