septastic/data_loader/src/main.rs

98 lines
2.6 KiB
Rust

use std::{clone, env, fs::File, io::{Read, Write}, path::{Path, PathBuf}, sync::{Arc, Mutex}, thread, time::Duration};
use dotenv::dotenv;
use serde::{Deserialize, Serialize};
use sqlx::postgres::PgPoolOptions;
use env_logger::{Builder, Env};
use log::{error, info, warn};
#[tokio::main]
async fn main() -> ::anyhow::Result<()> {
dotenv().ok();
let env = Env::new().filter_or("RUST_LOG", "data_loader=info");
Builder::from_env(env).init();
let mut file = File::open("config.yaml")?;
let mut file_contents = String::new();
file.read_to_string(&mut file_contents);
let config_file = serde_yaml::from_str::<Config>(file_contents.as_str());
let database_url = std::env::var("DATABASE_URL").expect("Database URL");
let pool = PgPoolOptions::new()
.max_connections(5)
.connect(&database_url)
.await?;
let mut tx = pool.begin().await?;
tx.commit().await?;
pool.close().await;
Ok(())
}
#[derive(Serialize, Deserialize, PartialEq, Debug)]
struct Config {
gtfs_zips: Vec<String>
}
struct GtfsFile {
pub url: String,
pub hash: Option<String>
}
struct GtfsPullServiceState {
pub gtfs_files: Vec<GtfsFile>,
pub tmp_dir: PathBuf
}
pub struct GtfsPullService {
state: Arc<Mutex<GtfsPullServiceState>>
}
impl GtfsPullService {
const UPDATE_SECONDS: u64 = 3600*24;
pub fn new(config: Config) -> Self {
Self {
state: Arc::new(Mutex::new(
GtfsPullServiceState {
gtfs_files: config.gtfs_zips.iter().map(|f| { GtfsFile { url: f.clone(), hash: None} }).collect(),
tmp_dir: env::temp_dir()
}
))
}
}
pub fn start(&self) {
let cloned_state = Arc::clone(&self.state);
thread::spawn(move || {
loop {
let recloned_state = Arc::clone(&cloned_state);
let res = Self::update_gtfs_data(recloned_state);
match res {
Err(err) => {
error!("{}", err);
}
_ => {}
}
thread::sleep(Duration::from_secs(Self::UPDATE_SECONDS));
}
});
}
pub fn update_gtfs_data(state: Arc<Mutex<GtfsPullServiceState>>) -> anyhow::Result<()> {
let l_state = state.lock().unwrap();
for gtfs_file in l_state.gtfs_files.iter() {
let resp = reqwest::blocking::get(gtfs_file.url.clone())?;
}
Ok(())
}
}