Feat/import huggingface

src/server/src/controllers/workspaces/files.rs

@@ -14,10 +14,15 @@ use actix_web::{web, HttpRequest, HttpResponse};
 
 use actix_multipart::Multipart;
 use actix_web::Error;
+use flate2::read::GzDecoder;
+use futures::StreamExt;
 use futures_util::TryStreamExt as _;
+use std::io::Read;
 use std::io::Write;
 use std::path::{Path, PathBuf};
 
+const BUFFER_SIZE_THRESHOLD: usize = 262144; // 256kb
+
 pub async fn get(
     req: HttpRequest,
     query: web::Query<ImgResize>,
@@ -79,6 +84,139 @@ pub async fn add(req: HttpRequest, payload: Multipart) -> Result<HttpResponse, O
     }))
 }
 
+pub async fn add_stream(
+    req: HttpRequest,
+    mut payload: web::Payload,
+) -> Result<HttpResponse, OxenHttpError> {
+    // whether if file chunk is compressed
+    let is_gzip_encoded = req
+        .headers()
+        .get("Content-Encoding")
+        .map_or(false, |v| v == "gzip");
+
+    let filename = req
+        .headers()
+        .get("X-Filename")
+        .and_then(|h| h.to_str().ok())
+        .ok_or_else(|| OxenHttpError::BadRequest("Missing X-Filename header".into()))?;
+
+    let total_size = req
+        .headers()
+        .get("X-Total-Size")
+        .and_then(|h| h.to_str().ok())
+        .and_then(|s| s.parse::<u64>().ok())
+        .ok_or_else(|| {
+            OxenHttpError::BadRequest("Missing or invalid X-Total-Size header".into())
+        })?;
+
+    let app_data = app_data(&req)?;
+    let namespace = path_param(&req, "namespace")?;
+    let repo_name = path_param(&req, "repo_name")?;
+    let workspace_id = path_param(&req, "workspace_id")?;
+    let repo = get_repo(&app_data.path, namespace, &repo_name)?;
+    let directory = PathBuf::from(path_param(&req, "path")?);
+
+    let workspace = repositories::workspaces::get(&repo, &workspace_id)?;
+
+    log::debug!("workspace::files::add_stream Got uploaded file name: {filename:?}");
+
+    let workspace_dir = workspace.dir();
+
+    log::debug!("workspace::files::add_stream Got workspace dir: {workspace_dir:?}");
+
+    let full_dir = workspace_dir.join(directory);
+
+    log::debug!("workspace::files::add_stream Got full dir: {full_dir:?}");
+
+    if !full_dir.exists() {
+        std::fs::create_dir_all(&full_dir)?;
+    }
+
+    let filepath = full_dir.join(filename);
+
+    log::debug!("workspace::files::add_stream Got filepath: {:?}", filepath);
+
+    let mut files = vec![];
+
+    let bytes_written = if filepath.exists() {
+        std::fs::metadata(&filepath)?.len()
+    } else {
+        0
+    };
+
+    // Log progress every 5MB
+    if bytes_written % (10 * 1024 * 1024) == 0 {
+        log::debug!(
+            "workspace::files::add_stream file upload progress: {:.1}% ({}/{} bytes)",
+            (bytes_written as f64 / total_size as f64) * 100.0,
+            bytes_written,
+            total_size
+        );
+    }
+
+    let mut buffer = web::BytesMut::new();
+
+    while let Some(chunk) = payload.next().await {
+        let chunk = chunk.map_err(|_| OxenHttpError::BadRequest("Error reading payload".into()))?;
+
+        // check if received eof signal
+        if chunk.len() == 1 && chunk[0] == 0 {
+            // validate file size match
+            if bytes_written == total_size {
+                log::info!("add_stream upload completed: {} bytes", total_size);
+
+                files.push(filepath.clone());
+
+                let path = repositories::workspaces::files::add(&workspace, filepath)?;
+                log::debug!("add_stream ✅ success! staged file {:?}", path);
+
+                return Ok(HttpResponse::Ok().json(FilePathsResponse {
+                    status: StatusMessage::resource_created(),
+                    paths: files,
+                }));
+            } else {
+                log::error!(
+                    "Upload stream incomplete. Expected {} bytes but received {} bytes",
+                    total_size,
+                    bytes_written
+                );
+                return Ok(HttpResponse::InternalServerError()
+                    .json(StatusMessage::internal_server_error()));
+            }
+        } else {
+            // not eof, save stream to file
+            let processed_chunk = if is_gzip_encoded {
+                let mut decoder = GzDecoder::new(&chunk[..]);
+                let mut decompressed = Vec::new();
+                decoder.read_to_end(&mut decompressed).map_err(|e| {
+                    OxenHttpError::BadRequest(
+                        format!("Failed to decompress gzip data: {}", e).into(),
+                    )
+                })?;
+                decompressed
+            } else {
+                chunk.to_vec()
+            };
+            buffer.extend_from_slice(&processed_chunk);
+
+            if buffer.len() > BUFFER_SIZE_THRESHOLD {
+                save_stream(&filepath, buffer.split().freeze().to_vec()).await?;
+            }
+        }
+    }
+
+    if !buffer.is_empty() {
+        save_stream(&filepath, buffer.freeze().to_vec()).await?;
+    }
+
+    files.push(filepath.clone());
+
+    Ok(HttpResponse::Ok().json(FilePathsResponse {
+        status: StatusMessage::resource_created(),
+        paths: files,
+    }))
+}
+
 pub async fn delete(req: HttpRequest) -> Result<HttpResponse, OxenHttpError> {
     let app_data = app_data(&req)?;
     let namespace = path_param(&req, "namespace")?;
@@ -169,3 +307,26 @@ async fn save_parts(
 
     Ok(files)
 }
+
+async fn save_stream(filepath: &PathBuf, chunk: Vec<u8>) -> Result<&PathBuf, Error> {
+    log::debug!(
+        "workspace::files::save_stream writing {} bytes to file",
+        chunk.len()
+    );
+
+    let filepath_cpy = filepath.clone();
+
+    let mut file = web::block(move || {
+        std::fs::OpenOptions::new()
+            .create(true)
+            .append(true)
+            .open(filepath_cpy)
+    })
+    .await??;
+
+    log::debug!("workspace::files::save_stream is writing file");
+
+    web::block(move || file.write_all(&chunk).map(|_| file)).await??;
+
+    Ok(filepath)
+}

src/server/src/services/workspaces.rs

@@ -31,6 +31,10 @@ pub fn workspace() -> Scope {
                     "/files/{path:.*}",
                     web::post().to(controllers::workspaces::files::add),
                 )
+                .route(
+                    "/file_stream/{path:.*}",
+                    web::post().to(controllers::workspaces::files::add_stream),
+                )
                 .route(
                     "/files/{path:.*}",
                     web::delete().to(controllers::workspaces::files::delete),