use rust_embed instead

smallcloudai · Dec 19, 2024 · 4dd1db0 · 4dd1db0
1 parent beac40c
commit 4dd1db0
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 55 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -110,3 +110,4 @@ resvg = "0.44.0"
 async-tar = "0.5.0"
 git2 = "0.19.0"
 process-wrap = { version = "8.0.2", features = ["tokio1"] }
+rust-embed = "8.5.0"
diff --git a/build.rs b/build.rs
@@ -1,50 +1,4 @@
-use std::env;
-use std::fs::{self, read_dir, File};
-use std::io::Write;
-
 
 fn main() -> shadow_rs::SdResult<()> {
-    let assets_dir = "assets/integrations";
-    let out_dir = env::var("OUT_DIR").unwrap();
-    let output_file_path = format!("{}/available_icons.rs", out_dir);
-    let mut output_file = File::create(&output_file_path).expect("Failed to create output file");
-
-    writeln!(
-        output_file,
-        "use std::collections::HashMap;\n\npub fn get_available_icons() -> HashMap<&'static str, &'static [u8]> {{\n    let mut icons = HashMap::new();"
-    )
-        .expect("Failed to write to output file");
-
-    for entry in read_dir(assets_dir).expect("Failed to read assets directory") {
-        let entry = entry.expect("Failed to read directory entry");
-        let path = entry.path();
-
-        if path.extension().and_then(|ext| ext.to_str()) == Some("png") {
-            let image_data = fs::read(&path).expect("Failed to read image file");
-            let file_stem = path
-                .file_stem()
-                .and_then(|stem| stem.to_str())
-                .expect("Failed to get file stem");
-
-            let constant_name = format!("{}_ICON_BYTES", file_stem.to_uppercase());
-
-            writeln!(
-                output_file,
-                "    pub const {}: &[u8] = &{:?};",
-                constant_name, image_data
-            )
-                .expect("Failed to write constant definition");
-
-            writeln!(
-                output_file,
-                "    icons.insert(\"{}.png\", {});",
-                file_stem, constant_name
-            )
-                .expect("Failed to write HashMap entry");
-        }
-    }
-
-    writeln!(output_file, "    icons\n}}").expect("Failed to write closing brace");
-
     shadow_rs::new()
 }
diff --git a/src/http/routers/v1/v1_integrations.rs b/src/http/routers/v1/v1_integrations.rs
@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::path::PathBuf;
 use std::sync::Arc;
 use axum::Extension;
@@ -8,12 +9,12 @@ use tokio::sync::RwLock as ARwLock;
 use regex::Regex;
 use axum::extract::Path;
 use axum::extract::Query;
-
-
+use rust_embed::RustEmbed;
 use crate::custom_error::ScratchError;
 use crate::global_context::GlobalContext;
 use crate::integrations::setting_up_integrations::split_path_into_project_and_integration;
 
+
 pub async fn handle_v1_integrations(
     Extension(gcx): Extension<Arc<ARwLock<GlobalContext>>>,
     _: hyper::body::Bytes,
@@ -130,29 +131,29 @@ pub async fn handle_v1_integration_save(
        .unwrap())
 }
 
-mod generated {
-    include!(concat!(env!("OUT_DIR"), "/available_icons.rs"));
-}
+#[derive(RustEmbed)]
+#[folder = "assets/integrations/"]
+struct IntegrationAsset;
 
 pub async fn handle_v1_integration_icon(
     Path(icon_name): Path<String>,
 ) -> axum::response::Result<Response<Body>, ScratchError> {
-    let icons = generated::get_available_icons();
     let sanitized_icon_name = icon_name
         .split('/').last()
         .map(|x| x.replace("_TEMPLATE", "")).ok_or(
         ScratchError::new(StatusCode::BAD_REQUEST, "invalid file name".to_string())
     )?;
-    if let Some(icon_bytes) = icons.get(sanitized_icon_name.as_str()) {
+    if let Some(icon_bytes) = IntegrationAsset::get(&sanitized_icon_name).map(|file| file.data) {
         return Ok(Response::builder()
             .status(StatusCode::OK)
             .header("Content-Type", "image/png")
             .header("Content-Disposition", "inline")
-            .body(Body::from(*icon_bytes))
+            .body(Body::from(icon_bytes))
             .unwrap());
     }
-    Err(ScratchError::new(StatusCode::NOT_FOUND, "icon not found".to_string()))
+    Err(ScratchError::new(StatusCode::NOT_FOUND, format!("icon {} not found", sanitized_icon_name)))
 }
+
 // Define a structure to match query parameters
 #[derive(Deserialize)]
 pub struct HTTPIntegrationDeleteQueryParams {