cosmos_utils/

lib.rs

use azure_core::{error::Error as AzureError, prelude::IfMatchCondition, HttpClient, StatusCode};
use azure_cosmos::prelude::*;
use azure_storage::prelude::*;
use azure_storage_blobs::prelude::*;
use futures::StreamExt;
use lazy_static::lazy_static;
use once_cell::sync::OnceCell;
use serde::{de::DeserializeOwned, Serialize};
use std::sync::Arc;
use tokio::time::{sleep, Duration};
use uuid::Uuid;
use warp::{filters::multipart::FormData, Buf as OtherBuf};

/// Re-export `CosmosEntity` which must be implemented by every model.
pub use azure_cosmos::CosmosEntity;

mod cosmos_saga;
pub use cosmos_saga::{
    CosmosSaga, ErasedCosmosEntity, ErasedCosmosEntityClone, UpcastErasedSerdeSerialize,
};

/// This struct keeps the static state of cosmos DB, it must be initialized before any cosmos
/// functions are called.
#[derive(Debug)]
pub struct CosmosState {
    pub cosmos_account: String,
    pub cosmos_database: String,
    pub cosmos_master_key: String,
    pub storage_account: Option<String>,
    pub storage_master_key: Option<String>,
    pub image_storage_container: Option<String>,
}

static COSMOS_STATE: OnceCell<CosmosState> = OnceCell::new();

lazy_static! {
    static ref HTTP_CLIENT: Arc<Box<dyn HttpClient>> = Arc::new(Box::new(reqwest::Client::new()));
}

pub fn set_state(state: CosmosState) {
    COSMOS_STATE.set(state).unwrap();
}

type CosmosError = CosmosErrorStruct;

#[derive(Debug)]
pub struct CosmosErrorStruct {
    pub err: String,
    pub kind: CosmosErrorKind,
}

impl std::fmt::Display for CosmosErrorStruct {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        fmt.write_fmt(format_args!("kind: {}, err: {}", self.kind, self.err))?;
        Ok(())
    }
}

#[derive(Debug)]
pub enum CosmosErrorKind {
    PreconditionFailed,
    NotFound,
    BadRequest,
    InternalError,
    Conflict,
    BlobError,
    ModificationError(warp::Rejection),
}

impl std::fmt::Display for CosmosErrorKind {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            CosmosErrorKind::PreconditionFailed => {
                fmt.write_str("PreconditionFailed")?;
            }
            CosmosErrorKind::NotFound => {
                fmt.write_str("NotFound")?;
            }
            CosmosErrorKind::BadRequest => {
                fmt.write_str("BadRequest")?;
            }
            CosmosErrorKind::InternalError => {
                fmt.write_str("InternalError")?;
            }
            CosmosErrorKind::Conflict => {
                fmt.write_str("Conflict")?;
            }
            CosmosErrorKind::BlobError => {
                fmt.write_str("BlobError")?;
            }
            CosmosErrorKind::ModificationError(rej) => {
                fmt.write_fmt(format_args!("ModificationError({:?})", rej))?;
            }
        }
        Ok(())
    }
}

/// Marker trait for all types that implement `CosmosEntity` and `Serialize` and `Clone`.
/// Necessary in order to fullfille the requirements for Cosmos DB objects.
/// This trait should be implemented automatically for any type that needs it.
pub trait CosmosObject: CosmosEntity + Serialize + Clone + Send + Sync + 'static {}
impl<T> CosmosObject for T where T: CosmosEntity + Serialize + Clone + Send + Sync + 'static {}

impl warp::reject::Reject for CosmosErrorStruct {}

/// Utility function that returns a closure that converts whatever error into a Reject error.
/// Usage is: function_call_which_returns_non_reject_result().map_err(into_cosmos_error("custom error message"))?;
fn into_cosmos_error<M: ToString>(msg: M) -> impl FnOnce(AzureError) -> CosmosError {
    move |e: AzureError| new_cosmos_error(msg.to_string(), e)
}

/// Creates a new cosmos error from a given error
fn new_cosmos_error<M: ToString>(msg: M, err: AzureError) -> CosmosError {
    let kind = match err.as_http_error() {
        Some(err) => match err.status() {
            StatusCode::PreconditionFailed => CosmosErrorKind::PreconditionFailed,
            StatusCode::NotFound => CosmosErrorKind::NotFound,
            StatusCode::BadRequest => CosmosErrorKind::BadRequest,
            StatusCode::Conflict => CosmosErrorKind::Conflict,
            _ => CosmosErrorKind::InternalError,
        },
        None => CosmosErrorKind::InternalError,
    };
    // We get more informative error messages if we use the debug printing rather than display
    // printing.
    let msg = format!("{} : {:?}", msg.to_string(), err);
    CosmosErrorStruct { kind, err: msg }
}

fn new_cosmos_error_internal<M: ToString>(msg: M) -> CosmosError {
    CosmosErrorStruct {
        kind: CosmosErrorKind::InternalError,
        err: msg.to_string(),
    }
}

/// Creates a new cosmos error from a given error and the error kind
fn new_cosmos_error_kind<E: ToString>(err: E, kind: CosmosErrorKind) -> CosmosError {
    let err = err.to_string();
    CosmosErrorStruct { kind, err }
}

/// Intermediate type for making a seemless conversation between a [&str; 1] and a `PartitionKey`
/// This is needed because the previous signature of a lot of functions required a slice of
/// PartitionKeys however the interface for CosmosDB was changed to require a single PartitionKey
/// breaking all of the previous signatures.
/// Without this intermediate type the codebase would require a lot of rewriting the signatures.
pub struct PartitionKeyIntermediate(String);

impl From<[&str; 1]> for PartitionKeyIntermediate {
    fn from(from: [&str; 1]) -> Self {
        Self(from.first().unwrap().to_string())
    }
}

impl From<[&String; 1]> for PartitionKeyIntermediate {
    fn from(from: [&String; 1]) -> Self {
        Self(from.first().unwrap().to_string())
    }
}

impl From<[&&String; 1]> for PartitionKeyIntermediate {
    fn from(from: [&&String; 1]) -> Self {
        Self(from.first().unwrap().to_string())
    }
}

impl From<[&(); 1]> for PartitionKeyIntermediate {
    fn from(_: [&(); 1]) -> Self {
        Self(String::default())
    }
}

impl From<&[(); 1]> for PartitionKeyIntermediate {
    fn from(_: &[(); 1]) -> Self {
        Self(String::default())
    }
}

impl From<[(); 1]> for PartitionKeyIntermediate {
    fn from(_: [(); 1]) -> Self {
        Self(String::default())
    }
}

impl From<String> for PartitionKeyIntermediate {
    fn from(from: String) -> Self {
        Self(from)
    }
}

impl From<PartitionKeyIntermediate> for String {
    fn from(from: PartitionKeyIntermediate) -> String {
        from.0
    }
}

async fn insert_internal<D: CosmosObject, P: Into<PartitionKeyIntermediate>, C: ToString>(
    collection_name: C,
    pk: P,
    document: &D,
    etag: Option<&String>,
    upsert: bool,
) -> Result<String, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);

    let pk: String = pk.into().into();

    let c = collection_client
        .create_document(document.clone())
        .is_upsert(upsert)
        .partition_key(&pk)
        .map_err(|e| new_cosmos_error_internal(format!("Could not create document: {}", e)))?;
    let c = match etag {
        Some(etag) => c.if_match_condition(IfMatchCondition::Match(etag.to_string())),
        None => c,
    };

    let resp = retry_loop(MAX_RETRY_LOOPS, || async {
        match c.clone().into_future().await {
            Ok(t) => Ok(t),
            Err(err) => Err(RetryLoopError::Permanent(new_cosmos_error(
                "Cosmos db error",
                err,
            ))),
        }
    })
    .await?;
    let etag = resp.etag;
    Ok(etag)
}

/// Insert a document into the cosmos database and returning an etag from the response if
/// successful
pub async fn insert<D: CosmosObject, P: Into<PartitionKeyIntermediate>, C: ToString>(
    collection_name: C,
    pk: P,
    document: &D,
    etag: Option<&String>,
) -> Result<String, CosmosError> {
    insert_internal(collection_name, pk, document, etag, false).await
}

/// Upsert a document into the cosmos database and returning an etag from the response if
/// successful
pub async fn upsert<
    D: CosmosEntity + Serialize + Clone + Send + Sync + 'static,
    P: Into<PartitionKeyIntermediate>,
    C: ToString,
>(
    collection_name: C,
    pk: P,
    document: &D,
    etag: Option<&String>,
) -> Result<String, CosmosError> {
    insert_internal(collection_name, pk, document, etag, true).await
}

/// Returns a specific document from the cosmos DB together with a corresponding etag
pub async fn get<
    D: DeserializeOwned + Send,
    P: Into<PartitionKeyIntermediate>,
    C: ToString,
    S: ToString,
>(
    collection_name: C,
    pk: P,
    document_id: S,
) -> Result<(D, String), CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();
    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(|_| new_cosmos_error_internal("Could not create document client"))?;
    let resp = match document_client
        .get_document()
        .into_future()
        .await
        .map_err(into_cosmos_error("Could not get document"))?
    {
        GetDocumentResponse::Found(resp) => resp,
        GetDocumentResponse::NotFound(resp) => {
            return Err(new_cosmos_error_kind(
                format!("Document not found: {:?}", resp),
                CosmosErrorKind::NotFound,
            ));
        }
    };
    let doc: D = resp.document.document;
    let etag = resp.etag;
    Ok((doc, etag))
}

/// Modifies a document in cosmos by applying `transform` async closure on the existing document and then
/// inserting the returned document and returning both the transformed document, the old
/// document and the etag if successful.
/// If the transform closure fails then no insertion is performed and the error it fails with
/// is returned
pub async fn modify_async_get_old<
    D: CosmosObject + DeserializeOwned + Clone,
    P: Into<PartitionKeyIntermediate>,
    F: Fn(D) -> Fut,
    C: ToString,
    S: ToString,
    Fut: futures::Future<Output = Result<D, warp::Rejection>>,
>(
    collection_name: C,
    pk: P,
    document_id: S,
    transform: F,
) -> Result<(D, D, String), CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();
    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(into_cosmos_error("Could not get document client"))?;
    let (doc, old_doc, etag) = retry_loop(MAX_RETRY_LOOPS, || async {
        let resp = match document_client
            .clone()
            .get_document()
            .into_future()
            .await
            .map_err(|e| RetryLoopError::Permanent(new_cosmos_error("Could not get document", e)))?
        {
            GetDocumentResponse::Found(resp) => resp,
            GetDocumentResponse::NotFound(resp) => {
                return Err(RetryLoopError::Permanent(new_cosmos_error_kind(
                    format!("Document not found: {:?}", resp),
                    CosmosErrorKind::NotFound,
                )));
            }
        };
        let doc: D = resp.document.document;
        let old_doc = doc.clone();

        // Perform changes to the document
        let doc = transform(doc).await.map_err(|e| {
            RetryLoopError::Permanent(new_cosmos_error_kind(
                format!("Modification not possible: {:?}", e),
                CosmosErrorKind::InternalError,
            ))
        })?;
        let c = collection_client
            .create_document(doc.clone())
            .is_upsert(true)
            .if_match_condition(IfMatchCondition::Match(resp.etag))
            .partition_key(&pk)
            .map_err(|e| {
                RetryLoopError::Permanent(new_cosmos_error(
                    format!("Could not create document after modification: {}", e),
                    e,
                ))
            })?;

        match c.into_future().await {
            Ok(resp) => Result::Ok::<_, RetryLoopError<CosmosError>>((doc, old_doc, resp.etag)),
            Err(err) => {
                let err = new_cosmos_error("Cosmos db error", err);
                //NOTE: 412 - PreconditionFailed means the document has been edited between read and write so it
                //means we need to retry the entire read/write block
                match err.kind {
                    CosmosErrorKind::PreconditionFailed => Err(RetryLoopError::Transient(err)),
                    _ => Err(RetryLoopError::Permanent(err)),
                }
            }
        }
    })
    .await?;
    Ok((doc, old_doc, etag))
}

/// Modifies a document in cosmos by applying `transform` async closure on the existing document and then
/// inserting the returned document and returning the transformed document if successful
/// If the transform closure fails then no insertion is performed and the error it fails with
/// is returned
pub async fn modify_async<
    D: CosmosObject + DeserializeOwned,
    P: Into<PartitionKeyIntermediate>,
    F: Fn(D) -> Fut,
    C: ToString,
    S: ToString,
    Fut: futures::Future<Output = Result<D, warp::Rejection>>,
>(
    collection_name: C,
    pk: P,
    document_id: S,
    transform: F,
) -> Result<D, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();
    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(|e| {
            new_cosmos_error_internal(format!("Could not create document client: {}", e))
        })?;
    let doc = retry_loop(MAX_RETRY_LOOPS, || async {
        let resp = match document_client
            .get_document()
            .into_future()
            .await
            .map_err(|e| RetryLoopError::Permanent(new_cosmos_error("Could not get document", e)))?
        {
            GetDocumentResponse::Found(resp) => resp,
            GetDocumentResponse::NotFound(resp) => {
                return Err(RetryLoopError::Permanent(new_cosmos_error_kind(
                    format!("Document not found: {:?}", resp),
                    CosmosErrorKind::NotFound,
                )));
            }
        };
        let doc: D = resp.document.document;

        // Perform changes to the document
        let doc = transform(doc).await.map_err(|e| {
            RetryLoopError::Permanent(new_cosmos_error_kind(
                format!("Modification not possible: {:?}", e),
                CosmosErrorKind::InternalError,
            ))
        })?;
        let c = collection_client
            .create_document(doc.clone())
            .is_upsert(true)
            .if_match_condition(IfMatchCondition::Match(resp.etag))
            .partition_key(&pk)
            .map_err(|e| {
                RetryLoopError::Permanent(new_cosmos_error(
                    format!("Could not create document after modification: {}", e),
                    e,
                ))
            })?;

        match c.into_future().await {
            Ok(_) => Result::Ok::<_, RetryLoopError<CosmosError>>(doc),
            Err(err) => {
                let err = new_cosmos_error("Cosmos db error", err);
                //NOTE: 412 means the document has been edited between read and write so it
                //means we need to retry the entire read/write block
                match err.kind {
                    CosmosErrorKind::PreconditionFailed => Err(RetryLoopError::Transient(err)),
                    _ => Err(RetryLoopError::Permanent(err)),
                }
            }
        }
    })
    .await?;
    Ok(doc)
}

#[derive(Debug, Clone)]
pub enum ModifyReturn<D> {
    Replace(D),
    DontReplace(D),
}

impl<D> ModifyReturn<D> {
    /// No risk of panicing, gets the inner value
    pub fn unwrap(self) -> D {
        match self {
            ModifyReturn::Replace(d) => d,
            ModifyReturn::DontReplace(d) => d,
        }
    }
}

/// Modifies a document in cosmos by applying `transform` closure on the existing document and
/// either inserting the returned value or just returning it to the caller
pub async fn maybe_modify<
    D: CosmosObject + DeserializeOwned + std::fmt::Debug,
    P: Into<PartitionKeyIntermediate>,
    F: Fn(D) -> Result<ModifyReturn<D>, warp::Rejection>,
    C: ToString,
    S: ToString,
>(
    collection_name: C,
    pk: P,
    document_id: S,
    transform: F,
) -> Result<ModifyReturn<D>, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();
    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);

    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(|e| {
            new_cosmos_error_internal(format!("Could not create document client: {}", e))
        })?;

    let doc = retry_loop(MAX_RETRY_LOOPS, || async {
        let resp = match document_client
            .get_document()
            .into_future()
            .await
            .map_err(|e| RetryLoopError::Permanent(new_cosmos_error("Could not get document", e)))?
        {
            GetDocumentResponse::Found(resp) => resp,
            GetDocumentResponse::NotFound(resp) => {
                return Err(RetryLoopError::Permanent(new_cosmos_error_kind(
                    format!("Document not found: {:?}", resp),
                    CosmosErrorKind::NotFound,
                )));
            }
        };
        let doc: D = resp.document.document;

        // Perform changes to the document
        let doc = transform(doc).map_err(|e| {
            RetryLoopError::Permanent(new_cosmos_error_kind(
                format!("Modification error: {:?}", e),
                CosmosErrorKind::InternalError,
            ))
        })?;

        let doc = match doc {
            ModifyReturn::Replace(doc) => doc,
            ModifyReturn::DontReplace(doc) => return Ok(ModifyReturn::DontReplace(doc)),
        };

        let c = collection_client
            .create_document(doc.clone())
            .is_upsert(true)
            .if_match_condition(IfMatchCondition::Match(resp.etag))
            .partition_key(&pk)
            .map_err(|e| {
                RetryLoopError::Permanent(new_cosmos_error(
                    format!("Could not create document after modification: {}", e),
                    e,
                ))
            })?;

        match c.into_future().await {
            Ok(_) => Result::Ok::<_, RetryLoopError<CosmosError>>(ModifyReturn::Replace(doc)),
            Err(err) => {
                let err = new_cosmos_error("Cosmos db error", err);
                //NOTE: 412 means the document has been edited between read and write so it
                //means we need to retry the entire read/write block
                match err.kind {
                    CosmosErrorKind::PreconditionFailed => Err(RetryLoopError::Transient(err)),
                    _ => Err(RetryLoopError::Permanent(err)),
                }
            }
        }
    })
    .await?;
    Ok(doc)
}

/// Modify in cosmos but does not retry on failure which is useful when one needs to pass in a
/// `FnOnce` closure
pub async fn modify_no_retry<
    D: CosmosObject + DeserializeOwned + std::fmt::Debug,
    P: Into<PartitionKeyIntermediate>,
    F: FnOnce(D) -> Result<D, warp::Rejection>,
    C: ToString,
    S: ToString,
>(
    collection_name: C,
    pk: P,
    document_id: S,
    transform: F,
) -> Result<D, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();
    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(|e| {
            new_cosmos_error_internal(format!("Could not create document client: {}", e))
        })?;
    let resp = match document_client
        .get_document()
        .into_future()
        .await
        .map_err(|e| new_cosmos_error("Could not get document", e))?
    {
        GetDocumentResponse::Found(resp) => resp,
        GetDocumentResponse::NotFound(resp) => {
            return Err(new_cosmos_error_kind(
                format!("Document not found: {:?}", resp),
                CosmosErrorKind::NotFound,
            ));
        }
    };
    let doc: D = resp.document.document;

    // Perform changes to the document
    let doc = transform(doc).map_err(|e| {
        new_cosmos_error_kind(
            format!("Modification error: {:?}", e),
            CosmosErrorKind::InternalError,
        )
    })?;
    let c = collection_client
        .create_document(doc.clone())
        .is_upsert(true)
        .if_match_condition(IfMatchCondition::Match(resp.etag))
        .partition_key(&pk)
        .map_err(|e| {
            new_cosmos_error(
                format!("Could not create document after modification: {}", e),
                e,
            )
        })?;

    let doc = match c.into_future().await {
        Ok(_) => doc,
        Err(err) => {
            return Err(new_cosmos_error("Cosmos db error", err));
        }
    };
    Ok(doc)
}

/// Modifies a document in cosmos by applying `transform` closure on the existing document and then
/// inserting the returned document and returning the transformed document if successful
/// If the transform closure fails then no insertion is performed and the error it fails with
/// is returned
pub async fn modify<
    D: CosmosObject + DeserializeOwned + std::fmt::Debug,
    P: Into<PartitionKeyIntermediate>,
    F: Fn(D) -> Result<D, warp::Rejection>,
    C: ToString,
    S: ToString,
>(
    collection_name: C,
    pk: P,
    document_id: S,
    transform: F,
) -> Result<D, CosmosError> {
    modify_async(collection_name, pk, document_id, |d| async { transform(d) }).await
}

pub async fn delete<C: ToString, S: ToString, P: Into<PartitionKeyIntermediate>>(
    collection_name: C,
    pk: P,
    document_id: S,
    etag: Option<String>,
) -> Result<(), CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let document_id = document_id.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token"))?;

    let pk: String = pk.into().into();

    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name);
    let document_client = collection_client
        .clone()
        .document_client(document_id, &pk)
        .map_err(|e| {
            new_cosmos_error_internal(format!("Could not create document client: {}", e))
        })?;

    let del_doc = document_client.delete_document();
    if let Some(etag) = etag {
        del_doc
            .if_match_condition(IfMatchCondition::Match(etag.to_string()))
            .into_future()
            .await
            .map_err(into_cosmos_error("Could not delete document"))?;
    } else {
        del_doc
            .into_future()
            .await
            .map_err(into_cosmos_error("Could not delete document"))?;
    }
    Ok(())
}

pub async fn query_crosspartition_etag<
    D: DeserializeOwned + Send + Sync,
    P: Into<PartitionKeyIntermediate>,
    C: ToString,
>(
    collection_name: C,
    pk: P,
    query: String,
    max_count: i32,
    cross_partition: bool,
) -> Result<Vec<(D, String)>, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    let collection_name = collection_name.to_string();
    let authorization_token = AuthorizationToken::primary_from_base64(&state.cosmos_master_key)
        .map_err(|_| new_cosmos_error_internal("Could not get authorization token."))?;
    let pk: String = pk.into().into();

    let client = CosmosClient::new(state.cosmos_account.to_string(), authorization_token);
    let database_client = client.database_client(&state.cosmos_database);
    let collection_client = database_client.collection_client(collection_name.clone());
    let mut documents: Vec<(D, String)> = vec![];

    let q = Query::new(query.clone());
    let mut query_documents_builder = collection_client
        .query_documents(q)
        .max_item_count(max_count);

    if cross_partition {
        query_documents_builder = query_documents_builder.query_cross_partition(true);
    } else {
        query_documents_builder = query_documents_builder.partition_key(&pk).map_err(|e| {
            new_cosmos_error_internal(format!(
                "Could not set partition key for query document builder: {}",
                e
            ))
        })?;
    }

    let mut query_documents_stream = query_documents_builder.into_stream();
    while let Some(query_documents_response) = query_documents_stream.next().await {
        let query_documents_response =
            query_documents_response.map_err(into_cosmos_error(format!(
                "Could not get query documents response: {} {} {:?}.",
                collection_name, query, pk
            )))?;
        let mut fetched_documents: Vec<(D, String)> = query_documents_response
            .results
            .into_iter()
            .map(|(document, attributes): (D, _)| {
                Ok((
                    document,
                    attributes
                        .ok_or_else(|| new_cosmos_error_internal("No etag returned by query"))?
                        .etag()
                        .to_string(),
                ))
            })
            .collect::<Result<Vec<_>, _>>()?;
        documents.append(&mut fetched_documents);
    }

    Ok(documents)
}

pub async fn query_crosspartition<
    D: DeserializeOwned + Send + Sync,
    P: Into<PartitionKeyIntermediate>,
    C: ToString,
>(
    collection_name: C,
    pk: P,
    query: String,
    max_count: i32,
    cross_partition: bool,
) -> Result<Vec<D>, CosmosError> {
    let v =
        query_crosspartition_etag(collection_name, pk, query, max_count, cross_partition).await?;
    Ok(v.into_iter().map(|(d, _)| d).collect())
}

pub async fn query<
    D: DeserializeOwned + Send + Sync,
    P: Into<PartitionKeyIntermediate>,
    C: ToString,
>(
    collection_name: C,
    pk: P,
    query: String,
    max_count: i32,
) -> Result<Vec<D>, CosmosError> {
    query_crosspartition(collection_name, pk, query, max_count, false).await
}

/// Gets the data from a blob storage
pub async fn get_blob<S: AsRef<str>>(
    id: S,
    storage_container: &str,
) -> Result<Vec<u8>, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    if state.storage_account.is_none() || state.storage_master_key.is_none() {
        return Err(new_cosmos_error_internal(
            "Cosmos storage values are not initialized",
        ));
    }

    let blob_client = StorageClient::new_access_key(
        state.storage_account.as_ref().unwrap(),
        state.storage_master_key.as_ref().unwrap(),
    )
    .container_client(storage_container)
    .blob_client(id.as_ref());

    let mut data = Vec::new();
    let mut blob_stream = blob_client.get().into_stream();

    while let Some(blob_response) = blob_stream.next().await {
        let blob_response =
            blob_response.map_err(into_cosmos_error("Could not get blob response"))?;
        match blob_response.data.collect().await {
            Ok(r) => data.extend_from_slice(&r[..]),
            Err(err) => {
                return Err(new_cosmos_error_kind(
                    format!("Could not get blob from storage account {:?}", err),
                    CosmosErrorKind::BlobError,
                ));
            }
        }
    }
    Ok(data)
}

/// Uploads form data to the blob storage and returns the blob id
pub async fn upload_blob(
    mut f: FormData,
    blob_type: &str,
    expected_content_type: &str,
    storage_container: &str,
) -> Result<String, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    if state.storage_account.is_none() || state.storage_master_key.is_none() {
        return Err(new_cosmos_error_internal(
            "Cosmos storage values are not initialized",
        ));
    }
    while let Some(r) = f.next().await {
        match r {
            Ok(part) => {
                if part.name() == blob_type {
                    if let Some(g) = part.content_type() {
                        let content_type = String::from(g);
                        if content_type.starts_with(&expected_content_type) {
                            match part.filename() {
                                Some(n) => {
                                    let pos = n.find('.');
                                    if let Some(pos) = pos {
                                        let ext = String::from(n);
                                        let ext = &ext[pos..];

                                        //FIXME(Jonathan): Here we are using a Vec<u8> when we probably want to
                                        //use a bytes::Bytes. Additionally we are copying the buffer
                                        //when moving or referencing is prerferrable.
                                        //The reason we're doing it this way is because there are
                                        //currently some versioning issues with azure
                                        let mut buf: Vec<u8> = Vec::with_capacity(64);
                                        let mut s = part.stream();
                                        while let Some(r) = s.next().await {
                                            match r {
                                                Ok(mut b) => {
                                                    buf.extend(b.copy_to_bytes(b.remaining()));
                                                }
                                                Err(err) => {
                                                    return Err(new_cosmos_error_kind(
                                                        format!(
                                                            "Error getting {} data: {:?}",
                                                            blob_type, err
                                                        ),
                                                        CosmosErrorKind::BlobError,
                                                    ));
                                                }
                                            }
                                        }

                                        let mut blob_id = Uuid::new_v4().to_string();

                                        // Add extension to blob id.
                                        blob_id.push_str(ext);

                                        let blob_client = StorageClient::new_access_key(
                                            state.storage_account.as_ref().unwrap(),
                                            state.storage_master_key.as_ref().unwrap(),
                                        )
                                        .container_client(storage_container)
                                        .blob_client(&blob_id);

                                        // Helps preventing spurious data to be uploaded.
                                        let digest: Hash = md5::compute(&buf[..]).into();
                                        {
                                            match blob_client
                                                .put_block_blob(buf)
                                                .content_type(content_type)
                                                .hash(digest)
                                                .into_future()
                                                .await
                                            {
                                                Ok(r) => r,
                                                Err(err) => {
                                                    return Err(new_cosmos_error_kind(
                                                                format!("Could not add blob to storage account {:?}", err),
                                                                CosmosErrorKind::BlobError));
                                                }
                                            };
                                        }

                                        // Set blob id.
                                        return Ok(blob_id);
                                    } else {
                                        return Err(new_cosmos_error_kind(
                                            "Could not get filename for blob.",
                                            CosmosErrorKind::BlobError,
                                        ));
                                    }
                                }
                                None => {
                                    return Err(new_cosmos_error_kind(
                                        "Could not get filename for blob.",
                                        CosmosErrorKind::BlobError,
                                    ));
                                }
                            }
                        } else {
                            return Err(new_cosmos_error_kind(
                                format!(
                                    "Blob does not have a {} content-type.",
                                    expected_content_type
                                ),
                                CosmosErrorKind::BlobError,
                            ));
                        }
                    }
                }
            }
            Err(err) => {
                return Err(new_cosmos_error_kind(
                    format!("Error getting multipart data {:?}", err),
                    CosmosErrorKind::BlobError,
                ));
            }
        }
    }
    Err(new_cosmos_error_kind(
        "No blob provided",
        CosmosErrorKind::BlobError,
    ))
}

/// Uploads a new form data image to the blob storage and returns the image_id
pub async fn upload_image(f: FormData) -> Result<String, CosmosError> {
    let state = COSMOS_STATE
        .get()
        .ok_or_else(|| new_cosmos_error_internal("Cosmos state not initialized"))?;
    if state.image_storage_container.is_none() {
        return Err(new_cosmos_error_internal(
            "Cosmos image storage container is not initialized",
        ));
    }
    upload_blob(
        f,
        "image",
        "image",
        state.image_storage_container.as_ref().unwrap(),
    )
    .await
}

/// The default value for amount of retry loops we do
const MAX_RETRY_LOOPS: usize = 4;
/// `retry_loop` is utilized in order to combat transient errors. A closure which generates a future
/// is used to generate the same future and run this untill completion, if the future returns success
/// with r then Ok(r) is returned. If the future returns an error then exponential backoff is tried
/// until eventually a max amount of tries is reached at which point the function returns the
/// latest error. The function runs at least once.
/// See here for microsoft documentation https://docs.microsoft.com/en-us/azure/azure-sql/database/troubleshoot-common-connectivity-issues#principles-for-retry
const RETRY_MAX_RANDOM: u64 = 5000;
const RETRY_START_WAIT: u64 = 5000;
pub enum RetryLoopError<E> {
    Permanent(E),
    Transient(E),
}
pub async fn retry_loop<F, A, R, E>(tries: usize, mut f: F) -> Result<R, E>
where
    F: FnMut() -> A,
    A: std::future::Future<Output = Result<R, RetryLoopError<E>>>,
{
    let mut counter = 0;
    let mut wait = RETRY_START_WAIT;
    loop {
        match f().await {
            Ok(r) => {
                return Ok(r);
            }
            Err(e) => {
                match e {
                    RetryLoopError::Permanent(e) => return Err(e),
                    RetryLoopError::Transient(e) => {
                        counter += 1;
                        let random_wait: u64 = rand::random::<u64>() % RETRY_MAX_RANDOM;
                        // Wait for an exponential backoff time + some random wait
                        sleep(Duration::from_millis(wait) + Duration::from_millis(random_wait))
                            .await;
                        if counter >= tries {
                            return Err(e);
                        }
                        //Exponential backoff
                        wait *= 2;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod util_tests {
    use super::*;
    use std::time::Instant;

    #[tokio::test]
    async fn retry_loop_test() {
        let calls = std::cell::RefCell::new(vec![]);
        retry_loop(8, || async {
            calls.borrow_mut().push(Instant::now());
            if calls.borrow().len() >= 4 {
                Ok(())
            } else {
                Err(RetryLoopError::Transient(()))
            }
        })
        .await
        .unwrap();
        let mut calls = calls.borrow_mut();
        assert_eq!(calls.len(), 4);
        let t = calls.pop().unwrap().elapsed().as_millis();
        assert!(t == 0);
        let t = calls.pop().unwrap().elapsed().as_millis();
        assert!((200..=1200).contains(&t));
        let t = calls.pop().unwrap().elapsed().as_millis();
        assert!((300..=2200).contains(&t));
        let t = calls.pop().unwrap().elapsed().as_millis();
        assert!((350..=3200).contains(&t));
    }

    #[tokio::test]
    #[ignore]
    // Ignored since it requires quite a bit of time to retry several times
    async fn rety_loop_failure() {
        let calls = std::cell::RefCell::new(vec![]);
        let result = retry_loop(8, || async {
            calls.borrow_mut().push(Instant::now());
            if 1 == 2 {
                return Ok(());
            }
            Err(RetryLoopError::Transient(()))
        })
        .await;
        if result.is_ok() {
            panic!();
        }
        let calls = calls.borrow();
        assert_eq!(calls.len(), 8);
    }
}