Artifact 1e0198fa610dfca4069956285ab338f79692ff32dcf751fb7c717c960d4eec17:

• File offload/src/main.rs — part of check-in [b33171b487] at 2018-06-01 10:35:23 on branch master — update script, multiple schemas, rust offloader (user: arcade@b1t.name, size: 2969) [annotate] [blame] [check-ins using]

extern crate config;
extern crate postgres;
extern crate redis;
extern crate regex;

use postgres::{Connection, TlsMode};
use redis::Commands;
use regex::Regex;
use std::collections::HashMap;
use std::time::Duration;
use std::thread;

fn main() {
    loop {
        let mut settings = config::Config::default();
        settings.merge(config::File::with_name("offload")).expect("Can't read configuration file");

        let reddb = redis::Client::open("redis://127.0.0.1/").expect("Can't connect to the database");
        let red = reddb.get_connection().expect("Can't initialize new connection");

        let re = Regex::new("^([0-9A-Z]+)_([a-z./0-9-]+)_([0-9.]+)$").expect("Can't parse regexp");
        let conn = Connection::connect(settings.get_str("pg").expect("Postgres connection absent in config"), TlsMode::None).expect("Can't connect to postgres");

        let schemas = settings.get_array("schemas").expect("Schema list not found in config").into_iter().map(|value| config::Value::into_str(value).expect("We require string here"));
        for schema in schemas {
            let data_key = schema.to_owned() + "_counter_pending";
            let cache_key = schema.to_owned() + "_counter_pending_now";

            if red.exists(&cache_key).expect("Can't query cache existance") {
            } else if red.exists(&data_key).expect("Can't query data existance") {
                let _ : bool = red.rename_nx(&data_key, &cache_key).expect("Can't readd unflushed hash");
            } else {
                continue;
            }

            let stats : HashMap<String, i16> = red.hgetall(&cache_key).expect("Can't query for stored stat keys");

            //println!("# {:?}", stats);

            let trans = conn.transaction().expect("Can't start transaction");
            let stmt = trans.prepare(&format!("select {}.merge_counter($1::text, $2::text, ($3::text)::inet, $4::smallint);", &schema)).expect("Can't prepare statement");
            for (client, count) in stats {
                //println!("# {:?}: {:?}", &client, &count);
                let cap = re.captures(&client).expect("Client match failed");
                //let addr = IpAddr::V4(cap[3].parse().expect("Can't parse IP"));
                //println!("insert into x values({:?}, {:?}, {:?}, {:?});", &cap[1], &cap[2], &cap[3], &count);
                stmt.execute(&[&cap[1].to_string(), &cap[2].to_string(), &cap[3].to_string(), &count]).expect("Can't execute prepared");
                //trans.query(&format!("select {}.merge_counter($1::text, $2::text, $3::inet, $4::smallint);", &schema), &[&cap[1].to_string(), &cap[2].to_string(), &cap[3].to_string(), &count]).expect("Can't prepare statement");
            }
            trans.commit().expect("Can't commit transaction");
            red.del(cache_key).expect("Can't remove stale key")
        }
        thread::sleep(Duration::new(settings.get("delay").expect("Delay specification absent in config"), 0));
    }
}