From 9aae3faab351b16f5c07445dbaeb34b0e0935be6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 May 2021 00:54:04 -0400 Subject: [PATCH] more common pieces --- .env | 4 ++-- src/buckets/file_hashing.rs | 29 +++++++++++++++++++++++++++ src/buckets/mod.rs | 1 + src/tts_download_job.rs | 40 ++++++++----------------------------- src/w2l_download_job.rs | 34 +++++-------------------------- 5 files changed, 45 insertions(+), 63 deletions(-) create mode 100644 src/buckets/file_hashing.rs diff --git a/.env b/.env index 2651903499..4730549117 100644 --- a/.env +++ b/.env @@ -11,5 +11,5 @@ COOKIE_DOMAIN=.jungle.horse # Bucket config (non-secret) REGION_NAME=us-east1 -BUCKET_NAME=dev-vocodes-private-uploads -#BUCKET_NAME=vocodes-audio-uploads +TTS_DOWNLOAD_BUCKET_NAME=dev-vocodes-private-uploads +TTS_DOWNLOAD_BUCKET_ROOT=/user_uploaded_tts_models diff --git a/src/buckets/file_hashing.rs b/src/buckets/file_hashing.rs new file mode 100644 index 0000000000..13fefee0a5 --- /dev/null +++ b/src/buckets/file_hashing.rs @@ -0,0 +1,29 @@ +use crate::util::anyhow_result::AnyhowResult; +use data_encoding::HEXLOWER_PERMISSIVE; +use ring::digest::{Context, Digest, SHA256}; +use std::fs::File; +use std::io::{Read, BufReader}; + +pub fn get_file_hash(filename: &str) -> AnyhowResult { + let input = File::open(filename)?; + let reader = BufReader::new(input); + let digest = sha256_digest(reader)?; + + let hash = HEXLOWER_PERMISSIVE.encode(digest.as_ref()); + Ok(hash) +} + +fn sha256_digest(mut reader: R) -> AnyhowResult { + let mut context = Context::new(&SHA256); + let mut buffer = [0; 1024]; + + loop { + let count = reader.read(&mut buffer)?; + if count == 0 { + break; + } + context.update(&buffer[..count]); + } + + Ok(context.finish()) +} diff --git a/src/buckets/mod.rs b/src/buckets/mod.rs index 90dd7c3a20..477fc55a78 100644 --- a/src/buckets/mod.rs +++ b/src/buckets/mod.rs @@ -1,2 +1,3 @@ pub mod bucket_client; pub mod bucket_paths; +pub mod file_hashing; diff --git a/src/tts_download_job.rs b/src/tts_download_job.rs index 6f3d774fa5..149fb8d7c0 100644 --- a/src/tts_download_job.rs +++ b/src/tts_download_job.rs @@ -12,6 +12,8 @@ pub mod util; use anyhow::anyhow; use chrono::Utc; use crate::buckets::bucket_client::BucketClient; +use crate::buckets::bucket_paths::hash_to_bucket_path; +use crate::buckets::file_hashing::get_file_hash; use crate::job::job_queries::TtsUploadJobRecord; use crate::job::job_queries::insert_tts_model; use crate::job::job_queries::mark_tts_upload_job_done; @@ -48,6 +50,8 @@ struct Downloader { pub bucket_client: BucketClient, // Command to run pub download_script: String, + // Root to store TTS results + pub bucket_root_tts_model_uploads: String, } #[tokio::main] @@ -77,7 +81,7 @@ async fn main() -> AnyhowResult<()> { &secret_key, ®ion_name, &bucket_name, - Some(&bucket_root), + None, )?; let temp_directory = easyenv::get_env_string_or_default( @@ -109,6 +113,7 @@ async fn main() -> AnyhowResult<()> { mysql_pool, bucket_client, download_script, + bucket_root_tts_model_uploads: bucket_root.to_string(), }; main_loop(downloader).await; @@ -217,31 +222,6 @@ async fn download_file(downloader: &Downloader, Ok(filename) } - -fn sha256_digest(mut reader: R) -> AnyhowResult { - let mut context = Context::new(&SHA256); - let mut buffer = [0; 1024]; - - loop { - let count = reader.read(&mut buffer)?; - if count == 0 { - break; - } - context.update(&buffer[..count]); - } - - Ok(context.finish()) -} - -fn get_file_hash(filename: &str) -> AnyhowResult { - let input = File::open(filename)?; - let reader = BufReader::new(input); - let digest = sha256_digest(reader)?; - - let hash = HEXLOWER_PERMISSIVE.encode(digest.as_ref()); - Ok(hash) -} - async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> AnyhowResult<()> { // TODO: 1. Mark processing. // TODO: 2. Download. (DONE) @@ -260,13 +240,9 @@ async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> Anyho info!("File hash: {}", private_bucket_hash); // NB: /.../a/b/c/d/abcdefg.bin - let object_name = format!( - "/user_uploaded_tts_models/{}/{}/{}/{}.bin", - &private_bucket_hash[0..1], - &private_bucket_hash[1..2], - &private_bucket_hash[2..3], + let object_name = hash_to_bucket_path( &private_bucket_hash, - ); + Some(&downloader.bucket_root_tts_model_uploads))?; info!("Destination bucket path: {}", object_name); diff --git a/src/w2l_download_job.rs b/src/w2l_download_job.rs index 4974baaf06..26b98b23f6 100644 --- a/src/w2l_download_job.rs +++ b/src/w2l_download_job.rs @@ -12,6 +12,7 @@ pub mod util; use anyhow::anyhow; use chrono::Utc; use crate::buckets::bucket_client::BucketClient; +use crate::buckets::file_hashing::get_file_hash; use crate::job::job_queries::TtsUploadJobRecord; use crate::job::job_queries::insert_tts_model; use crate::job::job_queries::mark_tts_upload_job_done; @@ -20,17 +21,17 @@ use crate::job::job_queries::query_tts_upload_job_records; use crate::util::anyhow_result::AnyhowResult; use crate::util::filesystem::check_directory_exists; use crate::util::random_token::random_token; +use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; use log::{warn, info}; +use ring::digest::{Context, Digest, SHA256}; use sqlx::MySqlPool; use sqlx::mysql::MySqlPoolOptions; +use std::fs::File; +use std::io::{BufReader, Read}; use std::path::{PathBuf, Path}; use std::process::Command; use std::time::Duration; use tempdir::TempDir; -use std::fs::File; -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"; @@ -216,31 +217,6 @@ async fn download_file(downloader: &Downloader, Ok(filename) } - -fn sha256_digest(mut reader: R) -> AnyhowResult { - let mut context = Context::new(&SHA256); - let mut buffer = [0; 1024]; - - loop { - let count = reader.read(&mut buffer)?; - if count == 0 { - break; - } - context.update(&buffer[..count]); - } - - Ok(context.finish()) -} - -fn get_file_hash(filename: &str) -> AnyhowResult { - let input = File::open(filename)?; - let reader = BufReader::new(input); - let digest = sha256_digest(reader)?; - - let hash = HEXLOWER_PERMISSIVE.encode(digest.as_ref()); - Ok(hash) -} - async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> AnyhowResult<()> { // TODO: 1. Mark processing. // TODO: 2. Download. (DONE)