<commit_before>use std::{collections::BTreeSet, mem};

use syn::{
    braced,
    parse::{Parse, ParseStream},
    spanned::Spanned,
    visit::Visit,
    Attribute, Field, Ident, Lifetime, Token, Type,
};

use super::{
    attribute::{Meta, MetaNameValue},
    request::{RequestField, RequestFieldKind, RequestLifetimes},
    response::{ResponseField, ResponseFieldKind},
    Api, Metadata, Request, Response,
};

mod kw {
    use syn::custom_keyword;

    custom_keyword!(error);
    custom_keyword!(request);
    custom_keyword!(response);
}

impl Parse for Api {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let metadata: Metadata = input.parse()?;

        let attributes = input.call(Attribute::parse_outer)?;
        let (request, attributes) = if input.peek(kw::request) {
            let request = parse_request(input, attributes)?;
            let attributes = input.call(Attribute::parse_outer)?;

            (Some(request), attributes)
        } else {
            (None, attributes)
        };

        let response = if input.peek(kw::response) {
            Some(parse_response(input, attributes)?)
        } else if !attributes.is_empty() {
            return Err(syn::Error::new_spanned(
                &attributes[0],
                "attributes are not supported on the error type",
            ));
        } else {
            None
        };

        // TODO: Use `bool::then` when MSRV >= 1.50
        let error_ty = if input.peek(kw::error) {
            let _: kw::error = input.parse()?;
            let _: Token![:] = input.parse()?;

            Some(input.parse()?)
        } else {
            None
        };

        if let Some(req) = &request {
            let newtype_body_field = req.newtype_body_field();
            if metadata.method == "GET" && (req.has_body_fields() || newtype_body_field.is_some()) {
                let mut combined_error: Option<syn::Error> = None;
                let mut add_error = |field| {
                    let error =
                        syn::Error::new_spanned(field, "GET endpoints can't have body fields");
                    if let Some(combined_error_ref) = &mut combined_error {
                        combined_error_ref.combine(error);
                    } else {
                        combined_error = Some(error);
                    }
                };

                for field in req.body_fields() {
                    add_error(field);
                }

                if let Some(field) = newtype_body_field {
                    add_error(field);
                }

                return Err(combined_error.unwrap());
            }
        }

        Ok(Self { metadata, request, response, error_ty })
    }
}

fn parse_request(input: ParseStream<'_>, attributes: Vec<Attribute>) -> syn::Result<Request> {
    let request_kw: kw::request = input.parse()?;
    let _: Token![:] = input.parse()?;
    let fields;
    braced!(fields in input);

    let mut newtype_body_field = None;
    let mut query_map_field = None;
    let mut lifetimes = RequestLifetimes::default();

    let fields: Vec<_> = fields
        .parse_terminated::<Field, Token![,]>(Field::parse_named)?
        .into_iter()
        .map(|mut field| {
            let mut field_kind = None;
            let mut header = None;

            for attr in mem::take(&mut field.attrs) {
                let meta = match Meta::from_attribute(&attr)? {
                    Some(m) => m,
                    None => {
                        field.attrs.push(attr);
                        continue;
                    }
                };

                if field_kind.is_some() {
                    return Err(syn::Error::new_spanned(
                        attr,
                        "There can only be one field kind attribute",
                    ));
                }

                field_kind = Some(match meta {
                    Meta::Word(ident) => match &ident.to_string()[..] {
                        attr @ "body" | attr @ "raw_body" => req_res_meta_word(
                            attr,
                            &field,
                            &mut newtype_body_field,
                            RequestFieldKind::NewtypeBody,
                            RequestFieldKind::NewtypeRawBody,
                        )?,
                        "path" => RequestFieldKind::Path,
                        "query" => RequestFieldKind::Query,
                        "query_map" => {
                            if let Some(f) = &query_map_field {
                                let mut error = syn::Error::new_spanned(
                                    field,
                                    "There can only be one query map field",
                                );
                                error.combine(syn::Error::new_spanned(
                                    f,
                                    "Previous query map field",
                                ));
                                return Err(error);
                            }

                            query_map_field = Some(field.clone());
                            RequestFieldKind::QueryMap
                        }
                        _ => {
                            return Err(syn::Error::new_spanned(
                                ident,
                                "Invalid #[ruma_api] argument, expected one of \
                                     `body`, `path`, `query`, `query_map`",
                            ));
                        }
                    },
                    Meta::NameValue(MetaNameValue { name, value }) => {
                        req_res_name_value(name, value, &mut header, RequestFieldKind::Header)?
                    }
                });
            }

            match field_kind.unwrap_or(RequestFieldKind::Body) {
                RequestFieldKind::Header => {
                    collect_lifetime_idents(&mut lifetimes.header, &field.ty)
                }
                RequestFieldKind::Body => collect_lifetime_idents(&mut lifetimes.body, &field.ty),
                RequestFieldKind::NewtypeBody => {
                    collect_lifetime_idents(&mut lifetimes.body, &field.ty)
                }
                RequestFieldKind::NewtypeRawBody => {
                    collect_lifetime_idents(&mut lifetimes.body, &field.ty)
                }
                RequestFieldKind::Path => collect_lifetime_idents(&mut lifetimes.path, &field.ty),
                RequestFieldKind::Query => collect_lifetime_idents(&mut lifetimes.query, &field.ty),
                RequestFieldKind::QueryMap => {
                    collect_lifetime_idents(&mut lifetimes.query, &field.ty)
                }
            }

            Ok(RequestField::new(field_kind.unwrap_or(RequestFieldKind::Body), field, header))
        })
        .collect::<syn::Result<_>>()?;

    if newtype_body_field.is_some() && fields.iter().any(|f| f.is_body()) {
        // TODO: highlight conflicting fields,
        return Err(syn::Error::new_spanned(
            request_kw,
            "Can't have both a newtype body field and regular body fields",
        ));
    }

    if query_map_field.is_some() && fields.iter().any(|f| f.is_query()) {
        return Err(syn::Error::new_spanned(
            // TODO: raw,
            request_kw,
            "Can't have both a query map field and regular query fields",
        ));
    }

    // TODO when/if `&[(&str, &str)]` is supported remove this
    if query_map_field.is_some() && !lifetimes.query.is_empty() {
        return Err(syn::Error::new_spanned(
            request_kw,
            "Lifetimes are not allowed for query_map fields",
        ));
    }

    Ok(Request { attributes, fields, lifetimes })
}

fn parse_response(input: ParseStream<'_>, attributes: Vec<Attribute>) -> syn::Result<Response> {
    let response_kw: kw::response = input.parse()?;
    let _: Token![:] = input.parse()?;
    let fields;
    braced!(fields in input);

    let mut newtype_body_field = None;

    let fields: Vec<_> = fields
        .parse_terminated::<Field, Token![,]>(Field::parse_named)?
        .into_iter()
        .map(|mut field| {
            if has_lifetime(&field.ty) {
                return Err(syn::Error::new(
                    field.ident.span(),
                    "Lifetimes on Response fields cannot be supported until GAT are stable",
                ));
            }

            let mut field_kind = None;
            let mut header = None;

            for attr in mem::take(&mut field.attrs) {
                let meta = match Meta::from_attribute(&attr)? {
                    Some(m) => m,
                    None => {
                        field.attrs.push(attr);
                        continue;
                    }
                };

                if field_kind.is_some() {
                    return Err(syn::Error::new_spanned(
                        attr,
                        "There can only be one field kind attribute",
                    ));
                }

                field_kind = Some(match meta {
                    Meta::Word(ident) => match &ident.to_string()[..] {
                        s @ "body" | s @ "raw_body" => req_res_meta_word(
                            s,
                            &field,
                            &mut newtype_body_field,
                            ResponseFieldKind::NewtypeBody,
                            ResponseFieldKind::NewtypeRawBody,
                        )?,
                        _ => {
                            return Err(syn::Error::new_spanned(
                                ident,
                                "Invalid #[ruma_api] argument with value, expected `body`",
                            ));
                        }
                    },
                    Meta::NameValue(MetaNameValue { name, value }) => {
                        req_res_name_value(name, value, &mut header, ResponseFieldKind::Header)?
                    }
                });
            }

            Ok(match field_kind.unwrap_or(ResponseFieldKind::Body) {
                ResponseFieldKind::Body => ResponseField::Body(field),
                ResponseFieldKind::Header => {
                    ResponseField::Header(field, header.expect("missing header name"))
                }
                ResponseFieldKind::NewtypeBody => ResponseField::NewtypeBody(field),
                ResponseFieldKind::NewtypeRawBody => ResponseField::NewtypeRawBody(field),
            })
        })
        .collect::<syn::Result<_>>()?;

    if newtype_body_field.is_some() && fields.iter().any(|f| f.is_body()) {
        // TODO: highlight conflicting fields,
        return Err(syn::Error::new_spanned(
            response_kw,
            "Can't have both a newtype body field and regular body fields",
        ));
    }

    Ok(Response { attributes, fields })
}

fn has_lifetime(ty: &Type) -> bool {
    let mut lifetimes = BTreeSet::new();
    collect_lifetime_idents(&mut lifetimes, ty);
    !lifetimes.is_empty()
}

fn collect_lifetime_idents(lifetimes: &mut BTreeSet<Lifetime>, ty: &Type) {
    struct Visitor<'lt>(&'lt mut BTreeSet<Lifetime>);
    impl<'ast> Visit<'ast> for Visitor<'_> {
        fn visit_lifetime(&mut self, lt: &'ast Lifetime) {
            self.0.insert(lt.clone());
        }
    }

    Visitor(lifetimes).visit_type(ty)
}

fn req_res_meta_word<T>(
    attr_kind: &str,
    field: &syn::Field,
    newtype_body_field: &mut Option<syn::Field>,
    body_field_kind: T,
    raw_field_kind: T,
) -> syn::Result<T> {
    if let Some(f) = &newtype_body_field {
        let mut error = syn::Error::new_spanned(field, "There can only be one newtype body field");
        error.combine(syn::Error::new_spanned(f, "Previous newtype body field"));
        return Err(error);
    }

    *newtype_body_field = Some(field.clone());
    Ok(match attr_kind {
        "body" => body_field_kind,
        "raw_body" => raw_field_kind,
        _ => unreachable!(),
    })
}

fn req_res_name_value<T>(
    name: Ident,
    value: Ident,
    header: &mut Option<Ident>,
    field_kind: T,
) -> syn::Result<T> {
    if name != "header" {
        return Err(syn::Error::new_spanned(
            name,
            "Invalid #[ruma_api] argument with value, expected `header`",
        ));
    }

    *header = Some(value);
    Ok(field_kind)
}
<commit_msg>Fix confusing variable naming in ruma-api-macros<commit_after>use std::{collections::BTreeSet, mem};

use syn::{
    braced,
    parse::{Parse, ParseStream},
    spanned::Spanned,
    visit::Visit,
    Attribute, Field, Ident, Lifetime, Token, Type,
};

use super::{
    attribute::{Meta, MetaNameValue},
    request::{RequestField, RequestFieldKind, RequestLifetimes},
    response::{ResponseField, ResponseFieldKind},
    Api, Metadata, Request, Response,
};

mod kw {
    use syn::custom_keyword;

    custom_keyword!(error);
    custom_keyword!(request);
    custom_keyword!(response);
}

impl Parse for Api {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let metadata: Metadata = input.parse()?;

        let req_attrs = input.call(Attribute::parse_outer)?;
        let (request, attributes) = if input.peek(kw::request) {
            let request = parse_request(input, req_attrs)?;
            let after_req_attrs = input.call(Attribute::parse_outer)?;

            (Some(request), after_req_attrs)
        } else {
            // There was no `request` field so the attributes are for `response`
            (None, req_attrs)
        };

        let response = if input.peek(kw::response) {
            Some(parse_response(input, attributes)?)
        } else if !attributes.is_empty() {
            return Err(syn::Error::new_spanned(
                &attributes[0],
                "attributes are not supported on the error type",
            ));
        } else {
            None
        };

        // TODO: Use `bool::then` when MSRV >= 1.50
        let error_ty = if input.peek(kw::error) {
            let _: kw::error = input.parse()?;
            let _: Token![:] = input.parse()?;

            Some(input.parse()?)
        } else {
            None
        };

        if let Some(req) = &request {
            let newtype_body_field = req.newtype_body_field();
            if metadata.method == "GET" && (req.has_body_fields() || newtype_body_field.is_some()) {
                let mut combined_error: Option<syn::Error> = None;
                let mut add_error = |field| {
                    let error =
                        syn::Error::new_spanned(field, "GET endpoints can't have body fields");
                    if let Some(combined_error_ref) = &mut combined_error {
                        combined_error_ref.combine(error);
                    } else {
                        combined_error = Some(error);
                    }
                };

                for field in req.body_fields() {
                    add_error(field);
                }

                if let Some(field) = newtype_body_field {
                    add_error(field);
                }

                return Err(combined_error.unwrap());
            }
        }

        Ok(Self { metadata, request, response, error_ty })
    }
}

fn parse_request(input: ParseStream<'_>, attributes: Vec<Attribute>) -> syn::Result<Request> {
    let request_kw: kw::request = input.parse()?;
    let _: Token![:] = input.parse()?;
    let fields;
    braced!(fields in input);

    let mut newtype_body_field = None;
    let mut query_map_field = None;
    let mut lifetimes = RequestLifetimes::default();

    let fields: Vec<_> = fields
        .parse_terminated::<Field, Token![,]>(Field::parse_named)?
        .into_iter()
        .map(|mut field| {
            let mut field_kind = None;
            let mut header = None;

            for attr in mem::take(&mut field.attrs) {
                let meta = match Meta::from_attribute(&attr)? {
                    Some(m) => m,
                    None => {
                        field.attrs.push(attr);
                        continue;
                    }
                };

                if field_kind.is_some() {
                    return Err(syn::Error::new_spanned(
                        attr,
                        "There can only be one field kind attribute",
                    ));
                }

                field_kind = Some(match meta {
                    Meta::Word(ident) => match &ident.to_string()[..] {
                        attr @ "body" | attr @ "raw_body" => req_res_meta_word(
                            attr,
                            &field,
                            &mut newtype_body_field,
                            RequestFieldKind::NewtypeBody,
                            RequestFieldKind::NewtypeRawBody,
                        )?,
                        "path" => RequestFieldKind::Path,
                        "query" => RequestFieldKind::Query,
                        "query_map" => {
                            if let Some(f) = &query_map_field {
                                let mut error = syn::Error::new_spanned(
                                    field,
                                    "There can only be one query map field",
                                );
                                error.combine(syn::Error::new_spanned(
                                    f,
                                    "Previous query map field",
                                ));
                                return Err(error);
                            }

                            query_map_field = Some(field.clone());
                            RequestFieldKind::QueryMap
                        }
                        _ => {
                            return Err(syn::Error::new_spanned(
                                ident,
                                "Invalid #[ruma_api] argument, expected one of \
                                     `body`, `path`, `query`, `query_map`",
                            ));
                        }
                    },
                    Meta::NameValue(MetaNameValue { name, value }) => {
                        req_res_name_value(name, value, &mut header, RequestFieldKind::Header)?
                    }
                });
            }

            match field_kind.unwrap_or(RequestFieldKind::Body) {
                RequestFieldKind::Header => {
                    collect_lifetime_idents(&mut lifetimes.header, &field.ty)
                }
                RequestFieldKind::Body => collect_lifetime_idents(&mut lifetimes.body, &field.ty),
                RequestFieldKind::NewtypeBody => {
                    collect_lifetime_idents(&mut lifetimes.body, &field.ty)
                }
                RequestFieldKind::NewtypeRawBody => {
                    collect_lifetime_idents(&mut lifetimes.body, &field.ty)
                }
                RequestFieldKind::Path => collect_lifetime_idents(&mut lifetimes.path, &field.ty),
                RequestFieldKind::Query => collect_lifetime_idents(&mut lifetimes.query, &field.ty),
                RequestFieldKind::QueryMap => {
                    collect_lifetime_idents(&mut lifetimes.query, &field.ty)
                }
            }

            Ok(RequestField::new(field_kind.unwrap_or(RequestFieldKind::Body), field, header))
        })
        .collect::<syn::Result<_>>()?;

    if newtype_body_field.is_some() && fields.iter().any(|f| f.is_body()) {
        // TODO: highlight conflicting fields,
        return Err(syn::Error::new_spanned(
            request_kw,
            "Can't have both a newtype body field and regular body fields",
        ));
    }

    if query_map_field.is_some() && fields.iter().any(|f| f.is_query()) {
        return Err(syn::Error::new_spanned(
            // TODO: raw,
            request_kw,
            "Can't have both a query map field and regular query fields",
        ));
    }

    // TODO when/if `&[(&str, &str)]` is supported remove this
    if query_map_field.is_some() && !lifetimes.query.is_empty() {
        return Err(syn::Error::new_spanned(
            request_kw,
            "Lifetimes are not allowed for query_map fields",
        ));
    }

    Ok(Request { attributes, fields, lifetimes })
}

fn parse_response(input: ParseStream<'_>, attributes: Vec<Attribute>) -> syn::Result<Response> {
    let response_kw: kw::response = input.parse()?;
    let _: Token![:] = input.parse()?;
    let fields;
    braced!(fields in input);

    let mut newtype_body_field = None;

    let fields: Vec<_> = fields
        .parse_terminated::<Field, Token![,]>(Field::parse_named)?
        .into_iter()
        .map(|mut field| {
            if has_lifetime(&field.ty) {
                return Err(syn::Error::new(
                    field.ident.span(),
                    "Lifetimes on Response fields cannot be supported until GAT are stable",
                ));
            }

            let mut field_kind = None;
            let mut header = None;

            for attr in mem::take(&mut field.attrs) {
                let meta = match Meta::from_attribute(&attr)? {
                    Some(m) => m,
                    None => {
                        field.attrs.push(attr);
                        continue;
                    }
                };

                if field_kind.is_some() {
                    return Err(syn::Error::new_spanned(
                        attr,
                        "There can only be one field kind attribute",
                    ));
                }

                field_kind = Some(match meta {
                    Meta::Word(ident) => match &ident.to_string()[..] {
                        s @ "body" | s @ "raw_body" => req_res_meta_word(
                            s,
                            &field,
                            &mut newtype_body_field,
                            ResponseFieldKind::NewtypeBody,
                            ResponseFieldKind::NewtypeRawBody,
                        )?,
                        _ => {
                            return Err(syn::Error::new_spanned(
                                ident,
                                "Invalid #[ruma_api] argument with value, expected `body`",
                            ));
                        }
                    },
                    Meta::NameValue(MetaNameValue { name, value }) => {
                        req_res_name_value(name, value, &mut header, ResponseFieldKind::Header)?
                    }
                });
            }

            Ok(match field_kind.unwrap_or(ResponseFieldKind::Body) {
                ResponseFieldKind::Body => ResponseField::Body(field),
                ResponseFieldKind::Header => {
                    ResponseField::Header(field, header.expect("missing header name"))
                }
                ResponseFieldKind::NewtypeBody => ResponseField::NewtypeBody(field),
                ResponseFieldKind::NewtypeRawBody => ResponseField::NewtypeRawBody(field),
            })
        })
        .collect::<syn::Result<_>>()?;

    if newtype_body_field.is_some() && fields.iter().any(|f| f.is_body()) {
        // TODO: highlight conflicting fields,
        return Err(syn::Error::new_spanned(
            response_kw,
            "Can't have both a newtype body field and regular body fields",
        ));
    }

    Ok(Response { attributes, fields })
}

fn has_lifetime(ty: &Type) -> bool {
    let mut lifetimes = BTreeSet::new();
    collect_lifetime_idents(&mut lifetimes, ty);
    !lifetimes.is_empty()
}

fn collect_lifetime_idents(lifetimes: &mut BTreeSet<Lifetime>, ty: &Type) {
    struct Visitor<'lt>(&'lt mut BTreeSet<Lifetime>);
    impl<'ast> Visit<'ast> for Visitor<'_> {
        fn visit_lifetime(&mut self, lt: &'ast Lifetime) {
            self.0.insert(lt.clone());
        }
    }

    Visitor(lifetimes).visit_type(ty)
}

fn req_res_meta_word<T>(
    attr_kind: &str,
    field: &syn::Field,
    newtype_body_field: &mut Option<syn::Field>,
    body_field_kind: T,
    raw_field_kind: T,
) -> syn::Result<T> {
    if let Some(f) = &newtype_body_field {
        let mut error = syn::Error::new_spanned(field, "There can only be one newtype body field");
        error.combine(syn::Error::new_spanned(f, "Previous newtype body field"));
        return Err(error);
    }

    *newtype_body_field = Some(field.clone());
    Ok(match attr_kind {
        "body" => body_field_kind,
        "raw_body" => raw_field_kind,
        _ => unreachable!(),
    })
}

fn req_res_name_value<T>(
    name: Ident,
    value: Ident,
    header: &mut Option<Ident>,
    field_kind: T,
) -> syn::Result<T> {
    if name != "header" {
        return Err(syn::Error::new_spanned(
            name,
            "Invalid #[ruma_api] argument with value, expected `header`",
        ));
    }

    *header = Some(value);
    Ok(field_kind)
}
<|endoftext|>


<commit_before>//! Methods associated with document parameters and calling documents

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use common::{
    eyre::{bail, Result},
    indexmap::IndexMap,
    itertools::Itertools,
};
use node_address::Address;
use node_pointer::{resolve, resolve_mut};
use node_validate::Validator;
use path_utils::lexiclean::Lexiclean;
use stencila_schema::{InlineContent, Node, Parameter};

use crate::{document::Document, When};

impl Document {
    /// Get the parameters of the document
    ///
    /// Returns all parameters within the root node of the document as a map, indexed by the
    /// parameter name, containing a tuple of the parameter `id`, [`Address`] as well as the [`Parameter`] itself.
    ///
    /// Used in `Document::call` and when compiling a `Call` node so that the `Call` inherits the parameters of
    /// the document as it's own `arguments`.
    pub async fn params(&mut self) -> Result<IndexMap<String, (String, Address, Parameter)>> {
        // Assemble the document to ensure its `addresses` are up to date
        self.assemble(When::Never, When::Never, When::Never).await?;

        // Collect parameters from addresses
        let addresses = self.addresses.read().await;
        let root = &*self.root.read().await;
        let params = addresses
            .iter()
            .filter_map(|(id, address)| {
                if let Ok(pointer) = resolve(root, Some(address.clone()), Some(id.clone())) {
                    if let Some(InlineContent::Parameter(param)) = pointer.as_inline() {
                        return Some((
                            param.name.clone(),
                            (id.clone(), address.clone(), param.clone()),
                        ));
                    }
                }
                None
            })
            .collect();

        Ok(params)
    }

    /// Call the document with arguments
    ///
    /// This function is similar to `Document::execute`, and indeed calls that function, but first
    /// sets the value of parameters to the `Node` with the same name in `args`.
    ///
    /// # Arguments
    ///
    /// - `args`: A map of the arguments to call the document with
    pub async fn call(&mut self, args: HashMap<String, Node>) -> Result<()> {
        let mut params = self.params().await?;

        {
            let root = &mut *self.root.write().await;
            for (name, value) in args {
                if let Some((id, address, param)) = params.remove(&name) {
                    if let Some(validator) = param.validator.as_deref() {
                        match validator.validate(&value) {
                            Ok(..) => {
                                if let Ok(mut pointer) = resolve_mut(root, Some(address), Some(id))
                                {
                                    if let Some(InlineContent::Parameter(param)) =
                                        pointer.as_inline_mut()
                                    {
                                        param.value = Some(Box::new(value));
                                    }
                                }
                            }
                            Err(error) => bail!(
                                "While attempting to parse document parameter `{}`: {}",
                                name,
                                error
                            ),
                        }
                    }
                } else {
                    bail!("Document does not have a parameter named `{}`", name)
                }
            }
        }

        self.execute(When::Never, None, None, None).await?;

        Ok(())
    }

    /// Call the document with arguments having string values
    ///
    /// This calls the `Document::call` function but first parses each string value into
    /// a `Node` using the corresponding validator.
    ///
    /// Used to call a document with strings taken from the command line.
    ///
    /// # Arguments
    ///
    /// - `args`: A map of the arguments to call the document with
    pub async fn call_with_strings(&mut self, args: HashMap<String, String>) -> Result<()> {
        let mut params = self.params().await?;
        let mut args_parsed = HashMap::new();
        for (name, value) in args {
            if let Some((_id, _address, param)) = params.remove(&name) {
                if let Some(validator) = param.validator.as_deref() {
                    match validator.parse(&value) {
                        Ok(value) => {
                            args_parsed.insert(name, value);
                        }
                        Err(error) => bail!(
                            "While attempting to parse document parameter `{}`: {}",
                            name,
                            error
                        ),
                    }
                }
            } else {
                bail!("Document does not have a parameter named `{}`", name)
            }
        }

        self.call(args_parsed).await?;

        Ok(())
    }

    /// Call the document with arguments parsed from a path
    ///
    /// This calls the `Document::call_with_strings` function but first extracts arguments
    /// from URL path by matching `/<arg>/` segments with `/$<param>/` (or `.$<param>.`)
    /// segments within the document's filesystem path.
    ///
    /// Used by the server and by the site generator to render a document with a parameterized
    /// path. Also exposed on the command line e.g. `stencila run 2020/west` would run either
    /// `./$year/$region.md` or `./$year.$region.docx` files.
    ///
    /// Will return an error if the `path` does not match the path of the document.
    ///
    /// # Arguments
    ///
    /// - `path`: The path to be called e.g. `2020/west` for a document with filesystem path
    ///           of `./$year/$region.md` or `./$year.$region.md`.
    ///
    /// - `dir`: The directory that the `path` is relative to (defaults to current working directory)
    pub async fn call_with_path(&mut self, path: &Path, dir: Option<&Path>) -> Result<()> {
        let dir = dir.map_or_else(
            || std::env::current_dir().expect("Directory not supplied and unable to get CWD"),
            PathBuf::from,
        );
        let path = dir.join(path).lexiclean();

        if !Self::is_matching_path(&path, &self.path) {
            bail!(
                "Unable to call document at `{}` with path `{}` because paths do not match",
                self.path.to_string_lossy(),
                path.to_string_lossy()
            );
        }

        let mut args = HashMap::new();
        let arg_values = path.components().collect_vec();
        for (index, component) in Self::path_parts(&self.path).iter().enumerate() {
            if let Some(name) = component.strip_prefix('$') {
                args.insert(
                    name.to_string(),
                    arg_values[index].as_os_str().to_string_lossy().to_string(),
                );
            }
        }

        self.call_with_strings(args).await?;

        Ok(())
    }

    /// Does a filesystem path contain parameters?
    pub fn path_has_parameters(path: &Path) -> bool {
        for component in path.components() {
            let component = component.as_os_str().to_string_lossy();
            if component.starts_with('$') || component.contains(".$") {
                return false;
            }
        }
        false
    }

    /// Split a filesystem path into parts, including parts separated by dots
    /// in directory of file names
    fn path_parts(path: &Path) -> Vec<String> {
        let mut parts = Vec::new();
        for component in path.components() {
            let string = component.as_os_str().to_string_lossy();
            for part in string.split('.') {
                parts.push(part.to_string());
            }
        }
        parts
    }

    /// Find a filesystem path that matches a, possibly parametrized, path
    ///
    /// Searches the `dir` for a document with a path (parameterized, or not) that matches
    /// `path`. Returns the first file, in lexical order of path, that matches the path.
    ///
    /// Used when opening documents from the command line or by the server.
    ///
    /// # Arguments
    ///
    /// - `path`: The path possibly containing arguments and lacking an extension e.g. `2020/west`
    ///
    /// - `dir`: The directory that the `path` is relative to (defaults to current working directory)
    pub fn find_matching_path(path: &Path, dir: Option<&Path>) -> Result<PathBuf> {
        if path.exists() {
            return Ok(path.to_path_buf());
        } else if path.is_absolute() {
            bail!("Path does not exist: {}", path.to_string_lossy())
        }

        let dir = dir.map_or_else(
            || std::env::current_dir().expect("Directory not supplied and unable to get CWD"),
            PathBuf::from,
        );

        let walker = ignore::WalkBuilder::new(&dir)
            .sort_by_file_path(Path::cmp)
            .build();
        for entry in walker.filter_map(Result::ok) {
            let entry_path = entry.path();
            let rel_path = entry_path.strip_prefix(&dir).unwrap_or(entry_path);
            if entry_path.is_file() && Self::is_matching_path(path, rel_path) {
                return Ok(entry_path.to_path_buf());
            }
        }

        bail!(
            "Unable to find a file matching path: {}",
            path.to_string_lossy()
        );
    }

    /// Match a path to a, possibly parametrized, filesystem path
    ///
    /// Used in `Document::find_path` to find a file system path matching a path.
    pub fn is_matching_path<P1: AsRef<Path>, P2: AsRef<Path>>(path: P1, file: P2) -> bool {
        let path = path.as_ref();
        let file = file.as_ref();

        // Equal as is? (note that Rust removes trailing slashes in paths)
        if path == file {
            return true;
        }

        // Equal to file stem (i.e. if drop extension)
        let file = match file.file_stem() {
            Some(stem) => file
                .components()
                .take(file.components().count() - 1)
                .collect::<PathBuf>()
                .join(stem),
            None => file.to_path_buf(),
        };
        if path == file {
            return true;
        }

        // Split path into components
        let path_parts = path.components();

        // Split the file into parts (includes dots)
        let file_parts = Self::path_parts(&file);

        // Equal number of parts?
        if path_parts.count() != file_parts.len() {
            return false;
        }

        // Do the `path_parts` match the `file_parts` allowing for `$<name>` "wildcards"
        for (index, component) in path.components().enumerate() {
            if file_parts[index].starts_with('$') {
                continue;
            }

            let path_part = component.as_os_str().to_string_lossy();
            if path_part != file_parts[index] {
                return false;
            }
        }

        true
    }
}

#[cfg(test)]
mod tests {
    use common::tokio;

    use crate::document::Document;

    #[tokio::test]
    async fn matches_path() {
        // Match on equal
        assert!(Document::is_matching_path("file.ext", "file.ext"));
        assert!(Document::is_matching_path("dir/file.ext", "dir/file.ext"));
        assert!(Document::is_matching_path("dir/file.ext/", "dir/file.ext"));

        // Match on stem
        assert!(Document::is_matching_path("file", "file.ext"));
        assert!(Document::is_matching_path("dir/file", "dir/file.ext"));
        assert!(Document::is_matching_path("dir/file/", "dir/file.ext"));

        // Match on parameters
        assert!(Document::is_matching_path("2002", "$year.ext"));
        assert!(Document::is_matching_path("2002/west", "$year.$region.ext"));
        assert!(Document::is_matching_path("2002/west", "$year/$region.ext"));
        assert!(Document::is_matching_path(
            "2002/west/",
            "$year/$region.ext"
        ));
        assert!(Document::is_matching_path(
            "2002/west/habitat/pelagic",
            "$year.$region.habitat.$habitat.ext"
        ));
        assert!(Document::is_matching_path(
            "2002/west/habitat/pelagic",
            "$year/$region/habitat/$habitat.ext"
        ));
        assert!(Document::is_matching_path(
            "2002/region/west/pelagic",
            "$year/region/$region.$habitat.ext"
        ));
        assert!(Document::is_matching_path(
            "2002/region/west/pelagic",
            "$year.region.$region/$habitat.ext"
        ));
        assert!(Document::is_matching_path(
            "2002/region/west/habitat/pelagic/",
            "$year.region.$region.habitat.$habitat.ext"
        ));

        // Non matches
        assert!(!Document::is_matching_path("a", "b"));
        assert!(!Document::is_matching_path("a", "b.ext"));
        assert!(!Document::is_matching_path("a/c", "b/$c"));
        assert!(!Document::is_matching_path("a/b/c/d", "a/$b/c"));
    }
}
<commit_msg>feat(Documents): Add `enumerate_urls` function<commit_after>//! Methods associated with document parameters and calling documents

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use common::{
    eyre::{bail, Result},
    indexmap::IndexMap,
    itertools::Itertools,
    serde_json, tracing,
};
use node_address::Address;
use node_pointer::{resolve, resolve_mut};
use node_validate::Validator;
use path_utils::lexiclean::Lexiclean;
use stencila_schema::{EnumValidator, InlineContent, Node, Parameter, ValidatorTypes};

use crate::{document::Document, When};

impl Document {
    /// Get the parameters of the document
    ///
    /// Returns all parameters within the root node of the document as a map, indexed by the
    /// parameter name, containing a tuple of the parameter `id`, [`Address`] as well as the [`Parameter`] itself.
    ///
    /// Used in `Document::call` and when compiling a `Call` node so that the `Call` inherits the parameters of
    /// the document as it's own `arguments`.
    pub async fn params(&mut self) -> Result<IndexMap<String, (String, Address, Parameter)>> {
        // Assemble the document to ensure its `addresses` are up to date
        self.assemble(When::Never, When::Never, When::Never).await?;

        // Collect parameters from addresses
        let addresses = self.addresses.read().await;
        let root = &*self.root.read().await;
        let params = addresses
            .iter()
            .filter_map(|(id, address)| {
                if let Ok(pointer) = resolve(root, Some(address.clone()), Some(id.clone())) {
                    if let Some(InlineContent::Parameter(param)) = pointer.as_inline() {
                        return Some((
                            param.name.clone(),
                            (id.clone(), address.clone(), param.clone()),
                        ));
                    }
                }
                None
            })
            .collect();

        Ok(params)
    }

    /// Call the document with arguments
    ///
    /// This function is similar to `Document::execute`, and indeed calls that function, but first
    /// sets the value of parameters to the `Node` with the same name in `args`.
    ///
    /// # Arguments
    ///
    /// - `args`: A map of the arguments to call the document with
    pub async fn call(&mut self, args: HashMap<String, Node>) -> Result<()> {
        let mut params = self.params().await?;

        {
            let root = &mut *self.root.write().await;
            for (name, value) in args {
                if let Some((id, address, param)) = params.remove(&name) {
                    if let Some(validator) = param.validator.as_deref() {
                        match validator.validate(&value) {
                            Ok(..) => {
                                if let Ok(mut pointer) = resolve_mut(root, Some(address), Some(id))
                                {
                                    if let Some(InlineContent::Parameter(param)) =
                                        pointer.as_inline_mut()
                                    {
                                        param.value = Some(Box::new(value));
                                    }
                                }
                            }
                            Err(error) => bail!(
                                "While attempting to parse document parameter `{}`: {}",
                                name,
                                error
                            ),
                        }
                    }
                } else {
                    bail!("Document does not have a parameter named `{}`", name)
                }
            }
        }

        self.execute(When::Never, None, None, None).await?;

        Ok(())
    }

    /// Call the document with arguments having string values
    ///
    /// This calls the `Document::call` function but first parses each string value into
    /// a `Node` using the corresponding validator.
    ///
    /// Used to call a document with strings taken from the command line.
    ///
    /// # Arguments
    ///
    /// - `args`: A map of the arguments to call the document with
    pub async fn call_with_strings(&mut self, args: HashMap<String, String>) -> Result<()> {
        let mut params = self.params().await?;
        let mut args_parsed = HashMap::new();
        for (name, value) in args {
            if let Some((_id, _address, param)) = params.remove(&name) {
                if let Some(validator) = param.validator.as_deref() {
                    match validator.parse(&value) {
                        Ok(value) => {
                            args_parsed.insert(name, value);
                        }
                        Err(error) => bail!(
                            "While attempting to parse document parameter `{}`: {}",
                            name,
                            error
                        ),
                    }
                }
            } else {
                bail!("Document does not have a parameter named `{}`", name)
            }
        }

        self.call(args_parsed).await?;

        Ok(())
    }

    /// Call the document with arguments parsed from a path
    ///
    /// This calls the `Document::call_with_strings` function but first extracts arguments
    /// from URL path by matching `/<arg>/` segments with `/$<param>/` (or `.$<param>.`)
    /// segments within the document's filesystem path.
    ///
    /// Used by the server and by the site generator to render a document with a parameterized
    /// path. Also exposed on the command line e.g. `stencila run 2020/west` would run either
    /// `./$year/$region.md` or `./$year.$region.docx` files.
    ///
    /// Will return an error if the `path` does not match the path of the document.
    ///
    /// # Arguments
    ///
    /// - `path`: The path to be called e.g. `2020/west` for a document with filesystem path
    ///           of `./$year/$region.md` or `./$year.$region.md`.
    ///
    /// - `dir`: The directory that the `path` is relative to (defaults to current working directory)
    pub async fn call_with_path(&mut self, path: &Path, dir: Option<&Path>) -> Result<()> {
        let dir = dir.map_or_else(
            || std::env::current_dir().expect("Directory not supplied and unable to get CWD"),
            PathBuf::from,
        );
        let path = dir.join(path).lexiclean();

        if !Self::path_matches(&path, &self.path) {
            bail!(
                "Unable to call document at `{}` with path `{}` because paths do not match",
                self.path.to_string_lossy(),
                path.to_string_lossy()
            );
        }

        let mut args = HashMap::new();
        let arg_values = path.components().collect_vec();
        for (index, component) in Self::path_segments(&self.path).iter().enumerate() {
            if let Some(name) = component.strip_prefix('$') {
                args.insert(
                    name.to_string(),
                    arg_values[index].as_os_str().to_string_lossy().to_string(),
                );
            }
        }

        self.call_with_strings(args).await?;

        Ok(())
    }

    /// Get a list of possible URL paths for a document
    ///
    /// When a document has a parameterized path that involves only parameters with `EnumValidator`s,
    /// tis function will return all possible combinations of paths for those parameters. If a document does not
    /// have any parameters, or one of the parameters in the path is not an `EnumValidator`, then
    /// only one path (`self.path`) will be returned.
    pub async fn enumerate_urls(&mut self, expand: bool) -> Result<Vec<String>> {
        let params = self.params().await?;

        let path_segments = Self::path_segments(&self.path);
        let path_params = Self::path_params(&self.path);

        // Collect the values of all enum parameters
        let param_values: HashMap<String, &Vec<Node>> = path_params.into_iter().filter_map(|name| match params.get(&name) {
            Some((.., param)) => {
                match param.validator.as_deref() {
                    Some(ValidatorTypes::EnumValidator(EnumValidator{values,..})) => {
                        match !values.is_empty() {
                            true => Some((name, values)),
                            false => None
                        }
                    },
                    _ => None
                }
            },
            None => {
                tracing::warn!("Parameter `${}` is used in path for document `{}` but is not represented by a `Parameter` in the content of the document", name, self.path.to_string_lossy());
                None
            }
        }).collect();

        // If there is not a set of values for all of the docs parameters then return the
        // expanded (forward slashes only) URL path
        if !expand || param_values.len() != params.len() {
            return Ok(vec![path_segments.join("/")]);
        }

        // Expand the segments ito URL paths using the enum values
        let mut urls: Vec<String> = Vec::new();
        for segment in path_segments {
            // A parameter segment so expand paths for that segment
            if let Some(param) = segment.strip_prefix('$') {
                if let Some(values) = param_values.get(param) {
                    urls = urls
                        .iter()
                        .flat_map(|path| {
                            values.iter().map(|node| {
                                [
                                    path,
                                    "/",
                                    &match node {
                                        Node::Boolean(value) => value.to_string(),
                                        Node::Integer(value) => value.to_string(),
                                        Node::Number(value) => value.0.to_string(),
                                        Node::String(value) => value.to_string(),
                                        _ => serde_json::to_string(node)
                                            .unwrap_or_else(|_| segment.clone()),
                                    },
                                ]
                                .concat()
                            })
                        })
                        .collect();
                    continue;
                }
            }

            // Not a parameter segment (or one that for some reason has no values) so just add to
            // each of the existing paths
            for path in urls.iter_mut() {
                path.push('/');
                path.push_str(&segment)
            }
        }

        Ok(urls)
    }

    /// Does a filesystem path contain parameters?
    pub fn path_has_parameters(path: &Path) -> bool {
        for component in path.components() {
            let component = component.as_os_str().to_string_lossy();
            if component.starts_with('$') || component.contains(".$") {
                return false;
            }
        }
        false
    }

    /// Split a filesystem path into segments, including segments separated by dots
    /// in directory of file names
    fn path_segments(path: &Path) -> Vec<String> {
        let mut parts = Vec::new();
        for component in path.components() {
            let string = component.as_os_str().to_string_lossy();
            for part in string.split('.') {
                parts.push(part.to_string());
            }
        }
        parts
    }

    /// Get the name of the parameters is a document's filesystem path
    fn path_params(path: &Path) -> Vec<String> {
        Self::path_segments(path)
            .iter()
            .filter_map(|part| part.strip_prefix('$').map(String::from))
            .collect_vec()
    }

    /// Match a URL path to a, possibly parametrized, filesystem path
    ///
    /// Used in `Document::find_path` to find a file system path matching a path.
    pub fn path_matches<P1: AsRef<Path>, P2: AsRef<Path>>(path: P1, file: P2) -> bool {
        let path = path.as_ref();
        let file = file.as_ref();

        // Equal as is? (note that Rust removes trailing slashes in paths)
        if path == file {
            return true;
        }

        // Equal to file stem (i.e. if drop extension)
        let file = match file.file_stem() {
            Some(stem) => file
                .components()
                .take(file.components().count() - 1)
                .collect::<PathBuf>()
                .join(stem),
            None => file.to_path_buf(),
        };
        if path == file {
            return true;
        }

        // Split path into segments
        let path_segments = path.components();

        // Split the file into segments (includes dots)
        let file_segments = Self::path_segments(&file);

        // Equal number of segments?
        if path_segments.count() != file_segments.len() {
            return false;
        }

        // Do the `path_segments` match the `file_segments` allowing for `$<name>` "wildcards"
        for (index, component) in path.components().enumerate() {
            if file_segments[index].starts_with('$') {
                continue;
            }

            let path_part = component.as_os_str().to_string_lossy();
            if path_part != file_segments[index] {
                return false;
            }
        }

        true
    }

    /// Find a filesystem path that matches a, possibly parametrized, path
    ///
    /// Searches the `dir` for a document with a path (parameterized, or not) that matches
    /// `path`. Returns the first file, in lexical order of path, that matches the path.
    ///
    /// Used when opening documents from the command line or by the server.
    ///
    /// # Arguments
    ///
    /// - `path`: The path possibly containing arguments and lacking an extension e.g. `2020/west`
    ///
    /// - `dir`: The directory that the `path` is relative to (defaults to current working directory)
    pub fn find_matching_path(path: &Path, dir: Option<&Path>) -> Result<PathBuf> {
        if path.exists() {
            return Ok(path.to_path_buf());
        } else if path.is_absolute() {
            bail!("Path does not exist: {}", path.to_string_lossy())
        }

        let dir = dir.map_or_else(
            || std::env::current_dir().expect("Directory not supplied and unable to get CWD"),
            PathBuf::from,
        );

        let walker = ignore::WalkBuilder::new(&dir)
            .sort_by_file_path(Path::cmp)
            .build();
        for entry in walker.filter_map(Result::ok) {
            let entry_path = entry.path();
            let rel_path = entry_path.strip_prefix(&dir).unwrap_or(entry_path);
            if entry_path.is_file() && Self::path_matches(path, rel_path) {
                return Ok(entry_path.to_path_buf());
            }
        }

        bail!(
            "Unable to find a file matching path: {}",
            path.to_string_lossy()
        );
    }
}

#[cfg(test)]
mod tests {
    use common::tokio;

    use crate::document::Document;

    #[tokio::test]
    async fn matches_path() {
        // Match on equal
        assert!(Document::path_matches("file.ext", "file.ext"));
        assert!(Document::path_matches("dir/file.ext", "dir/file.ext"));
        assert!(Document::path_matches("dir/file.ext/", "dir/file.ext"));

        // Match on stem
        assert!(Document::path_matches("file", "file.ext"));
        assert!(Document::path_matches("dir/file", "dir/file.ext"));
        assert!(Document::path_matches("dir/file/", "dir/file.ext"));

        // Match on parameters
        assert!(Document::path_matches("2002", "$year.ext"));
        assert!(Document::path_matches("2002/west", "$year.$region.ext"));
        assert!(Document::path_matches("2002/west", "$year/$region.ext"));
        assert!(Document::path_matches("2002/west/", "$year/$region.ext"));
        assert!(Document::path_matches(
            "2002/west/habitat/pelagic",
            "$year.$region.habitat.$habitat.ext"
        ));
        assert!(Document::path_matches(
            "2002/west/habitat/pelagic",
            "$year/$region/habitat/$habitat.ext"
        ));
        assert!(Document::path_matches(
            "2002/region/west/pelagic",
            "$year/region/$region.$habitat.ext"
        ));
        assert!(Document::path_matches(
            "2002/region/west/pelagic",
            "$year.region.$region/$habitat.ext"
        ));
        assert!(Document::path_matches(
            "2002/region/west/habitat/pelagic/",
            "$year.region.$region.habitat.$habitat.ext"
        ));

        // Non matches
        assert!(!Document::path_matches("a", "b"));
        assert!(!Document::path_matches("a", "b.ext"));
        assert!(!Document::path_matches("a/c", "b/$c"));
        assert!(!Document::path_matches("a/b/c/d", "a/$b/c"));
    }
}
<|endoftext|>


<commit_before>use conduit::RequestExt;
use diesel::prelude::*;
use diesel::r2d2::{self, ConnectionManager, CustomizeConnection};
use parking_lot::{ReentrantMutex, ReentrantMutexGuard};
use prometheus::Histogram;
use std::sync::Arc;
use std::{ops::Deref, time::Duration};
use thiserror::Error;
use url::Url;

use crate::config;
use crate::middleware::app::RequestApp;

#[derive(Clone)]
pub enum DieselPool {
    Pool {
        pool: r2d2::Pool<ConnectionManager<PgConnection>>,
        time_to_obtain_connection_metric: Histogram,
    },
    Test(Arc<ReentrantMutex<PgConnection>>),
}

impl DieselPool {
    pub(crate) fn new(
        url: &str,
        config: &config::DatabasePools,
        r2d2_config: r2d2::Builder<ConnectionManager<PgConnection>>,
        time_to_obtain_connection_metric: Histogram,
    ) -> Result<DieselPool, PoolError> {
        let manager = ConnectionManager::new(connection_url(config, url));

        // For crates.io we want the behavior of creating a database pool to be slightly different
        // than the defaults of R2D2: the library's build() method assumes its consumers always
        // need a database connection to operate, so it blocks creating a pool until a minimum
        // number of connections is available.
        //
        // crates.io can actually operate in a limited capacity without a database connections,
        // especially by serving download requests to our users. Because of that we don't want to
        // block indefinitely waiting for a connection: we instead need to wait for a bit (to avoid
        // serving errors for the first connections until the pool is initialized) and if we can't
        // establish any connection continue booting up the application. The database pool will
        // automatically be marked as unhealthy and the rest of the application will adapt.
        let pool = DieselPool::Pool {
            pool: r2d2_config.build_unchecked(manager),
            time_to_obtain_connection_metric,
        };
        match pool.wait_until_healthy(Duration::from_secs(5)) {
            Ok(()) => {}
            Err(PoolError::UnhealthyPool) => {}
            Err(err) => return Err(err),
        }

        Ok(pool)
    }

    pub(crate) fn new_test(config: &config::DatabasePools, url: &str) -> DieselPool {
        let conn = PgConnection::establish(&connection_url(config, url))
            .expect("failed to establish connection");
        conn.begin_test_transaction()
            .expect("failed to begin test transaction");
        DieselPool::Test(Arc::new(ReentrantMutex::new(conn)))
    }

    pub fn get(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        match self {
            DieselPool::Pool {
                pool,
                time_to_obtain_connection_metric,
            } => time_to_obtain_connection_metric.observe_closure_duration(|| {
                if let Some(conn) = pool.try_get() {
                    Ok(DieselPooledConn::Pool(conn))
                } else if !self.is_healthy() {
                    Err(PoolError::UnhealthyPool)
                } else {
                    Ok(DieselPooledConn::Pool(pool.get()?))
                }
            }),
            DieselPool::Test(conn) => Ok(DieselPooledConn::Test(conn.lock())),
        }
    }

    pub fn state(&self) -> PoolState {
        match self {
            DieselPool::Pool { pool, .. } => {
                let state = pool.state();
                PoolState {
                    connections: state.connections,
                    idle_connections: state.idle_connections,
                }
            }
            DieselPool::Test(_) => PoolState {
                connections: 0,
                idle_connections: 0,
            },
        }
    }

    pub fn wait_until_healthy(&self, timeout: Duration) -> Result<(), PoolError> {
        match self {
            DieselPool::Pool { pool, .. } => match pool.get_timeout(timeout) {
                Ok(_) => Ok(()),
                Err(_) if !self.is_healthy() => Err(PoolError::UnhealthyPool),
                Err(err) => Err(PoolError::R2D2(err)),
            },
            DieselPool::Test(_) => Ok(()),
        }
    }

    fn is_healthy(&self) -> bool {
        self.state().connections > 0
    }
}

#[derive(Debug, Copy, Clone)]
pub struct PoolState {
    pub connections: u32,
    pub idle_connections: u32,
}

pub enum DieselPooledConn<'a> {
    Pool(r2d2::PooledConnection<ConnectionManager<PgConnection>>),
    Test(ReentrantMutexGuard<'a, PgConnection>),
}

impl Deref for DieselPooledConn<'_> {
    type Target = PgConnection;

    fn deref(&self) -> &Self::Target {
        match self {
            DieselPooledConn::Pool(conn) => conn.deref(),
            DieselPooledConn::Test(conn) => conn.deref(),
        }
    }
}

pub fn oneoff_connection_with_config(
    config: &config::DatabasePools,
) -> ConnectionResult<PgConnection> {
    let url = connection_url(config, &config.primary.url);
    PgConnection::establish(&url)
}

pub fn oneoff_connection() -> ConnectionResult<PgConnection> {
    let config = config::DatabasePools::full_from_environment(&config::Base::from_environment());
    oneoff_connection_with_config(&config)
}

pub fn connection_url(config: &config::DatabasePools, url: &str) -> String {
    let mut url = Url::parse(url).expect("Invalid database URL");

    if config.enforce_tls {
        maybe_append_url_param(&mut url, "sslmode", "require");
    }

    // Configure the time it takes for diesel to return an error when there is full packet loss
    // between the application and the database.
    maybe_append_url_param(
        &mut url,
        "tcp_user_timeout",
        &config.tcp_timeout_ms.to_string(),
    );

    url.into()
}

fn maybe_append_url_param(url: &mut Url, key: &str, value: &str) {
    if !url.query_pairs().any(|(k, _)| k == key) {
        url.query_pairs_mut().append_pair(key, value);
    }
}

pub trait RequestTransaction {
    /// Obtain a read/write database connection from the primary pool
    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError>;

    /// Obtain a readonly database connection from the replica pool
    ///
    /// If there is no replica pool, the primary pool is used instead.
    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError>;
}

impl<T: RequestExt + ?Sized> RequestTransaction for T {
    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        self.app().primary_database.get()
    }

    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        match &self.app().read_only_replica_database {
            Some(pool) => pool.get(),
            None => self.app().primary_database.get(),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ConnectionConfig {
    pub statement_timeout: u64,
    pub read_only: bool,
}

impl CustomizeConnection<PgConnection, r2d2::Error> for ConnectionConfig {
    fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), r2d2::Error> {
        use diesel::sql_query;

        sql_query(format!(
            "SET statement_timeout = {}",
            self.statement_timeout * 1000
        ))
        .execute(conn)
        .map_err(r2d2::Error::QueryError)?;
        if self.read_only {
            sql_query("SET default_transaction_read_only = 't'")
                .execute(conn)
                .map_err(r2d2::Error::QueryError)?;
        }
        Ok(())
    }
}

#[cfg(test)]
pub(crate) fn test_conn() -> PgConnection {
    let conn = PgConnection::establish(&crate::env("TEST_DATABASE_URL")).unwrap();
    conn.begin_test_transaction().unwrap();
    conn
}

#[derive(Debug, Error)]
pub enum PoolError {
    #[error(transparent)]
    R2D2(#[from] r2d2::PoolError),
    #[error("unhealthy database pool")]
    UnhealthyPool,
}
<commit_msg>db: Add primary database fallback to `db_read_only()` function<commit_after>use conduit::RequestExt;
use diesel::prelude::*;
use diesel::r2d2::{self, ConnectionManager, CustomizeConnection};
use parking_lot::{ReentrantMutex, ReentrantMutexGuard};
use prometheus::Histogram;
use std::sync::Arc;
use std::{ops::Deref, time::Duration};
use thiserror::Error;
use url::Url;

use crate::config;
use crate::middleware::app::RequestApp;

#[derive(Clone)]
pub enum DieselPool {
    Pool {
        pool: r2d2::Pool<ConnectionManager<PgConnection>>,
        time_to_obtain_connection_metric: Histogram,
    },
    Test(Arc<ReentrantMutex<PgConnection>>),
}

impl DieselPool {
    pub(crate) fn new(
        url: &str,
        config: &config::DatabasePools,
        r2d2_config: r2d2::Builder<ConnectionManager<PgConnection>>,
        time_to_obtain_connection_metric: Histogram,
    ) -> Result<DieselPool, PoolError> {
        let manager = ConnectionManager::new(connection_url(config, url));

        // For crates.io we want the behavior of creating a database pool to be slightly different
        // than the defaults of R2D2: the library's build() method assumes its consumers always
        // need a database connection to operate, so it blocks creating a pool until a minimum
        // number of connections is available.
        //
        // crates.io can actually operate in a limited capacity without a database connections,
        // especially by serving download requests to our users. Because of that we don't want to
        // block indefinitely waiting for a connection: we instead need to wait for a bit (to avoid
        // serving errors for the first connections until the pool is initialized) and if we can't
        // establish any connection continue booting up the application. The database pool will
        // automatically be marked as unhealthy and the rest of the application will adapt.
        let pool = DieselPool::Pool {
            pool: r2d2_config.build_unchecked(manager),
            time_to_obtain_connection_metric,
        };
        match pool.wait_until_healthy(Duration::from_secs(5)) {
            Ok(()) => {}
            Err(PoolError::UnhealthyPool) => {}
            Err(err) => return Err(err),
        }

        Ok(pool)
    }

    pub(crate) fn new_test(config: &config::DatabasePools, url: &str) -> DieselPool {
        let conn = PgConnection::establish(&connection_url(config, url))
            .expect("failed to establish connection");
        conn.begin_test_transaction()
            .expect("failed to begin test transaction");
        DieselPool::Test(Arc::new(ReentrantMutex::new(conn)))
    }

    pub fn get(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        match self {
            DieselPool::Pool {
                pool,
                time_to_obtain_connection_metric,
            } => time_to_obtain_connection_metric.observe_closure_duration(|| {
                if let Some(conn) = pool.try_get() {
                    Ok(DieselPooledConn::Pool(conn))
                } else if !self.is_healthy() {
                    Err(PoolError::UnhealthyPool)
                } else {
                    Ok(DieselPooledConn::Pool(pool.get()?))
                }
            }),
            DieselPool::Test(conn) => Ok(DieselPooledConn::Test(conn.lock())),
        }
    }

    pub fn state(&self) -> PoolState {
        match self {
            DieselPool::Pool { pool, .. } => {
                let state = pool.state();
                PoolState {
                    connections: state.connections,
                    idle_connections: state.idle_connections,
                }
            }
            DieselPool::Test(_) => PoolState {
                connections: 0,
                idle_connections: 0,
            },
        }
    }

    pub fn wait_until_healthy(&self, timeout: Duration) -> Result<(), PoolError> {
        match self {
            DieselPool::Pool { pool, .. } => match pool.get_timeout(timeout) {
                Ok(_) => Ok(()),
                Err(_) if !self.is_healthy() => Err(PoolError::UnhealthyPool),
                Err(err) => Err(PoolError::R2D2(err)),
            },
            DieselPool::Test(_) => Ok(()),
        }
    }

    fn is_healthy(&self) -> bool {
        self.state().connections > 0
    }
}

#[derive(Debug, Copy, Clone)]
pub struct PoolState {
    pub connections: u32,
    pub idle_connections: u32,
}

pub enum DieselPooledConn<'a> {
    Pool(r2d2::PooledConnection<ConnectionManager<PgConnection>>),
    Test(ReentrantMutexGuard<'a, PgConnection>),
}

impl Deref for DieselPooledConn<'_> {
    type Target = PgConnection;

    fn deref(&self) -> &Self::Target {
        match self {
            DieselPooledConn::Pool(conn) => conn.deref(),
            DieselPooledConn::Test(conn) => conn.deref(),
        }
    }
}

pub fn oneoff_connection_with_config(
    config: &config::DatabasePools,
) -> ConnectionResult<PgConnection> {
    let url = connection_url(config, &config.primary.url);
    PgConnection::establish(&url)
}

pub fn oneoff_connection() -> ConnectionResult<PgConnection> {
    let config = config::DatabasePools::full_from_environment(&config::Base::from_environment());
    oneoff_connection_with_config(&config)
}

pub fn connection_url(config: &config::DatabasePools, url: &str) -> String {
    let mut url = Url::parse(url).expect("Invalid database URL");

    if config.enforce_tls {
        maybe_append_url_param(&mut url, "sslmode", "require");
    }

    // Configure the time it takes for diesel to return an error when there is full packet loss
    // between the application and the database.
    maybe_append_url_param(
        &mut url,
        "tcp_user_timeout",
        &config.tcp_timeout_ms.to_string(),
    );

    url.into()
}

fn maybe_append_url_param(url: &mut Url, key: &str, value: &str) {
    if !url.query_pairs().any(|(k, _)| k == key) {
        url.query_pairs_mut().append_pair(key, value);
    }
}

pub trait RequestTransaction {
    /// Obtain a read/write database connection from the primary pool
    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError>;

    /// Obtain a readonly database connection from the replica pool
    ///
    /// If the replica pool is disabled or unavailable, the primary pool is used instead.
    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError>;
}

impl<T: RequestExt + ?Sized> RequestTransaction for T {
    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        self.app().primary_database.get()
    }

    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError> {
        let read_only_pool = self.app().read_only_replica_database.as_ref();
        match read_only_pool.map(|pool| pool.get()) {
            // Replica is available
            Some(Ok(connection)) => Ok(connection),

            // Replica is not available, but primary might be available
            Some(Err(PoolError::UnhealthyPool)) => self.app().primary_database.get(),

            // Replica failed
            Some(Err(error)) => Err(error),

            // Replica is disabled, but primary might be available
            None => self.app().primary_database.get(),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ConnectionConfig {
    pub statement_timeout: u64,
    pub read_only: bool,
}

impl CustomizeConnection<PgConnection, r2d2::Error> for ConnectionConfig {
    fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), r2d2::Error> {
        use diesel::sql_query;

        sql_query(format!(
            "SET statement_timeout = {}",
            self.statement_timeout * 1000
        ))
        .execute(conn)
        .map_err(r2d2::Error::QueryError)?;
        if self.read_only {
            sql_query("SET default_transaction_read_only = 't'")
                .execute(conn)
                .map_err(r2d2::Error::QueryError)?;
        }
        Ok(())
    }
}

#[cfg(test)]
pub(crate) fn test_conn() -> PgConnection {
    let conn = PgConnection::establish(&crate::env("TEST_DATABASE_URL")).unwrap();
    conn.begin_test_transaction().unwrap();
    conn
}

#[derive(Debug, Error)]
pub enum PoolError {
    #[error(transparent)]
    R2D2(#[from] r2d2::PoolError),
    #[error("unhealthy database pool")]
    UnhealthyPool,
}
<|endoftext|>


<commit_before>use ffi;
use std::os::raw::c_int;
use std::ptr;
use types::*;
use ffi::MaybeError;
use device::DeviceInfo;

// Input
// -----
/// Representation of an input midi port
#[allow(missing_copy_implementations)]
pub struct InputPort {
    stream: *const ffi::PortMidiStream,
    device: ffi::PmDeviceId,
    buffer_size: i32, // TODO: replace with an usize
}

const EVENT_BUFFER_SIZE: usize = 128; // replace this with a parameter to InputPort new
impl InputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(device: DeviceInfo, buffer_size: i32) -> Result<InputPort> {
        if device.is_output() {
            return Err(Error::Invalid);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenInput(&raw_stream as *const *const _,
                              device.id(),
                              ptr::null(), // *inputDriverInfo, not needed for normal operation
                              EVENT_BUFFER_SIZE as c_int,
                              ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                              ptr::null()) // time_info, a pointer passed to the time procedure
        }));

        Ok(InputPort {
            stream: raw_stream,
            device: device.id(),
            buffer_size: EVENT_BUFFER_SIZE as c_int,
        })
    }

    pub fn read_n(&mut self, cnt: usize) -> Result<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > EVENT_BUFFER_SIZE {
            EVENT_BUFFER_SIZE as i32
        } else {
            cnt as i32
        };
        let mut event_buffer = [ffi::PmEvent::default(); EVENT_BUFFER_SIZE];
        let res = unsafe { ffi::Pm_Read(self.stream, event_buffer.as_mut_ptr(), read_cnt) };
        if res < 0 {
            let err = ffi::PmError::try_from(res).unwrap();
            // TODO: Return the error
            println!("error: {:?}", err);
            return Ok(None);
        } else if res == 0 {
            Ok(None)
        } else {
            // remove mutability and replace return value
            let events = (0..res as usize)
                             .map(|i| MidiEvent::from(event_buffer[i].clone()))
                             .collect::<Vec<MidiEvent>>();
            Ok(Some(events))
        }
    }

    /// Reads a single `MidiEvent` if one is avaible
    ///
    /// A `Result` of `None` means no event was available.
    ///
    /// See the PortMidi documentation for information on how it deals with input
    /// overflows
    /// TODO: call `read_n`
    pub fn read(&mut self) -> Result<Option<MidiEvent>> {
        match self.read_n(1) {
            Ok(Some(mut vec)) => Ok(vec.pop()),
            Ok(_) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// `poll` tests if there is input available, either returing a bool or an error
    pub fn poll(&self) -> Result<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            err @ _ => Err(Error::PortMidi(err)),
        }
    }
}
impl Drop for InputPort {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}

// Output
// ------
/// Representation of an output midi port
///
#[allow(missing_copy_implementations)]
pub struct OutputPort {
    stream: *const ffi::PortMidiStream,
    device: ffi::PmDeviceId,
    buffer_size: i32, // TODO: replace with an usize
}
impl OutputPort {
    /// Construct a new `OutputPort` for `input_device`
    pub fn new(device: DeviceInfo, buffer_size: i32) -> Result<OutputPort> {
        if device.is_input() {
            return Err(Error::Invalid);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenOutput(&raw_stream as *const *const _,
                               device.id(),
                               ptr::null(), // *inputDriverInfo, not needed for normal operation
                               EVENT_BUFFER_SIZE as c_int,
                               ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                               ptr::null(), // time_info, a pointer passed to the time procedure
                               0) //latency
        }));

        Ok(OutputPort {
            stream: raw_stream,
            device: device.id(),
            buffer_size: EVENT_BUFFER_SIZE as c_int,
        })
    }

    /// Terminates outgoing messages immediately
    ///
    /// The caller should immediately close the output port, this may
    /// result in transmission of a partial midi message. Note, not all platforms
    /// support abort.
    pub fn abort(&mut self) -> Result<()> {
        Result::from(unsafe { ffi::Pm_Abort(self.stream) })
    }

    /// Write a single `MidiEvent`
    pub fn write_event(&mut self, midi_event: MidiEvent) -> Result<()> {
        let event = midi_event.into();
        Result::from(unsafe { ffi::Pm_Write(self.stream, &event, 1) })
    }

    /// Write a single `MidiMessage` immediately
    pub fn write_message(&mut self, midi_message: MidiMessage) -> Result<()> {
        let message = midi_message.into();
        Result::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, message) })
    }
}
impl Drop for OutputPort {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
<commit_msg>Use more appropriate error types<commit_after>use ffi;
use std::os::raw::c_int;
use std::ptr;
use types::*;
use ffi::MaybeError;
use device::DeviceInfo;

// Input
// -----
/// Representation of an input midi port
#[allow(missing_copy_implementations)]
pub struct InputPort {
    stream: *const ffi::PortMidiStream,
    device: ffi::PmDeviceId,
    buffer_size: i32, // TODO: replace with an usize
}

const EVENT_BUFFER_SIZE: usize = 128; // replace this with a parameter to InputPort new
impl InputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(device: DeviceInfo, buffer_size: i32) -> Result<InputPort> {
        if device.is_output() {
            return Err(Error::NotAnInputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenInput(&raw_stream as *const *const _,
                              device.id(),
                              ptr::null(), // *inputDriverInfo, not needed for normal operation
                              EVENT_BUFFER_SIZE as c_int,
                              ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                              ptr::null()) // time_info, a pointer passed to the time procedure
        }));

        Ok(InputPort {
            stream: raw_stream,
            device: device.id(),
            buffer_size: EVENT_BUFFER_SIZE as c_int,
        })
    }

    pub fn read_n(&mut self, cnt: usize) -> Result<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > EVENT_BUFFER_SIZE {
            EVENT_BUFFER_SIZE as i32
        } else {
            cnt as i32
        };
        let mut event_buffer = [ffi::PmEvent::default(); EVENT_BUFFER_SIZE];
        let res = unsafe { ffi::Pm_Read(self.stream, event_buffer.as_mut_ptr(), read_cnt) };
        if res < 0 {
            let err = ffi::PmError::try_from(res).unwrap();
            // TODO: Return the error
            println!("error: {:?}", err);
            return Ok(None);
        } else if res == 0 {
            Ok(None)
        } else {
            // remove mutability and replace return value
            let events = (0..res as usize)
                             .map(|i| MidiEvent::from(event_buffer[i].clone()))
                             .collect::<Vec<MidiEvent>>();
            Ok(Some(events))
        }
    }

    /// Reads a single `MidiEvent` if one is avaible
    ///
    /// A `Result` of `None` means no event was available.
    ///
    /// See the PortMidi documentation for information on how it deals with input
    /// overflows
    /// TODO: call `read_n`
    pub fn read(&mut self) -> Result<Option<MidiEvent>> {
        match self.read_n(1) {
            Ok(Some(mut vec)) => Ok(vec.pop()),
            Ok(_) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// `poll` tests if there is input available, either returing a bool or an error
    pub fn poll(&self) -> Result<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            err @ _ => Err(Error::PortMidi(err)),
        }
    }
}
impl Drop for InputPort {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}

// Output
// ------
/// Representation of an output midi port
///
#[allow(missing_copy_implementations)]
pub struct OutputPort {
    stream: *const ffi::PortMidiStream,
    device: ffi::PmDeviceId,
    buffer_size: i32, // TODO: replace with an usize
}
impl OutputPort {
    /// Construct a new `OutputPort` for `input_device`
    pub fn new(device: DeviceInfo, buffer_size: i32) -> Result<OutputPort> {
        if device.is_input() {
            return Err(Error::NotAnOutputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenOutput(&raw_stream as *const *const _,
                               device.id(),
                               ptr::null(), // *inputDriverInfo, not needed for normal operation
                               EVENT_BUFFER_SIZE as c_int,
                               ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                               ptr::null(), // time_info, a pointer passed to the time procedure
                               0) //latency
        }));

        Ok(OutputPort {
            stream: raw_stream,
            device: device.id(),
            buffer_size: EVENT_BUFFER_SIZE as c_int,
        })
    }

    /// Terminates outgoing messages immediately
    ///
    /// The caller should immediately close the output port, this may
    /// result in transmission of a partial midi message. Note, not all platforms
    /// support abort.
    pub fn abort(&mut self) -> Result<()> {
        Result::from(unsafe { ffi::Pm_Abort(self.stream) })
    }

    /// Write a single `MidiEvent`
    pub fn write_event(&mut self, midi_event: MidiEvent) -> Result<()> {
        let event = midi_event.into();
        Result::from(unsafe { ffi::Pm_Write(self.stream, &event, 1) })
    }

    /// Write a single `MidiMessage` immediately
    pub fn write_message(&mut self, midi_message: MidiMessage) -> Result<()> {
        let message = midi_message.into();
        Result::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, message) })
    }
}
impl Drop for OutputPort {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
<|endoftext|>