Policy Architecture (Azure#204)

* Start policy architecture

* Request and Response wrapper types

* Options

Author: rylev

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/etag.rs renamed to sdk/core/src/models/etag.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/lease.rs renamed to sdk/core/src/models/lease.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/core/src/models/mod.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/core/src/stored_access_policy.rs renamed to sdk/core/src/models/stored_access_policy.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);