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

use ruru::AnyObject;
trait Wrap {
    fn wrap(self) -> AnyObject;
}

#[allow(unused_variables)]
pub mod storeid {
    use std::path::PathBuf;

    use ruru::{Class, Object, AnyObject, Boolean, RString, NilClass};

    use libimagstore::storeid::StoreId;

    wrappable_struct!(StoreId, StoreIdWrapper, STOREID_WRAPPER);
    class!(RStoreId);

    use store::Wrap;
    impl Wrap for StoreId {
        fn wrap(self) -> AnyObject {
            Class::from_existing("RStoreId").wrap_data(self, &*STOREID_WRAPPER)
        }
    }

    methods!(
        RStoreId,
        itself,

        fn r_storeid_new(base: RString, id: RString) -> AnyObject {
            let base = match base.map(|b| b.to_string()).map(PathBuf::from) {
                Ok(base) => base,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            let id = match id.map(|id| id.to_string()).map(PathBuf::from) {
                Ok(id) => id,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            match StoreId::new(Some(base), id) {
                Ok(sid) => Class::from_existing("RStoreId").wrap_data(sid, &*STOREID_WRAPPER),
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            }
        }

        fn r_storeid_new_baseless(id: RString) -> AnyObject {
            let id = match id.map(|id| id.to_string()).map(PathBuf::from) {
                Ok(id) => id,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            match StoreId::new(None, id) {
                Ok(sid) => Class::from_existing("RStoreId").wrap_data(sid, &*STOREID_WRAPPER),
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            }
        }

        fn r_storeid_without_base() -> RStoreId {
            let withoutbase : StoreId = itself.get_data(&*STOREID_WRAPPER).clone().without_base();
            Class::from_existing("RStoreId").wrap_data(withoutbase, &*STOREID_WRAPPER)
        }

        fn r_storeid_with_base(base: RString) -> AnyObject {
            let base : PathBuf = match base.map(|b| b.to_string()).map(PathBuf::from) {
                Ok(pb) => pb,
                Err(e) => {
                    // TODO: Exception!
                    error!("Error: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };
            let withoutbase : StoreId = itself.get_data(&*STOREID_WRAPPER).clone().with_base(base);
            Class::from_existing("RStoreId").wrap_data(withoutbase, &*STOREID_WRAPPER)
        }

        fn r_storeid_into_pathbuf() -> AnyObject {
            itself.get_data(&*STOREID_WRAPPER)
                .clone()
                .into_pathbuf()
                // TODO: No unwraps
                .map(|pb| pb.to_str().map(String::from).unwrap())
                .as_ref()
                .map(|s| AnyObject::from(RString::new(s).value()))
                // TODO: Exception!
                .unwrap_or(AnyObject::from(NilClass::new().value()))
        }

        fn r_storeid_exists() -> Boolean {
            Boolean::new(itself.get_data(&*STOREID_WRAPPER).exists())
        }

        fn r_storeid_to_str() -> AnyObject {
            itself.get_data(&*STOREID_WRAPPER)
                .to_str()
                .as_ref()
                .map(|s| AnyObject::from(RString::new(s).value()))
                // TODO: Exception!
                .unwrap_or(AnyObject::from(NilClass::new().value()))
        }

        fn r_storeid_local() -> RString {
            let local = itself.get_data(&*STOREID_WRAPPER).local();
            let local = local.to_str().unwrap(); // TODO: No unwraps
            RString::new(local)
        }

    );

    pub fn setup() -> Class {
        let mut class = Class::new("RStoreId", None);
        class.define(|itself| {
            itself.def_self("new"          , r_storeid_new);
            itself.def_self("new_baseless" , r_storeid_new_baseless);

            itself.def("without_base"      , r_storeid_without_base);
            itself.def("with_base"         , r_storeid_with_base);
            itself.def("into_pathbuf"      , r_storeid_into_pathbuf);
            itself.def("exists"            , r_storeid_exists);
            itself.def("to_str"            , r_storeid_to_str);
            itself.def("local"             , r_storeid_local);
        });
        class
    }

}

#[allow(unused_variables)]
pub mod store {
    pub mod entry {
        use std::collections::BTreeMap;
        use std::error::Error;
        use std::ops::Deref;
        use std::ops::DerefMut;

        use ruru::{Class, Object, AnyObject, Boolean, RString, VM};

        use libimagstore::store::FileLockEntry as FLE;
        use libimagstore::store::EntryHeader;
        use libimagstore::store::EntryContent;
        use libimagstore::store::Entry;

        use ruby_utils::IntoToml;
        use toml_utils::IntoRuby;
        use store::Wrap;

        pub struct FLECustomWrapper(Box<FLE<'static>>);

        impl Deref for FLECustomWrapper {
            type Target = Box<FLE<'static>>;

            fn deref(&self) -> &Self::Target {
                &self.0
            }
        }

        impl DerefMut for FLECustomWrapper {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.0
            }
        }

        wrappable_struct!(FLECustomWrapper, FileLockEntryWrapper, FLE_WRAPPER);
        class!(RFileLockEntry);

        methods!(
            RFileLockEntry,
            itself,

            fn r_get_location() -> AnyObject {
                itself.get_data(&*FLE_WRAPPER).get_location().clone().wrap()
            }
        );

        wrappable_struct!(EntryHeader, EntryHeaderWrapper, ENTRY_HEADER_WRAPPER);
        class!(REntryHeader);

        impl Wrap for EntryHeader {
            fn wrap(self) -> AnyObject {
                Class::from_existing("REntryHeader").wrap_data(self, &*ENTRY_HEADER_WRAPPER)
            }
        }

        methods!(
            REntryHeader,
            itself,

            fn r_entry_header_new() -> AnyObject {
                EntryHeader::new().wrap()
            }

            fn r_entry_header_insert(spec: RString, obj: AnyObject) -> Boolean {
                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false);
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.
                let obj = obj.unwrap(); // possibly not safe... TODO

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).insert(&spec, obj.into_toml()) {
                    Ok(b) => Boolean::new(b),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false);
                    }
                }
            }

            fn r_entry_header_set(spec: RString, obj: AnyObject) -> AnyObject {
                use ruru::NilClass;

                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false).to_any_object();
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.
                let obj = obj.unwrap(); // possibly not safe... TODO

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).set(&spec, obj.into_toml()) {
                    Ok(Some(v)) => v.into_ruby(),
                    Ok(None)    => NilClass::new().to_any_object(),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false).to_any_object();
                    }
                }
            }

            fn r_entry_header_get(spec: RString) -> AnyObject {
                use ruru::NilClass;

                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false).to_any_object();
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).read(&spec) {
                    Ok(Some(v)) => v.into_ruby(),
                    Ok(None)    => NilClass::new().to_any_object(),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false).to_any_object();
                    }
                }
            }

        );

        wrappable_struct!(EntryContent, EntryContentWrapper, ENTRY_CONTENT_WRAPPER);
        class!(REntryContent);

        impl Wrap for EntryContent {
            fn wrap(self) -> AnyObject {
                Class::from_existing("REntryContent").wrap_data(self, &*ENTRY_CONTENT_WRAPPER)
            }
        }

        wrappable_struct!(Entry, EntryWrapper, ENTRY_WRAPPER);
        class!(REntry);

        pub fn setup_filelockentry() -> Class {
            let mut class = Class::new("RFileLockEntry", None);
            class
        }

        pub fn setup_entryheader() -> Class {
            let mut class = Class::new("REntryHeader", None);
            class.define(|itself| {
                itself.def("insert", r_entry_header_insert);
                itself.def("set"   , r_entry_header_set);
                itself.def("read"  , r_entry_header_get);
            });
            class
        }

        pub fn setup_entrycontent() -> Class {
            let string = Class::from_existing("String");
            let mut class = Class::new("REntryContent", Some(&string));
            class
        }
    }

}

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

use ruru::AnyObject;
trait Wrap {
    fn wrap(self) -> AnyObject;
}

#[allow(unused_variables)]
pub mod storeid {
    use std::path::PathBuf;

    use ruru::{Class, Object, AnyObject, Boolean, RString, NilClass};

    use libimagstore::storeid::StoreId;

    wrappable_struct!(StoreId, StoreIdWrapper, STOREID_WRAPPER);
    class!(RStoreId);

    use store::Wrap;
    impl Wrap for StoreId {
        fn wrap(self) -> AnyObject {
            Class::from_existing("RStoreId").wrap_data(self, &*STOREID_WRAPPER)
        }
    }

    methods!(
        RStoreId,
        itself,

        fn r_storeid_new(base: RString, id: RString) -> AnyObject {
            let base = match base.map(|b| b.to_string()).map(PathBuf::from) {
                Ok(base) => base,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            let id = match id.map(|id| id.to_string()).map(PathBuf::from) {
                Ok(id) => id,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            match StoreId::new(Some(base), id) {
                Ok(sid) => Class::from_existing("RStoreId").wrap_data(sid, &*STOREID_WRAPPER),
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            }
        }

        fn r_storeid_new_baseless(id: RString) -> AnyObject {
            let id = match id.map(|id| id.to_string()).map(PathBuf::from) {
                Ok(id) => id,
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };

            match StoreId::new(None, id) {
                Ok(sid) => Class::from_existing("RStoreId").wrap_data(sid, &*STOREID_WRAPPER),
                Err(e) => {
                    // TODO: Exception!
                    error!("Building StoreId object failed: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            }
        }

        fn r_storeid_without_base() -> RStoreId {
            let withoutbase : StoreId = itself.get_data(&*STOREID_WRAPPER).clone().without_base();
            Class::from_existing("RStoreId").wrap_data(withoutbase, &*STOREID_WRAPPER)
        }

        fn r_storeid_with_base(base: RString) -> AnyObject {
            let base : PathBuf = match base.map(|b| b.to_string()).map(PathBuf::from) {
                Ok(pb) => pb,
                Err(e) => {
                    // TODO: Exception!
                    error!("Error: {:?}", e);
                    return AnyObject::from(NilClass::new().value());
                },
            };
            let withoutbase : StoreId = itself.get_data(&*STOREID_WRAPPER).clone().with_base(base);
            Class::from_existing("RStoreId").wrap_data(withoutbase, &*STOREID_WRAPPER)
        }

        fn r_storeid_into_pathbuf() -> AnyObject {
            itself.get_data(&*STOREID_WRAPPER)
                .clone()
                .into_pathbuf()
                // TODO: No unwraps
                .map(|pb| pb.to_str().map(String::from).unwrap())
                .as_ref()
                .map(|s| AnyObject::from(RString::new(s).value()))
                // TODO: Exception!
                .unwrap_or(AnyObject::from(NilClass::new().value()))
        }

        fn r_storeid_exists() -> Boolean {
            Boolean::new(itself.get_data(&*STOREID_WRAPPER).exists())
        }

        fn r_storeid_to_str() -> AnyObject {
            itself.get_data(&*STOREID_WRAPPER)
                .to_str()
                .as_ref()
                .map(|s| AnyObject::from(RString::new(s).value()))
                // TODO: Exception!
                .unwrap_or(AnyObject::from(NilClass::new().value()))
        }

        fn r_storeid_local() -> RString {
            let local = itself.get_data(&*STOREID_WRAPPER).local();
            let local = local.to_str().unwrap(); // TODO: No unwraps
            RString::new(local)
        }

    );

    pub fn setup() -> Class {
        let mut class = Class::new("RStoreId", None);
        class.define(|itself| {
            itself.def_self("new"          , r_storeid_new);
            itself.def_self("new_baseless" , r_storeid_new_baseless);

            itself.def("without_base"      , r_storeid_without_base);
            itself.def("with_base"         , r_storeid_with_base);
            itself.def("into_pathbuf"      , r_storeid_into_pathbuf);
            itself.def("exists"            , r_storeid_exists);
            itself.def("to_str"            , r_storeid_to_str);
            itself.def("local"             , r_storeid_local);
        });
        class
    }

}

#[allow(unused_variables)]
pub mod store {
    pub mod entry {
        use std::collections::BTreeMap;
        use std::error::Error;
        use std::ops::Deref;
        use std::ops::DerefMut;

        use ruru::{Class, Object, AnyObject, Boolean, RString, VM};

        use libimagstore::store::FileLockEntry as FLE;
        use libimagstore::store::EntryHeader;
        use libimagstore::store::EntryContent;
        use libimagstore::store::Entry;

        use ruby_utils::IntoToml;
        use toml_utils::IntoRuby;
        use store::Wrap;

        pub struct FLECustomWrapper(Box<FLE<'static>>);

        impl Deref for FLECustomWrapper {
            type Target = Box<FLE<'static>>;

            fn deref(&self) -> &Self::Target {
                &self.0
            }
        }

        impl DerefMut for FLECustomWrapper {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.0
            }
        }

        wrappable_struct!(FLECustomWrapper, FileLockEntryWrapper, FLE_WRAPPER);
        class!(RFileLockEntry);

        methods!(
            RFileLockEntry,
            itself,

            fn r_get_location() -> AnyObject {
                itself.get_data(&*FLE_WRAPPER).get_location().clone().wrap()
            }

            fn r_get_header() -> AnyObject {
                itself.get_data(&*FLE_WRAPPER).get_header().clone().wrap()
            }

        );

        wrappable_struct!(EntryHeader, EntryHeaderWrapper, ENTRY_HEADER_WRAPPER);
        class!(REntryHeader);

        impl Wrap for EntryHeader {
            fn wrap(self) -> AnyObject {
                Class::from_existing("REntryHeader").wrap_data(self, &*ENTRY_HEADER_WRAPPER)
            }
        }

        methods!(
            REntryHeader,
            itself,

            fn r_entry_header_new() -> AnyObject {
                EntryHeader::new().wrap()
            }

            fn r_entry_header_insert(spec: RString, obj: AnyObject) -> Boolean {
                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false);
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.
                let obj = obj.unwrap(); // possibly not safe... TODO

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).insert(&spec, obj.into_toml()) {
                    Ok(b) => Boolean::new(b),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false);
                    }
                }
            }

            fn r_entry_header_set(spec: RString, obj: AnyObject) -> AnyObject {
                use ruru::NilClass;

                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false).to_any_object();
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.
                let obj = obj.unwrap(); // possibly not safe... TODO

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).set(&spec, obj.into_toml()) {
                    Ok(Some(v)) => v.into_ruby(),
                    Ok(None)    => NilClass::new().to_any_object(),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false).to_any_object();
                    }
                }
            }

            fn r_entry_header_get(spec: RString) -> AnyObject {
                use ruru::NilClass;

                if let Err(ref error) = spec { // raise exception if "spec" is not a String
                    VM::raise(error.to_exception(), error.description());
                    return Boolean::new(false).to_any_object();
                }

                let spec = spec.unwrap().to_string(); // safe because of check above.

                match itself.get_data(&*ENTRY_HEADER_WRAPPER).read(&spec) {
                    Ok(Some(v)) => v.into_ruby(),
                    Ok(None)    => NilClass::new().to_any_object(),
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        return Boolean::new(false).to_any_object();
                    }
                }
            }

        );

        wrappable_struct!(EntryContent, EntryContentWrapper, ENTRY_CONTENT_WRAPPER);
        class!(REntryContent);

        impl Wrap for EntryContent {
            fn wrap(self) -> AnyObject {
                Class::from_existing("REntryContent").wrap_data(self, &*ENTRY_CONTENT_WRAPPER)
            }
        }

        wrappable_struct!(Entry, EntryWrapper, ENTRY_WRAPPER);
        class!(REntry);

        pub fn setup_filelockentry() -> Class {
            let mut class = Class::new("RFileLockEntry", None);
            class
        }

        pub fn setup_entryheader() -> Class {
            let mut class = Class::new("REntryHeader", None);
            class.define(|itself| {
                itself.def("insert", r_entry_header_insert);
                itself.def("set"   , r_entry_header_set);
                itself.def("read"  , r_entry_header_get);
            });
            class
        }

        pub fn setup_entrycontent() -> Class {
            let string = Class::from_existing("String");
            let mut class = Class::new("REntryContent", Some(&string));
            class
        }
    }

}

<|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::objects::{DocType, EsId};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        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();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .get(&url)
                              .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .put(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .post(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    Err(format!("cannot create index: {:?}", res))
                }
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        debug!("get last index: {base_index}/_aliases",
               base_index = get_main_index(doc_type, dataset));
        self.get(&format!("{base_index}/_aliases",
                          base_index = get_main_index(doc_type, dataset)))
            .map_err(|e| e.to_string())
            .and_then(|res| {
                match res.status {
                    StatusCode::Ok => {
                        let value: serde_json::Value = try!(res.read_response()
                                                               .map_err(|e| e.to_string()));
                        Ok(value.as_object()
                                .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                                .unwrap_or_else(|| {
                                    info!("no previous index to delete for type {} and dataset {}",
                                          doc_type,
                                          dataset);
                                    vec![]
                                }))
                    }
                    StatusCode::NotFound => {
                        info!("impossible to find alias {}, no last index to remove",
                              get_main_index(doc_type, dataset));
                        Ok(vec![])
                    }
                    _ => Err(format!("invalid elasticsearch response: {:?}", res)),
                }

            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));
        let main_index = get_main_index(doc_type, dataset);
        try!(self.alias(&main_index, &vec![index.clone()], &last_indexes));
        try!(self.alias("munin", &vec![main_index.to_string()], &vec![]));
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations = AliasOperations {
            actions: add_operations.chain(remove_operations).collect(),
        };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    error!("failed to change aliases for {}, es response: {:?}",
                           alias,
                           res);
                    Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
                }
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
                      .delete_index(&index)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            mut iter: I)
                            -> Result<u32, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut chunk = Vec::new();
        let mut nb = 0;
        loop {
            chunk.clear();
            let addr = match iter.next() {
                Some(a) => a,
                None => break,
            };
            //create the action and apply with_id to it if es_id returned Something
            let action = addr.es_id().into_iter().fold(Action::index(addr), |action, id| action.with_id(id));
            chunk.push(action);

            nb += 1;
            for addr in iter.by_ref().take(1000) {
                let action = addr.es_id().into_iter().fold(Action::index(addr), |action, id| action.with_id(id));
                chunk.push(action);
                nb += 1;
            }
            try!(self.es_client
                     .bulk(&chunk)
                     .with_index(&index)
                     .with_doc_type(T::doc_type())
                     .send());
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<u32, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<commit_msg>refacto of the bulk_insert method<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::objects::{DocType, EsId};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        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();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .get(&url)
                              .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .put(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .post(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    Err(format!("cannot create index: {:?}", res))
                }
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        debug!("get last index: {base_index}/_aliases",
               base_index = get_main_index(doc_type, dataset));
        self.get(&format!("{base_index}/_aliases",
                          base_index = get_main_index(doc_type, dataset)))
            .map_err(|e| e.to_string())
            .and_then(|res| {
                match res.status {
                    StatusCode::Ok => {
                        let value: serde_json::Value = try!(res.read_response()
                                                               .map_err(|e| e.to_string()));
                        Ok(value.as_object()
                                .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                                .unwrap_or_else(|| {
                                    info!("no previous index to delete for type {} and dataset {}",
                                          doc_type,
                                          dataset);
                                    vec![]
                                }))
                    }
                    StatusCode::NotFound => {
                        info!("impossible to find alias {}, no last index to remove",
                              get_main_index(doc_type, dataset));
                        Ok(vec![])
                    }
                    _ => Err(format!("invalid elasticsearch response: {:?}", res)),
                }

            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));
        let main_index = get_main_index(doc_type, dataset);
        try!(self.alias(&main_index, &vec![index.clone()], &last_indexes));
        try!(self.alias("munin", &vec![main_index.to_string()], &vec![]));
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations = AliasOperations {
            actions: add_operations.chain(remove_operations).collect(),
        };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    error!("failed to change aliases for {}, es response: {:?}",
                           alias,
                           res);
                    Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
                }
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
                      .delete_index(&index)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            iter: I)
                            -> Result<usize, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut nb = 0;
        let chunk_size = 1000;
        //fold is used for creating the action and optionally set the id of the object
        let mut actions = iter.map(|v| v.es_id()
                                        .into_iter()
                                        .fold(Action::index(v), |action, id| action.with_id(id)));
        loop {
            let chunk = actions.by_ref().take(chunk_size).collect::<Vec<_>>();
            nb += chunk.len();
            try!(self.es_client
                     .bulk(&chunk)
                     .with_index(&index)
                     .with_doc_type(T::doc_type())
                     .send());

            if chunk.len() < chunk_size {
                break;
            }
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<usize, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<|endoftext|>


<commit_before>//! Provides functions for maintaining database schema.
//!
//! A database migration always provides procedures to update the schema, as well as to revert
//! itself. Diesel's migrations are versioned, and run in order. Diesel also takes care of tracking
//! which migrations have already been run automatically. Your migrations don't need to be
//! idempotent, as Diesel will ensure no migration is run twice unless it has been reverted.
//!
//! Migrations should be placed in a `/migrations` directory at the root of your project (the same
//! directiry as `Cargo.toml`). When any of these functions are run, Diesel will search for the
//! migrations directory in the current directory and its parents, stopping when it finds the
//! directory containing `Cargo.toml`.
//!
//! Individual migrations should be a folder containing exactly two files, `up.sql` and `down.sql`.
//! `up.sql` will be used to run the migration, while `down.sql` will be used for reverting it. The
//! folder itself should have the structure `{version}_{migration_name}`. It is recommended that
//! you use the timestamp of creation for the version.
//!
//! ## Example
//!
//! ```text
//! # Directory Structure
//! - 20151219180527_create_users
//!     - up.sql
//!     - down.sql
//! - 20160107082941_create_posts
//!     - up.sql
//!     - down.sql
//! ```
//!
//! ```sql
//! -- 20151219180527_create_users/up.sql
//! CREATE TABLE users (
//!   id SERIAL PRIMARY KEY,
//!   name VARCHAR NOT NULL,
//!   hair_color VARCHAR
//! );
//! ```
//!
//! ```sql
//! -- 20151219180527_create_users/down.sql
//! DROP TABLE users;
//! ```
//!
//! ```sql
//! -- 20160107082941_create_posts/up.sql
//! CREATE TABLE posts (
//!   id SERIAL PRIMARY KEY,
//!   user_id INTEGER NOT NULL,
//!   title VARCHAR NOT NULL,
//!   body TEXT
//! );
//! ```
//!
//! ```sql
//! -- 20160107082941_create_posts/down.sql
//! DROP TABLE posts;
//! ```
mod migration;
mod migration_error;
mod schema;

pub use self::migration_error::*;

use ::expression::expression_methods::*;
use ::query_dsl::*;
use self::migration::*;
use self::migration_error::MigrationError::*;
use self::schema::NewMigration;
use self::schema::__diesel_schema_migrations::dsl::*;
use {Connection, QueryResult};

use std::collections::HashSet;
use std::env;
use std::path::{PathBuf, Path};

/// Runs all migrations that have not yet been run. This function will print all progress to
/// stdout. This function will return an `Err` if some error occurs reading the migrations, or if
/// any migration fails to run. Each migration is run in its own transaction, so some migrations
/// may be committed, even if a later migration fails to run.
///
/// It should be noted that this runs all migrations that have not already been run, regardless of
/// whether or not their version is later than the latest run migration. This is generally not a
/// problem, and eases the more common case of two developers generating independent migrations on
/// a branch. Whoever created the second one will eventually need to run the first when both
/// branches are merged.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn run_pending_migrations(conn: &Connection) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let already_run = try!(previously_run_migration_versions(conn));
    let migrations_dir = try!(find_migrations_directory());
    let all_migrations = try!(migrations_in_directory(&migrations_dir));
    let pending_migrations = all_migrations.into_iter().filter(|m| {
        !already_run.contains(m.version())
    });
    run_migrations(conn, pending_migrations)
}

/// Reverts the last migration that was run. This function will return an `Err` if no migrations
/// have ever been run.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn revert_latest_migration(conn: &Connection) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let latest_migration_version = try!(latest_run_migration_version(conn));
    revert_migration_with_version(conn, latest_migration_version)
}

/// Reverts the migration with the given version. This function will return an `Err` if a migration
/// with that version cannot be found, an error occurs reading it, or if an error occurs when
/// reverting it.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn revert_migration_with_version(conn: &Connection, ver: String) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let migrations_dir = try!(find_migrations_directory());
    let all_migrations = try!(migrations_in_directory(&migrations_dir));
    let migration_to_revert = all_migrations.into_iter().find(|m| {
        m.version() == ver
    });
    match migration_to_revert {
        Some(m) => revert_migration(&conn, m),
        None => Err(UnknownMigrationVersion(ver).into()),
    }
}

fn create_schema_migrations_table_if_needed(conn: &Connection) -> QueryResult<usize> {
    conn.execute("CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (
        version VARCHAR PRIMARY KEY NOT NULL,
        run_on TIMESTAMP NOT NULL DEFAULT NOW()
    )")
}

fn previously_run_migration_versions(conn: &Connection) -> QueryResult<HashSet<String>> {
    __diesel_schema_migrations.select(version)
        .load(&conn)
        .map(|r| r.collect())
}

fn latest_run_migration_version(conn: &Connection) -> QueryResult<String> {
    use ::expression::dsl::max;
    __diesel_schema_migrations.select(max(version))
        .first(&conn)
}

fn migrations_in_directory(path: &Path) -> Result<Vec<Box<Migration>>, MigrationError> {
    use self::migration::migration_from;

    try!(path.read_dir())
        .filter_map(|entry| {
            let entry = match entry {
                Ok(e) => e,
                Err(e) => return Some(Err(e.into())),
            };
            if !entry.file_name().to_string_lossy().starts_with(".") {
                Some(migration_from(entry.path()))
            } else {
                None
            }
        }).collect()
}

fn run_migrations<T>(conn: &Connection, migrations: T)
    -> Result<(), RunMigrationsError> where
        T: Iterator<Item=Box<Migration>>
{
    use ::query_builder::insert;

    for migration in migrations {
        try!(conn.transaction(|| {
            println!("Running migration {}", migration.version());
            try!(migration.run(conn));
            try!(insert(&NewMigration(migration.version()))
                 .into(__diesel_schema_migrations)
                 .execute(&conn));
            Ok(())
        }));
    }
    Ok(())
}

fn revert_migration(conn: &Connection, migration: Box<Migration>)
    -> Result<(), RunMigrationsError>
{
    use ::query_builder::delete;

    try!(conn.transaction(|| {
        println!("Rolling back migration {}", migration.version());
        try!(migration.revert(conn));
        let target = __diesel_schema_migrations.filter(version.eq(migration.version()));
        try!(delete(target).execute(&conn));
        Ok(())
    }));
    Ok(())
}

/// Returns the directory containing migrations. Will look at for
/// $PWD/migrations. If it is not found, it will search the parents
/// of the current directory, until it reaches the directory containing
/// `Cargo.toml`. Returns `MigrationError::MigrationDirectoryNotFound`
/// if no directory is found.
fn find_migrations_directory() -> Result<PathBuf, MigrationError> {
    search_for_migrations_directory(&try!(env::current_dir()))
}

fn search_for_migrations_directory(path: &Path) -> Result<PathBuf, MigrationError> {
    let migration_path = path.join("migrations");
    if migration_path.is_dir() {
        Ok(migration_path)
    } else if path.join("Cargo.toml").exists() {
        Err(MigrationError::MigrationDirectoryNotFound)
    } else {
        path.parent().map(search_for_migrations_directory)
            .unwrap_or(Err(MigrationError::MigrationDirectoryNotFound))
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use super::*;
    use super::find_migrations_directory;

    use self::tempdir::TempDir;
    use std::{env, fs};

    #[test]
    fn migration_directory_not_found_if_no_migration_dir_exists() {
        let dir = TempDir::new("diesel").unwrap();

        env::set_current_dir(dir.path()).unwrap();

        assert_eq!(Err(MigrationError::MigrationDirectoryNotFound),
            find_migrations_directory());
    }

    #[test]
    fn migration_directory_defaults_to_pwd_slash_migrations() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let migrations_path = temp_path.join("migrations");

        env::set_current_dir(temp_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();

        assert_eq!(Ok(migrations_path), find_migrations_directory());
    }

    #[test]
    fn migration_directory_checks_parents() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let migrations_path = temp_path.join("migrations");
        let child_path = temp_path.join("child");

        fs::create_dir(&child_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();
        env::set_current_dir(child_path).unwrap();

        assert_eq!(Ok(migrations_path), find_migrations_directory());
    }

    #[test]
    fn migration_directory_stops_at_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path();
        let migrations_path = temp_path.join("migrations");
        let child_path = temp_path.join("child");
        let grandchild_path = child_path.join("grandchild");

        fs::create_dir_all(&grandchild_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();
        fs::File::create(child_path.join("Cargo.toml")).unwrap();
        env::set_current_dir(grandchild_path).unwrap();

        assert_eq!(Err(MigrationError::MigrationDirectoryNotFound),
            find_migrations_directory());
    }
}
<commit_msg>Fix a typo in the migrations module documentation<commit_after>//! Provides functions for maintaining database schema.
//!
//! A database migration always provides procedures to update the schema, as well as to revert
//! itself. Diesel's migrations are versioned, and run in order. Diesel also takes care of tracking
//! which migrations have already been run automatically. Your migrations don't need to be
//! idempotent, as Diesel will ensure no migration is run twice unless it has been reverted.
//!
//! Migrations should be placed in a `/migrations` directory at the root of your project (the same
//! directory as `Cargo.toml`). When any of these functions are run, Diesel will search for the
//! migrations directory in the current directory and its parents, stopping when it finds the
//! directory containing `Cargo.toml`.
//!
//! Individual migrations should be a folder containing exactly two files, `up.sql` and `down.sql`.
//! `up.sql` will be used to run the migration, while `down.sql` will be used for reverting it. The
//! folder itself should have the structure `{version}_{migration_name}`. It is recommended that
//! you use the timestamp of creation for the version.
//!
//! ## Example
//!
//! ```text
//! # Directory Structure
//! - 20151219180527_create_users
//!     - up.sql
//!     - down.sql
//! - 20160107082941_create_posts
//!     - up.sql
//!     - down.sql
//! ```
//!
//! ```sql
//! -- 20151219180527_create_users/up.sql
//! CREATE TABLE users (
//!   id SERIAL PRIMARY KEY,
//!   name VARCHAR NOT NULL,
//!   hair_color VARCHAR
//! );
//! ```
//!
//! ```sql
//! -- 20151219180527_create_users/down.sql
//! DROP TABLE users;
//! ```
//!
//! ```sql
//! -- 20160107082941_create_posts/up.sql
//! CREATE TABLE posts (
//!   id SERIAL PRIMARY KEY,
//!   user_id INTEGER NOT NULL,
//!   title VARCHAR NOT NULL,
//!   body TEXT
//! );
//! ```
//!
//! ```sql
//! -- 20160107082941_create_posts/down.sql
//! DROP TABLE posts;
//! ```
mod migration;
mod migration_error;
mod schema;

pub use self::migration_error::*;

use ::expression::expression_methods::*;
use ::query_dsl::*;
use self::migration::*;
use self::migration_error::MigrationError::*;
use self::schema::NewMigration;
use self::schema::__diesel_schema_migrations::dsl::*;
use {Connection, QueryResult};

use std::collections::HashSet;
use std::env;
use std::path::{PathBuf, Path};

/// Runs all migrations that have not yet been run. This function will print all progress to
/// stdout. This function will return an `Err` if some error occurs reading the migrations, or if
/// any migration fails to run. Each migration is run in its own transaction, so some migrations
/// may be committed, even if a later migration fails to run.
///
/// It should be noted that this runs all migrations that have not already been run, regardless of
/// whether or not their version is later than the latest run migration. This is generally not a
/// problem, and eases the more common case of two developers generating independent migrations on
/// a branch. Whoever created the second one will eventually need to run the first when both
/// branches are merged.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn run_pending_migrations(conn: &Connection) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let already_run = try!(previously_run_migration_versions(conn));
    let migrations_dir = try!(find_migrations_directory());
    let all_migrations = try!(migrations_in_directory(&migrations_dir));
    let pending_migrations = all_migrations.into_iter().filter(|m| {
        !already_run.contains(m.version())
    });
    run_migrations(conn, pending_migrations)
}

/// Reverts the last migration that was run. This function will return an `Err` if no migrations
/// have ever been run.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn revert_latest_migration(conn: &Connection) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let latest_migration_version = try!(latest_run_migration_version(conn));
    revert_migration_with_version(conn, latest_migration_version)
}

/// Reverts the migration with the given version. This function will return an `Err` if a migration
/// with that version cannot be found, an error occurs reading it, or if an error occurs when
/// reverting it.
///
/// See the [module level documentation](index.html) for information on how migrations should be
/// structured, and where Diesel will look for them by default.
pub fn revert_migration_with_version(conn: &Connection, ver: String) -> Result<(), RunMigrationsError> {
    try!(create_schema_migrations_table_if_needed(conn));
    let migrations_dir = try!(find_migrations_directory());
    let all_migrations = try!(migrations_in_directory(&migrations_dir));
    let migration_to_revert = all_migrations.into_iter().find(|m| {
        m.version() == ver
    });
    match migration_to_revert {
        Some(m) => revert_migration(&conn, m),
        None => Err(UnknownMigrationVersion(ver).into()),
    }
}

fn create_schema_migrations_table_if_needed(conn: &Connection) -> QueryResult<usize> {
    conn.execute("CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (
        version VARCHAR PRIMARY KEY NOT NULL,
        run_on TIMESTAMP NOT NULL DEFAULT NOW()
    )")
}

fn previously_run_migration_versions(conn: &Connection) -> QueryResult<HashSet<String>> {
    __diesel_schema_migrations.select(version)
        .load(&conn)
        .map(|r| r.collect())
}

fn latest_run_migration_version(conn: &Connection) -> QueryResult<String> {
    use ::expression::dsl::max;
    __diesel_schema_migrations.select(max(version))
        .first(&conn)
}

fn migrations_in_directory(path: &Path) -> Result<Vec<Box<Migration>>, MigrationError> {
    use self::migration::migration_from;

    try!(path.read_dir())
        .filter_map(|entry| {
            let entry = match entry {
                Ok(e) => e,
                Err(e) => return Some(Err(e.into())),
            };
            if !entry.file_name().to_string_lossy().starts_with(".") {
                Some(migration_from(entry.path()))
            } else {
                None
            }
        }).collect()
}

fn run_migrations<T>(conn: &Connection, migrations: T)
    -> Result<(), RunMigrationsError> where
        T: Iterator<Item=Box<Migration>>
{
    use ::query_builder::insert;

    for migration in migrations {
        try!(conn.transaction(|| {
            println!("Running migration {}", migration.version());
            try!(migration.run(conn));
            try!(insert(&NewMigration(migration.version()))
                 .into(__diesel_schema_migrations)
                 .execute(&conn));
            Ok(())
        }));
    }
    Ok(())
}

fn revert_migration(conn: &Connection, migration: Box<Migration>)
    -> Result<(), RunMigrationsError>
{
    use ::query_builder::delete;

    try!(conn.transaction(|| {
        println!("Rolling back migration {}", migration.version());
        try!(migration.revert(conn));
        let target = __diesel_schema_migrations.filter(version.eq(migration.version()));
        try!(delete(target).execute(&conn));
        Ok(())
    }));
    Ok(())
}

/// Returns the directory containing migrations. Will look at for
/// $PWD/migrations. If it is not found, it will search the parents
/// of the current directory, until it reaches the directory containing
/// `Cargo.toml`. Returns `MigrationError::MigrationDirectoryNotFound`
/// if no directory is found.
fn find_migrations_directory() -> Result<PathBuf, MigrationError> {
    search_for_migrations_directory(&try!(env::current_dir()))
}

fn search_for_migrations_directory(path: &Path) -> Result<PathBuf, MigrationError> {
    let migration_path = path.join("migrations");
    if migration_path.is_dir() {
        Ok(migration_path)
    } else if path.join("Cargo.toml").exists() {
        Err(MigrationError::MigrationDirectoryNotFound)
    } else {
        path.parent().map(search_for_migrations_directory)
            .unwrap_or(Err(MigrationError::MigrationDirectoryNotFound))
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use super::*;
    use super::find_migrations_directory;

    use self::tempdir::TempDir;
    use std::{env, fs};

    #[test]
    fn migration_directory_not_found_if_no_migration_dir_exists() {
        let dir = TempDir::new("diesel").unwrap();

        env::set_current_dir(dir.path()).unwrap();

        assert_eq!(Err(MigrationError::MigrationDirectoryNotFound),
            find_migrations_directory());
    }

    #[test]
    fn migration_directory_defaults_to_pwd_slash_migrations() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let migrations_path = temp_path.join("migrations");

        env::set_current_dir(temp_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();

        assert_eq!(Ok(migrations_path), find_migrations_directory());
    }

    #[test]
    fn migration_directory_checks_parents() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let migrations_path = temp_path.join("migrations");
        let child_path = temp_path.join("child");

        fs::create_dir(&child_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();
        env::set_current_dir(child_path).unwrap();

        assert_eq!(Ok(migrations_path), find_migrations_directory());
    }

    #[test]
    fn migration_directory_stops_at_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path();
        let migrations_path = temp_path.join("migrations");
        let child_path = temp_path.join("child");
        let grandchild_path = child_path.join("grandchild");

        fs::create_dir_all(&grandchild_path).unwrap();
        fs::create_dir(&migrations_path).unwrap();
        fs::File::create(child_path.join("Cargo.toml")).unwrap();
        env::set_current_dir(grandchild_path).unwrap();

        assert_eq!(Err(MigrationError::MigrationDirectoryNotFound),
            find_migrations_directory());
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, Write};

use super::{Deserialize, Error, Module, Serialize, VarUint32, VarUint7, Type};
use super::index_map::IndexMap;

const NAME_TYPE_MODULE: u8 = 0;
const NAME_TYPE_FUNCTION: u8 = 1;
const NAME_TYPE_LOCAL: u8 = 2;

/// Debug name information.
#[derive(Clone, Debug, PartialEq)]
pub enum NameSection {
    /// Module name section.
    Module(ModuleNameSection),

    /// Function name section.
    Function(FunctionNameSection),

    /// Local name section.
    Local(LocalNameSection),

    /// Name section is unparsed.
    Unparsed {
        /// The numeric identifier for this name section type.
        name_type: u8,
        /// The contents of this name section, unparsed.
        name_payload: Vec<u8>,
    },
}

impl NameSection {
    /// Deserialize a name section.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<NameSection, Error> {
        let name_type: u8 = VarUint7::deserialize(rdr)?.into();
        let name_payload_len: u32 = VarUint32::deserialize(rdr)?.into();
        let name_section = match name_type {
            NAME_TYPE_MODULE => NameSection::Module(ModuleNameSection::deserialize(rdr)?),
            NAME_TYPE_FUNCTION => NameSection::Function(FunctionNameSection::deserialize(module, rdr)?),
            NAME_TYPE_LOCAL => NameSection::Local(LocalNameSection::deserialize(module, rdr)?),
            _ => {
                let mut name_payload = vec![0u8; name_payload_len as usize];
                rdr.read_exact(&mut name_payload)?;
                NameSection::Unparsed {
                    name_type,
                    name_payload,
                }
            }
        };
        Ok(name_section)
    }
}

impl Serialize for NameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        let (name_type, name_payload) = match self {
            NameSection::Module(mod_name) => {
                let mut buffer = vec![];
                mod_name.serialize(&mut buffer)?;
                (NAME_TYPE_MODULE, buffer)
            }
            NameSection::Function(fn_names) => {
                let mut buffer = vec![];
                fn_names.serialize(&mut buffer)?;
                (NAME_TYPE_FUNCTION, buffer)
            }
            NameSection::Local(local_names) => {
                let mut buffer = vec![];
                local_names.serialize(&mut buffer)?;
                (NAME_TYPE_LOCAL, buffer)
            }
            NameSection::Unparsed {
                name_type,
                name_payload,
            } => (name_type, name_payload),
        };
        VarUint7::from(name_type).serialize(wtr)?;
        VarUint32::from(name_payload.len()).serialize(wtr)?;
        wtr.write_all(&name_payload)?;
        Ok(())
    }
}

/// The name of this module.
#[derive(Clone, Debug, PartialEq)]
pub struct ModuleNameSection {
    name: String,
}

impl ModuleNameSection {
    /// Create a new module name section with the specified name.
    pub fn new<S: Into<String>>(name: S) -> ModuleNameSection {
        ModuleNameSection { name: name.into() }
    }

    /// The name of this module.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The name of this module (mutable).
    pub fn name_mut(&mut self) -> &mut String {
        &mut self.name
    }
}

impl Serialize for ModuleNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.name.serialize(wtr)
    }
}

impl Deserialize for ModuleNameSection {
    type Error = Error;

    fn deserialize<R: Read>(rdr: &mut R) -> Result<ModuleNameSection, Error> {
        let name = String::deserialize(rdr)?;
        Ok(ModuleNameSection { name })
    }
}

/// The names of the functions in this module.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FunctionNameSection {
    names: NameMap,
}

impl FunctionNameSection {
    /// A map from function indices to names.
    pub fn names(&self) -> &NameMap {
        &self.names
    }

    /// A map from function indices to names (mutable).
    pub fn names_mut(&mut self) -> &mut NameMap {
        &mut self.names
    }

    /// Deserialize names, making sure that all names correspond to functions.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<FunctionNameSection, Error> {
        let names = IndexMap::deserialize(module.functions_space(), rdr)?;
        Ok(FunctionNameSection { names })
    }
}

impl Serialize for FunctionNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.names.serialize(wtr)
    }
}

/// The names of the local variables in this module's functions.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct LocalNameSection {
    local_names: IndexMap<NameMap>,
}

impl LocalNameSection {
    /// A map from function indices to a map from variables indices to names.
    pub fn local_names(&self) -> &IndexMap<NameMap> {
        &self.local_names
    }

    /// A map from function indices to a map from variables indices to names
    /// (mutable).
    pub fn local_names_mut(&mut self) -> &mut IndexMap<NameMap> {
        &mut self.local_names
    }

    /// Deserialize names, making sure that all names correspond to local
    /// variables.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<LocalNameSection, Error> {
        let funcs = module.function_section().ok_or_else(|| {
            Error::Other("cannot deserialize local names without a function section")
        })?;
        let max_entry_space = funcs.entries().len();

        let max_signature_args = module
            .type_section()
            .map(|ts|
                ts.types()
                    .iter()
                    .map(|x| { let Type::Function(ref func) = *x; func.params().len() })
                    .max()
                    .unwrap_or(0))
            .unwrap_or(0);

        let max_locals = module
            .code_section()
            .map(|cs| cs.bodies().iter().map(|f| f.locals().len()).max().unwrap_or(0))
            .unwrap_or(0);

        let max_space = max_signature_args + max_locals;

        let deserialize_locals = |_: u32, rdr: &mut R| IndexMap::deserialize(max_space, rdr);

        let local_names = IndexMap::deserialize_with(
            max_entry_space,
            &deserialize_locals,
            rdr,
        )?;
        Ok(LocalNameSection { local_names })
    }}

impl Serialize for LocalNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.local_names.serialize(wtr)
    }
}

/// A map from indices to names.
pub type NameMap = IndexMap<String>;

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use super::*;
    use super::super::{deserialize_file, Section};

    // A helper funtion for the tests. Serialize a section, deserialize it,
    // and make sure it matches the original.
    fn serialize_test(original: NameSection) -> Vec<u8> {
        let mut buffer = vec![];
        original
            .serialize(&mut buffer)
            .expect("serialize error");
        buffer
    }

    #[test]
    fn serialize_module_name() {
        let original = NameSection::Module(ModuleNameSection::new("my_mod"));
        serialize_test(original.clone());
    }

    #[test]
    fn serialize_function_names() {
        let mut sect = FunctionNameSection::default();
        sect.names_mut().insert(0, "hello_world".to_string());
        serialize_test(NameSection::Function(sect));
    }

    #[test]
    fn serialize_local_names() {
        let mut sect = LocalNameSection::default();
        let mut locals = NameMap::default();
        locals.insert(0, "msg".to_string());
        sect.local_names_mut().insert(0, locals);
        serialize_test(NameSection::Local(sect));
    }

    #[test]
    fn serialize_and_deserialize_unparsed() {
        let original = NameSection::Unparsed {
            // A made-up name section type which is unlikely to be allocated
            // soon, in order to allow us to test `Unparsed`.
            name_type: 120,
            name_payload: vec![0u8, 1, 2],
        };
        serialize_test(original.clone());
    }

    #[test]
    fn handles_real_wasm_file() {
        let module = deserialize_file("./res/cases/v1/names.wasm").expect("Should be deserialized");

        // Find the custom section containing function names.
        //
        // TODO: Add a real API for this.
        let mut function_names_opt = None;
        for s in module.sections() {
            match *s {
                Section::Custom(ref cs) if cs.name() == "name" => {
                    let mut cursor = Cursor::new(cs.payload());
                    let section = NameSection::deserialize(&module, &mut cursor)
                        .expect("could not parse name section");
                    if let NameSection::Function(function_names) = section {
                        function_names_opt = Some(function_names);
                    }
                }
                _ => {}
            }
        }
        let function_names = function_names_opt.expect("no name section found");

        assert_eq!(
            function_names.names().get(0).expect("name 0 missing"),
            "double"
        );

        // TODO: Find a test file with module and local names.
    }
}
<commit_msg>remove obsolete test<commit_after>use std::io::{Read, Write};

use super::{Deserialize, Error, Module, Serialize, VarUint32, VarUint7, Type};
use super::index_map::IndexMap;

const NAME_TYPE_MODULE: u8 = 0;
const NAME_TYPE_FUNCTION: u8 = 1;
const NAME_TYPE_LOCAL: u8 = 2;

/// Debug name information.
#[derive(Clone, Debug, PartialEq)]
pub enum NameSection {
    /// Module name section.
    Module(ModuleNameSection),

    /// Function name section.
    Function(FunctionNameSection),

    /// Local name section.
    Local(LocalNameSection),

    /// Name section is unparsed.
    Unparsed {
        /// The numeric identifier for this name section type.
        name_type: u8,
        /// The contents of this name section, unparsed.
        name_payload: Vec<u8>,
    },
}

impl NameSection {
    /// Deserialize a name section.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<NameSection, Error> {
        let name_type: u8 = VarUint7::deserialize(rdr)?.into();
        let name_payload_len: u32 = VarUint32::deserialize(rdr)?.into();
        let name_section = match name_type {
            NAME_TYPE_MODULE => NameSection::Module(ModuleNameSection::deserialize(rdr)?),
            NAME_TYPE_FUNCTION => NameSection::Function(FunctionNameSection::deserialize(module, rdr)?),
            NAME_TYPE_LOCAL => NameSection::Local(LocalNameSection::deserialize(module, rdr)?),
            _ => {
                let mut name_payload = vec![0u8; name_payload_len as usize];
                rdr.read_exact(&mut name_payload)?;
                NameSection::Unparsed {
                    name_type,
                    name_payload,
                }
            }
        };
        Ok(name_section)
    }
}

impl Serialize for NameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        let (name_type, name_payload) = match self {
            NameSection::Module(mod_name) => {
                let mut buffer = vec![];
                mod_name.serialize(&mut buffer)?;
                (NAME_TYPE_MODULE, buffer)
            }
            NameSection::Function(fn_names) => {
                let mut buffer = vec![];
                fn_names.serialize(&mut buffer)?;
                (NAME_TYPE_FUNCTION, buffer)
            }
            NameSection::Local(local_names) => {
                let mut buffer = vec![];
                local_names.serialize(&mut buffer)?;
                (NAME_TYPE_LOCAL, buffer)
            }
            NameSection::Unparsed {
                name_type,
                name_payload,
            } => (name_type, name_payload),
        };
        VarUint7::from(name_type).serialize(wtr)?;
        VarUint32::from(name_payload.len()).serialize(wtr)?;
        wtr.write_all(&name_payload)?;
        Ok(())
    }
}

/// The name of this module.
#[derive(Clone, Debug, PartialEq)]
pub struct ModuleNameSection {
    name: String,
}

impl ModuleNameSection {
    /// Create a new module name section with the specified name.
    pub fn new<S: Into<String>>(name: S) -> ModuleNameSection {
        ModuleNameSection { name: name.into() }
    }

    /// The name of this module.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The name of this module (mutable).
    pub fn name_mut(&mut self) -> &mut String {
        &mut self.name
    }
}

impl Serialize for ModuleNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.name.serialize(wtr)
    }
}

impl Deserialize for ModuleNameSection {
    type Error = Error;

    fn deserialize<R: Read>(rdr: &mut R) -> Result<ModuleNameSection, Error> {
        let name = String::deserialize(rdr)?;
        Ok(ModuleNameSection { name })
    }
}

/// The names of the functions in this module.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FunctionNameSection {
    names: NameMap,
}

impl FunctionNameSection {
    /// A map from function indices to names.
    pub fn names(&self) -> &NameMap {
        &self.names
    }

    /// A map from function indices to names (mutable).
    pub fn names_mut(&mut self) -> &mut NameMap {
        &mut self.names
    }

    /// Deserialize names, making sure that all names correspond to functions.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<FunctionNameSection, Error> {
        let names = IndexMap::deserialize(module.functions_space(), rdr)?;
        Ok(FunctionNameSection { names })
    }
}

impl Serialize for FunctionNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.names.serialize(wtr)
    }
}

/// The names of the local variables in this module's functions.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct LocalNameSection {
    local_names: IndexMap<NameMap>,
}

impl LocalNameSection {
    /// A map from function indices to a map from variables indices to names.
    pub fn local_names(&self) -> &IndexMap<NameMap> {
        &self.local_names
    }

    /// A map from function indices to a map from variables indices to names
    /// (mutable).
    pub fn local_names_mut(&mut self) -> &mut IndexMap<NameMap> {
        &mut self.local_names
    }

    /// Deserialize names, making sure that all names correspond to local
    /// variables.
    pub fn deserialize<R: Read>(
        module: &Module,
        rdr: &mut R,
    ) -> Result<LocalNameSection, Error> {
        let funcs = module.function_section().ok_or_else(|| {
            Error::Other("cannot deserialize local names without a function section")
        })?;
        let max_entry_space = funcs.entries().len();

        let max_signature_args = module
            .type_section()
            .map(|ts|
                ts.types()
                    .iter()
                    .map(|x| { let Type::Function(ref func) = *x; func.params().len() })
                    .max()
                    .unwrap_or(0))
            .unwrap_or(0);

        let max_locals = module
            .code_section()
            .map(|cs| cs.bodies().iter().map(|f| f.locals().len()).max().unwrap_or(0))
            .unwrap_or(0);

        let max_space = max_signature_args + max_locals;

        let deserialize_locals = |_: u32, rdr: &mut R| IndexMap::deserialize(max_space, rdr);

        let local_names = IndexMap::deserialize_with(
            max_entry_space,
            &deserialize_locals,
            rdr,
        )?;
        Ok(LocalNameSection { local_names })
    }}

impl Serialize for LocalNameSection {
    type Error = Error;

    fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
        self.local_names.serialize(wtr)
    }
}

/// A map from indices to names.
pub type NameMap = IndexMap<String>;

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use super::*;
    use super::super::{deserialize_file, Section};

    // A helper funtion for the tests. Serialize a section, deserialize it,
    // and make sure it matches the original.
    fn serialize_test(original: NameSection) -> Vec<u8> {
        let mut buffer = vec![];
        original
            .serialize(&mut buffer)
            .expect("serialize error");
        buffer
    }

    #[test]
    fn serialize_module_name() {
        let original = NameSection::Module(ModuleNameSection::new("my_mod"));
        serialize_test(original.clone());
    }

    #[test]
    fn serialize_function_names() {
        let mut sect = FunctionNameSection::default();
        sect.names_mut().insert(0, "hello_world".to_string());
        serialize_test(NameSection::Function(sect));
    }

    #[test]
    fn serialize_local_names() {
        let mut sect = LocalNameSection::default();
        let mut locals = NameMap::default();
        locals.insert(0, "msg".to_string());
        sect.local_names_mut().insert(0, locals);
        serialize_test(NameSection::Local(sect));
    }

    #[test]
    fn serialize_and_deserialize_unparsed() {
        let original = NameSection::Unparsed {
            // A made-up name section type which is unlikely to be allocated
            // soon, in order to allow us to test `Unparsed`.
            name_type: 120,
            name_payload: vec![0u8, 1, 2],
        };
        serialize_test(original.clone());
    }
}
<|endoftext|>