<commit_before>//!Router asigns handlers to paths and resolves them per request

#[cfg(test)]
use http::method;
use http::method::{ Method, Get, Post, Put, Delete };
use http::server::request::{AbsolutePath};
use regex::Regex;
use std::collections::hashmap::HashMap;
use request::Request;
use response::Response;
use middleware::{ Middleware, Action, Halt, Continue };
use nickel_error::NickelError;

/// A Route is the basic data structure that stores both the path
/// and the handler that gets executed for the route.
/// The path can contain variable pattern such as `user/:userid/invoices`
pub struct Route {
    pub path: String,
    pub method: Method,
    pub handler: fn(request: &Request, response: &mut Response),
    pub variables: HashMap<String, uint>,
    matcher: Regex
}

impl Clone for Route {
    fn clone(&self) -> Route {
        Route {
            path: self.path.clone(),
            method: self.method.clone(),
            handler: self.handler,
            matcher: self.matcher.clone(),
            variables: self.variables.clone()
        }
    }
}

/// A RouteResult is what the router returns when `match_route` is called.
/// It contains the matched `route` and also a `params` property holding
/// a HashMap with the keys being the variable names and the value being the
/// evaluated string
pub struct RouteResult<'a> {
    pub route: &'a Route,
    pub params: HashMap<String, String>
}

/// The PathUtils collects some small helper methods that operate on the path
struct PathUtils;

// matches named variables (e.g. :userid)
static REGEX_VAR_SEQ: Regex                 = regex!(r":([,a-zA-Z0-9_-]*)");
static VAR_SEQ:&'static str                 = "[,a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_SLASH:&'static str      = "[,/a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_CAPTURE:&'static str    = "([,a-zA-Z0-9%_-]*)";
// matches request params (e.g. ?foo=true&bar=false)
static REGEX_PARAM_SEQ:&'static str         = "(\\?[a-zA-Z0-9%_=&-]*)?";
static REGEX_START:&'static str             = "^";
static REGEX_END:&'static str               = "$";

impl PathUtils {
    fn create_regex (route_path: &str) -> Regex {

        let updated_path = route_path.to_string()
                                     // first mark all double wildcards for replacement. We can't directly replace them
                                     // since the replacement does contain the * symbol as well, which would get overwritten
                                     // with the next replace call
                                     .replace("**", "___DOUBLE_WILDCARD___")
                                     // then replace the regular wildcard symbols (*) with the appropriate regex
                                     .replace("*", VAR_SEQ)
                                     // now replace the previously marked double wild cards (**)
                                     .replace("___DOUBLE_WILDCARD___", VAR_SEQ_WITH_SLASH);

        // then replace the variable symbols (:variable) with the appropriate regex
        let result = REGEX_START.to_string()
                                .append(REGEX_VAR_SEQ.replace_all(updated_path.as_slice(), VAR_SEQ_WITH_CAPTURE).as_slice())
                                .append(REGEX_PARAM_SEQ)
                                .append(REGEX_END);

        match Regex::new(result.as_slice()) {
            Ok(re) => re,
            Err(err) => fail!("{}", err)
        }
    }

    fn get_variable_info (route_path: &str) -> HashMap<String, uint> {
        REGEX_VAR_SEQ.captures_iter(route_path)
             .enumerate()
             .map(|(i, matched)| (matched.at(1).to_string(), i))
             .collect()
    }
}

/// The Router's job is it to hold routes and to resolve them later against
/// concrete URLs. The router is also a regular middleware and needs to be
/// added to the middleware stack with `server.utilize(router)`.

#[deriving(Clone)]
pub struct Router{
    routes: Vec<Route>,
}

impl Router {
    pub fn new () -> Router {
        Router {
            routes: Vec::new()
        }
    }

    /// Registers a handler to be used for a specific GET request.
    /// Handlers are assigned to paths and paths are allowed to contain
    /// variables and wildcards.
    ///
    /// # Example without variables and wildcards
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user");
    /// };
    /// router.get("/user", handler);
    /// ```
    /// # Example with variables
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     let text = format!("This is user: {}", request.params.get(&"userid".to_string()));
    ///     response.send(text.as_slice());
    /// };
    /// router.get("/user/:userid", handler);
    /// ```
    /// # Example with simple wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 but not /user/extended/list/4711");
    /// };
    /// router.get("/user/*/:userid", handler);
    /// ```
    /// # Example with double wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 and also /user/extended/list/4711");
    /// };
    /// router.get("/user/**/:userid", handler);
    /// ```
    pub fn get(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Get, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific POST request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/post/request");
    /// };
    /// router.post("/a/post/request", handler);
    /// ```
    /// Take a look at `get()` for a more detailed description.
    pub fn post(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Post, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific PUT request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/put/request");
    /// };
    /// router.put("/a/put/request", handler);
    /// ```
    /// Take a look at `get(..)` for a more detailed description.
    pub fn put(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Put, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific DELETE request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a DELETE request to /a/delete/request");
    /// };
    /// router.delete("/a/delete/request", handler);
    /// ```
    /// Take a look at `get(...)` for a more detailed description.
    pub fn delete(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Delete, String::from_str(uri), handler);
    }

    pub fn add_route (&mut self, method: Method, path: String, handler: fn(request: &Request, response: &mut Response)) -> () {
        let matcher = PathUtils::create_regex(path.as_slice());
        let variable_infos = PathUtils::get_variable_info(path.as_slice());
        let route = Route {
            path: path,
            method: method,
            matcher: matcher,
            handler: handler,
            variables: variable_infos
        };
        self.routes.push(route);
    }

    pub fn match_route(&self, method: Method, path: String) -> Option<RouteResult> {
        self.routes.iter().find(|item| item.method == method && item.matcher.is_match(path.as_slice()))
            .map(|route| {
                let map = match route.matcher.captures(path.as_slice()) {
                    Some(captures) => {
                        route.variables.iter().map(|(name, pos)| {
                            (name.to_string(), captures.at(pos + 1).to_string())
                        }).collect()
                    },
                    None => HashMap::new(),
                };
                RouteResult {
                    route: route,
                    params: map
                }
            })
    }
}

impl Middleware for Router {
    fn invoke (&self, req: &mut Request, res: &mut Response) -> Result<Action, NickelError> {
        match &req.origin.request_uri {
            &AbsolutePath(ref url) => {
                match self.match_route(req.origin.method.clone(), url.clone()) {
                    Some(route_result) => {
                        res.origin.status = ::http::status::Ok;
                        req.params = route_result.params.clone();
                        (route_result.route.handler)(req, res);
                        Ok(Halt)
                    },
                    None => Ok(Continue)
                }
            },
            _ => Ok(Continue)
        }
    }
}

#[test]
fn creates_map_with_var_variable_infos () {
    let map = PathUtils::get_variable_info("foo/:uid/bar/:groupid");

    assert_eq!(map.len(), 2);
    assert_eq!(map.get(&"uid".to_string()), &0);
    assert_eq!(map.get(&"groupid".to_string()), &1);
}

#[test]
fn creates_regex_with_captures () {
    let regex = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let caps = regex.captures("foo/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/*/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/**/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/another/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");
}

#[test]
fn creates_valid_regex_for_routes () {
    let regex1 = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let regex2 = PathUtils::create_regex("foo/*/bar");
    let regex3 = PathUtils::create_regex("foo/**/bar");

    assert_eq!(regex1.is_match("foo/4711/bar/5490"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=true&bar=false"), false);
    assert_eq!(regex1.is_match("foo/4711/bar/test%20spacing"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5281?foo=test%20spacing&bar=false"), true);

    assert_eq!(regex2.is_match("foo/4711/bar"), true);
    assert_eq!(regex2.is_match("foo/4711/barr"), false);
    assert_eq!(regex2.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex2.is_match("foo/4711/4712/bar"), false);
    assert_eq!(regex2.is_match("foo/4711/4712/bar?foo=true&bar=false"), false);

    assert_eq!(regex3.is_match("foo/4711/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar?foo=true&bar=false"), true);

    //ensure that this works with commas too
    assert_eq!(regex1.is_match("foo/4711/bar/5490,1234"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490,1234?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=1,2,3&bar=false"), false);
}

#[test]
fn can_match_var_routes () {
    let route_store = &mut Router::new();

    fn handler (_request: &Request, response: &mut Response) -> () {
        let _ = response.origin.write("hello from foo".as_bytes());
    };

    route_store.add_route(method::Get, "/foo/:userid".to_string(), handler);
    route_store.add_route(method::Get, "/bar".to_string(), handler);

    let route_result = route_store.match_route(method::Get, "/foo/4711".to_string()).unwrap();
    let route = route_result.route;

    assert_eq!(route_result.params.get(&"userid".to_string()), &"4711".to_string());

    //assert the route has identified the variable
    assert_eq!(route.variables.len(), 1);
    assert_eq!(route.variables.get(&"userid".to_string()), &0);

    let route_result = route_store.match_route(method::Get, "/bar/4711".to_string());
    assert!(route_result.is_none());

    let route_result = route_store.match_route(method::Get, "/foo".to_string());
    assert!(route_result.is_none());

    //ensure that this will work with commas too
    let route_result = route_store.match_route(method::Get, "/foo/123,456".to_string());
    assert!(route_result.is_some());

    let route_result = route_result.unwrap();
    assert_eq!(route_result.params.get(&"userid".to_string()), &"123,456".to_string());

    //ensure that this will work with spacing too
    let route_result = route_store.match_route(method::Get, "/foo/John%20Doe".to_string());
    assert!(route_result.is_some());

    let route_result = route_result.unwrap();
    assert_eq!(route_result.params.get(&"userid".to_string()), &"John%20Doe".to_string());
}
<commit_msg>Remove multiple appends.<commit_after>//!Router asigns handlers to paths and resolves them per request

#[cfg(test)]
use http::method;
use http::method::{ Method, Get, Post, Put, Delete };
use http::server::request::{AbsolutePath};
use regex::Regex;
use std::collections::hashmap::HashMap;
use request::Request;
use response::Response;
use middleware::{ Middleware, Action, Halt, Continue };
use nickel_error::NickelError;

/// A Route is the basic data structure that stores both the path
/// and the handler that gets executed for the route.
/// The path can contain variable pattern such as `user/:userid/invoices`
pub struct Route {
    pub path: String,
    pub method: Method,
    pub handler: fn(request: &Request, response: &mut Response),
    pub variables: HashMap<String, uint>,
    matcher: Regex
}

impl Clone for Route {
    fn clone(&self) -> Route {
        Route {
            path: self.path.clone(),
            method: self.method.clone(),
            handler: self.handler,
            matcher: self.matcher.clone(),
            variables: self.variables.clone()
        }
    }
}

/// A RouteResult is what the router returns when `match_route` is called.
/// It contains the matched `route` and also a `params` property holding
/// a HashMap with the keys being the variable names and the value being the
/// evaluated string
pub struct RouteResult<'a> {
    pub route: &'a Route,
    pub params: HashMap<String, String>
}

/// The PathUtils collects some small helper methods that operate on the path
struct PathUtils;

// matches named variables (e.g. :userid)
static REGEX_VAR_SEQ: Regex                 = regex!(r":([,a-zA-Z0-9_-]*)");
static VAR_SEQ:&'static str                 = "[,a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_SLASH:&'static str      = "[,/a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_CAPTURE:&'static str    = "([,a-zA-Z0-9%_-]*)";
// matches request params (e.g. ?foo=true&bar=false)
static REGEX_PARAM_SEQ:&'static str         = "(\\?[a-zA-Z0-9%_=&-]*)?";
static REGEX_START:&'static str             = "^";
static REGEX_END:&'static str               = "$";

impl PathUtils {
    fn create_regex (route_path: &str) -> Regex {

        let updated_path = route_path.to_string()
                                     // first mark all double wildcards for replacement. We can't directly replace them
                                     // since the replacement does contain the * symbol as well, which would get overwritten
                                     // with the next replace call
                                     .replace("**", "___DOUBLE_WILDCARD___")
                                     // then replace the regular wildcard symbols (*) with the appropriate regex
                                     .replace("*", VAR_SEQ)
                                     // now replace the previously marked double wild cards (**)
                                     .replace("___DOUBLE_WILDCARD___", VAR_SEQ_WITH_SLASH);

        // then replace the variable symbols (:variable) with the appropriate regex
        let result = [REGEX_START,
                      REGEX_VAR_SEQ.replace_all(updated_path.as_slice(),
                                                VAR_SEQ_WITH_CAPTURE)
                                   .as_slice(),
                      REGEX_PARAM_SEQ,
                      REGEX_END].concat();

        match Regex::new(result.as_slice()) {
            Ok(re) => re,
            Err(err) => fail!("{}", err)
        }
    }

    fn get_variable_info (route_path: &str) -> HashMap<String, uint> {
        REGEX_VAR_SEQ.captures_iter(route_path)
             .enumerate()
             .map(|(i, matched)| (matched.at(1).to_string(), i))
             .collect()
    }
}

/// The Router's job is it to hold routes and to resolve them later against
/// concrete URLs. The router is also a regular middleware and needs to be
/// added to the middleware stack with `server.utilize(router)`.

#[deriving(Clone)]
pub struct Router{
    routes: Vec<Route>,
}

impl Router {
    pub fn new () -> Router {
        Router {
            routes: Vec::new()
        }
    }

    /// Registers a handler to be used for a specific GET request.
    /// Handlers are assigned to paths and paths are allowed to contain
    /// variables and wildcards.
    ///
    /// # Example without variables and wildcards
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user");
    /// };
    /// router.get("/user", handler);
    /// ```
    /// # Example with variables
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     let text = format!("This is user: {}", request.params.get(&"userid".to_string()));
    ///     response.send(text.as_slice());
    /// };
    /// router.get("/user/:userid", handler);
    /// ```
    /// # Example with simple wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 but not /user/extended/list/4711");
    /// };
    /// router.get("/user/*/:userid", handler);
    /// ```
    /// # Example with double wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 and also /user/extended/list/4711");
    /// };
    /// router.get("/user/**/:userid", handler);
    /// ```
    pub fn get(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Get, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific POST request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/post/request");
    /// };
    /// router.post("/a/post/request", handler);
    /// ```
    /// Take a look at `get()` for a more detailed description.
    pub fn post(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Post, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific PUT request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/put/request");
    /// };
    /// router.put("/a/put/request", handler);
    /// ```
    /// Take a look at `get(..)` for a more detailed description.
    pub fn put(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Put, String::from_str(uri), handler);
    }

    /// Registers a handler to be used for a specific DELETE request.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches a DELETE request to /a/delete/request");
    /// };
    /// router.delete("/a/delete/request", handler);
    /// ```
    /// Take a look at `get(...)` for a more detailed description.
    pub fn delete(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.add_route(Delete, String::from_str(uri), handler);
    }

    pub fn add_route (&mut self, method: Method, path: String, handler: fn(request: &Request, response: &mut Response)) -> () {
        let matcher = PathUtils::create_regex(path.as_slice());
        let variable_infos = PathUtils::get_variable_info(path.as_slice());
        let route = Route {
            path: path,
            method: method,
            matcher: matcher,
            handler: handler,
            variables: variable_infos
        };
        self.routes.push(route);
    }

    pub fn match_route(&self, method: Method, path: String) -> Option<RouteResult> {
        self.routes.iter().find(|item| item.method == method && item.matcher.is_match(path.as_slice()))
            .map(|route| {
                let map = match route.matcher.captures(path.as_slice()) {
                    Some(captures) => {
                        route.variables.iter().map(|(name, pos)| {
                            (name.to_string(), captures.at(pos + 1).to_string())
                        }).collect()
                    },
                    None => HashMap::new(),
                };
                RouteResult {
                    route: route,
                    params: map
                }
            })
    }
}

impl Middleware for Router {
    fn invoke (&self, req: &mut Request, res: &mut Response) -> Result<Action, NickelError> {
        match &req.origin.request_uri {
            &AbsolutePath(ref url) => {
                match self.match_route(req.origin.method.clone(), url.clone()) {
                    Some(route_result) => {
                        res.origin.status = ::http::status::Ok;
                        req.params = route_result.params.clone();
                        (route_result.route.handler)(req, res);
                        Ok(Halt)
                    },
                    None => Ok(Continue)
                }
            },
            _ => Ok(Continue)
        }
    }
}

#[test]
fn creates_map_with_var_variable_infos () {
    let map = PathUtils::get_variable_info("foo/:uid/bar/:groupid");

    assert_eq!(map.len(), 2);
    assert_eq!(map.get(&"uid".to_string()), &0);
    assert_eq!(map.get(&"groupid".to_string()), &1);
}

#[test]
fn creates_regex_with_captures () {
    let regex = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let caps = regex.captures("foo/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/*/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/**/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/another/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");
}

#[test]
fn creates_valid_regex_for_routes () {
    let regex1 = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let regex2 = PathUtils::create_regex("foo/*/bar");
    let regex3 = PathUtils::create_regex("foo/**/bar");

    assert_eq!(regex1.is_match("foo/4711/bar/5490"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=true&bar=false"), false);
    assert_eq!(regex1.is_match("foo/4711/bar/test%20spacing"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5281?foo=test%20spacing&bar=false"), true);

    assert_eq!(regex2.is_match("foo/4711/bar"), true);
    assert_eq!(regex2.is_match("foo/4711/barr"), false);
    assert_eq!(regex2.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex2.is_match("foo/4711/4712/bar"), false);
    assert_eq!(regex2.is_match("foo/4711/4712/bar?foo=true&bar=false"), false);

    assert_eq!(regex3.is_match("foo/4711/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar?foo=true&bar=false"), true);

    //ensure that this works with commas too
    assert_eq!(regex1.is_match("foo/4711/bar/5490,1234"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490,1234?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=1,2,3&bar=false"), false);
}

#[test]
fn can_match_var_routes () {
    let route_store = &mut Router::new();

    fn handler (_request: &Request, response: &mut Response) -> () {
        let _ = response.origin.write("hello from foo".as_bytes());
    };

    route_store.add_route(method::Get, "/foo/:userid".to_string(), handler);
    route_store.add_route(method::Get, "/bar".to_string(), handler);

    let route_result = route_store.match_route(method::Get, "/foo/4711".to_string()).unwrap();
    let route = route_result.route;

    assert_eq!(route_result.params.get(&"userid".to_string()), &"4711".to_string());

    //assert the route has identified the variable
    assert_eq!(route.variables.len(), 1);
    assert_eq!(route.variables.get(&"userid".to_string()), &0);

    let route_result = route_store.match_route(method::Get, "/bar/4711".to_string());
    assert!(route_result.is_none());

    let route_result = route_store.match_route(method::Get, "/foo".to_string());
    assert!(route_result.is_none());

    //ensure that this will work with commas too
    let route_result = route_store.match_route(method::Get, "/foo/123,456".to_string());
    assert!(route_result.is_some());

    let route_result = route_result.unwrap();
    assert_eq!(route_result.params.get(&"userid".to_string()), &"123,456".to_string());

    //ensure that this will work with spacing too
    let route_result = route_store.match_route(method::Get, "/foo/John%20Doe".to_string());
    assert!(route_result.is_some());

    let route_result = route_result.unwrap();
    assert_eq!(route_result.params.get(&"userid".to_string()), &"John%20Doe".to_string());
}
<|endoftext|>


<commit_before>// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.

use std::time::Duration;

use super::config::{ReadableDuration, ReadableSize};
use super::time::Limiter;
use super::timer::GLOBAL_TIMER_HANDLE;

use cpu_time::ThreadTime;
use futures::compat::Future01CompatExt;

const CPU_LIMITER_REFILL_DURATION: Duration = Duration::from_millis(50);

// Quota limiter allows users to obtain stable performance by increasing the
// completion time of tasks through restrictions of different metrics.
#[derive(Debug)]
pub struct QuotaLimiter {
    cputime_limiter: Limiter,
    write_bandwidth_limiter: Limiter,
    read_bandwidth_limiter: Limiter,
    max_delay_duration: Duration,
}

// Throttle must be consumed in quota limiter.
pub struct Sample {
    read_bytes: usize,
    write_bytes: usize,
    cpu_time: Duration,
    enable_cpu_limit: bool,
}

impl<'a> Sample {
    pub fn add_read_bytes(&mut self, bytes: usize) {
        self.read_bytes += bytes;
    }

    pub fn add_write_bytes(&mut self, bytes: usize) {
        self.write_bytes += bytes;
    }

    // Record the cpu time in the lifetime. Use this function inside code block.
    // If `cputime_limiter` is not enabled, guard will do nothing when dropped.
    pub fn observe_cpu(&'a mut self) -> CpuObserveGuard<'a> {
        if self.enable_cpu_limit {
            CpuObserveGuard {
                timer: Some(ThreadTime::now()),
                sample: self,
            }
        } else {
            CpuObserveGuard {
                timer: None,
                sample: self,
            }
        }
    }

    fn add_cpu_time(&mut self, time: Duration) {
        self.cpu_time += time;
    }
}

pub struct CpuObserveGuard<'a> {
    timer: Option<ThreadTime>,
    sample: &'a mut Sample,
}

impl<'a> Drop for CpuObserveGuard<'a> {
    fn drop(&mut self) {
        if let Some(timer) = self.timer {
            self.sample.add_cpu_time(timer.elapsed());
        }
    }
}

impl Default for QuotaLimiter {
    fn default() -> Self {
        Self {
            cputime_limiter: Limiter::new(f64::INFINITY),
            write_bandwidth_limiter: Limiter::new(f64::INFINITY),
            read_bandwidth_limiter: Limiter::new(f64::INFINITY),
            max_delay_duration: Duration::ZERO,
        }
    }
}

impl QuotaLimiter {
    // 1000 millicpu equals to 1vCPU, 0 means unlimited
    pub fn new(
        cpu_quota: usize,
        write_bandwidth: ReadableSize,
        read_bandwidth: ReadableSize,
        max_delay_duration: ReadableDuration,
    ) -> Self {
        let cputime_limiter = if cpu_quota == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::builder(cpu_quota as f64 * 1000_f64)
                .refill(CPU_LIMITER_REFILL_DURATION)
                .build()
        };

        let write_bandwidth_limiter = if write_bandwidth.0 == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::new(write_bandwidth.0 as f64)
        };

        let read_bandwidth_limiter = if read_bandwidth.0 == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::new(read_bandwidth.0 as f64)
        };

        let max_delay_duration = max_delay_duration.0;

        Self {
            cputime_limiter,
            write_bandwidth_limiter,
            read_bandwidth_limiter,
            max_delay_duration,
        }
    }

    // To generate a sampler.
    pub fn new_sample(&self) -> Sample {
        Sample {
            read_bytes: 0,
            write_bytes: 0,
            cpu_time: Duration::ZERO,
            enable_cpu_limit: !self.cputime_limiter.speed_limit().is_infinite(),
        }
    }

    // To consume a sampler and return delayed duration.
    // If the sampler is null, the speed limiter will just return ZERO.
    pub async fn async_consume(&self, sample: Sample) -> Duration {
        let cpu_dur = if sample.cpu_time > Duration::ZERO {
            self.cputime_limiter
                .consume_duration(sample.cpu_time.as_micros() as usize)
        } else {
            Duration::ZERO
        };

        let w_bw_dur = if sample.write_bytes > 0 {
            self.write_bandwidth_limiter
                .consume_duration(sample.write_bytes)
        } else {
            Duration::ZERO
        };

        let r_bw_dur = if sample.read_bytes > 0 {
            self.read_bandwidth_limiter
                .consume_duration(sample.read_bytes)
        } else {
            Duration::ZERO
        };

        let mut exec_delay = std::cmp::max(cpu_dur, std::cmp::max(w_bw_dur, r_bw_dur));
        if !self.max_delay_duration.is_zero() {
            exec_delay = std::cmp::min(self.max_delay_duration, exec_delay);
        };

        if !exec_delay.is_zero() {
            GLOBAL_TIMER_HANDLE
                .delay(std::time::Instant::now() + exec_delay)
                .compat()
                .await
                .unwrap();
        }

        exec_delay
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::executor::block_on;

    #[test]
    fn test_quota_limiter() {
        // refill duration = 50ms
        // bucket capacity = 50
        let quota_limiter = QuotaLimiter::new(
            1000,
            ReadableSize::kb(1),
            ReadableSize::kb(1),
            ReadableDuration::millis(0),
        );

        let thread_start_time = ThreadTime::now();

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(20));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::ZERO);

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(30));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::from_millis(50));

        std::thread::sleep(Duration::from_millis(10));

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(30));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        // should less 20+30+30
        assert!(should_delay < Duration::from_millis(80));

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(200));
        sample.add_write_bytes(256);
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::from_millis(250));

        // ThreadTime elapsed time is not long.
        assert!(thread_start_time.elapsed() < Duration::from_millis(50));
    }
}
<commit_msg>Quota limiter: increase refill duration to avoid much context switch cost (#12264)<commit_after>// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.

use std::time::Duration;

use super::config::{ReadableDuration, ReadableSize};
use super::time::Limiter;
use super::timer::GLOBAL_TIMER_HANDLE;

use cpu_time::ThreadTime;
use futures::compat::Future01CompatExt;

// TODO: This value is fixed based on experience of AWS 4vCPU TPC-C bench test.
// It's better to use a universal approach.
const CPU_LIMITER_REFILL_DURATION: Duration = Duration::from_millis(100);

// Quota limiter allows users to obtain stable performance by increasing the
// completion time of tasks through restrictions of different metrics.
#[derive(Debug)]
pub struct QuotaLimiter {
    cputime_limiter: Limiter,
    write_bandwidth_limiter: Limiter,
    read_bandwidth_limiter: Limiter,
    max_delay_duration: Duration,
}

// Throttle must be consumed in quota limiter.
pub struct Sample {
    read_bytes: usize,
    write_bytes: usize,
    cpu_time: Duration,
    enable_cpu_limit: bool,
}

impl<'a> Sample {
    pub fn add_read_bytes(&mut self, bytes: usize) {
        self.read_bytes += bytes;
    }

    pub fn add_write_bytes(&mut self, bytes: usize) {
        self.write_bytes += bytes;
    }

    // Record the cpu time in the lifetime. Use this function inside code block.
    // If `cputime_limiter` is not enabled, guard will do nothing when dropped.
    pub fn observe_cpu(&'a mut self) -> CpuObserveGuard<'a> {
        if self.enable_cpu_limit {
            CpuObserveGuard {
                timer: Some(ThreadTime::now()),
                sample: self,
            }
        } else {
            CpuObserveGuard {
                timer: None,
                sample: self,
            }
        }
    }

    fn add_cpu_time(&mut self, time: Duration) {
        self.cpu_time += time;
    }
}

pub struct CpuObserveGuard<'a> {
    timer: Option<ThreadTime>,
    sample: &'a mut Sample,
}

impl<'a> Drop for CpuObserveGuard<'a> {
    fn drop(&mut self) {
        if let Some(timer) = self.timer {
            self.sample.add_cpu_time(timer.elapsed());
        }
    }
}

impl Default for QuotaLimiter {
    fn default() -> Self {
        Self {
            cputime_limiter: Limiter::new(f64::INFINITY),
            write_bandwidth_limiter: Limiter::new(f64::INFINITY),
            read_bandwidth_limiter: Limiter::new(f64::INFINITY),
            max_delay_duration: Duration::ZERO,
        }
    }
}

impl QuotaLimiter {
    // 1000 millicpu equals to 1vCPU, 0 means unlimited
    pub fn new(
        cpu_quota: usize,
        write_bandwidth: ReadableSize,
        read_bandwidth: ReadableSize,
        max_delay_duration: ReadableDuration,
    ) -> Self {
        let cputime_limiter = if cpu_quota == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::builder(cpu_quota as f64 * 1000_f64)
                .refill(CPU_LIMITER_REFILL_DURATION)
                .build()
        };

        let write_bandwidth_limiter = if write_bandwidth.0 == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::new(write_bandwidth.0 as f64)
        };

        let read_bandwidth_limiter = if read_bandwidth.0 == 0 {
            Limiter::new(f64::INFINITY)
        } else {
            Limiter::new(read_bandwidth.0 as f64)
        };

        let max_delay_duration = max_delay_duration.0;

        Self {
            cputime_limiter,
            write_bandwidth_limiter,
            read_bandwidth_limiter,
            max_delay_duration,
        }
    }

    // To generate a sampler.
    pub fn new_sample(&self) -> Sample {
        Sample {
            read_bytes: 0,
            write_bytes: 0,
            cpu_time: Duration::ZERO,
            enable_cpu_limit: !self.cputime_limiter.speed_limit().is_infinite(),
        }
    }

    // To consume a sampler and return delayed duration.
    // If the sampler is null, the speed limiter will just return ZERO.
    pub async fn async_consume(&self, sample: Sample) -> Duration {
        let cpu_dur = if sample.cpu_time > Duration::ZERO {
            self.cputime_limiter
                .consume_duration(sample.cpu_time.as_micros() as usize)
        } else {
            Duration::ZERO
        };

        let w_bw_dur = if sample.write_bytes > 0 {
            self.write_bandwidth_limiter
                .consume_duration(sample.write_bytes)
        } else {
            Duration::ZERO
        };

        let r_bw_dur = if sample.read_bytes > 0 {
            self.read_bandwidth_limiter
                .consume_duration(sample.read_bytes)
        } else {
            Duration::ZERO
        };

        let mut exec_delay = std::cmp::max(cpu_dur, std::cmp::max(w_bw_dur, r_bw_dur));
        if !self.max_delay_duration.is_zero() {
            exec_delay = std::cmp::min(self.max_delay_duration, exec_delay);
        };

        if !exec_delay.is_zero() {
            GLOBAL_TIMER_HANDLE
                .delay(std::time::Instant::now() + exec_delay)
                .compat()
                .await
                .unwrap();
        }

        exec_delay
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::executor::block_on;

    #[test]
    fn test_quota_limiter() {
        // refill duration = 100ms
        // bucket capacity = 100
        let quota_limiter = QuotaLimiter::new(
            1000,
            ReadableSize::kb(1),
            ReadableSize::kb(1),
            ReadableDuration::millis(0),
        );

        let thread_start_time = ThreadTime::now();

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(60));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::ZERO);

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(40));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::from_millis(100));

        std::thread::sleep(Duration::from_millis(10));

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(30));
        let should_delay = block_on(quota_limiter.async_consume(sample));
        // should less 60+40+30
        assert!(should_delay < Duration::from_millis(130));

        let mut sample = quota_limiter.new_sample();
        sample.add_cpu_time(Duration::from_millis(200));
        sample.add_write_bytes(256);
        let should_delay = block_on(quota_limiter.async_consume(sample));
        assert_eq!(should_delay, Duration::from_millis(250));

        // ThreadTime elapsed time is not long.
        assert!(thread_start_time.elapsed() < Duration::from_millis(50));
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<i32>, tab: Option<i32>) -> Result<String, String> {
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get tab {} of window {}", t_idx, w_idx),
        None => format!("tell application \"Safari\" to get URL of document {}", w_idx),
      }
    },
    None => "tell application \"Safari\" to get URL of document 1".to_string(),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Result<Vec<String>, String> {
  let script = include_str!("scripts/list-open-tabs.scpt");
  let output = run_applescript(&script);
  if output.status.success() {
    Ok(output.stdout
      .trim()
      .split(", ")
      .map(|url| urls::tidy_url(url))
      .filter(|url| url != "favorites://")
      .filter(|url| url != "://missing value")
      .collect()
    )
  } else {
    error!("Unexpected error from osascript: {:?}", output.stderr);
  }
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}
<commit_msg>Split out the parsing code for easier testing<commit_after>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<i32>, tab: Option<i32>) -> Result<String, String> {
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get tab {} of window {}", t_idx, w_idx),
        None => format!("tell application \"Safari\" to get URL of document {}", w_idx),
      }
    },
    None => "tell application \"Safari\" to get URL of document 1".to_string(),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Result<Vec<String>, String> {
  let script = include_str!("scripts/list-open-tabs.scpt");
  let output = run_applescript(&script);
  if output.status.success() {
    Ok(parse_list_open_tabs_output(&output.stdout))
  } else {
    error!("Unexpected error from osascript: {:?}", output.stderr);
  }
}

/// Given AppleScript output from list-open-tabs.scpt, return a list of URLs.
fn parse_list_open_tabs_output(stdout: &str) -> Vec<String> {
  stdout
    .trim()
    .split(", ")
    .map(|url| urls::tidy_url(url))
    .filter(|url| url != "favorites://")
    .filter(|url| url != "://missing value")
    .collect()
}

macro_rules! parse_list_open_tabs_output_tests {
  ($($name:ident: $value:expr,)*) => {
    $(
      #[test]
      fn $name() {
        let (input, expected) = $value;
        assert_eq!(expected, parse_list_open_tabs_output(input));
      }
    )*
  }
}

parse_list_open_tabs_output_tests! {

  // I don't have a good way to put an empty vector on the right-hand side of
  // this comparison, so no test checks for nothing (yet).

  favorites_with_url: (
    "favorites://, http://example.org",
    vec!["http://example.org"]
  ),

  url_with_missing_value: (
    "://missing value, https://www.example.net",
    vec!["https://www.example.net"]
  ),

  single_url: (
    "http://foo_bar.com",
    vec!["http://foo_bar.com"]
  ),

  multiple_urls: (
    "http://example.org, https://www.example.net, http://test.co.uk",
    vec!["http://example.org", "https://www.example.net", "http://test.co.uk"]
  ),

  with_extra_whitespace: (
    "    http://space.org, https://www.nasa.gov   ",
    vec!["http://space.org", "https://www.nasa.gov"]
  ),

  applies_tidy_url_transform: (
    "https://mobile.twitter.com",
    vec!["https://twitter.com"]
  ),
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}
<|endoftext|>


<commit_before>use nom::{alphanumeric, multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use caseless_tag::*;
use parser::ConditionExpression;
use parser::{fieldlist, unsigned_number, statement_terminator};

use condition::*;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<String>,
    having: String, // XXX(malte): should this be an arbitrary expr?
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
enum OrderType {
    //   OrderAscending,
    //   OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    order_type: OrderType,
    order_cols: Vec<String>, // TODO(malte): can this be an arbitrary expr?
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    table: String,
    distinct: bool,
    fields: Vec<String>,
    where_clause: Option<ConditionExpression>,
    group_by: Option<GroupByClause>,
    order: Option<OrderClause>,
    limit: Option<LimitClause>,
}

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    chain!(
        space ~
        caseless_tag!("limit") ~
        space ~
        limit_val: unsigned_number ~
        space? ~
        offset_val: opt!(chain!(
                caseless_tag!("offset") ~
                space ~
                val: unsigned_number,
                || { val }
            )
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    })
);

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    dbg_dmp!(chain!(
        space ~
        caseless_tag!("where") ~
        space ~
        cond: condition_expr,
        || { cond }
    ))
);

named!(table_reference<&[u8], &str>,
    chain!(
        table: map_res!(alphanumeric, str::from_utf8) ~
        opt!(
            chain!(
                space ~
                caseless_tag!("as") ~
                space ~
                alias: map_res!(alphanumeric, str::from_utf8),
                || { println!("got alias: {} -> {}", table, alias); alias }
            )
        ),
        || { table }
    )
);

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    dbg_dmp!(chain!(
        caseless_tag!("select") ~
        space ~
        distinct: opt!(caseless_tag!("distinct")) ~
        space? ~
        fields: fieldlist ~
        delimited!(multispace, caseless_tag!("from"), multispace) ~
        table: table_reference ~
        cond: opt!(where_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                table: String::from(table),
                distinct: distinct.is_some(),
                fields: fields.iter().map(|s| String::from(*s)).collect(),
                where_clause: cond,
                group_by: None,
                order: None,
                limit: limit,
            }
        }
    ))
);

#[cfg(test)]
mod tests {
    use super::*;
    use parser::{ConditionBase, ConditionExpression, ConditionTree};

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        assert_eq!(selection(qstring_sem.as_bytes()).unwrap(),
                   selection(qstring_linebreak.as_bytes()).unwrap());
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("email"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("ContactInfo"),
                       fields: vec!["ALL".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       table: String::from("PaperTag"),
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("PaperTag"),
                       distinct: true,
                       fields: vec!["tag".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: String::from("and"),
        }));
        println!("res: {:#?}", res);
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("PaperStorage"),
                       fields: vec!["infoJson".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }
}
<commit_msg>Make SelectStatement fields public.<commit_after>use nom::{alphanumeric, multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use caseless_tag::*;
use parser::ConditionExpression;
use parser::{fieldlist, unsigned_number, statement_terminator};

use condition::*;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<String>,
    having: String, // XXX(malte): should this be an arbitrary expr?
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
enum OrderType {
    //   OrderAscending,
    //   OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    order_type: OrderType,
    order_cols: Vec<String>, // TODO(malte): can this be an arbitrary expr?
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    pub table: String,
    pub distinct: bool,
    pub fields: Vec<String>,
    pub where_clause: Option<ConditionExpression>,
    pub group_by: Option<GroupByClause>,
    pub order: Option<OrderClause>,
    pub limit: Option<LimitClause>,
}

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    chain!(
        space ~
        caseless_tag!("limit") ~
        space ~
        limit_val: unsigned_number ~
        space? ~
        offset_val: opt!(chain!(
                caseless_tag!("offset") ~
                space ~
                val: unsigned_number,
                || { val }
            )
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    })
);

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    dbg_dmp!(chain!(
        space ~
        caseless_tag!("where") ~
        space ~
        cond: condition_expr,
        || { cond }
    ))
);

named!(table_reference<&[u8], &str>,
    chain!(
        table: map_res!(alphanumeric, str::from_utf8) ~
        opt!(
            chain!(
                space ~
                caseless_tag!("as") ~
                space ~
                alias: map_res!(alphanumeric, str::from_utf8),
                || { println!("got alias: {} -> {}", table, alias); alias }
            )
        ),
        || { table }
    )
);

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    dbg_dmp!(chain!(
        caseless_tag!("select") ~
        space ~
        distinct: opt!(caseless_tag!("distinct")) ~
        space? ~
        fields: fieldlist ~
        delimited!(multispace, caseless_tag!("from"), multispace) ~
        table: table_reference ~
        cond: opt!(where_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                table: String::from(table),
                distinct: distinct.is_some(),
                fields: fields.iter().map(|s| String::from(*s)).collect(),
                where_clause: cond,
                group_by: None,
                order: None,
                limit: limit,
            }
        }
    ))
);

#[cfg(test)]
mod tests {
    use super::*;
    use parser::{ConditionBase, ConditionExpression, ConditionTree};

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("users"),
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        assert_eq!(selection(qstring_sem.as_bytes()).unwrap(),
                   selection(qstring_linebreak.as_bytes()).unwrap());
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("email"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("ContactInfo"),
                       fields: vec!["ALL".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       table: String::from("PaperTag"),
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("PaperTag"),
                       distinct: true,
                       fields: vec!["tag".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Expr(ConditionBase::Field(String::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Expr(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: String::from("and"),
        }));
        println!("res: {:#?}", res);
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       table: String::from("PaperStorage"),
                       fields: vec!["infoJson".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }
}
<|endoftext|>