bucket client (but needs tokio runtime!)

src/tts_download_job.rs

@@ -10,6 +10,7 @@ pub mod util;
 
 use anyhow::anyhow;
 use chrono::Utc;
+use crate::util::bucket_client::BucketClient;
 use crate::job::job_queries::TtsUploadJobRecord;
 use crate::job::job_queries::insert_tts_model;
 use crate::job::job_queries::mark_tts_upload_job_done;
@@ -21,7 +22,7 @@ use crate::util::random_token::random_token;
 use log::{warn, info};
 use sqlx::MySqlPool;
 use sqlx::mysql::MySqlPoolOptions;
-use std::path::PathBuf;
+use std::path::{PathBuf, Path};
 use std::process::Command;
 use std::time::Duration;
 use tempdir::TempDir;
@@ -30,12 +31,19 @@ use std::io::{BufReader, Read};
 use ring::digest::{Context, Digest, SHA256};
 use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE};
 
+// Buckets
+const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
+const ENV_SECRET_KEY : &'static str = "SECRET_KEY";
+const ENV_REGION_NAME : &'static str = "REGION_NAME";
+const ENV_BUCKET_NAME : &'static str = "BUCKET_NAME";
+
 const DEFAULT_RUST_LOG: &'static str = "debug,actix_web=info";
 const DEFAULT_TEMP_DIR: &'static str = "/tmp";
 
 struct Downloader {
   pub download_temp_directory: PathBuf,
   pub mysql_pool: MySqlPool,
+  pub bucket_client: BucketClient,
   // Command to run
   pub download_script: String,
 }
@@ -44,6 +52,8 @@ struct Downloader {
 async fn main() -> AnyhowResult<()> {
   easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG));
 
+  let _ = dotenv::from_filename(".env-secrets").ok();
+
   info!("Obtaining hostname...");
 
   let server_hostname = hostname::get()
@@ -53,6 +63,19 @@ async fn main() -> AnyhowResult<()> {
 
   info!("Hostname: {}", &server_hostname);
 
+  // Bucket stuff
+  let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?;
+  let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?;
+  let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?;
+  let bucket_name = easyenv::get_env_string_required(ENV_BUCKET_NAME)?;
+
+  let bucket_client = BucketClient::create(
+    &access_key,
+    &secret_key,
+    &region_name,
+    &bucket_name,
+  )?;
+
   let temp_directory = easyenv::get_env_string_or_default(
     "DOWNLOAD_TEMP_DIR",
     DEFAULT_TEMP_DIR);
@@ -80,6 +103,7 @@ async fn main() -> AnyhowResult<()> {
   let downloader = Downloader {
     download_temp_directory: temp_directory,
     mysql_pool,
+    bucket_client,
     download_script,
   };
 
@@ -231,6 +255,19 @@ async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> Anyho
 
   info!("File hash: {}", private_bucket_hash);
 
+  let object_name = format!(
+    "user_uploaded_tts_models/{}/{}/{}/{}.bin",
+    &private_bucket_hash[0..1],
+    &private_bucket_hash[1..2],
+    &private_bucket_hash[2..3],
+    &private_bucket_hash,
+  );
+
+  info!("Destination bucket path: {}", object_name);
+
+  let file_path = PathBuf::from(download_filename);
+  downloader.bucket_client.upload_filename(&object_name, &file_path).await?;
+
   info!("Saving model record...");
   let id = insert_tts_model(&downloader.mysql_pool, job, &private_bucket_hash).await?;
 

src/util/bucket_client.rs

@@ -0,0 +1,106 @@
+use anyhow::bail;
+use log::info;
+use s3::bucket::Bucket;
+use s3::creds::Credentials;
+use s3::region::Region;
+use std::path::{PathBuf, Path};
+use std::str::FromStr;
+use std::fs::File;
+use std::io::Read;
+
+#[derive(Clone)]
+pub struct BucketClient {
+  bucket: Bucket,
+}
+
+impl BucketClient {
+
+  pub fn create(
+    access_key: &str,
+    secret_key: &str,
+    region_name: &str,
+    bucket_name: &str) -> anyhow::Result<Self>
+  {
+    let credentials = Credentials::new(
+      Some(&access_key),
+      Some(&secret_key),
+      None,
+      None,
+      None,
+    )?;
+    // NB: The GCS buckets aren't supported by default.
+    let region = Region::Custom {
+      region: region_name.to_owned(),
+      endpoint: "https://storage.googleapis.com".to_owned(),
+    };
+    let bucket = Bucket::new(&bucket_name, region, credentials)?;
+
+    Ok(Self {
+      bucket,
+    })
+  }
+
+  pub async fn upload_file(&self, object_name: &str, bytes: &[u8]) -> anyhow::Result<()> {
+    info!("Filename for bucket: {}", object_name);
+
+    let (_, code) = self.bucket.put_object(object_name, bytes).await?;
+
+    info!("upload code: {}", code);
+
+    Ok(())
+  }
+
+  pub fn upload_file_blocking(&self, object_name: &str, bytes: &[u8]) -> anyhow::Result<()> {
+    info!("Filename for bucket: {}", object_name);
+
+    let (_, code) = self.bucket.put_object_blocking(object_name, bytes)?;
+
+    info!("upload code: {}", code);
+
+    Ok(())
+  }
+
+  pub async fn upload_filename(&self, object_name: &str, filename: &Path) -> anyhow::Result<()> {
+    // TODO: does a newer version of this crate handle streaming/buffering file contents?
+    let mut file = File::open(filename)?;
+    let mut buffer : Vec<u8> = Vec::new();
+    file.read_to_end(&mut buffer)?;
+
+    self.upload_file(object_name, &buffer).await
+  }
+
+  pub fn upload_filename_blocking(&self, object_name: &str, filename: &Path) -> anyhow::Result<()> {
+    // TODO: does a newer version of this crate handle streaming/buffering file contents?
+    let mut file = File::open(filename)?;
+    let mut buffer : Vec<u8> = Vec::new();
+    file.read_to_end(&mut buffer)?;
+
+    self.upload_file_blocking(object_name, &buffer)
+  }
+
+  pub async fn download_file(&self, path: &str) -> anyhow::Result<Vec<u8>> {
+    info!("downloading from bucket: {}", path);
+    let (bytes, code) = self.bucket.get_object(path).await?;
+
+    match code {
+      404 => bail!("File not found in bucket: {}", path),
+      _ => {},
+    }
+
+    info!("download code: {}", code);
+    Ok(bytes)
+  }
+
+  pub fn download_file_blocking(&self, path: &str) -> anyhow::Result<Vec<u8>> {
+    info!("downloading from bucket: {}", path);
+    let (bytes, code) = self.bucket.get_object_blocking(path)?;
+
+    match code {
+      404 => bail!("File not found in bucket: {}", path),
+      _ => {},
+    }
+
+    info!("download code: {}", code);
+    Ok(bytes)
+  }
+}

src/util/mod.rs

@@ -1,4 +1,5 @@
 pub mod anyhow_result;
+pub mod bucket_client;
 pub mod cookie_manager;
 pub mod filesystem;
 pub mod ip_address;