<commit_before>use irc::client::prelude::*;
use antidote::RwLock;

use time;
use std::time::Duration;
use chrono::NaiveDateTime;
use humantime::format_duration;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

pub mod database;
use self::database::Database;

#[derive(PluginName)]
pub struct Tell<T: Database> {
    tells: RwLock<T>,
}

impl<T: Database> Tell<T> {
    pub fn new(db: T) -> Tell<T> {
        Tell {
            tells: RwLock::new(db),
        }
    }

    fn tell_command(
        &self,
        client: &IrcClient,
        command: PluginCommand,
    ) -> Result<String, TellError> {
        if command.tokens.len() < 2 {
            return Ok(self.invalid_command().to_owned());
        }

        let mut online = Vec::new();

        let receivers = command.tokens[0].split(',').filter(|&s| !s.is_empty());
        let sender = command.source;

        let mut no_receiver = true;
        for receiver in receivers {
            if receiver.eq_ignore_ascii_case(client.current_nickname())
                || receiver.eq_ignore_ascii_case(&sender)
            {
                online.push(receiver);
                continue;
            }

            let channels = client
                .list_channels()
                .expect("The irc crate should not be compiled with the \"nochanlists\" feature");

            if let Some(_) = channels
                .iter()
                .map(|channel| client.list_users(&channel))
                .map(|option| {
                    option.and_then(|users| {
                        users
                            .into_iter()
                            .find(|user| user.get_nickname().eq_ignore_ascii_case(&receiver))
                    })
                })
                .find(|option| option.is_some())
            {
                online.push(receiver);
                continue;
            }

            let tm = time::now().to_timespec();
            let message = command.tokens[1..].join(" ");
            let tell = database::NewTellMessage {
                sender: &sender,
                receiver: &receiver.to_lowercase(),
                time: NaiveDateTime::from_timestamp(tm.sec, 0u32),
                message: &message,
            };

            debug!("Saving tell for {:?}", receiver);
            self.tells.write().insert_tell(&tell)?;
            no_receiver = false;
        }

        Ok(if no_receiver {
            format!("Invalid receiver.")
        } else {
            match online.len() {
                0 => format!("Got it!"),
                1 => format!("{} is currently online.", online[0]),
                _ => format!("{} are currently online.", online.join(", ")),
            }
        })
    }

    fn on_namelist(&self, client: &IrcClient, channel: &str) -> Result<(), FrippyError> {
        let receivers = self.tells.read()
            .get_receivers()
            .context(FrippyErrorKind::Tell)?;

        if let Some(users) = client.list_users(channel) {
            debug!("Outstanding tells for {:?}", receivers);

            for receiver in users
                .iter()
                .map(|u| u.get_nickname())
                .filter(|u| receivers.iter().any(|r| r == &u.to_lowercase()))
            {
                self.send_tells(client, receiver)?;
            }

            Ok(())
        } else {
            Ok(())
        }
    }

    fn send_tells(&self, client: &IrcClient, receiver: &str) -> Result<(), FrippyError> {
        if client.current_nickname() == receiver {
            return Ok(());
        }

        let mut tells = self.tells.write();

        let tell_messages = match tells.get_tells(&receiver.to_lowercase()) {
            Ok(t) => t,
            Err(e) => {
                // This warning only occurs if frippy is built without a database
                #[allow(unreachable_patterns)]
                return match e.kind() {
                    ErrorKind::NotFound => Ok(()),
                    _ => Err(e.context(FrippyErrorKind::Tell))?,
                };
            }
        };

        for tell in tell_messages {
            let now = Duration::new(time::now().to_timespec().sec as u64, 0);
            let dur = now - Duration::new(tell.time.timestamp() as u64, 0);
            let human_dur = format_duration(dur);

            client
                .send_notice(
                    receiver,
                    &format!(
                        "Tell from {} {} ago: {}",
                        tell.sender, human_dur, tell.message
                    ),
                )
                .context(FrippyErrorKind::Connection)?;

            debug!(
                "Sent {:?} from {:?} to {:?}",
                tell.message, tell.sender, receiver
            );
        }

        tells
            .delete_tells(&receiver.to_lowercase())
            .context(FrippyErrorKind::Tell)?;

        Ok(())
    }

    fn invalid_command(&self) -> &str {
        "Incorrect Command. \
         Send \"tell help\" for help."
    }

    fn help(&self) -> &str {
        "usage: {} tell user message\r\n\
         example: {0} tell Foobar Hello!"
    }
}

impl<T: Database> Plugin for Tell<T> {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        let res = match message.command {
            Command::JOIN(_, _, _) => self.send_tells(client, message.source_nickname().unwrap()),
            Command::NICK(ref nick) => self.send_tells(client, nick),
            Command::Response(resp, ref chan_info, _) => {
                if resp == Response::RPL_NAMREPLY {
                    debug!("NAMREPLY info: {:?}", chan_info);

                    self.on_namelist(client, &chan_info[chan_info.len() - 1])
                } else {
                    Ok(())
                }
            }
            _ => Ok(()),
        };

        match res {
            Ok(_) => ExecutionStatus::Done,
            Err(e) => ExecutionStatus::Err(e),
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Tell should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            return Ok(client
                .send_notice(&command.source, &self.invalid_command())
                .context(FrippyErrorKind::Connection)?);
        }

        let sender = command.source.to_owned();

        Ok(match command.tokens[0].as_ref() {
            "help" => client
                .send_notice(&command.source, &self.help())
                .context(FrippyErrorKind::Connection)
                .into(),
            _ => match self.tell_command(client, command) {
                Ok(msg) => client
                    .send_notice(&sender, &msg)
                    .context(FrippyErrorKind::Connection),
                Err(e) => client
                    .send_notice(&sender, &e.to_string())
                    .context(FrippyErrorKind::Connection)
                    .into(),
            },
        }?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from("This Plugin does not implement any commands."))
    }
}

use std::fmt;
impl<T: Database> fmt::Debug for Tell<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Tell {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "TellError"]
    pub enum ErrorKind {
        /// Not found command error
        #[fail(display = "Tell was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<commit_msg>Fix tell error message<commit_after>use irc::client::prelude::*;
use antidote::RwLock;

use time;
use std::time::Duration;
use chrono::NaiveDateTime;
use humantime::format_duration;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

pub mod database;
use self::database::Database;

#[derive(PluginName)]
pub struct Tell<T: Database> {
    tells: RwLock<T>,
}

impl<T: Database> Tell<T> {
    pub fn new(db: T) -> Tell<T> {
        Tell {
            tells: RwLock::new(db),
        }
    }

    fn tell_command(
        &self,
        client: &IrcClient,
        command: PluginCommand,
    ) -> Result<String, TellError> {
        if command.tokens.len() < 2 {
            return Ok(self.invalid_command().to_owned());
        }

        let mut online = Vec::new();

        let receivers = command.tokens[0].split(',').filter(|&s| !s.is_empty());
        let sender = command.source;

        let mut no_receiver = true;
        for receiver in receivers {
            if receiver.eq_ignore_ascii_case(client.current_nickname())
                || receiver.eq_ignore_ascii_case(&sender)
            {
                online.push(receiver);
                continue;
            }

            let channels = client
                .list_channels()
                .expect("The irc crate should not be compiled with the \"nochanlists\" feature");

            if let Some(_) = channels
                .iter()
                .map(|channel| client.list_users(&channel))
                .map(|option| {
                    option.and_then(|users| {
                        users
                            .into_iter()
                            .find(|user| user.get_nickname().eq_ignore_ascii_case(&receiver))
                    })
                })
                .find(|option| option.is_some())
            {
                online.push(receiver);
                continue;
            }

            let tm = time::now().to_timespec();
            let message = command.tokens[1..].join(" ");
            let tell = database::NewTellMessage {
                sender: &sender,
                receiver: &receiver.to_lowercase(),
                time: NaiveDateTime::from_timestamp(tm.sec, 0u32),
                message: &message,
            };

            debug!("Saving tell for {:?}", receiver);
            self.tells.write().insert_tell(&tell)?;
            no_receiver = false;
        }

        Ok(if no_receiver && online.is_empty() {
            format!("Invalid receiver.")
        } else {
            match online.len() {
                0 => format!("Got it!"),
                1 => format!("{} is currently online.", online[0]),
                _ => format!("{} are currently online.", online.join(", ")),
            }
        })
    }

    fn on_namelist(&self, client: &IrcClient, channel: &str) -> Result<(), FrippyError> {
        let receivers = self.tells.read()
            .get_receivers()
            .context(FrippyErrorKind::Tell)?;

        if let Some(users) = client.list_users(channel) {
            debug!("Outstanding tells for {:?}", receivers);

            for receiver in users
                .iter()
                .map(|u| u.get_nickname())
                .filter(|u| receivers.iter().any(|r| r == &u.to_lowercase()))
            {
                self.send_tells(client, receiver)?;
            }

            Ok(())
        } else {
            Ok(())
        }
    }

    fn send_tells(&self, client: &IrcClient, receiver: &str) -> Result<(), FrippyError> {
        if client.current_nickname() == receiver {
            return Ok(());
        }

        let mut tells = self.tells.write();

        let tell_messages = match tells.get_tells(&receiver.to_lowercase()) {
            Ok(t) => t,
            Err(e) => {
                // This warning only occurs if frippy is built without a database
                #[allow(unreachable_patterns)]
                return match e.kind() {
                    ErrorKind::NotFound => Ok(()),
                    _ => Err(e.context(FrippyErrorKind::Tell))?,
                };
            }
        };

        for tell in tell_messages {
            let now = Duration::new(time::now().to_timespec().sec as u64, 0);
            let dur = now - Duration::new(tell.time.timestamp() as u64, 0);
            let human_dur = format_duration(dur);

            client
                .send_notice(
                    receiver,
                    &format!(
                        "Tell from {} {} ago: {}",
                        tell.sender, human_dur, tell.message
                    ),
                )
                .context(FrippyErrorKind::Connection)?;

            debug!(
                "Sent {:?} from {:?} to {:?}",
                tell.message, tell.sender, receiver
            );
        }

        tells
            .delete_tells(&receiver.to_lowercase())
            .context(FrippyErrorKind::Tell)?;

        Ok(())
    }

    fn invalid_command(&self) -> &str {
        "Incorrect Command. \
         Send \"tell help\" for help."
    }

    fn help(&self) -> &str {
        "usage: {} tell user message\r\n\
         example: {0} tell Foobar Hello!"
    }
}

impl<T: Database> Plugin for Tell<T> {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        let res = match message.command {
            Command::JOIN(_, _, _) => self.send_tells(client, message.source_nickname().unwrap()),
            Command::NICK(ref nick) => self.send_tells(client, nick),
            Command::Response(resp, ref chan_info, _) => {
                if resp == Response::RPL_NAMREPLY {
                    debug!("NAMREPLY info: {:?}", chan_info);

                    self.on_namelist(client, &chan_info[chan_info.len() - 1])
                } else {
                    Ok(())
                }
            }
            _ => Ok(()),
        };

        match res {
            Ok(_) => ExecutionStatus::Done,
            Err(e) => ExecutionStatus::Err(e),
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Tell should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            return Ok(client
                .send_notice(&command.source, &self.invalid_command())
                .context(FrippyErrorKind::Connection)?);
        }

        let sender = command.source.to_owned();

        Ok(match command.tokens[0].as_ref() {
            "help" => client
                .send_notice(&command.source, &self.help())
                .context(FrippyErrorKind::Connection)
                .into(),
            _ => match self.tell_command(client, command) {
                Ok(msg) => client
                    .send_notice(&sender, &msg)
                    .context(FrippyErrorKind::Connection),
                Err(e) => client
                    .send_notice(&sender, &e.to_string())
                    .context(FrippyErrorKind::Connection)
                    .into(),
            },
        }?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from("This Plugin does not implement any commands."))
    }
}

use std::fmt;
impl<T: Database> fmt::Debug for Tell<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Tell {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "TellError"]
    pub enum ErrorKind {
        /// Not found command error
        #[fail(display = "Tell was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<|endoftext|>


<commit_before>extern crate herder;
extern crate serde_json;

use herder::api::v1::{
    Account,
    Application,
    Attachment,
    Card,
    Context,
    Error,
    Instance,
    Mention,
    Notification,
    Relationship,
    Report,
    Results,
    Status,
    Tag
};

#[test]
fn account_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn application_deserialize_from_json() {
    let ojson = r#"{
                    "name": "MYNAME",
                    "website": "MYURI"
                 }"#;
    let application: Application = serde_json::from_str(ojson).unwrap();
    assert_eq!(application, Application {
                        name: String::from("MYNAME"),
                        website: Some(String::from("MYURI"))
    });
}

#[test]
fn attachment_deserialize_from_json() {
    let ojson = r#"{
                    "id": 1234,
                    "type": "video",
                    "url": "MYURL",
                    "remote_url": "MYURL",
                    "preview_url": "MYURL",
                    "text_url": "MYURL"
                 }"#;
    let attachment: Attachment = serde_json::from_str(ojson).unwrap();
    assert_eq!(attachment, Attachment {
                        attachment_id: 1234,
                        attachment_type: String::from("video"),
                        url: String::from("MYURL"),
                        remote_url: String::from("MYURL"),
                        preview_url: String::from("MYURL"),
                        text_url: String::from("MYURL")
    });
}

#[test]
fn card_deserialize_from_json() {
    let ojson = r#"{
                    "url": "MYURL",
                    "title": "My Title",
                    "description": "My description",
                    "image": "MYIMG.png"
                 }"#;
    let card: Card = serde_json::from_str(ojson).unwrap();
    assert_eq!(card, Card {
                        url: String::from("MYURL"),
                        title: String::from("My Title"),
                        description: String::from("My description"),
                        image: String::from("MYIMG.png")
    });
}

#[test]
fn context_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn error_deserialize_from_json() {
    let ojson = r#"{
                    "error": "Error message"
                 }"#;
    let error: Error = serde_json::from_str(ojson).unwrap();
    assert_eq!(error, Error {
                        error: String::from("Error message")
    });
}

#[test]
fn instance_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn mention_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn notification_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn relationship_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn report_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn results_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn status_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn tag_deserialize_from_json() {
    unimplemented!();
}
<commit_msg>test Report and Tag entities<commit_after>extern crate herder;
extern crate serde_json;

use herder::api::v1::{
    Account,
    Application,
    Attachment,
    Card,
    Context,
    Error,
    Instance,
    Mention,
    Notification,
    Relationship,
    Report,
    Results,
    Status,
    Tag
};

#[test]
fn account_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn application_deserialize_from_json() {
    let ojson = r#"{
                    "name": "MYNAME",
                    "website": "MYURI"
                 }"#;
    let application: Application = serde_json::from_str(ojson).unwrap();
    assert_eq!(application, Application {
                        name: String::from("MYNAME"),
                        website: Some(String::from("MYURI"))
    });
}

#[test]
fn attachment_deserialize_from_json() {
    let ojson = r#"{
                    "id": 1234,
                    "type": "video",
                    "url": "MYURL",
                    "remote_url": "MYURL",
                    "preview_url": "MYURL",
                    "text_url": "MYURL"
                 }"#;
    let attachment: Attachment = serde_json::from_str(ojson).unwrap();
    assert_eq!(attachment, Attachment {
                        attachment_id: 1234,
                        attachment_type: String::from("video"),
                        url: String::from("MYURL"),
                        remote_url: String::from("MYURL"),
                        preview_url: String::from("MYURL"),
                        text_url: String::from("MYURL")
    });
}

#[test]
fn card_deserialize_from_json() {
    let ojson = r#"{
                    "url": "MYURL",
                    "title": "My Title",
                    "description": "My description",
                    "image": "MYIMG.png"
                 }"#;
    let card: Card = serde_json::from_str(ojson).unwrap();
    assert_eq!(card, Card {
                        url: String::from("MYURL"),
                        title: String::from("My Title"),
                        description: String::from("My description"),
                        image: String::from("MYIMG.png")
    });
}

#[test]
fn context_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn error_deserialize_from_json() {
    let ojson = r#"{
                    "error": "Error message"
                 }"#;
    let error: Error = serde_json::from_str(ojson).unwrap();
    assert_eq!(error, Error {
                        error: String::from("Error message")
    });
}

#[test]
fn instance_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn mention_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn notification_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn relationship_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn report_deserialize_from_json() {
    let ojson = r#"{
                    "id": 1234,
                    "action_taken": "Some action taken"
                 }"#;
    let report: Report = serde_json::from_str(ojson).unwrap();
    assert_eq!(report, Report {
                        report_id: 1234,
                        action_taken: String::from("Some action taken")
    });
}

#[test]
fn results_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn status_deserialize_from_json() {
    unimplemented!();
}

#[test]
fn tag_deserialize_from_json() {
    let ojson = r#"{
                    "name": "rustlang",
                    "url": "MYURL"
                 }"#;
    let tag: Tag = serde_json::from_str(ojson).unwrap();
    assert_eq!(tag, Tag {
                        name: String::from("rustlang"),
                        url: String::from("MYURL")
    });
}
<|endoftext|>


<commit_before>
extern crate serde_json;
extern crate postgres;

use std::fs::File;
use std::io::{Write, BufWriter};
use std::collections::HashMap;

use self::serde_json::Value;
use self::postgres::Connection;

type CvName = String;

pub type MiscExtRange = String;

pub type GeneUniquename = String;
pub type TermUniquename = String;
pub type GeneName = String;
pub type TypeName = String;
pub type GeneProduct = String;

pub type TermName = String;
pub type TermId = String;
pub type TermDef = String;

pub type Evidence = String;
pub type With = String;
pub type Residue = String;
pub type Condition = String;
pub type Qualifier = String;

pub type TypeInteractionAnnotationMap =
    HashMap<TypeName, Vec<InteractionAnnotation>>;

pub type UniquenameGeneMap =
    HashMap<GeneUniquename, GeneDetails>;

pub type TranscriptUniquename = String;

pub type UniquenameTranscriptMap =
    HashMap<TranscriptUniquename, TranscriptDetails>;

pub type GenotypeUniquename = String;
pub type AlleleUniquename = String;

pub type UniquenameAlleleMap = HashMap<AlleleUniquename, AlleleShort>;
pub type UniquenameGenotypeMap = HashMap<GenotypeUniquename, GenotypeShort>;

pub type IdGeneMap = HashMap<GeneUniquename, GeneDetails>;
pub type IdGeneShortMap = HashMap<GeneUniquename, GeneShort>;
pub type IdTermDetailsMap = HashMap<TermId, Rc<TermDetails>>;
pub type IdReferenceMap = HashMap<TermId, ReferenceDetails>;

include!(concat!(env!("OUT_DIR"), "/data_serde.rs"));

impl WebData {
    fn get_genes(&self) -> &IdGeneMap {
        &self.genes
    }
    fn get_terms(&self) -> &IdTermDetailsMap {
        &self.terms
    }

    fn write_reference_details(&self, output_dir: &str) {
        for (reference_uniquename, reference_details) in &self.references {
            let s = serde_json::to_string(&reference_details).unwrap();
            let file_name = String::new() + &output_dir + "/reference/" + &reference_uniquename + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_gene_details(&self, output_dir: &str) {
        for (gene_uniquename, gene_details) in &self.genes {
            let s = serde_json::to_string(&gene_details).unwrap();
            let file_name = String::new() + &output_dir + "/gene/" + &gene_uniquename + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_gene_summary(&self, output_dir: &str, organism_genus_species: &str) {
        let gene_summaries =
            self.gene_summaries.values()
            .filter(|gene_short| {
                let gene_uniquename = &gene_short.uniquename;
                let gene_details = self.get_genes().get(gene_uniquename).unwrap();
                let feature_org_genus_species = String::new() +
                    &gene_details.organism.genus + "_" + &gene_details.organism.species;
                feature_org_genus_species == organism_genus_species
            })
            .map(|gene_short| gene_short.clone())
            .collect::<Vec<GeneShort>>();
        let s = serde_json::to_string(&gene_summaries).unwrap();
        let file_name = String::new() + &output_dir + "/gene_summaries.json";
        let f = File::create(file_name).expect("Unable to open file");
        let mut writer = BufWriter::new(&f);
        writer.write_all(s.as_bytes()).expect("Unable to write!");
    }

    fn write_terms(&self, output_dir: &str) {
        for (termid, term_details) in &self.terms {
            let s = serde_json::to_string(&term_details).unwrap();
            let file_name = String::new() + &output_dir + "/term/" + &termid + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_metadata(&self, output_dir: &str) {
        let s = serde_json::to_string(&self.metadata).unwrap();
        let file_name = String::new() + &output_dir + "/metadata.json";
        let f = File::create(file_name).expect("Unable to open file");
        let mut writer = BufWriter::new(&f);
        writer.write_all(s.as_bytes()).expect("Unable to write!");
    }

    pub fn write(&self, output_dir: &str, organism_genus_species: &str) {
        self.write_reference_details(output_dir);
        self.write_gene_details(output_dir);
        self.write_gene_summary(output_dir, organism_genus_species);
        self.write_terms(output_dir);
        self.write_metadata(output_dir);

        println!("wrote {} genes", self.get_genes().len());
        println!("wrote {} terms", self.get_terms().len());
    }

    pub fn store_jsonb(&self, conn: &Connection) {
        let trans = conn.transaction().unwrap();

        for (uniquename, gene_details) in &self.genes {
            // FIXME - how do we get a Valie from gene_details? - this is dumb:
            let gene_details_json = serde_json::to_string(&gene_details).unwrap();
            let serde_value: Value = serde_json::from_str(&gene_details_json).unwrap();

            trans.execute("INSERT INTO web_json.gene (uniquename, data) values ($1, $2)",
                          &[&uniquename, &serde_value]).unwrap();
        }
        for (uniquename, ref_details) in &self.references {
            // FIXME - this is dumb:
            let reference_details_json = serde_json::to_string(&ref_details).unwrap();
            let serde_value: Value = serde_json::from_str(&reference_details_json).unwrap();

            trans.execute("INSERT INTO web_json.reference (uniquename, data) values ($1, $2)",
                          &[&uniquename, &serde_value]).unwrap();
        }
        for (termid, term_details) in &self.terms {
            // FIXME - this is dumb:
            let term_details_json = serde_json::to_string(&term_details).unwrap();
            let serde_value: Value = serde_json::from_str(&term_details_json).unwrap();

            trans.execute("INSERT INTO web_json.term (termid, data) values ($1, $2)",
                         &[&termid, &serde_value]).unwrap();
        }

        trans.execute("CREATE INDEX gene_jsonb_idx ON web_json.gene USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX gene_jsonb_name_idx ON web_json.gene USING gin ((data->>'name') gin_trgm_ops);", &[]).unwrap();
        trans.execute("CREATE INDEX term_jsonb_idx ON web_json.term USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX term_jsonb_name_idx ON web_json.term USING gin ((data->>'name') gin_trgm_ops);", &[]).unwrap();
        trans.execute("CREATE INDEX reference_jsonb_idx ON web_json.reference USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX reference_jsonb_title_idx ON web_json.reference USING gin ((data->>'title') gin_trgm_ops);", &[]).unwrap();

        trans.commit().unwrap();
    }
}
<commit_msg>Make JSONB serialisation code less dumb<commit_after>extern crate serde_json;
extern crate postgres;

use std::fs::File;
use std::io::{Write, BufWriter};
use std::collections::HashMap;

use self::postgres::Connection;

type CvName = String;

pub type MiscExtRange = String;

pub type GeneUniquename = String;
pub type TermUniquename = String;
pub type GeneName = String;
pub type TypeName = String;
pub type GeneProduct = String;

pub type TermName = String;
pub type TermId = String;
pub type TermDef = String;

pub type Evidence = String;
pub type With = String;
pub type Residue = String;
pub type Condition = String;
pub type Qualifier = String;

pub type TypeInteractionAnnotationMap =
    HashMap<TypeName, Vec<InteractionAnnotation>>;

pub type UniquenameGeneMap =
    HashMap<GeneUniquename, GeneDetails>;

pub type TranscriptUniquename = String;

pub type UniquenameTranscriptMap =
    HashMap<TranscriptUniquename, TranscriptDetails>;

pub type GenotypeUniquename = String;
pub type AlleleUniquename = String;

pub type UniquenameAlleleMap = HashMap<AlleleUniquename, AlleleShort>;
pub type UniquenameGenotypeMap = HashMap<GenotypeUniquename, GenotypeShort>;

pub type IdGeneMap = HashMap<GeneUniquename, GeneDetails>;
pub type IdGeneShortMap = HashMap<GeneUniquename, GeneShort>;
pub type IdTermDetailsMap = HashMap<TermId, Rc<TermDetails>>;
pub type IdReferenceMap = HashMap<TermId, ReferenceDetails>;

include!(concat!(env!("OUT_DIR"), "/data_serde.rs"));

impl WebData {
    fn get_genes(&self) -> &IdGeneMap {
        &self.genes
    }
    fn get_terms(&self) -> &IdTermDetailsMap {
        &self.terms
    }

    fn write_reference_details(&self, output_dir: &str) {
        for (reference_uniquename, reference_details) in &self.references {
            let s = serde_json::to_string(&reference_details).unwrap();
            let file_name = String::new() + &output_dir + "/reference/" + &reference_uniquename + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_gene_details(&self, output_dir: &str) {
        for (gene_uniquename, gene_details) in &self.genes {
            let s = serde_json::to_string(&gene_details).unwrap();
            let file_name = String::new() + &output_dir + "/gene/" + &gene_uniquename + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_gene_summary(&self, output_dir: &str, organism_genus_species: &str) {
        let gene_summaries =
            self.gene_summaries.values()
            .filter(|gene_short| {
                let gene_uniquename = &gene_short.uniquename;
                let gene_details = self.get_genes().get(gene_uniquename).unwrap();
                let feature_org_genus_species = String::new() +
                    &gene_details.organism.genus + "_" + &gene_details.organism.species;
                feature_org_genus_species == organism_genus_species
            })
            .map(|gene_short| gene_short.clone())
            .collect::<Vec<GeneShort>>();
        let s = serde_json::to_string(&gene_summaries).unwrap();
        let file_name = String::new() + &output_dir + "/gene_summaries.json";
        let f = File::create(file_name).expect("Unable to open file");
        let mut writer = BufWriter::new(&f);
        writer.write_all(s.as_bytes()).expect("Unable to write!");
    }

    fn write_terms(&self, output_dir: &str) {
        for (termid, term_details) in &self.terms {
            let s = serde_json::to_string(&term_details).unwrap();
            let file_name = String::new() + &output_dir + "/term/" + &termid + ".json";
            let f = File::create(file_name).expect("Unable to open file");
            let mut writer = BufWriter::new(&f);
            writer.write_all(s.as_bytes()).expect("Unable to write!");
        }
    }

    fn write_metadata(&self, output_dir: &str) {
        let s = serde_json::to_string(&self.metadata).unwrap();
        let file_name = String::new() + &output_dir + "/metadata.json";
        let f = File::create(file_name).expect("Unable to open file");
        let mut writer = BufWriter::new(&f);
        writer.write_all(s.as_bytes()).expect("Unable to write!");
    }

    pub fn write(&self, output_dir: &str, organism_genus_species: &str) {
        self.write_reference_details(output_dir);
        self.write_gene_details(output_dir);
        self.write_gene_summary(output_dir, organism_genus_species);
        self.write_terms(output_dir);
        self.write_metadata(output_dir);

        println!("wrote {} genes", self.get_genes().len());
        println!("wrote {} terms", self.get_terms().len());
    }

    pub fn store_jsonb(&self, conn: &Connection) {
        let trans = conn.transaction().unwrap();

        for (uniquename, gene_details) in &self.genes {
            let serde_value = serde_json::value::to_value(&gene_details);
            trans.execute("INSERT INTO web_json.gene (uniquename, data) values ($1, $2)",
                          &[&uniquename, &serde_value]).unwrap();
        }
        for (uniquename, ref_details) in &self.references {
            let serde_value = serde_json::value::to_value(&ref_details);
            trans.execute("INSERT INTO web_json.reference (uniquename, data) values ($1, $2)",
                          &[&uniquename, &serde_value]).unwrap();
        }
        for (termid, term_details) in &self.terms {
            let serde_value = serde_json::value::to_value(&term_details);
            trans.execute("INSERT INTO web_json.term (termid, data) values ($1, $2)",
                         &[&termid, &serde_value]).unwrap();
        }

        trans.execute("CREATE INDEX gene_jsonb_idx ON web_json.gene USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX gene_jsonb_name_idx ON web_json.gene USING gin ((data->>'name') gin_trgm_ops);", &[]).unwrap();
        trans.execute("CREATE INDEX term_jsonb_idx ON web_json.term USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX term_jsonb_name_idx ON web_json.term USING gin ((data->>'name') gin_trgm_ops);", &[]).unwrap();
        trans.execute("CREATE INDEX reference_jsonb_idx ON web_json.reference USING gin (data jsonb_path_ops)", &[]).unwrap();
        trans.execute("CREATE INDEX reference_jsonb_title_idx ON web_json.reference USING gin ((data->>'title') gin_trgm_ops);", &[]).unwrap();

        trans.commit().unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program 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
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    let boost_addr = rs_q::build_term("_type", "addr").with_boost(1000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_match_query = rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                         "zip_codes.prefix".to_string()],
                                                    q.to_string())
        .with_boost(100)
        .build();

    let mut should_query = vec![boost_addr, boost_match_query];
    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(300)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    use rs_es::query::{CombinationMinimumShouldMatch, MinimumShouldMatch};
    let min_match = |min, threshold| {
        CombinationMinimumShouldMatch::new(MinimumShouldMatch::from(min),
                                           MinimumShouldMatch::from(threshold))
    };

    let minimum_should_match = match match_type {
        MatchType::Prefix => {
            // when we search [0, 2] fields, we need to match 100% of them
            // when we search [3, 4] fields, we need to match 90% of them
            // when we search [5, +] fields, we only need to match 70% on them
            // this way if the input has more fields than the documents
            // (additional information like the country, the region, ...)
            // they can be ignored
            MinimumShouldMatch::from(vec![min_match(0, 90f64), min_match(4, 70f64)])
        }
        // for fuzzy search we lower our expectation and we accept 50% of token match
        MatchType::Fuzzy => MinimumShouldMatch::from(50f64),
    };

    let mut must = vec![rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                     "zip_codes.prefix".to_string()],
                                                q.to_string())
                            .with_minimum_should_match(minimum_should_match)
                            .build()];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let filter = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn query(q: &str,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&["munin"])
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape)));
    if results.is_empty() {
        query(&q,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape))
    } else {
        Ok(results)
    }
}
<commit_msg>no error if 2 tokens<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program 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
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    let boost_addr = rs_q::build_term("_type", "addr").with_boost(1000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_match_query = rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                         "zip_codes.prefix".to_string()],
                                                    q.to_string())
        .with_boost(100)
        .build();

    let mut should_query = vec![boost_addr, boost_match_query];
    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(300)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    use rs_es::query::{CombinationMinimumShouldMatch, MinimumShouldMatch};
    let min_match = |min, threshold| {
        CombinationMinimumShouldMatch::new(MinimumShouldMatch::from(min),
                                           MinimumShouldMatch::from(threshold))
    };

    let minimum_should_match = match match_type {
        MatchType::Prefix => {
            // when we search [0, 2] fields, we need to match 100% of them
            // when we search [3, 4] fields, we need to match 90% of them
            // when we search [5, +] fields, we only need to match 70% on them
            // this way if the input has more fields than the documents
            // (additional information like the country, the region, ...)
            // they can be ignored
            MinimumShouldMatch::from(vec![min_match(2, 90f64), min_match(4, 70f64)])
        }
        // for fuzzy search we lower our expectation and we accept 50% of token match
        MatchType::Fuzzy => MinimumShouldMatch::from(50f64),
    };

    let mut must = vec![rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                     "zip_codes.prefix".to_string()],
                                                q.to_string())
                            .with_minimum_should_match(minimum_should_match)
                            .build()];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let filter = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn query(q: &str,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&["munin"])
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape)));
    if results.is_empty() {
        query(&q,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape))
    } else {
        Ok(results)
    }
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::fs;

use super::{git2, GitStatuses, GitError, GitReference, GitBranch, FileStatus};

pub struct GitRepo {
    repo: git2::Repository,
    path: PathBuf,
}

unsafe impl Send for GitRepo {}
unsafe impl Sync for GitRepo {}

impl GitRepo {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, GitError> {
        let path_ref = path.as_ref();
        let repo = git2::Repository::open(path_ref)
            .map_err(|_| GitError::OpenRepo)?;

        Ok(Self {
               repo: repo,
               path: path_ref.to_owned(),
           })
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn statuses(&self) -> Result<GitStatuses, GitError> {
        // self.repo.graph_ahead_behind

        let mut opts = git2::StatusOptions::new();

        opts.include_ignored(false)
            .include_untracked(true)
            .recurse_untracked_dirs(true)
            .include_unreadable_as_untracked(true)
            .disable_pathspec_match(true)
            .exclude_submodules(true);

        let statuses = self.repo
            .statuses(Some(&mut opts))
            .map_err(|_| GitError::Status)?;

        Ok(GitStatuses::new(statuses))
    }

    pub fn reset(&self) -> Result<GitReference, GitError> {
        let head = self.repo.head().map_err(|_| GitError::Reset)?;
        let obj = head.peel(git2::ObjectType::Any)
            .map_err(|_| GitError::Reset)?;

        let mut builder = git2::build::CheckoutBuilder::new();
        let options = builder
            .remove_untracked(true) // this is ignored for a reset :()
            .progress(|path, a, b| {
                          if path == None {
                              return;
                          }

                          println!("{:?} {:?} {:?}", path, a, b)
                      });

        self.repo
            .reset(&obj, git2::ResetType::Hard, Some(options))
            .map_err(|_| GitError::Reset)?;

        Ok(GitReference::new(&head))
    }

    pub fn checkout(&self, branch_name: &str) -> Result<(), GitError> {
        let branch_type = match branch_name.find("origin/") {
            Some(_) => git2::BranchType::Remote,
            None => git2::BranchType::Local,
        };

        let branch = self.repo
            .find_branch(branch_name, branch_type)
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        let obj = branch
            .get()
            .peel(git2::ObjectType::Commit)
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        // This needs to get non-hacky.
        // Should only set head and reset if not already pointing to head.
        match branch_type {
            git2::BranchType::Local => (),
            git2::BranchType::Remote => {
                self.repo.set_head_detached(obj.id()).expect("wut");
                self.repo
                    .reset(&obj, git2::ResetType::Hard, None)
                    .expect("wuufttttt");
                //let branch_str = "refs/heads/topic/ARTC-233";
                //self.repo.set_head(&branch_str).expect("wut");
                return Ok(());
            }
        }

        let mut opts = git2::build::CheckoutBuilder::new();

        self.repo
            .checkout_tree(&obj, Some(&mut opts))
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        let branch_str = format!("refs/heads/{}", branch_name);
        println!("- checking out '{}'", &branch_str);
        let branch_ref = match self.repo.find_reference(&branch_str) {
            Ok(r) => r,
            Err(e) => {
                println!("- checkout error: {}", e.message());
                return Err(GitError::Checkout(GitBranch::from(branch_type)));
            }
        };

        self.repo
            .set_head(branch_ref.name().unwrap())
            .expect("Error setting head");

        println!(" {} {}",
                 match branch_type {
                     git2::BranchType::Local => " [Local]",
                     git2::BranchType::Remote => "[Remote]",
                 },
                 branch_name);

        Ok(())
    }

    pub fn remove_untracked(&self) -> Result<(), GitError> {
        let statuses = self.statuses()?;
        let iter = statuses
            .iter()
            .filter(|x| match *x.status() {
                        FileStatus::New => true,
                        _ => false,
                    });

        // TODO: Finish this nonsense
        for entry in iter {
            let p = self.path.join(entry.path());

            // The whole file/directory distinction might be useless.
            // If a untracked file is removed from an untracked directory, should also
            // remove now empty directory?
            if p.is_file() {
                fs::remove_file(p)
                    .map_err(|_| GitError::RemoveUntracked)?;
            } else if p.is_dir() {
                fs::remove_dir_all(p)
                    .map_err(|_| GitError::RemoveUntracked)?;
            }
        }

        Ok(())
    }

    pub fn state(&self) -> RepoState {
        RepoState::from(self.repo.state())
    }
}

#[derive(Debug)]
pub enum RepoState {
    Clean,
    Merge,
    Revert,
    RevertSequence,
    CherryPick,
    CherryPickSequence,
    Bisect,
    Rebase,
    RebaseInteractive,
    RebaseMerge,
    ApplyMailbox,
    ApplyMailboxOrRebase,
}

impl From<git2::RepositoryState> for RepoState {
    fn from(state: git2::RepositoryState) -> Self {
        use git2::RepositoryState;

        match state {
            RepositoryState::Clean => RepoState::Clean,
            RepositoryState::Merge => RepoState::Merge,
            RepositoryState::Revert => RepoState::Revert,
            RepositoryState::RevertSequence => RepoState::RevertSequence,
            RepositoryState::CherryPick => RepoState::CherryPick,
            RepositoryState::CherryPickSequence => RepoState::CherryPickSequence,
            RepositoryState::Bisect => RepoState::Bisect,
            RepositoryState::Rebase => RepoState::Rebase,
            RepositoryState::RebaseInteractive => RepoState::RebaseInteractive,
            RepositoryState::RebaseMerge => RepoState::RebaseMerge,
            RepositoryState::ApplyMailbox => RepoState::ApplyMailbox,
            RepositoryState::ApplyMailboxOrRebase => RepoState::ApplyMailboxOrRebase,
        }
    }
}
<commit_msg>Only checkout remote when not already pointing to remote<commit_after>use std::path::{Path, PathBuf};
use std::fs;

use super::{git2, GitStatuses, GitError, GitReference, GitBranch, FileStatus};

pub struct GitRepo {
    repo: git2::Repository,
    path: PathBuf,
}

unsafe impl Send for GitRepo {}
unsafe impl Sync for GitRepo {}

impl GitRepo {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, GitError> {
        let path_ref = path.as_ref();
        let repo = git2::Repository::open(path_ref)
            .map_err(|_| GitError::OpenRepo)?;

        Ok(Self {
               repo: repo,
               path: path_ref.to_owned(),
           })
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn statuses(&self) -> Result<GitStatuses, GitError> {
        // self.repo.graph_ahead_behind

        let mut opts = git2::StatusOptions::new();

        opts.include_ignored(false)
            .include_untracked(true)
            .recurse_untracked_dirs(true)
            .include_unreadable_as_untracked(true)
            .disable_pathspec_match(true)
            .exclude_submodules(true);

        let statuses = self.repo
            .statuses(Some(&mut opts))
            .map_err(|_| GitError::Status)?;

        Ok(GitStatuses::new(statuses))
    }

    pub fn reset(&self) -> Result<GitReference, GitError> {
        let head = self.repo.head().map_err(|_| GitError::Reset)?;
        let obj = head.peel(git2::ObjectType::Any)
            .map_err(|_| GitError::Reset)?;

        let mut builder = git2::build::CheckoutBuilder::new();
        let options = builder
            .remove_untracked(true) // this is ignored for a reset :()
            .progress(|path, a, b| {
                          if path == None {
                              return;
                          }

                          println!("{:?} {:?} {:?}", path, a, b)
                      });

        self.repo
            .reset(&obj, git2::ResetType::Hard, Some(options))
            .map_err(|_| GitError::Reset)?;

        Ok(GitReference::new(&head))
    }

    pub fn checkout(&self, branch_name: &str) -> Result<(), GitError> {
        let branch_type = match branch_name.find("origin/") {
            Some(_) => git2::BranchType::Remote,
            None => git2::BranchType::Local,
        };

        let branch = self.repo
            .find_branch(branch_name, branch_type)
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        let obj = branch
            .get()
            .peel(git2::ObjectType::Commit)
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        // This needs to get non-hacky.
        // Should only set head and reset if not already pointing to head.
        match branch_type {
            git2::BranchType::Local => (),
            git2::BranchType::Remote => {
                let head_id = self.repo
                    .head()
                    .expect("Could not resolve head")
                    .peel(git2::ObjectType::Any)
                    .expect("Could not get head ref")
                    .id();

                if head_id == obj.id() {
                    println!("bailing out");
                    return Ok(());
                }

                self.repo.set_head_detached(obj.id()).expect("wut");
                self.repo
                    .reset(&obj, git2::ResetType::Hard, None)
                    .expect("wuufttttt");
                //let branch_str = "refs/heads/topic/ARTC-233";
                //self.repo.set_head(&branch_str).expect("wut");
                return Ok(());
            }
        }

        let mut opts = git2::build::CheckoutBuilder::new();

        self.repo
            .checkout_tree(&obj, Some(&mut opts))
            .map_err(|_| GitError::Checkout(GitBranch::from(branch_type)))?;

        let branch_str = format!("refs/heads/{}", branch_name);
        println!("- checking out '{}'", &branch_str);
        let branch_ref = match self.repo.find_reference(&branch_str) {
            Ok(r) => r,
            Err(e) => {
                println!("- checkout error: {}", e.message());
                return Err(GitError::Checkout(GitBranch::from(branch_type)));
            }
        };

        self.repo
            .set_head(branch_ref.name().unwrap())
            .expect("Error setting head");

        println!(" {} {}",
                 match branch_type {
                     git2::BranchType::Local => " [Local]",
                     git2::BranchType::Remote => "[Remote]",
                 },
                 branch_name);

        Ok(())
    }

    pub fn remove_untracked(&self) -> Result<(), GitError> {
        let statuses = self.statuses()?;
        let iter = statuses
            .iter()
            .filter(|x| match *x.status() {
                        FileStatus::New => true,
                        _ => false,
                    });

        // TODO: Finish this nonsense
        for entry in iter {
            let p = self.path.join(entry.path());

            // The whole file/directory distinction might be useless.
            // If a untracked file is removed from an untracked directory, should also
            // remove now empty directory?
            if p.is_file() {
                fs::remove_file(p)
                    .map_err(|_| GitError::RemoveUntracked)?;
            } else if p.is_dir() {
                fs::remove_dir_all(p)
                    .map_err(|_| GitError::RemoveUntracked)?;
            }
        }

        Ok(())
    }

    pub fn state(&self) -> RepoState {
        RepoState::from(self.repo.state())
    }
}

#[derive(Debug)]
pub enum RepoState {
    Clean,
    Merge,
    Revert,
    RevertSequence,
    CherryPick,
    CherryPickSequence,
    Bisect,
    Rebase,
    RebaseInteractive,
    RebaseMerge,
    ApplyMailbox,
    ApplyMailboxOrRebase,
}

impl From<git2::RepositoryState> for RepoState {
    fn from(state: git2::RepositoryState) -> Self {
        use git2::RepositoryState;

        match state {
            RepositoryState::Clean => RepoState::Clean,
            RepositoryState::Merge => RepoState::Merge,
            RepositoryState::Revert => RepoState::Revert,
            RepositoryState::RevertSequence => RepoState::RevertSequence,
            RepositoryState::CherryPick => RepoState::CherryPick,
            RepositoryState::CherryPickSequence => RepoState::CherryPickSequence,
            RepositoryState::Bisect => RepoState::Bisect,
            RepositoryState::Rebase => RepoState::Rebase,
            RepositoryState::RebaseInteractive => RepoState::RebaseInteractive,
            RepositoryState::RebaseMerge => RepoState::RebaseMerge,
            RepositoryState::ApplyMailbox => RepoState::ApplyMailbox,
            RepositoryState::ApplyMailboxOrRebase => RepoState::ApplyMailboxOrRebase,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::iter::FromIterator;
use std::sync::Arc;

use hab_core::users;
use hab_net;
use hab_net::socket::DEFAULT_CONTEXT;
use protocol::{message, jobsrv};
use zmq;

use config::Config;
use error::{Error, Result};
use feat;
use heartbeat::{HeartbeatCli, HeartbeatMgr};
use log_forwarder::LogForwarder;
use runner::{studio, RunnerCli, RunnerMgr};

enum State {
    Ready,
    Busy,
}

impl Default for State {
    fn default() -> State {
        State::Ready
    }
}

pub struct Server {
    config: Arc<Config>,
    /// Dealer Socket connected to JobSrv
    fe_sock: zmq::Socket,
    hb_cli: HeartbeatCli,
    runner_cli: RunnerCli,
    state: State,
    msg: zmq::Message,
    net_ident: Arc<String>,
}

impl Server {
    pub fn new(config: Config) -> Result<Self> {
        let net_ident = hab_net::socket::srv_ident();
        let fe_sock = (**DEFAULT_CONTEXT).as_mut().socket(zmq::DEALER)?;
        let hb_cli = HeartbeatCli::new(net_ident.clone());
        let runner_cli = RunnerCli::new();
        fe_sock.set_identity(net_ident.as_bytes())?;
        Ok(Server {
            config: Arc::new(config),
            fe_sock: fe_sock,
            hb_cli: hb_cli,
            runner_cli: runner_cli,
            state: State::default(),
            msg: zmq::Message::new()?,
            net_ident: Arc::new(net_ident),
        })
    }

    pub fn run(&mut self) -> Result<()> {
        if self.config.auth_token.is_empty() {
            error!(
                "ERROR: No 'auth_token' config value specified which prevents the \
                   worker from download fetching signing keys."
            );
            return Err(Error::NoAuthTokenError);
        }
        if self.config.network_interface.is_some() && self.config.network_gateway.is_none() {
            error!(
                "ERROR: No 'network_gateway' config value specfied when 'network_interface' \
                   was provided. Both must be present to work correctly."
            );
            return Err(Error::NoNetworkGatewayError);
        }
        if self.config.network_gateway.is_some() && self.config.network_interface.is_some() {
            error!(
                "ERROR: No 'network_interface' config value specfied when 'network_gateway' \
                   was provided. Both must be present to work correctly."
            );
            return Err(Error::NoNetworkInterfaceError);
        }
        init_users()?;
        self.enable_features_from_config();

        HeartbeatMgr::start(&self.config, (&*self.net_ident).clone())?;
        RunnerMgr::start(self.config.clone(), self.net_ident.clone())?;
        LogForwarder::start(&self.config)?;
        self.hb_cli.connect()?;
        self.runner_cli.connect()?;
        for (_, queue, _) in self.config.jobsrv_addrs() {
            println!("Connecting to job queue, {}", queue);
            self.fe_sock.connect(&queue)?;
        }

        let mut fe_msg = false;
        let mut runner_msg = false;
        info!("builder-worker is ready to go.");
        loop {
            {
                let mut items = [
                    self.fe_sock.as_poll_item(1),
                    self.runner_cli.as_poll_item(1),
                ];
                zmq::poll(&mut items, -1)?;
                if items[0].get_revents() & zmq::POLLIN > 0 {
                    fe_msg = true;
                }
                if items[1].get_revents() & zmq::POLLIN > 0 {
                    runner_msg = true;
                }
            }
            if runner_msg {
                {
                    let reply = self.runner_cli.recv_complete()?;
                    self.fe_sock.send(reply, 0)?;
                }
                self.set_ready()?;
                runner_msg = false;
            }
            if fe_msg {
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive empty msg
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive Command msg

                let wc = message::decode::<jobsrv::WorkerCommand>(&self.msg)?;
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive Job msg

                match self.state {
                    State::Ready => {
                        match wc.get_op() {
                            jobsrv::WorkerOperation::StartJob => self.start_job()?,
                            jobsrv::WorkerOperation::CancelJob => {
                                warn!("Received unexpected Cancel for Ready worker")
                            }
                        }
                    }
                    State::Busy => {
                        match wc.get_op() {
                            jobsrv::WorkerOperation::StartJob => self.reject_job()?,
                            jobsrv::WorkerOperation::CancelJob => self.cancel_job()?,
                        }
                    }
                }
                fe_msg = false;
            }
        }
    }

    fn start_job(&mut self) -> Result<()> {
        self.runner_cli.start_job(&self.msg)?;
        {
            let reply = self.runner_cli.recv_ack()?;
            self.fe_sock.send(reply, 0)?;
        }
        self.set_busy()?;
        Ok(())
    }

    fn cancel_job(&mut self) -> Result<()> {
        self.runner_cli.cancel_job(&self.msg)?;
        {
            let reply = self.runner_cli.recv_ack()?;
            self.fe_sock.send(reply, 0)?;
        }
        Ok(())
    }

    fn reject_job(&mut self) -> Result<()> {
        let mut reply = message::decode::<jobsrv::Job>(&self.msg)?;
        reply.set_state(jobsrv::JobState::Rejected);
        self.fe_sock.send(&message::encode(&reply)?, 0)?;
        Ok(())
    }

    fn set_busy(&mut self) -> Result<()> {
        self.hb_cli.set_busy()?;
        self.state = State::Busy;
        Ok(())
    }

    fn set_ready(&mut self) -> Result<()> {
        self.hb_cli.set_ready()?;
        self.state = State::Ready;
        Ok(())
    }

    fn enable_features_from_config(&self) {
        let features: HashMap<_, _> = HashMap::from_iter(vec![("LIST", feat::List)]);
        let features_enabled = self.config.features_enabled.split(",").map(|f| {
            f.trim().to_uppercase()
        });
        for key in features_enabled {
            if features.contains_key(key.as_str()) {
                info!("Enabling feature: {}", key);
                feat::enable(features.get(key.as_str()).unwrap().clone());
            }
        }

        if feat::is_enabled(feat::List) {
            println!("Listing possible feature flags: {:?}", features.keys());
            println!("Enable features by populating 'features_enabled' in config");
        }
    }
}

pub fn run(config: Config) -> Result<()> {
    Server::new(config)?.run()
}

fn init_users() -> Result<()> {
    let uid = users::get_uid_by_name(studio::STUDIO_USER).ok_or(
        Error::NoStudioUser,
    )?;
    let gid = users::get_gid_by_name(studio::STUDIO_GROUP).ok_or(
        Error::NoStudioGroup,
    )?;
    let mut home = studio::STUDIO_HOME.lock().unwrap();
    *home = users::get_home_for_user(studio::STUDIO_USER).ok_or(
        Error::NoStudioGroup,
    )?;
    studio::set_studio_uid(uid);
    studio::set_studio_gid(gid);
    Ok(())
}
<commit_msg>Fix validation of worker config on startup<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::iter::FromIterator;
use std::sync::Arc;

use hab_core::users;
use hab_net;
use hab_net::socket::DEFAULT_CONTEXT;
use protocol::{message, jobsrv};
use zmq;

use config::Config;
use error::{Error, Result};
use feat;
use heartbeat::{HeartbeatCli, HeartbeatMgr};
use log_forwarder::LogForwarder;
use runner::{studio, RunnerCli, RunnerMgr};

enum State {
    Ready,
    Busy,
}

impl Default for State {
    fn default() -> State {
        State::Ready
    }
}

pub struct Server {
    config: Arc<Config>,
    /// Dealer Socket connected to JobSrv
    fe_sock: zmq::Socket,
    hb_cli: HeartbeatCli,
    runner_cli: RunnerCli,
    state: State,
    msg: zmq::Message,
    net_ident: Arc<String>,
}

impl Server {
    pub fn new(config: Config) -> Result<Self> {
        let net_ident = hab_net::socket::srv_ident();
        let fe_sock = (**DEFAULT_CONTEXT).as_mut().socket(zmq::DEALER)?;
        let hb_cli = HeartbeatCli::new(net_ident.clone());
        let runner_cli = RunnerCli::new();
        fe_sock.set_identity(net_ident.as_bytes())?;
        Ok(Server {
            config: Arc::new(config),
            fe_sock: fe_sock,
            hb_cli: hb_cli,
            runner_cli: runner_cli,
            state: State::default(),
            msg: zmq::Message::new()?,
            net_ident: Arc::new(net_ident),
        })
    }

    pub fn run(&mut self) -> Result<()> {
        if self.config.auth_token.is_empty() {
            error!(
                "ERROR: No 'auth_token' config value specified which prevents the \
                   worker from download fetching signing keys."
            );
            return Err(Error::NoAuthTokenError);
        }
        if self.config.network_interface.is_some() && self.config.network_gateway.is_none() {
            error!(
                "ERROR: No 'network_gateway' config value specfied when 'network_interface' \
                   was provided. Both must be present to work correctly."
            );
            return Err(Error::NoNetworkGatewayError);
        }
        if self.config.network_gateway.is_some() && self.config.network_interface.is_none() {
            error!(
                "ERROR: No 'network_interface' config value specfied when 'network_gateway' \
                   was provided. Both must be present to work correctly."
            );
            return Err(Error::NoNetworkInterfaceError);
        }
        init_users()?;
        self.enable_features_from_config();

        HeartbeatMgr::start(&self.config, (&*self.net_ident).clone())?;
        RunnerMgr::start(self.config.clone(), self.net_ident.clone())?;
        LogForwarder::start(&self.config)?;
        self.hb_cli.connect()?;
        self.runner_cli.connect()?;
        for (_, queue, _) in self.config.jobsrv_addrs() {
            println!("Connecting to job queue, {}", queue);
            self.fe_sock.connect(&queue)?;
        }

        let mut fe_msg = false;
        let mut runner_msg = false;
        info!("builder-worker is ready to go.");
        loop {
            {
                let mut items = [
                    self.fe_sock.as_poll_item(1),
                    self.runner_cli.as_poll_item(1),
                ];
                zmq::poll(&mut items, -1)?;
                if items[0].get_revents() & zmq::POLLIN > 0 {
                    fe_msg = true;
                }
                if items[1].get_revents() & zmq::POLLIN > 0 {
                    runner_msg = true;
                }
            }
            if runner_msg {
                {
                    let reply = self.runner_cli.recv_complete()?;
                    self.fe_sock.send(reply, 0)?;
                }
                self.set_ready()?;
                runner_msg = false;
            }
            if fe_msg {
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive empty msg
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive Command msg

                let wc = message::decode::<jobsrv::WorkerCommand>(&self.msg)?;
                self.fe_sock.recv(&mut self.msg, 0)?; // Receive Job msg

                match self.state {
                    State::Ready => {
                        match wc.get_op() {
                            jobsrv::WorkerOperation::StartJob => self.start_job()?,
                            jobsrv::WorkerOperation::CancelJob => {
                                warn!("Received unexpected Cancel for Ready worker")
                            }
                        }
                    }
                    State::Busy => {
                        match wc.get_op() {
                            jobsrv::WorkerOperation::StartJob => self.reject_job()?,
                            jobsrv::WorkerOperation::CancelJob => self.cancel_job()?,
                        }
                    }
                }
                fe_msg = false;
            }
        }
    }

    fn start_job(&mut self) -> Result<()> {
        self.runner_cli.start_job(&self.msg)?;
        {
            let reply = self.runner_cli.recv_ack()?;
            self.fe_sock.send(reply, 0)?;
        }
        self.set_busy()?;
        Ok(())
    }

    fn cancel_job(&mut self) -> Result<()> {
        self.runner_cli.cancel_job(&self.msg)?;
        {
            let reply = self.runner_cli.recv_ack()?;
            self.fe_sock.send(reply, 0)?;
        }
        Ok(())
    }

    fn reject_job(&mut self) -> Result<()> {
        let mut reply = message::decode::<jobsrv::Job>(&self.msg)?;
        reply.set_state(jobsrv::JobState::Rejected);
        self.fe_sock.send(&message::encode(&reply)?, 0)?;
        Ok(())
    }

    fn set_busy(&mut self) -> Result<()> {
        self.hb_cli.set_busy()?;
        self.state = State::Busy;
        Ok(())
    }

    fn set_ready(&mut self) -> Result<()> {
        self.hb_cli.set_ready()?;
        self.state = State::Ready;
        Ok(())
    }

    fn enable_features_from_config(&self) {
        let features: HashMap<_, _> = HashMap::from_iter(vec![("LIST", feat::List)]);
        let features_enabled = self.config.features_enabled.split(",").map(|f| {
            f.trim().to_uppercase()
        });
        for key in features_enabled {
            if features.contains_key(key.as_str()) {
                info!("Enabling feature: {}", key);
                feat::enable(features.get(key.as_str()).unwrap().clone());
            }
        }

        if feat::is_enabled(feat::List) {
            println!("Listing possible feature flags: {:?}", features.keys());
            println!("Enable features by populating 'features_enabled' in config");
        }
    }
}

pub fn run(config: Config) -> Result<()> {
    Server::new(config)?.run()
}

fn init_users() -> Result<()> {
    let uid = users::get_uid_by_name(studio::STUDIO_USER).ok_or(
        Error::NoStudioUser,
    )?;
    let gid = users::get_gid_by_name(studio::STUDIO_GROUP).ok_or(
        Error::NoStudioGroup,
    )?;
    let mut home = studio::STUDIO_HOME.lock().unwrap();
    *home = users::get_home_for_user(studio::STUDIO_USER).ok_or(
        Error::NoStudioGroup,
    )?;
    studio::set_studio_uid(uid);
    studio::set_studio_gid(gid);
    Ok(())
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::js::{JS, JSRef, Temporary};
use dom::bindings::utils::{Reflectable, WindowProxyHandler};
use dom::document::{Document, DocumentHelpers};
use dom::window::Window;

use js::jsapi::JSObject;
use js::glue::{WrapperNew, CreateWrapperProxyHandler, ProxyTraps};
use js::rust::with_compartment;

use std::ptr;

#[allow(raw_pointer_deriving)]
#[jstraceable]
#[privatize]
pub struct BrowserContext {
    history: Vec<SessionHistoryEntry>,
    active_index: uint,
    window_proxy: *mut JSObject,
}

impl BrowserContext {
    pub fn new(document: JSRef<Document>) -> BrowserContext {
        let mut context = BrowserContext {
            history: vec!(SessionHistoryEntry::new(document)),
            active_index: 0,
            window_proxy: ptr::null_mut(),
        };
        context.create_window_proxy();
        context
    }

    pub fn active_document(&self) -> Temporary<Document> {
        Temporary::new(self.history[self.active_index].document.clone())
    }

    pub fn active_window(&self) -> Temporary<Window> {
        let doc = self.active_document().root();
        doc.r().window()
    }

    pub fn window_proxy(&self) -> *mut JSObject {
        assert!(self.window_proxy.is_not_null());
        self.window_proxy
    }

    fn create_window_proxy(&mut self) {
        let win = self.active_window().root();
        let page = win.page();
        let js_info = page.js_info();

        let WindowProxyHandler(handler) = js_info.as_ref().unwrap().dom_static.windowproxy_handler;
        assert!(handler.is_not_null());

        let parent = win.r().reflector().get_jsobject();
        let cx = js_info.as_ref().unwrap().js_context.ptr;
        let wrapper = with_compartment(cx, parent, || unsafe {
            WrapperNew(cx, parent, handler)
        });
        assert!(wrapper.is_not_null());
        self.window_proxy = wrapper;
    }
}

// This isn't a DOM struct, just a convenience struct
// without a reflector, so we don't mark this as #[dom_struct]
#[must_root]
#[privatize]
#[jstraceable]
pub struct SessionHistoryEntry {
    document: JS<Document>,
    children: Vec<BrowserContext>
}

impl SessionHistoryEntry {
    fn new(document: JSRef<Document>) -> SessionHistoryEntry {
        SessionHistoryEntry {
            document: JS::from_rooted(document),
            children: vec!()
        }
    }
}

static PROXY_HANDLER: ProxyTraps = ProxyTraps {
    getPropertyDescriptor: None,
    getOwnPropertyDescriptor: None,
    defineProperty: None,
    getOwnPropertyNames: None,
    delete_: None,
    enumerate: None,

    has: None,
    hasOwn: None,
    get: None,
    set: None,
    keys: None,
    iterate: None,

    call: None,
    construct: None,
    nativeCall: 0 as *const u8,
    hasInstance: None,
    typeOf: None,
    objectClassIs: None,
    obj_toString: None,
    fun_toString: None,
    //regexp_toShared: 0 as *u8,
    defaultValue: None,
    iteratorNext: None,
    finalize: None,
    getElementIfPresent: None,
    getPrototypeOf: None,
    trace: None
};

pub fn new_window_proxy_handler() -> WindowProxyHandler {
    unsafe {
        WindowProxyHandler(CreateWrapperProxyHandler(&PROXY_HANDLER))
    }
}
<commit_msg>Remove unsound Root::deref() call in BrowserContext::create_window_proxy.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::js::{JS, JSRef, Temporary};
use dom::bindings::utils::{Reflectable, WindowProxyHandler};
use dom::document::{Document, DocumentHelpers};
use dom::window::Window;

use js::jsapi::JSObject;
use js::glue::{WrapperNew, CreateWrapperProxyHandler, ProxyTraps};
use js::rust::with_compartment;

use std::ptr;

#[allow(raw_pointer_deriving)]
#[jstraceable]
#[privatize]
pub struct BrowserContext {
    history: Vec<SessionHistoryEntry>,
    active_index: uint,
    window_proxy: *mut JSObject,
}

impl BrowserContext {
    pub fn new(document: JSRef<Document>) -> BrowserContext {
        let mut context = BrowserContext {
            history: vec!(SessionHistoryEntry::new(document)),
            active_index: 0,
            window_proxy: ptr::null_mut(),
        };
        context.create_window_proxy();
        context
    }

    pub fn active_document(&self) -> Temporary<Document> {
        Temporary::new(self.history[self.active_index].document.clone())
    }

    pub fn active_window(&self) -> Temporary<Window> {
        let doc = self.active_document().root();
        doc.r().window()
    }

    pub fn window_proxy(&self) -> *mut JSObject {
        assert!(self.window_proxy.is_not_null());
        self.window_proxy
    }

    fn create_window_proxy(&mut self) {
        let win = self.active_window().root();
        let win = win.r();
        let page = win.page();
        let js_info = page.js_info();

        let WindowProxyHandler(handler) = js_info.as_ref().unwrap().dom_static.windowproxy_handler;
        assert!(handler.is_not_null());

        let parent = win.reflector().get_jsobject();
        let cx = js_info.as_ref().unwrap().js_context.ptr;
        let wrapper = with_compartment(cx, parent, || unsafe {
            WrapperNew(cx, parent, handler)
        });
        assert!(wrapper.is_not_null());
        self.window_proxy = wrapper;
    }
}

// This isn't a DOM struct, just a convenience struct
// without a reflector, so we don't mark this as #[dom_struct]
#[must_root]
#[privatize]
#[jstraceable]
pub struct SessionHistoryEntry {
    document: JS<Document>,
    children: Vec<BrowserContext>
}

impl SessionHistoryEntry {
    fn new(document: JSRef<Document>) -> SessionHistoryEntry {
        SessionHistoryEntry {
            document: JS::from_rooted(document),
            children: vec!()
        }
    }
}

static PROXY_HANDLER: ProxyTraps = ProxyTraps {
    getPropertyDescriptor: None,
    getOwnPropertyDescriptor: None,
    defineProperty: None,
    getOwnPropertyNames: None,
    delete_: None,
    enumerate: None,

    has: None,
    hasOwn: None,
    get: None,
    set: None,
    keys: None,
    iterate: None,

    call: None,
    construct: None,
    nativeCall: 0 as *const u8,
    hasInstance: None,
    typeOf: None,
    objectClassIs: None,
    obj_toString: None,
    fun_toString: None,
    //regexp_toShared: 0 as *u8,
    defaultValue: None,
    iteratorNext: None,
    finalize: None,
    getElementIfPresent: None,
    getPrototypeOf: None,
    trace: None
};

pub fn new_window_proxy_handler() -> WindowProxyHandler {
    unsafe {
        WindowProxyHandler(CreateWrapperProxyHandler(&PROXY_HANDLER))
    }
}
<|endoftext|>


<commit_before>use regex::Regex;

use support::project;

#[test]
fn migration_generate_creates_a_migration_with_the_proper_name() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("hello")
        .run();

    let expected_stdout = Regex::new("\
Creating migrations/\\d{14}_hello/up.sql
Creating migrations/\\d{14}_hello/down.sql\
        ").unwrap();
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert!(expected_stdout.is_match(result.stdout()));

    let migrations = p.migrations();
    assert_eq!(1, migrations.len());

    let migration = &migrations[0];
    assert_eq!("hello", migration.name());
    assert!(migration.path().join("up.sql").exists());
    assert!(migration.path().join("down.sql").exists());
}

#[test]
fn migration_generate_doesnt_require_database_url_to_be_set() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command_without_datatabase_url("migration")
        .arg("generate")
        .arg("hello")
        .run();

    assert!(result.is_success(), "Command failed: {:?}", result);
}

#[test]
fn migration_version_can_be_specified_on_creation() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("hello")
        .arg("--version=1234")
        .run();

    let expected_stdout = "\
Creating migrations/1234_hello/up.sql
Creating migrations/1234_hello/down.sql
";
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert_eq!(expected_stdout, result.stdout());

    assert!(p.has_file("migrations/1234_hello/up.sql"));
    assert!(p.has_file("migrations/1234_hello/down.sql"));
}

#[test]
fn migration_directory_can_be_specified_for_generate_by_command_line_arg() {
    let p = project("migration_name")
        .folder("foo")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("stuff")
        .arg("--version=12345")
        .arg("--migration-dir=foo")
        .run();

    let expected_stdout = "\
Creating foo/12345_stuff/up.sql
Creating foo/12345_stuff/down.sql
";
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert_eq!(expected_stdout, result.stdout());

    assert!(p.has_file("foo/12345_stuff/up.sql"));
    assert!(p.has_file("foo/12345_stuff/down.sql"));
}

#[test]
fn migration_directory_can_be_specified_for_generate_by_env_var() {
    let p = project("migration_name")
        .folder("bar")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("stuff")
        .arg("--version=12345")
        .env("MIGRATION_DIRECTORY", "bar")
        .run();

    let expected_stdout = "\
Creating bar/12345_stuff/up.sql
Creating bar/12345_stuff/down.sql
";
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert_eq!(expected_stdout, result.stdout());

    assert!(p.has_file("bar/12345_stuff/up.sql"));
    assert!(p.has_file("bar/12345_stuff/down.sql"));
}
<commit_msg>Fix some test failures on Windows<commit_after>use regex::Regex;

use support::project;

#[test]
fn migration_generate_creates_a_migration_with_the_proper_name() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("hello")
        .run();

    let expected_stdout = Regex::new("\
Creating migrations.\\d{14}_hello.up.sql
Creating migrations.\\d{14}_hello.down.sql\
        ").unwrap();
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert!(expected_stdout.is_match(result.stdout()));

    let migrations = p.migrations();
    assert_eq!(1, migrations.len());

    let migration = &migrations[0];
    assert_eq!("hello", migration.name());
    assert!(migration.path().join("up.sql").exists());
    assert!(migration.path().join("down.sql").exists());
}

#[test]
fn migration_generate_doesnt_require_database_url_to_be_set() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command_without_datatabase_url("migration")
        .arg("generate")
        .arg("hello")
        .run();

    assert!(result.is_success(), "Command failed: {:?}", result);
}

#[test]
fn migration_version_can_be_specified_on_creation() {
    let p = project("migration_name")
        .folder("migrations")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("hello")
        .arg("--version=1234")
        .run();

    let expected_stdout = Regex::new("\
Creating migrations.1234_hello.up.sql
Creating migrations.1234_hello.down.sql
").unwrap();
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert!(expected_stdout.is_match(result.stdout()));

    assert!(p.has_file("migrations/1234_hello/up.sql"));
    assert!(p.has_file("migrations/1234_hello/down.sql"));
}

#[test]
fn migration_directory_can_be_specified_for_generate_by_command_line_arg() {
    let p = project("migration_name")
        .folder("foo")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("stuff")
        .arg("--version=12345")
        .arg("--migration-dir=foo")
        .run();

    let expected_stdout = Regex::new("\
Creating foo.12345_stuff.up.sql
Creating foo.12345_stuff.down.sql
").unwrap();
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert!(expected_stdout.is_match(result.stdout()));

    assert!(p.has_file("foo/12345_stuff/up.sql"));
    assert!(p.has_file("foo/12345_stuff/down.sql"));
}

#[test]
fn migration_directory_can_be_specified_for_generate_by_env_var() {
    let p = project("migration_name")
        .folder("bar")
        .build();
    let result = p.command("migration")
        .arg("generate")
        .arg("stuff")
        .arg("--version=12345")
        .env("MIGRATION_DIRECTORY", "bar")
        .run();

    let expected_stdout = Regex::new("\
Creating bar.12345_stuff.up.sql
Creating bar.12345_stuff.down.sql
").unwrap();
    assert!(result.is_success(), "Command failed: {:?}", result);
    assert!(expected_stdout.is_match(result.stdout()));

    assert!(p.has_file("bar/12345_stuff/up.sql"));
    assert!(p.has_file("bar/12345_stuff/down.sql"));
}
<|endoftext|>