diff --git a/src/common_queries/firehose_publisher.rs b/src/common_queries/firehose_publisher.rs index 3fecd0a7da..e266218275 100644 --- a/src/common_queries/firehose_publisher.rs +++ b/src/common_queries/firehose_publisher.rs @@ -5,15 +5,48 @@ use log::{warn,info}; use sqlx::{MySqlPool}; use std::sync::Arc; use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDone; +#[derive(Debug, Clone, Copy)] enum FirehoseEvent { UserSignUp, + + TtsModelUploadStarted, + TtsModelUploadCompleted, + TtsInferenceStarted, + TtsInferenceCompleted, + + W2lTemplateUploadStarted, + W2lTemplateUploadCompleted, + W2lInferenceStarted, + W2lInferenceCompleted, + + TwitterMention, + TwitterRetweet, + DiscordJoin, + DiscordMessage, + TwitchSubscribe, + TwitchFollow, } impl FirehoseEvent { pub fn to_db_value(&self) -> &'static str { match self { FirehoseEvent::UserSignUp => "user_sign_up", + FirehoseEvent::TtsModelUploadStarted => "tts_model_upload_started", + FirehoseEvent::TtsModelUploadCompleted => "tts_model_upload_completed", + FirehoseEvent::TtsInferenceStarted => "tts_inference_started", + FirehoseEvent::TtsInferenceCompleted => "tts_inference_completed", + FirehoseEvent::W2lTemplateUploadStarted => "w2l_template_upload_started", + FirehoseEvent::W2lTemplateUploadCompleted => "w2l_template_upload_completed", + FirehoseEvent::W2lInferenceStarted => "w2l_inference_started", + FirehoseEvent::W2lInferenceCompleted => "w2l_inference_completed", + FirehoseEvent::TwitterMention => "twitter_mention", + FirehoseEvent::TwitterRetweet => "twitter_retweet", + FirehoseEvent::DiscordJoin => "discord_join", + FirehoseEvent::DiscordMessage => "discord_message", + FirehoseEvent::TwitchSubscribe => "twitch_subscribe", + FirehoseEvent::TwitchFollow => "twitch_follow", } } } @@ -29,6 +62,49 @@ impl FirehosePublisher { let token = random_prefix_crockford_token("EV", 32)?; let event_type = FirehoseEvent::UserSignUp; + let _record_id = self.insert( + event_type, + Some(user_token), + Some(user_token) + ).await?; + + Ok(()) + } + + pub async fn publish_w2l_template_upload_enqueue(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let token = random_prefix_crockford_token("EV", 32)?; + let event_type = FirehoseEvent::W2lTemplateUploadStarted; + + let _record_id = self.insert( + event_type, + Some(user_token), + Some(job_token) + ).await?; + + Ok(()) + } + + pub async fn publish_w2l_template_upload_finished(&self, user_token: &str, template_token: &str) -> AnyhowResult<()> { + let token = random_prefix_crockford_token("EV", 32)?; + let event_type = FirehoseEvent::W2lTemplateUploadCompleted; + + let _record_id = self.insert( + event_type, + Some(user_token), + Some(template_token) + ).await?; + + Ok(()) + } + + async fn insert( + &self, + event_type: FirehoseEvent, + user_token: Option<&str>, + entity_token: Option<&str> + ) -> AnyhowResult { + let token = random_prefix_crockford_token("EV", 32)?; + let query_result = sqlx::query!( r#" INSERT INTO firehose_entries @@ -41,11 +117,16 @@ SET token, event_type.to_db_value(), user_token, - user_token, + entity_token, ) .execute(&self.mysql_pool) .await; + let record_id = Self::handle_results(query_result)?; + Ok(record_id) + } + + fn handle_results(query_result: Result) -> AnyhowResult { let record_id = match query_result { Ok(res) => { res.last_insert_id() @@ -75,7 +156,6 @@ SET } }; - - Ok(()) + Ok(record_id) } } diff --git a/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs b/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs index 8682f86673..e9b505d266 100644 --- a/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs +++ b/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs @@ -200,6 +200,14 @@ SET info!("new w2l template upload job id: {}", record_id); + server_state.firehose_publisher.publish_w2l_template_upload_enqueue(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadW2lTemplateError::ServerError + })?; + + let response = UploadW2lTemplateSuccessResponse { success: true, job_token: job_token.to_string(), diff --git a/src/job_queries/w2l_download_job_queries.rs b/src/job_queries/w2l_download_job_queries.rs index a8ea1ee508..d3c782ea75 100644 --- a/src/job_queries/w2l_download_job_queries.rs +++ b/src/job_queries/w2l_download_job_queries.rs @@ -126,7 +126,7 @@ pub async fn insert_w2l_template(pool: &MySqlPool, frame_count: u64, fps: f32, duration_millis: u64) - -> AnyhowResult + -> AnyhowResult<(u64, String)> { let model_token = random_prefix_crockford_token("W2L_TPL:", 32)?; let updatable_slug = model_token.clone(); @@ -190,5 +190,5 @@ SET } }; - Ok(record_id) + Ok((record_id, model_token.clone())) } diff --git a/src/w2l_download_job.rs b/src/w2l_download_job.rs index d837779dd4..b731c82efd 100644 --- a/src/w2l_download_job.rs +++ b/src/w2l_download_job.rs @@ -10,6 +10,7 @@ pub mod buckets; pub mod job_queries; pub mod script_execution; +pub mod common_queries; pub mod util; use anyhow::anyhow; @@ -17,12 +18,17 @@ 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::common_queries::firehose_publisher::FirehosePublisher; use crate::job_queries::w2l_download_job_queries::W2lTemplateUploadJobRecord; use crate::job_queries::w2l_download_job_queries::insert_w2l_template; use crate::job_queries::w2l_download_job_queries::mark_w2l_template_upload_job_done; use crate::job_queries::w2l_download_job_queries::mark_w2l_template_upload_job_failure; use crate::job_queries::w2l_download_job_queries::query_w2l_template_upload_job_records; +use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand; +use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand; use crate::script_execution::google_drive_download_command::GoogleDriveDownloadCommand; +use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand; +use crate::script_execution::wav2lip_process_upload_command::Wav2LipPreprocessClient; use crate::util::anyhow_result::AnyhowResult; use crate::util::filesystem::check_directory_exists; use crate::util::filesystem::check_file_exists; @@ -38,10 +44,6 @@ use std::path::{PathBuf, Path}; use std::process::Command; use std::time::Duration; use tempdir::TempDir; -use crate::script_execution::wav2lip_process_upload_command::Wav2LipPreprocessClient; -use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand; -use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand; -use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -72,6 +74,8 @@ struct Downloader { pub private_bucket_client: BucketClient, pub public_bucket_client: BucketClient, + pub firehose_publisher: FirehosePublisher, + pub google_drive_downloader: GoogleDriveDownloadCommand, pub w2l_processor: Wav2LipPreprocessClient, pub ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand, @@ -165,6 +169,10 @@ async fn main() -> AnyhowResult<()> { .max_connections(5) .connect(&db_connection_string) .await?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + }; let downloader = Downloader { download_temp_directory: temp_directory, @@ -173,6 +181,7 @@ async fn main() -> AnyhowResult<()> { private_bucket_client, download_script, google_drive_downloader, + firehose_publisher, ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand {}, ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand {}, imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand {}, @@ -454,7 +463,7 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord) let template_type = if file_metadata.is_video { "video" } else { "image" }; info!("Saving model record..."); - let id = insert_w2l_template( + let (id, model_token) = insert_w2l_template( &downloader.mysql_pool, template_type, job, @@ -475,6 +484,14 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord) info!("Job {} complete success! Downloaded, processed, and uploaded. Saved model record: {}", job.id, id); + downloader.firehose_publisher.publish_w2l_template_upload_finished(&job.creator_user_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + mark_w2l_template_upload_job_done(&downloader.mysql_pool, job, true).await?; Ok(())