<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::ops::Deref;
use std::path::Path;
use std::path::PathBuf;

use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;
use std::path::Components;

use libimagerror::into::IntoError;

use error::StoreErrorKind as SEK;
use error::MapErrInto;
use store::Result;

/// The Index into the Store
#[derive(Debug, Clone, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    base: Option<PathBuf>,
    id:   PathBuf,
}

impl PartialEq for StoreId {
    fn eq(&self, other: &StoreId) -> bool {
        self.id == other.id
    }
}

impl StoreId {

    pub fn new(base: Option<PathBuf>, id: PathBuf) -> Result<StoreId> {
        StoreId::new_baseless(id).map(|mut sid| { sid.base = base; sid })
    }

    /// Try to create a StoreId object from a filesystem-absolute path.
    ///
    /// Automatically creates a StoreId object which has a `base` set to `store_part` if stripping
    /// the `store_part` from the `full_path` succeeded.
    ///
    /// Returns a `StoreErrorKind::StoreIdBuildFromFullPathError` if stripping failes.
    pub fn from_full_path<D>(store_part: &PathBuf, full_path: D) -> Result<StoreId>
        where D: Deref<Target = Path>
    {
        let p = try!(
            full_path.strip_prefix(store_part).map_err_into(SEK::StoreIdBuildFromFullPathError)
        );
        StoreId::new(Some(store_part.clone()), PathBuf::from(p))
    }

    pub fn new_baseless(id: PathBuf) -> Result<StoreId> {
        if id.is_absolute() {
            Err(SEK::StoreIdLocalPartAbsoluteError.into_error())
        } else {
            Ok(StoreId {
                base: None,
                id: id
            })
        }
    }

    pub fn without_base(mut self) -> StoreId {
        self.base = None;
        self
    }

    pub fn with_base(mut self, base: PathBuf) -> Self {
        self.base = Some(base);
        self
    }

    /// Transform the StoreId object into a PathBuf, error if the base of the StoreId is not
    /// specified.
    pub fn into_pathbuf(self) -> Result<PathBuf> {
        let mut base = try!(self.base.ok_or(SEK::StoreIdHasNoBaseError.into_error()));
        base.push(self.id);
        Ok(base)
    }

    pub fn exists(&self) -> Result<bool> {
        self.clone().into_pathbuf().map(|pb| pb.exists())
    }

    pub fn to_str(&self) -> Result<String> {
        self.base
            .as_ref()
            .cloned()
            .map(|mut base| { base.push(self.id.clone()); base })
            .unwrap_or_else(|| self.id.clone())
            .to_str()
            .map(String::from)
            .ok_or(SEK::StoreIdHandlingError.into_error())
    }

    /// Returns the components of the `id` part of the StoreId object.
    ///
    /// Can be used to check whether a StoreId points to an entry in a specific collection of
    /// StoreIds.
    pub fn components(&self) -> Components {
        self.id.components()
    }

    /// Get the _local_ part of a StoreId object, as in "the part from the store root to the entry".
    pub fn local(&self) -> &PathBuf {
        &self.id
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> Result<StoreId>;
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> Result<StoreId> {
        Ok(self)
    }
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> Result<StoreId> {
        StoreId::new_baseless(self)
    }
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;
            use $crate::store::Result;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(name);
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> Result<$crate::storeid::StoreId> {
                    StoreId::new(None, self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {
    use std::path::PathBuf;

    use storeid::StoreId;
    use storeid::IntoStoreId;
    use error::StoreErrorKind as SEK;

    module_entry_path_mod!("test");

    #[test]
    fn test_correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().unwrap().to_str().unwrap(), "test/test");
    }

    #[test]
    fn test_baseless_path() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());
        assert_eq!(id.unwrap(), StoreId {
            base: None,
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_base_path() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());
        assert_eq!(id.unwrap(), StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_adding_base_to_baseless_path() {
        let id = StoreId::new_baseless(PathBuf::from("test"));

        assert!(id.is_ok());
        let id = id.unwrap();

        assert_eq!(id, StoreId { base: None, id: PathBuf::from("test") });

        let id = id.with_base(PathBuf::from("/tmp/"));
        assert_eq!(id, StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_removing_base_from_base_path() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));

        assert!(id.is_ok());
        let id = id.unwrap();

        assert_eq!(id, StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });

        let id = id.without_base();
        assert_eq!(id, StoreId {
            base: None,
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_baseless_into_pathbuf_is_err() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());
        assert!(id.unwrap().into_pathbuf().is_err());
    }

    #[test]
    fn test_baseless_into_pathbuf_is_storeidhasnobaseerror() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());

        let pb = id.unwrap().into_pathbuf();
        assert!(pb.is_err());

        assert_eq!(pb.unwrap_err().err_type(), SEK::StoreIdHasNoBaseError);
    }

    #[test]
    fn test_basefull_into_pathbuf_is_ok() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());
        assert!(id.unwrap().into_pathbuf().is_ok());
    }

    #[test]
    fn test_basefull_into_pathbuf_is_correct() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());

        let pb = id.unwrap().into_pathbuf();
        assert!(pb.is_ok());

        assert_eq!(pb.unwrap(), PathBuf::from("/tmp/test"));
    }

}
<commit_msg>Add fn StoreId::local_push() for altering the local part of a StoreId<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::ops::Deref;
use std::path::Path;
use std::path::PathBuf;

use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;
use std::path::Components;

use libimagerror::into::IntoError;

use error::StoreErrorKind as SEK;
use error::MapErrInto;
use store::Result;

/// The Index into the Store
#[derive(Debug, Clone, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    base: Option<PathBuf>,
    id:   PathBuf,
}

impl PartialEq for StoreId {
    fn eq(&self, other: &StoreId) -> bool {
        self.id == other.id
    }
}

impl StoreId {

    pub fn new(base: Option<PathBuf>, id: PathBuf) -> Result<StoreId> {
        StoreId::new_baseless(id).map(|mut sid| { sid.base = base; sid })
    }

    /// Try to create a StoreId object from a filesystem-absolute path.
    ///
    /// Automatically creates a StoreId object which has a `base` set to `store_part` if stripping
    /// the `store_part` from the `full_path` succeeded.
    ///
    /// Returns a `StoreErrorKind::StoreIdBuildFromFullPathError` if stripping failes.
    pub fn from_full_path<D>(store_part: &PathBuf, full_path: D) -> Result<StoreId>
        where D: Deref<Target = Path>
    {
        let p = try!(
            full_path.strip_prefix(store_part).map_err_into(SEK::StoreIdBuildFromFullPathError)
        );
        StoreId::new(Some(store_part.clone()), PathBuf::from(p))
    }

    pub fn new_baseless(id: PathBuf) -> Result<StoreId> {
        if id.is_absolute() {
            Err(SEK::StoreIdLocalPartAbsoluteError.into_error())
        } else {
            Ok(StoreId {
                base: None,
                id: id
            })
        }
    }

    pub fn without_base(mut self) -> StoreId {
        self.base = None;
        self
    }

    pub fn with_base(mut self, base: PathBuf) -> Self {
        self.base = Some(base);
        self
    }

    /// Transform the StoreId object into a PathBuf, error if the base of the StoreId is not
    /// specified.
    pub fn into_pathbuf(self) -> Result<PathBuf> {
        let mut base = try!(self.base.ok_or(SEK::StoreIdHasNoBaseError.into_error()));
        base.push(self.id);
        Ok(base)
    }

    pub fn exists(&self) -> Result<bool> {
        self.clone().into_pathbuf().map(|pb| pb.exists())
    }

    pub fn to_str(&self) -> Result<String> {
        self.base
            .as_ref()
            .cloned()
            .map(|mut base| { base.push(self.id.clone()); base })
            .unwrap_or_else(|| self.id.clone())
            .to_str()
            .map(String::from)
            .ok_or(SEK::StoreIdHandlingError.into_error())
    }

    /// Returns the components of the `id` part of the StoreId object.
    ///
    /// Can be used to check whether a StoreId points to an entry in a specific collection of
    /// StoreIds.
    pub fn components(&self) -> Components {
        self.id.components()
    }

    /// Get the _local_ part of a StoreId object, as in "the part from the store root to the entry".
    pub fn local(&self) -> &PathBuf {
        &self.id
    }

    pub fn local_push<P: AsRef<Path>>(&mut self, path: P) {
        self.id.push(path)
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> Result<StoreId>;
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> Result<StoreId> {
        Ok(self)
    }
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> Result<StoreId> {
        StoreId::new_baseless(self)
    }
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;
            use $crate::store::Result;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(name);
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> Result<$crate::storeid::StoreId> {
                    StoreId::new(None, self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {
    use std::path::PathBuf;

    use storeid::StoreId;
    use storeid::IntoStoreId;
    use error::StoreErrorKind as SEK;

    module_entry_path_mod!("test");

    #[test]
    fn test_correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().unwrap().to_str().unwrap(), "test/test");
    }

    #[test]
    fn test_baseless_path() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());
        assert_eq!(id.unwrap(), StoreId {
            base: None,
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_base_path() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());
        assert_eq!(id.unwrap(), StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_adding_base_to_baseless_path() {
        let id = StoreId::new_baseless(PathBuf::from("test"));

        assert!(id.is_ok());
        let id = id.unwrap();

        assert_eq!(id, StoreId { base: None, id: PathBuf::from("test") });

        let id = id.with_base(PathBuf::from("/tmp/"));
        assert_eq!(id, StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_removing_base_from_base_path() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));

        assert!(id.is_ok());
        let id = id.unwrap();

        assert_eq!(id, StoreId {
            base: Some(PathBuf::from("/tmp/")),
            id: PathBuf::from("test")
        });

        let id = id.without_base();
        assert_eq!(id, StoreId {
            base: None,
            id: PathBuf::from("test")
        });
    }

    #[test]
    fn test_baseless_into_pathbuf_is_err() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());
        assert!(id.unwrap().into_pathbuf().is_err());
    }

    #[test]
    fn test_baseless_into_pathbuf_is_storeidhasnobaseerror() {
        let id = StoreId::new_baseless(PathBuf::from("test"));
        assert!(id.is_ok());

        let pb = id.unwrap().into_pathbuf();
        assert!(pb.is_err());

        assert_eq!(pb.unwrap_err().err_type(), SEK::StoreIdHasNoBaseError);
    }

    #[test]
    fn test_basefull_into_pathbuf_is_ok() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());
        assert!(id.unwrap().into_pathbuf().is_ok());
    }

    #[test]
    fn test_basefull_into_pathbuf_is_correct() {
        let id = StoreId::from_full_path(&PathBuf::from("/tmp/"), PathBuf::from("/tmp/test"));
        assert!(id.is_ok());

        let pb = id.unwrap().into_pathbuf();
        assert!(pb.is_ok());

        assert_eq!(pb.unwrap(), PathBuf::from("/tmp/test"));
    }

}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fmt;
use std::io::{BufRead, BufReader, Read, Write};

use client::Stats;
use error::MemcacheError;
use stream::Stream;
use value::{FromMemcacheValue, ToMemcacheValue};

#[derive(Default)]
pub struct Options {
    pub noreply: bool,
    pub exptime: u32,
    pub flags: u32,
}

enum StoreCommand {
    Set,
    Add,
    Replace,
    Append,
    Prepend,
}

impl fmt::Display for StoreCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StoreCommand::Set => write!(f, "set"),
            StoreCommand::Add => write!(f, "add"),
            StoreCommand::Replace => write!(f, "replace"),
            StoreCommand::Append => write!(f, "append"),
            StoreCommand::Prepend => write!(f, "prepend"),
        }
    }
}
pub struct AsciiProtocol<C: Read + Write + Sized> {
    pub reader: BufReader<C>,
}

impl AsciiProtocol<Stream> {
    fn store<V: ToMemcacheValue<Stream>>(
        &mut self,
        command: StoreCommand,
        key: &str,
        value: V,
        options: &Options,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }

        let mut header = format!(
            "{} {} {} {} {}",
            command,
            key,
            value.get_flags(),
            options.exptime,
            value.get_length()
        );
        if options.noreply {
            header += " noreply";
        }
        header += "\r\n";
        self.reader.get_mut().write(header.as_bytes())?;
        value.write_to(self.reader.get_mut())?;
        self.reader.get_mut().write(b"\r\n")?;
        self.reader.get_mut().flush()?;

        if options.noreply {
            return Ok(());
        }

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "STORED\r\n" {
            return Ok(());
        } else if s == "NOT_STORED\r\n" {
            return Ok(());
        } else {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
    }

    pub(super) fn version(&mut self) -> Result<String, MemcacheError> {
        self.reader.get_mut().write(b"version\r\n")?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if !s.starts_with("VERSION") {
            return Err(MemcacheError::ServerError(0));
        }
        let s = s.trim_start_matches("VERSION ");
        let s = s.trim_end_matches("\r\n");

        return Ok(s.to_string());
    }

    pub(super) fn flush(&mut self) -> Result<(), MemcacheError> {
        match self.reader.get_mut().write(b"flush_all\r\n") {
            Ok(_) => {}
            Err(err) => return Err(MemcacheError::from(err)),
        }
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s != "OK\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        return Ok(());
    }

    pub(super) fn flush_with_delay(&mut self, delay: u32) -> Result<(), MemcacheError> {
        write!(self.reader.get_mut(), "flush_all {}\r\n", delay)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s != "OK\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        return Ok(());
    }

    pub(super) fn get<V: FromMemcacheValue>(&mut self, key: &str) -> Result<Option<V>, MemcacheError> {
        write!(self.reader.get_mut(), "get {}\r\n", key)?;

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;

        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s.starts_with("END") {
            return Ok(None);
        } else if !s.starts_with("VALUE") {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
        if header.len() != 4 {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let key = header[1];
        if key != key {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        let flags = header[2].parse()?;
        let length = header[3].parse()?;

        let mut buffer = vec![0; length];
        self.reader.read_exact(buffer.as_mut_slice())?;

        // read the rest \r\n and END\r\n
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "END\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        return Ok(Some(FromMemcacheValue::from_memcache_value(buffer, flags)?));
    }

    pub(super) fn gets<V: FromMemcacheValue>(&mut self, keys: Vec<&str>) -> Result<HashMap<String, V>, MemcacheError> {
        write!(self.reader.get_mut(), "gets {}\r\n", keys.join(" "))?;

        let mut result: HashMap<String, V> = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("VALUE") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
            if header.len() != 5 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let key = header[1];
            if key != key {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
            let flags = header[2].parse()?;
            let length = header[3].parse()?;

            let mut buffer = vec![0; length];
            self.reader.read_exact(buffer.as_mut_slice())?;

            result.insert(key.to_string(), FromMemcacheValue::from_memcache_value(buffer, flags)?);

            // read the rest \r\n
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;
            if s != "\r\n" {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
        }

        return Ok(result);
    }

    pub(super) fn set<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Set, key, value, &options);
    }

    pub(super) fn add<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Add, key, value, &options);
    }

    pub(super) fn replace<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Replace, key, value, &options);
    }

    pub(super) fn append<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        return self.store(StoreCommand::Append, key, value, &Default::default());
    }

    pub(super) fn prepend<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        return self.store(StoreCommand::Prepend, key, value, &Default::default());
    }

    pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "delete {}\r\n", key)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "DELETED\r\n" {
            return Ok(true);
        } else if s == "NOT_FOUND\r\n" {
            return Ok(false);
        } else {
            return Err(MemcacheError::ClientError(String::from("invalid server response")));
        }
    }

    pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "incr {} {}\r\n", key, amount)?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "NOT_FOUND\r\n" {
            return Err(MemcacheError::from(1));
        } else {
            match s.trim_end_matches("\r\n").parse::<u64>() {
                Ok(n) => return Ok(n),
                Err(_) => return Err(MemcacheError::ClientError("invalid server response".into())),
            }
        }
    }

    pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "decr {} {}\r\n", key, amount)?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "NOT_FOUND\r\n" {
            return Err(MemcacheError::from(1));
        } else {
            match s.trim_end_matches("\r\n").parse::<u64>() {
                Ok(n) => return Ok(n),
                Err(_) => return Err(MemcacheError::ClientError("invalid server response".into())),
            }
        }
    }

    pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "touch {} {}\r\n", key, expiration)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "TOUCHED\r\n" {
            return Ok(true);
        } else if s == "NOT_FOUND\r\n" {
            return Ok(false);
        } else {
            return Err(MemcacheError::ClientError(String::from("invalid server response")));
        }
    }

    pub(super) fn stats(&mut self) -> Result<Stats, MemcacheError> {
        self.reader.get_mut().write(b"stats\r\n")?;
        self.reader.get_mut().flush()?;

        let mut result: Stats = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("STAT") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let stat: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
            if stat.len() < 3 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
            let key = stat[1];
            let value = s.trim_start_matches(format!("STAT {}", key).as_str());
            result.insert(key.into(), value.into());
        }

        return Ok(result);
    }
}

fn is_memcache_error(s: &str) -> bool {
    return s == "ERROR\r\n" || s.starts_with("CIENT_ERROR") || s.starts_with("SERVER_ERROR");
}
<commit_msg>:bug: Fix get command in ascii protocol validation<commit_after>use std::collections::HashMap;
use std::fmt;
use std::io::{BufRead, BufReader, Read, Write};

use client::Stats;
use error::MemcacheError;
use stream::Stream;
use value::{FromMemcacheValue, ToMemcacheValue};

#[derive(Default)]
pub struct Options {
    pub noreply: bool,
    pub exptime: u32,
    pub flags: u32,
}

enum StoreCommand {
    Set,
    Add,
    Replace,
    Append,
    Prepend,
}

impl fmt::Display for StoreCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StoreCommand::Set => write!(f, "set"),
            StoreCommand::Add => write!(f, "add"),
            StoreCommand::Replace => write!(f, "replace"),
            StoreCommand::Append => write!(f, "append"),
            StoreCommand::Prepend => write!(f, "prepend"),
        }
    }
}
pub struct AsciiProtocol<C: Read + Write + Sized> {
    pub reader: BufReader<C>,
}

impl AsciiProtocol<Stream> {
    fn store<V: ToMemcacheValue<Stream>>(
        &mut self,
        command: StoreCommand,
        key: &str,
        value: V,
        options: &Options,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }

        let mut header = format!(
            "{} {} {} {} {}",
            command,
            key,
            value.get_flags(),
            options.exptime,
            value.get_length()
        );
        if options.noreply {
            header += " noreply";
        }
        header += "\r\n";
        self.reader.get_mut().write(header.as_bytes())?;
        value.write_to(self.reader.get_mut())?;
        self.reader.get_mut().write(b"\r\n")?;
        self.reader.get_mut().flush()?;

        if options.noreply {
            return Ok(());
        }

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "STORED\r\n" {
            return Ok(());
        } else if s == "NOT_STORED\r\n" {
            return Ok(());
        } else {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
    }

    pub(super) fn version(&mut self) -> Result<String, MemcacheError> {
        self.reader.get_mut().write(b"version\r\n")?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if !s.starts_with("VERSION") {
            return Err(MemcacheError::ServerError(0));
        }
        let s = s.trim_start_matches("VERSION ");
        let s = s.trim_end_matches("\r\n");

        return Ok(s.to_string());
    }

    pub(super) fn flush(&mut self) -> Result<(), MemcacheError> {
        match self.reader.get_mut().write(b"flush_all\r\n") {
            Ok(_) => {}
            Err(err) => return Err(MemcacheError::from(err)),
        }
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s != "OK\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        return Ok(());
    }

    pub(super) fn flush_with_delay(&mut self, delay: u32) -> Result<(), MemcacheError> {
        write!(self.reader.get_mut(), "flush_all {}\r\n", delay)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s != "OK\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        return Ok(());
    }

    pub(super) fn get<V: FromMemcacheValue>(&mut self, key: &str) -> Result<Option<V>, MemcacheError> {
        write!(self.reader.get_mut(), "get {}\r\n", key)?;

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;

        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s.starts_with("END") {
            return Ok(None);
        } else if !s.starts_with("VALUE") {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
        if header.len() != 4 {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        if key != header[1] {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        let flags = header[2].parse()?;
        let length = header[3].parse()?;

        let mut buffer = vec![0; length];
        self.reader.read_exact(buffer.as_mut_slice())?;

        // read the rest \r\n and END\r\n
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "END\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        return Ok(Some(FromMemcacheValue::from_memcache_value(buffer, flags)?));
    }

    pub(super) fn gets<V: FromMemcacheValue>(&mut self, keys: Vec<&str>) -> Result<HashMap<String, V>, MemcacheError> {
        write!(self.reader.get_mut(), "gets {}\r\n", keys.join(" "))?;

        let mut result: HashMap<String, V> = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("VALUE") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
            if header.len() != 5 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let key = header[1];
            let flags = header[2].parse()?;
            let length = header[3].parse()?;

            let mut buffer = vec![0; length];
            self.reader.read_exact(buffer.as_mut_slice())?;

            result.insert(key.to_string(), FromMemcacheValue::from_memcache_value(buffer, flags)?);

            // read the rest \r\n
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;
            if s != "\r\n" {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
        }

        return Ok(result);
    }

    pub(super) fn set<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Set, key, value, &options);
    }

    pub(super) fn add<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Add, key, value, &options);
    }

    pub(super) fn replace<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options {
            exptime: expiration,
            ..Default::default()
        };
        return self.store(StoreCommand::Replace, key, value, &options);
    }

    pub(super) fn append<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        return self.store(StoreCommand::Append, key, value, &Default::default());
    }

    pub(super) fn prepend<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        return self.store(StoreCommand::Prepend, key, value, &Default::default());
    }

    pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "delete {}\r\n", key)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "DELETED\r\n" {
            return Ok(true);
        } else if s == "NOT_FOUND\r\n" {
            return Ok(false);
        } else {
            return Err(MemcacheError::ClientError(String::from("invalid server response")));
        }
    }

    pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "incr {} {}\r\n", key, amount)?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "NOT_FOUND\r\n" {
            return Err(MemcacheError::from(1));
        } else {
            match s.trim_end_matches("\r\n").parse::<u64>() {
                Ok(n) => return Ok(n),
                Err(_) => return Err(MemcacheError::ClientError("invalid server response".into())),
            }
        }
    }

    pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "decr {} {}\r\n", key, amount)?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "NOT_FOUND\r\n" {
            return Err(MemcacheError::from(1));
        } else {
            match s.trim_end_matches("\r\n").parse::<u64>() {
                Ok(n) => return Ok(n),
                Err(_) => return Err(MemcacheError::ClientError("invalid server response".into())),
            }
        }
    }

    pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        write!(self.reader.get_mut(), "touch {} {}\r\n", key, expiration)?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "TOUCHED\r\n" {
            return Ok(true);
        } else if s == "NOT_FOUND\r\n" {
            return Ok(false);
        } else {
            return Err(MemcacheError::ClientError(String::from("invalid server response")));
        }
    }

    pub(super) fn stats(&mut self) -> Result<Stats, MemcacheError> {
        self.reader.get_mut().write(b"stats\r\n")?;
        self.reader.get_mut().flush()?;

        let mut result: Stats = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("STAT") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let stat: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
            if stat.len() < 3 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
            let key = stat[1];
            let value = s.trim_start_matches(format!("STAT {}", key).as_str());
            result.insert(key.into(), value.into());
        }

        return Ok(result);
    }
}

fn is_memcache_error(s: &str) -> bool {
    return s == "ERROR\r\n" || s.starts_with("CIENT_ERROR") || s.starts_with("SERVER_ERROR");
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;
use std::kinds::marker;
use std::cell::Cell;
use green;

use {uvll, UvResult, Idle, Async, UvError};
use raw::{mod, Loop, Handle};
use queue::QueuePool;
use homing::HomeHandle;

tls!(local_loop: Cell<*mut EventLoop>)

pub struct EventLoop {
    uv_loop: Loop,
    pool: Option<Box<QueuePool>>,
}

pub struct BorrowedEventLoop {
    local: *mut EventLoop,
    marker1: marker::NoSend,
    marker2: marker::NoSync,
}

impl EventLoop {
    pub fn new() -> UvResult<EventLoop> {
        let mut uv_loop = try!(unsafe { Loop::new() });
        uv_loop.set_data(0 as *mut _);
        let pool = try!(QueuePool::new(&uv_loop));
        Ok(EventLoop {
            pool: Some(pool),
            uv_loop: uv_loop,
        })
    }

    /// Borrow a reference to the local event loop.
    ///
    /// If there is no local event loop, or the local event loop is already
    /// borrowed, then an error is returned.
    pub fn borrow() -> UvResult<BorrowedEventLoop> {
        let local = unsafe {
            local_loop.get(|local| {
                local.and_then(|c| {
                    match c.get() as uint {
                        0 => None,
                        n => { c.set(0 as *mut _); Some(n as *mut EventLoop) }
                    }
                })
            })
        };
        match local {
            Some(eloop) => Ok(BorrowedEventLoop {
                local: eloop,
                marker1: marker::NoSend,
                marker2: marker::NoSync,
            }),
            None => Err(UvError(uvll::UNKNOWN))
        }
    }

    /// Gain access to the underlying event loop.
    ///
    /// This method is unsafe as there is no guarantee that further safe methods
    /// called on the `Loop` will be valid as this event loop will deallocate it
    /// when it goes out of scope.
    pub unsafe fn uv_loop(&self) -> Loop { self.uv_loop }

    /// Create a new handle to this event loop.
    ///
    /// This handle can be used to return the current task home, assuming that
    /// it is a green task.
    pub fn make_handle(&mut self) -> HomeHandle {
        // It's understood by the homing code that the "local id" is just the
        // pointer of the local I/O factory cast to a uint.
        let id: uint = self as *mut _ as uint;
        HomeHandle::new(id, &mut **self.pool.get_mut_ref())
    }
}

impl green::EventLoop for EventLoop {
    fn run(&mut self) {
        let tls = Cell::new(self as *mut _);
        unsafe {
            local_loop.set(&tls, || {
                self.uv_loop.run(uvll::RUN_DEFAULT).unwrap();
            });
        }
    }

    fn callback(&mut self, f: proc()) {
        // Create a new idle handle, put the procedure into the custom data, and
        // then clean it up when the idle handle fires.
        unsafe {
            let mut idle = raw::Idle::new(&self.uv_loop).unwrap();
            idle.set_data(mem::transmute(box f));
            idle.start(onetime).unwrap();
        }

        extern fn onetime(handle: *mut uvll::uv_idle_t) {
            unsafe {
                let mut idle: raw::Idle = Handle::from_raw(handle);
                let f: Box<proc()> = mem::transmute(idle.get_data());
                idle.close_and_free();
                (*f)();
            }
        }
    }

    fn pausable_idle_callback(&mut self, cb: Box<green::Callback + Send>)
                              -> Box<green::PausableIdleCallback + Send> {
        box Idle::new_on(self, cb).unwrap()
            as Box<green::PausableIdleCallback + Send>
    }

    fn remote_callback(&mut self, f: Box<green::Callback + Send>)
                       -> Box<green::RemoteCallback + Send> {
        box Async::new_on(self, f).unwrap() as Box<green::RemoteCallback + Send>
    }

    fn has_active_io(&self) -> bool {
        true
        // self.uvio.loop_.get_blockers() > 0
    }
}

impl Drop for EventLoop {
    fn drop(&mut self) {
        use green::EventLoop;

        // Must first destroy the pool of handles before we destroy the loop
        // because otherwise the contained async handle will be destroyed after
        // the loop is free'd (use-after-free). We also must free the uv handle
        // after the loop has been closed because during the closing of the loop
        // the handle is required to be used apparently.
        //
        // Lastly, after we've closed the pool of handles we pump the event loop
        // one last time to run any closing callbacks to make sure the loop
        // shuts down cleanly.
        let mut handle = self.pool.get_ref().handle();
        drop(self.pool.take());
        self.run();

        unsafe {
            self.uv_loop.close().unwrap();
            handle.free();
            self.uv_loop.free();
        }
    }
}

impl Deref<EventLoop> for BorrowedEventLoop {
    fn deref<'a>(&'a self) -> &'a EventLoop { unsafe { &*self.local } }
}

impl DerefMut<EventLoop> for BorrowedEventLoop {
    fn deref_mut<'a>(&'a mut self) -> &'a mut EventLoop {
        unsafe { &mut *self.local }
    }
}

#[unsafe_destructor]
impl Drop for BorrowedEventLoop {
    fn drop(&mut self) {
        unsafe {
            local_loop.get(|l| l.unwrap().set(self.local))
        }
    }
}
<commit_msg>rustuv: Implement EventLoop::has_active_io<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;
use std::kinds::marker;
use std::cell::Cell;
use green;

use {uvll, UvResult, Idle, Async, UvError};
use raw::{mod, Loop, Handle};
use queue::QueuePool;
use homing::HomeHandle;

tls!(local_loop: Cell<*mut EventLoop>)

pub struct EventLoop {
    uv_loop: Loop,
    pool: Option<Box<QueuePool>>,
}

pub struct BorrowedEventLoop {
    local: *mut EventLoop,
    marker1: marker::NoSend,
    marker2: marker::NoSync,
}

impl EventLoop {
    pub fn new() -> UvResult<EventLoop> {
        let mut uv_loop = try!(unsafe { Loop::new() });
        uv_loop.set_data(0 as *mut _);
        let pool = try!(QueuePool::new(&uv_loop));
        Ok(EventLoop {
            pool: Some(pool),
            uv_loop: uv_loop,
        })
    }

    /// Borrow a reference to the local event loop.
    ///
    /// If there is no local event loop, or the local event loop is already
    /// borrowed, then an error is returned.
    pub fn borrow() -> UvResult<BorrowedEventLoop> {
        let local = unsafe {
            local_loop.get(|local| {
                local.and_then(|c| {
                    match c.get() as uint {
                        0 => None,
                        n => { c.set(0 as *mut _); Some(n as *mut EventLoop) }
                    }
                })
            })
        };
        match local {
            Some(eloop) => Ok(BorrowedEventLoop {
                local: eloop,
                marker1: marker::NoSend,
                marker2: marker::NoSync,
            }),
            None => Err(UvError(uvll::UNKNOWN))
        }
    }

    /// Gain access to the underlying event loop.
    ///
    /// This method is unsafe as there is no guarantee that further safe methods
    /// called on the `Loop` will be valid as this event loop will deallocate it
    /// when it goes out of scope.
    pub unsafe fn uv_loop(&self) -> Loop { self.uv_loop }

    /// Create a new handle to this event loop.
    ///
    /// This handle can be used to return the current task home, assuming that
    /// it is a green task.
    pub fn make_handle(&mut self) -> HomeHandle {
        // It's understood by the homing code that the "local id" is just the
        // pointer of the local I/O factory cast to a uint.
        let id: uint = self as *mut _ as uint;
        HomeHandle::new(id, &mut **self.pool.get_mut_ref())
    }
}

impl green::EventLoop for EventLoop {
    fn run(&mut self) {
        let tls = Cell::new(self as *mut _);
        unsafe {
            local_loop.set(&tls, || {
                self.uv_loop.run(uvll::RUN_DEFAULT).unwrap();
            });
        }
    }

    fn callback(&mut self, f: proc()) {
        // Create a new idle handle, put the procedure into the custom data, and
        // then clean it up when the idle handle fires.
        unsafe {
            let mut idle = raw::Idle::new(&self.uv_loop).unwrap();
            idle.set_data(mem::transmute(box f));
            idle.start(onetime).unwrap();
        }

        extern fn onetime(handle: *mut uvll::uv_idle_t) {
            unsafe {
                let mut idle: raw::Idle = Handle::from_raw(handle);
                let f: Box<proc()> = mem::transmute(idle.get_data());
                idle.close_and_free();
                (*f)();
            }
        }
    }

    fn pausable_idle_callback(&mut self, cb: Box<green::Callback + Send>)
                              -> Box<green::PausableIdleCallback + Send> {
        box Idle::new_on(self, cb).unwrap()
            as Box<green::PausableIdleCallback + Send>
    }

    fn remote_callback(&mut self, f: Box<green::Callback + Send>)
                       -> Box<green::RemoteCallback + Send> {
        box Async::new_on(self, f).unwrap() as Box<green::RemoteCallback + Send>
    }

    fn has_active_io(&self) -> bool {
        unsafe { self.uv_loop().get_data() as uint > 0 }
    }
}

impl Drop for EventLoop {
    fn drop(&mut self) {
        use green::EventLoop;

        // Must first destroy the pool of handles before we destroy the loop
        // because otherwise the contained async handle will be destroyed after
        // the loop is free'd (use-after-free). We also must free the uv handle
        // after the loop has been closed because during the closing of the loop
        // the handle is required to be used apparently.
        //
        // Lastly, after we've closed the pool of handles we pump the event loop
        // one last time to run any closing callbacks to make sure the loop
        // shuts down cleanly.
        let mut handle = self.pool.get_ref().handle();
        drop(self.pool.take());
        self.run();

        unsafe {
            self.uv_loop.close().unwrap();
            handle.free();
            self.uv_loop.free();
        }
    }
}

impl Deref<EventLoop> for BorrowedEventLoop {
    fn deref<'a>(&'a self) -> &'a EventLoop { unsafe { &*self.local } }
}

impl DerefMut<EventLoop> for BorrowedEventLoop {
    fn deref_mut<'a>(&'a mut self) -> &'a mut EventLoop {
        unsafe { &mut *self.local }
    }
}

#[unsafe_destructor]
impl Drop for BorrowedEventLoop {
    fn drop(&mut self) {
        unsafe {
            local_loop.get(|l| l.unwrap().set(self.local))
        }
    }
}
<|endoftext|>


<commit_before>use chrono::UTC;
use rustc_serialize::json::{self, Json, ToJson};

use params::*;

use rs_es::Client;
use rs_es::query::{Filter, Query};
use rs_es::units::JsonVal;
use rs_es::operations::search::{Sort, SortField, Order};
use rs_es::operations::index::IndexResult;
use rs_es::operations::mapping::*;
use rs_es::error::EsError;

extern crate searchspot;
use searchspot::terms::VectorOfTerms;

#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Talent {
  pub id:                 u32,
  pub accepted:           bool,
  pub work_roles:         Vec<String>,
  pub work_languages:     Vec<String>,
  pub work_experience:    String,
  pub work_locations:     Vec<String>,
  pub work_authorization: String,
  pub company_ids:        Vec<u32>,
  pub batch_starts_at:    String,
  pub batch_ends_at:      String,
  pub added_to_batch_at:  String,
  pub weight:             i32,
  pub blocked_companies:  Vec<u32>
}

impl ToJson for Talent {
  fn to_json(&self) -> Json {
    json::encode(&self).unwrap()
                       .to_json()
  }
}

impl Talent {
  /// Populate the ElasticSearch index with `self`.
  // I'm having problems with bulk actions. Let's wait for the next iteration.
  pub fn index(&self, mut es: &mut Client, index: &str) -> Result<IndexResult, EsError> {
    es.index(index, "talent")
      .with_doc(&self)
      .with_id(&*self.id.to_string())
      .send()
  }

  /// Query ElasticSearch on given `indexes` and `params` and return the IDs of
  /// the found talents.
  pub fn search(mut es: &mut Client, default_index: &str, params: &Map) -> Vec<u32> {
    let now   = UTC::now().to_rfc3339();
    let epoch = match params.find(&["epoch"]) {
      Some(epoch) => String::from_value(&epoch).unwrap_or(now),
      _           => now
    };

    let index: Vec<&str> = match params.find(&["index"]) {
      Some(&Value::String(ref index)) => vec![&index[..]],
      _ => vec![default_index]
    };

    let result = es.search_query()
                   .with_indexes(&*index)
                   .with_query(&Talent::search_filters(params, &*epoch))
                   .with_sort(&Talent::sorting_criteria())
                   .with_size(1000) // TODO
                   .send();

    match result {
      Ok(result) => {
        let mut results = result.hits.hits.into_iter()
                                          .map(|hit| hit.source.unwrap()["id"]
                                                               .as_u64()
                                                               .unwrap() as u32)
                                          .collect::<Vec<u32>>();
        results.dedup();
        results
      },
      Err(err) => {
        println!("{:?}", err);
        vec![]
      }
    }
  }

  /// Reset the given index. All the data will be destroyed and then the index
  /// will be created again. The map that will be used is hardcoded.
  pub fn reset_index(mut es: &mut Client, index: &str) -> Result<MappingResult, EsError> {
    let mapping = hashmap! {
      "talent" => hashmap! {
        "id" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "work_roles" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_languages" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_experience" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_locations" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_authorization" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "company_ids" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "accepted" => hashmap! {
          "type" => "boolean",
          "index" => "not_analyzed"
        },

        "batch_starts_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "batch_ends_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "added_to_batch_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "weight" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "blocked_companies" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        }
      }
    };

    es.delete_index(index);

    MappingOperation::new(&mut es, index, &mapping).send()
  }

  /// Return a `Vec<Filter>` with visibility criteria for the talents.
  /// The `epoch` must be given as `I64` (UNIX time in seconds) and is
  /// the range in which batches are searched.
  /// If `presented_talents` is provided, talents who match the IDs
  /// contained there skip the standard visibility criteria.
  ///
  /// Basically, the talents must be accepted into the platform and must be
  /// inside a living batch to match the visibility criteria.
  fn visibility_filters(epoch: &str, presented_talents: Vec<i32>) -> Vec<Filter> {
    let visibility_rules = Filter::build_bool()
                                  .with_must(
                                    vec![
                                      Filter::build_term("accepted", true)
                                             .build(),
                                      Filter::build_range("batch_starts_at")
                                             .with_lte(JsonVal::from(epoch))
                                             .with_format("dateOptionalTime")
                                             .build(),
                                      Filter::build_range("batch_ends_at")
                                             .with_gte(JsonVal::from(epoch))
                                             .with_format("dateOptionalTime")
                                             .build()
                                    ])
                                  .build();

    if presented_talents.len() > 0 { // preferred over !_.is_empty()
      let presented_talents_filters = Filter::build_bool()
                                             .with_must(
                                               vec![
                                                 <Filter as VectorOfTerms<i32>>::build_terms(
                                                   "ids", &presented_talents)
                                               ].into_iter()
                                                .flat_map(|x| x)
                                                .collect::<Vec<Filter>>())
                                             .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents_filters])
               .build()
      ]
    }
    else {
      vec![visibility_rules]
    }
  }

  /// Given parameters inside the query string mapped inside a `Map`,
  /// and the `epoch` (defined as UNIX time in seconds) for batches,
  /// return a `Query` for ElasticSearch.
  ///
  /// `VectorOfTerms` are ORred, while `Filter`s are ANDed.
  /// I.e.: given ["Fullstack", "DevOps"] as `work_roles`, found talents
  /// will present at least one of these roles), but both `work_roles`
  /// and `work_languages`, if provided, must not be empty.
  fn search_filters(params: &Map, epoch: &str) -> Query {
    let company_id = i32_vec_from_params!(params, "company_id");

    Query::build_filtered(Filter::build_bool()
                                 .with_must(
                                   vec![
                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_roles", &vec_from_params!(params, "work_roles")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_languages", &vec_from_params!(params, "work_languages")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_experience", &string_vec_from_params!(params, "work_experience")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                      "work_authorization", &string_vec_from_params!(params, "work_authorization")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_locations", &vec_from_params!(params, "work_locations")),

                                     Talent::visibility_filters(epoch,
                                       i32_vec_from_params!(params, "presented_talents"))
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .with_must_not(
                                   vec![
                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "company_ids", &company_id),

                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "blocked_companies", &company_id)
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .build())
          .build()
  }

  /// Return a `Sort` that makes values be sorted for given fields, descendently.
  fn sorting_criteria() -> Sort {
    Sort::new(
      vec![
        SortField::new("batch_starts_at",   Some(Order::Desc)).build(),
        SortField::new("weight",            Some(Order::Desc)).build(),
        SortField::new("added_to_batch_at", Some(Order::Desc)).build()
      ])
  }
}

#[cfg(test)]
mod tests {
  use chrono::*;

  use rs_es::Client;

  use params::*;

  extern crate searchspot;
  use searchspot::config::*;

  use resources::user::Talent;

  const CONFIG_FILE: &'static str = "examples/tests.toml";

  lazy_static! {
    static ref config: Config = Config::from_file(CONFIG_FILE.to_owned());
  }

  pub fn make_client() -> Client {
    Client::new(&*config.es.host, config.es.port)
  }

  macro_rules! epoch_from_year {
    ($year:expr) => {
      UTC.datetime_from_str(&format!("{}-01-01 12:00:00", $year),
        "%Y-%m-%d %H:%M:%S").unwrap().to_rfc3339()
    }
  }

  pub fn populate_index(mut client: &mut Client) -> bool {
    vec![
      Talent {
        id:                 1,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             -5,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 2,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 3,
        accepted:           false,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2007"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 4,
        accepted:           true,
        work_roles:         vec!["Fullstack".to_owned(), "DevOps".to_owned()],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization:  "yes".to_owned(),
        company_ids:        vec![6],
        batch_starts_at:    epoch_from_year!("2008"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             0,
        blocked_companies:  vec![]
      }
    ].iter()
     .map(|talent| talent.index(&mut client, &config.es.index)
                         .is_ok())
     .collect::<Vec<bool>>()
     .into_iter()
     .all(|result| result)
  }

  fn refresh_index(mut client: &mut Client) {
    client.refresh()
          .with_indexes(&[&config.es.index])
          .send()
          .unwrap();
  }

  #[test]
  fn test_search() {
    let mut client = make_client();
    assert!(Talent::reset_index(&mut client, &*config.es.index).is_ok());
    refresh_index(&mut client);
    assert!(populate_index(&mut client));
    refresh_index(&mut client);

    // no parameters are given
    {
      let results = Talent::search(&mut client, &*config.es.index, &Map::new());
      assert_eq!(vec![4, 2, 1], results);
    }

    // a non existing index is given
    {
      let mut map = Map::new();
      map.assign("index", Value::String("lololol".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // a date that doesn't match given indexes is given
    {
      let mut map = Map::new();
      map.assign("epoch", Value::String(epoch_from_year!("2040"))).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    {
      let mut map = Map::new();
      map.assign("work_roles[]", Value::String("Fullstack".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![4], results);
    }

    // filtering for given company_id
    {
      let mut map = Map::new();
      map.assign("company_id", Value::String("6".into())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![2, 1], results);
    }
  }
}
<commit_msg>resources/user: consider only the arrays in the query string<commit_after>use chrono::UTC;
use rustc_serialize::json::{self, Json, ToJson};

use params::*;

use rs_es::Client;
use rs_es::query::{Filter, Query};
use rs_es::units::JsonVal;
use rs_es::operations::search::{Sort, SortField, Order};
use rs_es::operations::index::IndexResult;
use rs_es::operations::mapping::*;
use rs_es::error::EsError;

extern crate searchspot;
use searchspot::terms::VectorOfTerms;

#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Talent {
  pub id:                 u32,
  pub accepted:           bool,
  pub work_roles:         Vec<String>,
  pub work_languages:     Vec<String>,
  pub work_experience:    String,
  pub work_locations:     Vec<String>,
  pub work_authorization: String,
  pub company_ids:        Vec<u32>,
  pub batch_starts_at:    String,
  pub batch_ends_at:      String,
  pub added_to_batch_at:  String,
  pub weight:             i32,
  pub blocked_companies:  Vec<u32>
}

impl ToJson for Talent {
  fn to_json(&self) -> Json {
    json::encode(&self).unwrap()
                       .to_json()
  }
}

impl Talent {
  /// Populate the ElasticSearch index with `self`.
  // I'm having problems with bulk actions. Let's wait for the next iteration.
  pub fn index(&self, mut es: &mut Client, index: &str) -> Result<IndexResult, EsError> {
    es.index(index, "talent")
      .with_doc(&self)
      .with_id(&*self.id.to_string())
      .send()
  }

  /// Query ElasticSearch on given `indexes` and `params` and return the IDs of
  /// the found talents.
  pub fn search(mut es: &mut Client, default_index: &str, params: &Map) -> Vec<u32> {
    let now   = UTC::now().to_rfc3339();
    let epoch = match params.find(&["epoch"]) {
      Some(epoch) => String::from_value(&epoch).unwrap_or(now),
      _           => now
    };

    let index: Vec<&str> = match params.find(&["index"]) {
      Some(&Value::String(ref index)) => vec![&index[..]],
      _ => vec![default_index]
    };

    let result = es.search_query()
                   .with_indexes(&*index)
                   .with_query(&Talent::search_filters(params, &*epoch))
                   .with_sort(&Talent::sorting_criteria())
                   .with_size(1000) // TODO
                   .send();

    match result {
      Ok(result) => {
        let mut results = result.hits.hits.into_iter()
                                          .map(|hit| hit.source.unwrap()["id"]
                                                               .as_u64()
                                                               .unwrap() as u32)
                                          .collect::<Vec<u32>>();
        results.dedup();
        results
      },
      Err(err) => {
        println!("{:?}", err);
        vec![]
      }
    }
  }

  /// Reset the given index. All the data will be destroyed and then the index
  /// will be created again. The map that will be used is hardcoded.
  pub fn reset_index(mut es: &mut Client, index: &str) -> Result<MappingResult, EsError> {
    let mapping = hashmap! {
      "talent" => hashmap! {
        "id" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "work_roles" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_languages" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_experience" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_locations" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "work_authorization" => hashmap! {
          "type" => "string",
          "index" => "not_analyzed"
        },

        "company_ids" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "accepted" => hashmap! {
          "type" => "boolean",
          "index" => "not_analyzed"
        },

        "batch_starts_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "batch_ends_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "added_to_batch_at" => hashmap! {
          "type" => "date",
          "format" => "dateOptionalTime",
          "index" => "not_analyzed"
        },

        "weight" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        },

        "blocked_companies" => hashmap! {
          "type" => "integer",
          "index" => "not_analyzed"
        }
      }
    };

    es.delete_index(index);

    MappingOperation::new(&mut es, index, &mapping).send()
  }

  /// Return a `Vec<Filter>` with visibility criteria for the talents.
  /// The `epoch` must be given as `I64` (UNIX time in seconds) and is
  /// the range in which batches are searched.
  /// If `presented_talents` is provided, talents who match the IDs
  /// contained there skip the standard visibility criteria.
  ///
  /// Basically, the talents must be accepted into the platform and must be
  /// inside a living batch to match the visibility criteria.
  fn visibility_filters(epoch: &str, presented_talents: Vec<i32>) -> Vec<Filter> {
    let visibility_rules = Filter::build_bool()
                                  .with_must(
                                    vec![
                                      Filter::build_term("accepted", true)
                                             .build(),
                                      Filter::build_range("batch_starts_at")
                                             .with_lte(JsonVal::from(epoch))
                                             .with_format("dateOptionalTime")
                                             .build(),
                                      Filter::build_range("batch_ends_at")
                                             .with_gte(JsonVal::from(epoch))
                                             .with_format("dateOptionalTime")
                                             .build()
                                    ])
                                  .build();

    if presented_talents.len() > 0 { // preferred over !_.is_empty()
      let presented_talents_filters = Filter::build_bool()
                                             .with_must(
                                               vec![
                                                 <Filter as VectorOfTerms<i32>>::build_terms(
                                                   "ids", &presented_talents)
                                               ].into_iter()
                                                .flat_map(|x| x)
                                                .collect::<Vec<Filter>>())
                                             .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents_filters])
               .build()
      ]
    }
    else {
      vec![visibility_rules]
    }
  }

  /// Given parameters inside the query string mapped inside a `Map`,
  /// and the `epoch` (defined as UNIX time in seconds) for batches,
  /// return a `Query` for ElasticSearch.
  ///
  /// `VectorOfTerms` are ORred, while `Filter`s are ANDed.
  /// I.e.: given ["Fullstack", "DevOps"] as `work_roles`, found talents
  /// will present at least one of these roles), but both `work_roles`
  /// and `work_languages`, if provided, must not be empty.
  fn search_filters(params: &Map, epoch: &str) -> Query {
    let company_id = i32_vec_from_params!(params, "company_id");

    Query::build_filtered(Filter::build_bool()
                                 .with_must(
                                   vec![
                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_roles", &vec_from_params!(params, "work_roles")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_languages", &vec_from_params!(params, "work_languages")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_experience", &vec_from_params!(params, "work_experience")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                      "work_authorization", &vec_from_params!(params, "work_authorization")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_locations", &vec_from_params!(params, "work_locations")),

                                     Talent::visibility_filters(epoch,
                                       i32_vec_from_params!(params, "presented_talents"))
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .with_must_not(
                                   vec![
                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "company_ids", &company_id),

                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "blocked_companies", &company_id)
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .build())
          .build()
  }

  /// Return a `Sort` that makes values be sorted for given fields, descendently.
  fn sorting_criteria() -> Sort {
    Sort::new(
      vec![
        SortField::new("batch_starts_at",   Some(Order::Desc)).build(),
        SortField::new("weight",            Some(Order::Desc)).build(),
        SortField::new("added_to_batch_at", Some(Order::Desc)).build()
      ])
  }
}

#[cfg(test)]
mod tests {
  use chrono::*;

  use rs_es::Client;

  use params::*;

  extern crate searchspot;
  use searchspot::config::*;

  use resources::user::Talent;

  const CONFIG_FILE: &'static str = "examples/tests.toml";

  lazy_static! {
    static ref config: Config = Config::from_file(CONFIG_FILE.to_owned());
  }

  pub fn make_client() -> Client {
    Client::new(&*config.es.host, config.es.port)
  }

  macro_rules! epoch_from_year {
    ($year:expr) => {
      UTC.datetime_from_str(&format!("{}-01-01 12:00:00", $year),
        "%Y-%m-%d %H:%M:%S").unwrap().to_rfc3339()
    }
  }

  pub fn populate_index(mut client: &mut Client) -> bool {
    vec![
      Talent {
        id:                 1,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             -5,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 2,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 3,
        accepted:           false,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2007"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 4,
        accepted:           true,
        work_roles:         vec!["Fullstack".to_owned(), "DevOps".to_owned()],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization:  "yes".to_owned(),
        company_ids:        vec![6],
        batch_starts_at:    epoch_from_year!("2008"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             0,
        blocked_companies:  vec![]
      }
    ].iter()
     .map(|talent| talent.index(&mut client, &config.es.index)
                         .is_ok())
     .collect::<Vec<bool>>()
     .into_iter()
     .all(|result| result)
  }

  fn refresh_index(mut client: &mut Client) {
    client.refresh()
          .with_indexes(&[&config.es.index])
          .send()
          .unwrap();
  }

  #[test]
  fn test_search() {
    let mut client = make_client();
    assert!(Talent::reset_index(&mut client, &*config.es.index).is_ok());
    refresh_index(&mut client);
    assert!(populate_index(&mut client));
    refresh_index(&mut client);

    // no parameters are given
    {
      let results = Talent::search(&mut client, &*config.es.index, &Map::new());
      assert_eq!(vec![4, 2, 1], results);
    }

    // a non existing index is given
    {
      let mut map = Map::new();
      map.assign("index", Value::String("lololol".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // a date that doesn't match given indexes is given
    {
      let mut map = Map::new();
      map.assign("epoch", Value::String(epoch_from_year!("2040"))).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    {
      let mut map = Map::new();
      map.assign("work_roles[]", Value::String("Fullstack".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![4], results);
    }

    // filtering for given company_id
    {
      let mut map = Map::new();
      map.assign("company_id", Value::String("6".into())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![2, 1], results);
    }
  }
}
<|endoftext|>