Policy Architecture (#204)

* Start policy architecture

* Request and Response wrapper types

* Options

sdk/core/src/headers/mod.rs

@@ -9,6 +9,9 @@ pub const MS_DATE: &str = "x-ms-date";
 
 pub trait AddAsHeader {
     fn add_as_header(&self, builder: Builder) -> Builder;
+    fn add_as_header2(&self, _request: &mut crate::Request) {
+        unimplemented!()
+    }
 }
 
 #[must_use]
@@ -27,6 +30,12 @@ pub fn add_optional_header<T: AddAsHeader>(item: &Option<T>, mut builder: Builde
     builder
 }
 
+pub fn add_optional_header2<T: AddAsHeader>(item: &Option<T>, request: &mut crate::Request) {
+    if let Some(item) = item {
+        item.add_as_header2(request);
+    }
+}
+
 #[must_use]
 pub fn add_mandatory_header<T: AddAsHeader>(item: &T, builder: Builder) -> Builder {
     item.add_as_header(builder)

sdk/core/src/http.rs

@@ -0,0 +1,44 @@
+pub struct Request {
+    inner: http::Request<bytes::Bytes>,
+}
+
+impl Request {
+    /// Get the inner http::Request object, replacing it
+    /// with an empty one.
+    /// Note: this method will soon be replaced
+    pub fn take_inner(&mut self) -> http::Request<bytes::Bytes> {
+        std::mem::replace(&mut self.inner, http::Request::new(bytes::Bytes::new()))
+    }
+
+    pub fn body<T: serde::Serialize>(
+        &mut self,
+        body: T,
+    ) -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
+        let b = self.inner.body_mut();
+        *b = crate::to_json(&body)?;
+        Ok(())
+    }
+}
+
+impl From<http::Request<bytes::Bytes>> for Request {
+    fn from(inner: http::Request<bytes::Bytes>) -> Self {
+        Self { inner }
+    }
+}
+
+pub struct Response {
+    inner: http::Response<bytes::Bytes>,
+}
+
+impl Response {
+    /// TODO: get rid of this
+    pub fn into_inner(self) -> http::Response<bytes::Bytes> {
+        self.inner
+    }
+}
+
+impl From<http::Response<bytes::Bytes>> for Response {
+    fn from(inner: http::Response<bytes::Bytes>) -> Self {
+        Self { inner }
+    }
+}

sdk/core/src/lib.rs

@@ -12,15 +12,15 @@ extern crate serde_derive;
 mod macros;
 
 pub mod errors;
-mod etag;
 pub mod headers;
+mod http;
 mod http_client;
 pub mod incompletevector;
-pub mod lease;
+mod models;
 pub mod parsing;
+mod policy;
 pub mod prelude;
 mod request_options;
-mod stored_access_policy;
 pub mod util;
 
 use chrono::{DateTime, Utc};
@@ -30,9 +30,11 @@ use oauth2::AccessToken;
 use std::fmt::Debug;
 use uuid::Uuid;
 
+pub use self::http::{Request, Response};
 pub use headers::AddAsHeader;
 pub use http_client::{to_json, HttpClient};
-pub use stored_access_policy::{StoredAccessPolicy, StoredAccessPolicyList};
+pub use models::*;
+pub use policy::*;
 
 pub type RequestId = Uuid;
 pub type SessionToken = String;

sdk/core/src/models/mod.rs

@@ -0,0 +1,5 @@
+pub(crate) mod etag;
+pub mod lease;
+mod stored_access_policy;
+
+pub use stored_access_policy::{StoredAccessPolicy, StoredAccessPolicyList};

sdk/core/src/policy.rs

@@ -0,0 +1,147 @@
+use crate::{Request, Response};
+
+use async_trait::async_trait;
+
+use std::error::Error;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+
+pub type PolicyResult<T> = Result<T, Box<dyn Error + Send + Sync>>;
+
+#[derive(Copy, Clone, Debug)]
+pub struct RetryOptions {
+    num_retries: usize,
+}
+
+impl RetryOptions {
+    pub fn new(num_retries: usize) -> Self {
+        Self { num_retries }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct RetryPolicy {
+    options: RetryOptions,
+}
+
+impl RetryPolicy {
+    pub fn new(options: RetryOptions) -> Self {
+        Self { options }
+    }
+}
+
+#[async_trait]
+impl Policy for RetryPolicy {
+    async fn send(
+        &self,
+        ctx: Context,
+        request: &mut Request,
+        next: &[Arc<dyn Policy>],
+    ) -> PolicyResult<Response> {
+        let retries = self.options.num_retries;
+        let mut last_result = next[0].send(ctx.clone(), request, &next[1..]).await;
+        loop {
+            if last_result.is_ok() || retries == 0 {
+                return last_result;
+            }
+
+            last_result = next[0].send(ctx.clone(), request, &next[1..]).await;
+        }
+    }
+}
+
+type BoxedFuture<T> = Box<dyn Future<Output = PolicyResult<T>> + Send>;
+type Transport = dyn Fn(Context, &mut Request) -> Pin<BoxedFuture<Response>> + Send;
+
+pub struct TransportOptions {
+    send: Box<Mutex<Transport>>,
+}
+
+impl TransportOptions {
+    pub fn new<F>(send: F) -> Self
+    where
+        F: Fn(Context, &mut Request) -> Pin<BoxedFuture<Response>> + Send + 'static,
+    {
+        Self {
+            send: Box::new(Mutex::new(send)),
+        }
+    }
+}
+
+impl std::fmt::Debug for TransportOptions {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("TransportOptions")
+    }
+}
+#[derive(Debug)]
+pub struct TransportPolicy {
+    options: TransportOptions,
+}
+
+impl TransportPolicy {
+    pub fn new(options: TransportOptions) -> Self {
+        Self { options }
+    }
+}
+
+#[async_trait::async_trait]
+impl Policy for TransportPolicy {
+    async fn send(
+        &self,
+        ctx: Context,
+        request: &mut Request,
+        next: &[Arc<dyn Policy>],
+    ) -> PolicyResult<Response> {
+        if !next.is_empty() {
+            panic!("Transport policy was not last policy")
+        }
+        let response = {
+            let transport = self.options.send.lock().unwrap();
+            (transport)(ctx, request)
+        };
+        Ok(response.await?)
+    }
+}
+
+#[derive(Clone)]
+pub struct Context {
+    // Temporary hack to make sure that Context is not initializeable
+    // Soon Context will have proper data fields
+    _priv: (),
+}
+
+impl Context {
+    pub fn new() -> Self {
+        Self { _priv: () }
+    }
+}
+
+#[async_trait::async_trait]
+pub trait Policy: Send + Sync + std::fmt::Debug {
+    async fn send(
+        &self,
+        ctx: Context,
+        request: &mut Request,
+        next: &[Arc<dyn Policy>],
+    ) -> PolicyResult<Response>;
+}
+
+#[derive(Debug, Clone)]
+pub struct Pipeline {
+    policies: Vec<Arc<dyn Policy>>,
+}
+
+impl Pipeline {
+    // TODO: how can we ensure that the transport policy is the last policy?
+    // Make this more idiot proof
+    pub fn new(policies: Vec<Arc<dyn Policy>>) -> Self {
+        Self { policies }
+    }
+
+    pub async fn send(&self, ctx: Context, mut request: Request) -> PolicyResult<Response> {
+        self.policies[0]
+            .send(ctx, &mut request, &self.policies[1..])
+            .await
+    }
+}

sdk/cosmos/examples/create_delete_database.rs

@@ -31,7 +31,11 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
     // authorization token at later time if you need, for example, to escalate the privileges for a
     // single operation.
     let http_client: Arc<Box<dyn HttpClient>> = Arc::new(Box::new(reqwest::Client::new()));
-    let client = CosmosClient::new(http_client, account, authorization_token);
+    let client = CosmosClient::new(
+        http_client.clone(),
+        account.clone(),
+        authorization_token.clone(),
+    );
 
     // The Cosmos' client exposes a lot of methods. This one lists the databases in the specified
     // account. Database do not implement Display but deref to &str so you can pass it to methods
@@ -40,7 +44,18 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
     let list_databases_response = client.list_databases().execute().await?;
     println!("list_databases_response = {:#?}", list_databases_response);
 
-    let db = client.create_database().execute(&database_name).await?;
+    let cosmos_client = CosmosClient::with_pipeline(
+        account,
+        authorization_token,
+        CosmosOptions::with_client(http_client),
+    );
+    let db = cosmos_client
+        .create_database(
+            azure_core::Context::new(),
+            &database_name,
+            azure_cosmos::operations::create_database::Options::new(),
+        )
+        .await?;
     println!("created database = {:#?}", db);
 
     // create collection!

sdk/cosmos/examples/document_00.rs

@@ -51,7 +51,11 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
     // Next we will create a Cosmos client. You need an authorization_token but you can later
     // change it if needed.
     let http_client: Arc<Box<dyn HttpClient>> = Arc::new(Box::new(reqwest::Client::new()));
-    let client = CosmosClient::new(http_client, account, authorization_token);
+    let client = CosmosClient::new(
+        http_client.clone(),
+        account.clone(),
+        authorization_token.clone(),
+    );
 
     // list_databases will give us the databases available in our account. If there is
     // an error (for example, the given key is not valid) you will receive a
@@ -65,10 +69,24 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         .into_iter()
         .find(|db| db.id == DATABASE);
 
+    let database_client = CosmosClient::with_pipeline(
+        account,
+        authorization_token,
+        CosmosOptions::with_client(http_client),
+    );
     // If the requested database is not found we create it.
     let database = match db {
         Some(db) => db,
-        None => client.create_database().execute(DATABASE).await?.database,
+        None => {
+            database_client
+                .create_database(
+                    azure_core::Context::new(),
+                    DATABASE,
+                    azure_cosmos::operations::create_database::Options::new(),
+                )
+                .await?
+                .database
+        }
     };
     println!("database == {:?}", database);
 

sdk/cosmos/src/clients/cosmos_client.rs

@@ -11,6 +11,7 @@ use ring::hmac;
 use url::form_urlencoded;
 
 use std::borrow::Cow;
+use std::convert::TryInto;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -18,13 +19,40 @@ const AZURE_VERSION: &str = "2018-12-31";
 const VERSION: &str = "1.0";
 const TIME_FORMAT: &str = "%a, %d %h %Y %T GMT";
 
+pub type Error = Box<dyn std::error::Error + Send + Sync>;
+use azure_core::*;
+
+#[derive(Debug, Clone)]
+enum TransportStack {
+    Client(Arc<Box<dyn HttpClient>>),
+    Pipeline(Pipeline),
+}
 /// A plain Cosmos client.
 #[derive(Debug, Clone)]
 pub struct CosmosClient {
-    http_client: Arc<Box<dyn HttpClient>>,
+    transport: TransportStack,
     auth_token: AuthorizationToken,
     cloud_location: CloudLocation,
 }
+/// TODO
+pub struct CosmosOptions {
+    retry: RetryOptions,
+    transport: TransportOptions,
+}
+
+impl CosmosOptions {
+    /// TODO
+    pub fn with_client(client: Arc<Box<dyn HttpClient>>) -> Self {
+        Self {
+            retry: RetryOptions::new(3),
+            transport: TransportOptions::new(move |_ctx, req| {
+                let client = client.clone();
+                let req = req.take_inner();
+                Box::pin(async move { Ok(client.execute_request(req).await?.into()) })
+            }),
+        }
+    }
+}
 
 impl CosmosClient {
     /// Create a new `CosmosClient` which connects to the account's instance in the public Azure cloud.
@@ -35,7 +63,7 @@ impl CosmosClient {
     ) -> Self {
         let cloud_location = CloudLocation::Public(account);
         Self {
-            http_client,
+            transport: TransportStack::Client(http_client),
             auth_token,
             cloud_location,
         }
@@ -49,7 +77,7 @@ impl CosmosClient {
     ) -> Self {
         let cloud_location = CloudLocation::China(account);
         Self {
-            http_client,
+            transport: TransportStack::Client(http_client),
             auth_token,
             cloud_location,
         }
@@ -64,7 +92,7 @@ impl CosmosClient {
     ) -> Self {
         let cloud_location = CloudLocation::Custom { account, uri };
         Self {
-            http_client,
+            transport: TransportStack::Client(http_client),
             auth_token,
             cloud_location,
         }
@@ -83,20 +111,58 @@ impl CosmosClient {
             uri,
         };
         Self {
-            http_client,
+            transport: TransportStack::Client(http_client),
             auth_token,
             cloud_location,
         }
     }
 
+    /// TODO
+    pub fn with_pipeline(
+        account: String, // TODO: this will eventually be a URL
+        auth_token: AuthorizationToken,
+        options: CosmosOptions,
+    ) -> Self {
+        use azure_core::*;
+        let mut policies = Vec::new();
+        let retry_policy = RetryPolicy::new(options.retry);
+        policies.push(Arc::new(retry_policy) as Arc<dyn Policy>);
+        let transport_policy = TransportPolicy::new(options.transport);
+        policies.push(Arc::new(transport_policy) as Arc<dyn Policy>);
+        let pipeline = Pipeline::new(policies);
+        Self {
+            transport: TransportStack::Pipeline(pipeline),
+            auth_token,
+            cloud_location: CloudLocation::Public(account),
+        }
+    }
+
     /// Set the auth token used
     pub fn auth_token(&mut self, auth_token: AuthorizationToken) {
         self.auth_token = auth_token;
     }
 
     /// Create a database
-    pub fn create_database(&self) -> requests::CreateDatabaseBuilder<'_> {
-        requests::CreateDatabaseBuilder::new(self)
+    pub async fn create_database<S: AsRef<str>>(
+        &self,
+        ctx: Context,
+        database_name: S,
+        options: crate::operations::create_database::Options,
+    ) -> Result<crate::operations::create_database::Response, Error> {
+        let mut request = self.prepare_request2("dbs", http::Method::POST, ResourceType::Databases);
+        options.decorate_request(&mut request, database_name.as_ref())?;
+        self.pipeline()
+            .unwrap()
+            .send(ctx, request)
+            .await?
+            .try_into()
+    }
+
+    fn pipeline(&self) -> Option<&Pipeline> {
+        match &self.transport {
+            TransportStack::Pipeline(p) => Some(p),
+            TransportStack::Client(_) => None,
+        }
     }
 
     /// List all databases
@@ -130,8 +196,22 @@ impl CosmosClient {
         self.prepare_request_with_signature(uri_path, http_method, &time, &auth)
     }
 
+    // Eventually this method will replace `prepare_request` fully
+    pub(crate) fn prepare_request2(
+        &self,
+        uri_path: &str,
+        http_method: http::Method,
+        resource_type: ResourceType,
+    ) -> Request {
+        let builder = self.prepare_request(uri_path, http_method, resource_type);
+        builder.body(bytes::Bytes::new()).unwrap().into()
+    }
+
     pub(crate) fn http_client(&self) -> &dyn HttpClient {
-        self.http_client.as_ref().as_ref()
+        match &self.transport {
+            TransportStack::Client(c) => c.as_ref().as_ref(),
+            TransportStack::Pipeline(_) => panic!("No client set"),
+        }
     }
 
     fn prepare_request_with_signature(

sdk/cosmos/src/clients/mod.rs

@@ -33,7 +33,7 @@ mod user_defined_function_client;
 
 pub use attachment_client::AttachmentClient;
 pub use collection_client::CollectionClient;
-pub use cosmos_client::CosmosClient;
+pub use cosmos_client::{CosmosClient, CosmosOptions};
 pub use database_client::DatabaseClient;
 pub use document_client::DocumentClient;
 pub use permission_client::PermissionClient;

sdk/cosmos/src/lib.rs

@@ -106,6 +106,7 @@ extern crate failure;
 extern crate azure_core;
 
 pub mod clients;
+pub mod operations;
 pub mod prelude;
 pub mod requests;
 pub mod resources;

sdk/cosmos/src/operations/create_database.rs

@@ -1,12 +1,48 @@
 use crate::headers::from_headers::*;
+use crate::prelude::*;
 use crate::resources::Database;
 use crate::{CosmosError, ResourceQuota};
 use azure_core::headers::{etag_from_headers, session_token_from_headers};
+use azure_core::{Request as HttpRequest, Response as HttpResponse};
 use chrono::{DateTime, Utc};
-use http::response::Response;
+
+#[derive(Debug, Clone)]
+pub struct Options {
+    consistency_level: Option<ConsistencyLevel>,
+}
+
+impl Options {
+    pub fn new() -> Self {
+        Self {
+            consistency_level: None,
+        }
+    }
+
+    setters! {
+        consistency_level: ConsistencyLevel => Some(consistency_level),
+    }
+}
+
+impl Options {
+    pub(crate) fn decorate_request(
+        &self,
+        request: &mut HttpRequest,
+        database_name: &str,
+    ) -> Result<(), CosmosError> {
+        #[derive(Serialize)]
+        struct CreateDatabaseRequest<'a> {
+            pub id: &'a str,
+        }
+        let req = CreateDatabaseRequest { id: database_name };
+
+        azure_core::headers::add_optional_header2(&self.consistency_level, request);
+        request.body(req)?;
+        Ok(())
+    }
+}
 
 #[derive(Debug, Clone, PartialEq, PartialOrd)]
-pub struct CreateDatabaseResponse {
+pub struct Response {
     pub database: Database,
     pub charge: f64,
     pub etag: String,
@@ -23,14 +59,15 @@ pub struct CreateDatabaseResponse {
     pub gateway_version: String,
 }
 
-impl std::convert::TryFrom<Response<bytes::Bytes>> for CreateDatabaseResponse {
+impl std::convert::TryFrom<HttpResponse> for Response {
     type Error = CosmosError;
 
-    fn try_from(response: Response<bytes::Bytes>) -> Result<Self, Self::Error> {
+    fn try_from(response: HttpResponse) -> Result<Self, Self::Error> {
+        let response = response.into_inner();
         let headers = response.headers();
         let body = response.body();
 
-        Ok(CreateDatabaseResponse {
+        Ok(Self {
             database: serde_json::from_slice(&body)?,
             charge: request_charge_from_headers(headers)?,
             etag: etag_from_headers(headers)?,

sdk/cosmos/src/operations/mod.rs

@@ -0,0 +1,5 @@
+//! TODO: Documentation
+
+#![allow(missing_docs)]
+
+pub mod create_database;

sdk/cosmos/src/prelude.rs

@@ -25,3 +25,5 @@ pub use crate::resources::document::*;
 pub use crate::resources::*;
 
 pub use permission::AuthorizationToken;
+
+pub use crate::operations::*;

sdk/cosmos/src/requests/create_database_builder.rs

@@ -1,70 +0,0 @@
-use crate::prelude::*;
-use crate::resources::ResourceType;
-use crate::responses::CreateDatabaseResponse;
-use azure_core::prelude::*;
-use http::StatusCode;
-use std::convert::TryInto;
-
-#[derive(Debug, Clone)]
-pub struct CreateDatabaseBuilder<'a> {
-    cosmos_client: &'a CosmosClient,
-    user_agent: Option<UserAgent<'a>>,
-    activity_id: Option<ActivityId<'a>>,
-    consistency_level: Option<ConsistencyLevel>,
-}
-
-impl<'a> CreateDatabaseBuilder<'a> {
-    pub(crate) fn new(cosmos_client: &'a CosmosClient) -> Self {
-        Self {
-            cosmos_client,
-            user_agent: None,
-            activity_id: None,
-            consistency_level: None,
-        }
-    }
-}
-
-impl<'a> CreateDatabaseBuilder<'a> {
-    setters! {
-        user_agent: &'a str => Some(UserAgent::new(user_agent)),
-        activity_id: &'a str => Some(ActivityId::new(activity_id)),
-        consistency_level: ConsistencyLevel => Some(consistency_level),
-    }
-}
-
-impl<'a> CreateDatabaseBuilder<'a> {
-    pub async fn execute<D: AsRef<str>>(
-        &self,
-        database_name: D,
-    ) -> Result<CreateDatabaseResponse, CosmosError> {
-        trace!("CreateDatabaseBuilder::execute called");
-
-        #[derive(Serialize, Debug)]
-        struct CreateDatabaseRequest<'a> {
-            pub id: &'a str,
-        }
-
-        let req = azure_core::to_json(&CreateDatabaseRequest {
-            id: database_name.as_ref(),
-        })?;
-
-        let request =
-            self.cosmos_client
-                .prepare_request("dbs", http::Method::POST, ResourceType::Databases);
-
-        let request = azure_core::headers::add_optional_header(&self.user_agent, request);
-        let request = azure_core::headers::add_optional_header(&self.activity_id, request);
-        let request = azure_core::headers::add_optional_header(&self.consistency_level, request);
-
-        let request = request.body(req)?; // todo: set content-length here and elsewhere without builders
-
-        debug!("create database request prepared == {:?}", request);
-
-        Ok(self
-            .cosmos_client
-            .http_client()
-            .execute_request_check_status(request, StatusCode::CREATED)
-            .await?
-            .try_into()?)
-    }
-}

sdk/cosmos/src/requests/mod.rs

@@ -7,7 +7,6 @@
 #![allow(missing_docs)]
 
 mod create_collection_builder;
-mod create_database_builder;
 mod create_document_builder;
 mod create_or_replace_trigger_builder;
 mod create_or_replace_user_defined_function_builder;
@@ -52,7 +51,6 @@ mod replace_stored_procedure_builder;
 mod replace_user_builder;
 
 pub use create_collection_builder::CreateCollectionBuilder;
-pub use create_database_builder::CreateDatabaseBuilder;
 pub use create_document_builder::CreateDocumentBuilder;
 pub use create_or_replace_trigger_builder::CreateOrReplaceTriggerBuilder;
 pub use create_or_replace_user_defined_function_builder::CreateOrReplaceUserDefinedFunctionBuilder;

sdk/cosmos/src/responses/mod.rs

@@ -3,7 +3,6 @@
 #![allow(missing_docs)]
 
 mod create_collection_response;
-mod create_database_response;
 mod create_document_response;
 mod create_permission_response;
 mod create_reference_attachment_response;
@@ -44,7 +43,6 @@ mod replace_reference_attachment_response;
 mod replace_stored_procedure_response;
 
 pub use create_collection_response::CreateCollectionResponse;
-pub use create_database_response::CreateDatabaseResponse;
 pub use create_document_response::CreateDocumentResponse;
 pub use create_permission_response::CreatePermissionResponse;
 pub use create_reference_attachment_response::CreateReferenceAttachmentResponse;

sdk/cosmos/tests/attachment_00.rs

@@ -36,8 +36,11 @@ async fn attachment() -> Result<(), CosmosError> {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/cosmos_collection.rs

@@ -12,8 +12,11 @@ async fn create_and_delete_collection() {
     let client = setup::initialize().unwrap();
 
     client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 
@@ -62,8 +65,11 @@ async fn replace_collection() {
     const COLLECTION_NAME: &str = "test-collection";
 
     client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/cosmos_database.rs

@@ -2,6 +2,8 @@
 
 mod setup;
 
+use azure_cosmos::prelude::*;
+
 #[tokio::test]
 async fn create_and_delete_database() {
     const DATABASE_NAME: &str = "cosmos-test-db-create-and-delete-database";
@@ -14,8 +16,11 @@ async fn create_and_delete_database() {
 
     // create a new database and check if the number of DBs increased
     let database = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
     let databases = client.list_databases().execute().await.unwrap();

sdk/cosmos/tests/cosmos_document.rs

@@ -31,8 +31,11 @@ async fn create_and_delete_document() {
     let client = setup::initialize().unwrap();
 
     client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 
@@ -118,8 +121,11 @@ async fn query_documents() {
     let client = setup::initialize().unwrap();
 
     client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
     let database_client = client.into_database_client(DATABASE_NAME);
@@ -190,8 +196,11 @@ async fn replace_document() {
     let client = setup::initialize().unwrap();
 
     client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
     let database_client = client.into_database_client(DATABASE_NAME);

sdk/cosmos/tests/permission.rs

@@ -16,8 +16,11 @@ async fn permissions() {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/permission_token_usage.rs

@@ -32,8 +32,11 @@ async fn permission_token_usage() {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/trigger.rs

@@ -44,8 +44,11 @@ async fn trigger() -> Result<(), CosmosError> {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/user.rs

@@ -1,5 +1,7 @@
 #![cfg(all(test, feature = "test_e2e"))]
 
+use azure_cosmos::prelude::*;
+
 mod setup;
 
 #[tokio::test]
@@ -12,8 +14,11 @@ async fn users() {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();
 

sdk/cosmos/tests/user_defined_function00.rs

@@ -28,8 +28,11 @@ async fn user_defined_function00() -> Result<(), CosmosError> {
 
     // create a temp database
     let _create_database_response = client
-        .create_database()
-        .execute(DATABASE_NAME)
+        .create_database(
+            azure_core::Context::new(),
+            DATABASE_NAME,
+            create_database::Options::new(),
+        )
         .await
         .unwrap();