Embed thumbnail and metadata into transcoded file

williamyang98 · Aug 15, 2024 · 39f386a · 39f386a
1 parent 81b9a2e
commit 39f386a
Show file tree

Hide file tree

Showing 16 changed files with 971 additions and 113 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -6,8 +6,8 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-actix-web = { version = "4.8.0" }
 actix-files = { version = "0.6.6" }
+actix-web = { version = "4.8.0" }
 clap = { version = "4.5.4", features = ["derive"] }
 dashmap = { version = "6.0.1" }
 derive_more = { version = "0.99.18" }
@@ -20,6 +20,7 @@ num-traits = { version = "0.2" }
 r2d2 = { version = "0.8.10" }
 r2d2_sqlite = { version = "0.24" }
 regex = { version = "1.10.5" }
+reqwest = { version = "0.12.5" }
 rusqlite = { version = "0.31", features = ["bundled"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0" }

diff --git a/scripts/cleanup.sh b/scripts/cleanup.sh
@@ -1,3 +1,2 @@
 #!/bin/sh
-rm ./index.db
 rm -rf ./data/
diff --git a/src/app.rs b/src/app.rs
@@ -2,6 +2,13 @@ use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Condvar};
 use thiserror::Error;
 use threadpool::ThreadPool;
+use dashmap::DashMap;
+use crate::{
+    database::{DatabasePool, VideoId, setup_database},
+    metadata::{MetadataCache, Metadata},
+    worker_download::{DownloadCache, DownloadState},
+    worker_transcode::{TranscodeCache, TranscodeKey, TranscodeState},
+};
 
 pub type WorkerThreadPool = Arc<Mutex<ThreadPool>>;
 pub type WorkerCacheEntry<T> = Arc<(Mutex<T>, Condvar)>;
@@ -61,3 +68,33 @@ impl AppConfig {
         Ok(())
     }
 }
+
+#[derive(Clone)]
+pub struct AppState {
+    pub app_config: Arc<AppConfig>,
+    pub db_pool: DatabasePool,
+    pub worker_thread_pool: WorkerThreadPool,
+    pub download_cache: DownloadCache,
+    pub transcode_cache: TranscodeCache,
+    pub metadata_cache: MetadataCache,
+}
+
+impl AppState {
+    pub fn new(app_config: AppConfig, total_transcode_threads: usize) -> Result<Self, Box<dyn std::error::Error>> {
+        let db_manager = r2d2_sqlite::SqliteConnectionManager::file(app_config.data.join("index.db"));
+        let db_pool = DatabasePool::new(db_manager)?;
+        setup_database(db_pool.get()?)?;
+        let worker_thread_pool: WorkerThreadPool = Arc::new(Mutex::new(ThreadPool::new(total_transcode_threads)));
+        let download_cache: DownloadCache = Arc::new(DashMap::<VideoId, WorkerCacheEntry<DownloadState>>::new());
+        let transcode_cache: TranscodeCache = Arc::new(DashMap::<TranscodeKey, WorkerCacheEntry<TranscodeState>>::new());
+        let metadata_cache: MetadataCache = Arc::new(DashMap::<VideoId, Arc<Metadata>>::new());
+        Ok(Self {
+            app_config: Arc::new(app_config),
+            db_pool, 
+            worker_thread_pool,
+            download_cache,
+            transcode_cache,
+            metadata_cache,
+        })
+    }
+}
diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs
@@ -138,6 +138,9 @@ const TIME_REGEX: &str = r"(?:\d+:)*\d+(?:\.\d+)?";
 
 #[derive(Clone,Copy,Debug,Default)]
 pub struct TranscodeProgress {
+    pub frame: Option<usize>,
+    pub fps: Option<f32>,
+    pub q_factor: Option<f32>,
     pub size_bytes: Option<usize>,
     pub total_time_transcoded: Option<Time>,
     pub speed_bits: Option<usize>,
@@ -160,7 +163,7 @@ pub enum ParsedStderrLine {
 pub fn parse_stderr_line(line: &str) -> Option<ParsedStderrLine> {
     lazy_static! {
         static ref PROGRESS_REGEX: Regex = Regex::new(format!(
-            r"size\s*=\s*(\d+)({0})\s+time\s*=\s*({1})\s+bitrate\s*=\s*({2})({3})\/s\s+speed\s*=\s*({2})\s*x",
+            r"frame\s*=\s*(\d+)\s+fps\s*=\s*({2})\s+q\s*=\s*({2})\s+size\s*=\s*(\d+)({0})\s+time\s*=\s*({1})\s+bitrate\s*=\s*({2})({3})\/s\s+speed\s*=\s*({2})\s*x",
             BYTES_REGEX, TIME_REGEX, FLOAT32_REGEX, BITS_LONG_REGEX,
         ).as_str()).unwrap();
         static ref SOURCE_INFO_REGEX: Regex = Regex::new(format!(
@@ -170,25 +173,31 @@ pub fn parse_stderr_line(line: &str) -> Option<ParsedStderrLine> {
     }
     let line = line.trim();
     if let Some(captures) = PROGRESS_REGEX.captures(line) {
+        let frame: Option<usize> = captures.get(1).and_then(|m| m.as_str().parse().ok());
+        let fps: Option<f32> = captures.get(2).and_then(|m| m.as_str().parse().ok());
+        let q_factor: Option<f32> = captures.get(3).and_then(|m| m.as_str().parse().ok());
         let size_bytes = {
-            let value: Option<u32> = captures.get(1).and_then(|m| m.as_str().parse().ok());
-            let unit: Option<SizeBytes> = captures.get(2).and_then(|m| m.as_str().try_into().ok());
+            let value: Option<u32> = captures.get(4).and_then(|m| m.as_str().parse().ok());
+            let unit: Option<SizeBytes> = captures.get(5).and_then(|m| m.as_str().try_into().ok());
             match (value, unit) {
                 (Some(value), Some(unit)) => Some(value as usize * unit.to_bytes()),
                 _ => None,
             }
         };
-        let total_time_transcoded: Option<Time> = captures.get(3).and_then(|m| Time::try_from_str(m.as_str()).ok());
+        let total_time_transcoded: Option<Time> = captures.get(6).and_then(|m| Time::try_from_str(m.as_str()).ok());
         let speed_bits = {
-            let value: Option<f32> = captures.get(4).and_then(|m| m.as_str().parse().ok());
-            let unit: Option<SizeBits> = captures.get(5).and_then(|m| SizeBits::try_from_long(m.as_str()));
+            let value: Option<f32> = captures.get(7).and_then(|m| m.as_str().parse().ok());
+            let unit: Option<SizeBits> = captures.get(8).and_then(|m| SizeBits::try_from_long(m.as_str()));
             match (value, unit) {
                 (Some(value), Some(unit)) => Some((value * unit.to_bits() as f32) as usize),
                 _ => None,
             }
         };
-        let speed_factor: Option<f32> = captures.get(6).and_then(|m| m.as_str().parse().ok());
+        let speed_factor: Option<f32> = captures.get(9).and_then(|m| m.as_str().parse().ok());
         let result = TranscodeProgress {
+            frame,
+            fps,
+            q_factor,
             size_bytes,
             total_time_transcoded,
             speed_bits,

diff --git a/src/lib.rs b/src/lib.rs
@@ -1,6 +1,7 @@
 pub mod app;
 pub mod database;
 pub mod ffmpeg;
+pub mod metadata;
 pub mod routes;
 pub mod util;
 pub mod worker_download;

diff --git a/src/main.rs b/src/main.rs
@@ -1,14 +1,8 @@
-use std::sync::{Arc, Mutex};
 use std::path::PathBuf;
 use actix_web::{middleware, web, App, HttpServer};
 use clap::Parser;
-use dashmap::DashMap;
-use threadpool::ThreadPool;
 use ytdlp_server::{
-    app::{AppConfig, WorkerThreadPool, WorkerCacheEntry},
-    database::{DatabasePool, VideoId, setup_database},
-    worker_download::{DownloadCache, DownloadState},
-    worker_transcode::{TranscodeCache, TranscodeKey, TranscodeState},
+    app::{AppConfig, AppState},
     routes,
 };
 
@@ -42,6 +36,9 @@ struct Args {
 #[actix_web::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args = Args::parse();
+    if std::env::var("RUST_LOG").is_err() {
+        std::env::set_var("RUST_LOG", "INFO");
+    }
     env_logger::init();
 
     let total_transcode_threads: usize = match args.total_transcode_threads {
@@ -56,21 +53,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     if let Some(path) = args.ytdlp_binary_path { app_config.ytdlp_binary = PathBuf::from(path); }
     if let Some(path) = args.ffmpeg_binary_path { app_config.ffmpeg_binary = PathBuf::from(path); }
     app_config.seed_directories()?;
-    let db_manager = r2d2_sqlite::SqliteConnectionManager::file(app_config.root.join("index.db"));
-    let db_pool = DatabasePool::new(db_manager)?;
-    setup_database(db_pool.get()?)?;
-    let download_cache: DownloadCache = Arc::new(DashMap::<VideoId, WorkerCacheEntry<DownloadState>>::new());
-    let transcode_cache: TranscodeCache = Arc::new(DashMap::<TranscodeKey, WorkerCacheEntry<TranscodeState>>::new());
-    let worker_thread_pool: WorkerThreadPool = Arc::new(Mutex::new(ThreadPool::new(total_transcode_threads)));
+    let app_state = AppState::new(app_config, total_transcode_threads)?;
     // start server
     const API_PREFIX: &str = "/api/v1";
     HttpServer::new(move || {
         App::new()
-            .app_data(app_config.clone())
-            .app_data(db_pool.clone())
-            .app_data(worker_thread_pool.clone())
-            .app_data(download_cache.clone())
-            .app_data(transcode_cache.clone())
+            .app_data(app_state.clone())
             .service(web::scope(API_PREFIX)
                 .service(routes::request_transcode)
                 .service(routes::delete_transcode)
@@ -82,6 +70,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .service(routes::get_download_state)
                 .service(routes::get_transcode_state)
                 .service(routes::get_download_link)
+                .service(routes::get_metadata)
             )
             .service(actix_files::Files::new("/data", "./data/").show_files_listing())
             .service(actix_files::Files::new("/", "./static/").index_file("index.html"))

diff --git a/src/metadata.rs b/src/metadata.rs
@@ -0,0 +1,74 @@
+use std::{collections::HashMap, sync::Arc};
+use dashmap::DashMap;
+use serde::{Serialize,Deserialize};
+use crate::database::VideoId;
+
+pub type MetadataCache = Arc<DashMap<VideoId, Arc<Metadata>>>;
+
+pub fn get_metadata_url(video_id: &str) -> String {
+    const URL: &str = "https://www.googleapis.com/youtube/v3/videos";
+    const PARTS: &str = "snippet,contentDetails";
+    const API_KEY: &str = "AIzaSyDkmFSz9gH9slSnonGjs8TZEjtAKS4e9cg";
+    format!("{URL}?part={PARTS}&id={video_id}&key={API_KEY}")
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct Thumbnail {
+    pub url: String,
+    pub width: usize,
+    pub height: usize,
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct ContentDetails {
+    pub duration: String,
+    pub dimension: String,
+    pub definition: String,
+    pub caption: String,
+    #[serde(rename="licensedContent")]
+    pub licensed_content: bool,
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct Snippet {
+    #[serde(rename="publishedAt")]
+    pub published_at: String,
+    #[serde(rename="channelId")]
+    pub channel_id: String,
+    pub title: String,
+    pub description: String,
+    pub thumbnails: HashMap<String, Thumbnail>,
+    #[serde(rename="channelTitle")]
+    pub channel_title: String,
+    pub tags: Vec<String>,
+    #[serde(rename="categoryId")]
+    pub category_id: String,
+
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct Item {
+    pub id: String,
+    pub etag: String,
+    pub kind: String,
+    pub snippet: Snippet,
+    #[serde(rename="contentDetails")]
+    pub content_details: ContentDetails,
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct PageInfo {
+    #[serde(rename="totalResults")]
+    pub total_results: usize,
+    #[serde(rename="resultsPerPage")]
+    pub results_per_page: usize,
+}
+
+#[derive(Clone,Debug,Deserialize,Serialize)]
+pub struct Metadata {
+    pub kind: String,
+    pub etag: String,
+    pub items: Vec<Item>,
+    #[serde(rename="pageInfo")]
+    pub page_info: PageInfo,
+}