fix: stream multipart uploads to avoid OOM on large files (#477)

* fix: stream multipart uploads to avoid OOM on large files

Replace buffered file read + build_multipart_body in build_http_request
with streaming build_multipart_stream using tokio_util::io::ReaderStream.

Memory usage drops from O(file_size) to O(64 KB) regardless of upload size.
Content-Length is pre-computed from file metadata so Google APIs still
receive the correct header without buffering.

Fixes #244

* refactor: improve error messages per review feedback

- Metadata error now says 'Failed to get metadata' instead of misleading
  'Failed to read upload file'
- File::open error in stream now includes the file path for easier debugging

* test: add Drive upload smoketest to CI

Uploads a small text file, verifies the response has a file ID,
then cleans up by deleting it. Validates the streaming multipart
upload path end-to-end against real Google APIs.

* fix(ci): use drive +upload helper for upload smoketest

The upload is via the +upload helper command, not files create --upload.
Also pipe stderr through tee so errors are visible in CI logs.

* revert: remove Drive upload smoketest (insufficient CI scopes)

---------

Co-authored-by: jpoehnelt-bot <jpoehnelt-bot@users.noreply.github.com>
This commit is contained in:
Justin Poehnelt
2026-03-13 14:56:33 -06:00
committed by GitHub
parent 835e1f16c0
commit 945ac91604
4 changed files with 160 additions and 6 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@googleworkspace/cli": patch
---
Stream multipart uploads to avoid OOM on large files. File content is now streamed in chunks via `ReaderStream` instead of being read entirely into memory, reducing memory usage from O(file_size) to O(64 KB).
Generated
+2
View File
@@ -867,6 +867,7 @@ dependencies = [
"anyhow",
"async-trait",
"base64",
"bytes",
"chrono",
"clap",
"crossterm",
@@ -888,6 +889,7 @@ dependencies = [
"tempfile",
"thiserror 2.0.18",
"tokio",
"tokio-util",
"yup-oauth2",
"zeroize",
]
+2
View File
@@ -47,6 +47,8 @@ thiserror = "2"
tokio = { version = "1", features = ["full"] }
yup-oauth2 = "12"
futures-util = "0.3"
tokio-util = { version = "0.7", features = ["io"] }
bytes = "1"
base64 = "0.22.1"
derive_builder = "0.20.2"
ratatui = "0.30.0"
+151 -6
View File
@@ -22,6 +22,7 @@ use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use anyhow::Context;
use futures_util::stream::TryStreamExt;
use futures_util::StreamExt;
use serde_json::{json, Map, Value};
use tokio::io::AsyncWriteExt;
@@ -183,20 +184,22 @@ async fn build_http_request(
if input.is_upload {
let upload_path = upload_path.expect("upload_path must be Some when is_upload is true");
let file_bytes = tokio::fs::read(upload_path).await.map_err(|e| {
let file_meta = tokio::fs::metadata(upload_path).await.map_err(|e| {
GwsError::Validation(format!(
"Failed to read upload file '{}': {}",
"Failed to get metadata for upload file '{}': {}",
upload_path, e
))
})?;
let file_size = file_meta.len();
request = request.query(&[("uploadType", "multipart")]);
let media_mime =
resolve_upload_mime(upload_content_type, Some(upload_path), &input.body);
let (multipart_body, content_type) =
build_multipart_body(&input.body, &file_bytes, &media_mime)?;
let (body, content_type, content_length) =
build_multipart_stream(&input.body, upload_path, file_size, &media_mime)?;
request = request.header("Content-Type", content_type);
request = request.body(multipart_body);
request = request.header("Content-Length", content_length);
request = request.body(body);
} else if let Some(ref body_val) = input.body {
request = request.header("Content-Type", "application/json");
request = request.json(body_val);
@@ -827,9 +830,75 @@ fn mime_from_extension(path: &str) -> Option<&'static str> {
}
}
/// Builds a multipart/related body for media upload requests.
/// Builds a streaming multipart/related body for media upload requests.
///
/// Instead of reading the entire file into memory, this streams the file in
/// chunks via `ReaderStream`, keeping memory usage at O(64 KB) regardless of
/// file size. The `Content-Length` is pre-computed from file metadata so Google
/// APIs still receive the correct header without buffering.
///
/// Returns `(body, content_type, content_length)`.
fn build_multipart_stream(
metadata: &Option<Value>,
file_path: &str,
file_size: u64,
media_mime: &str,
) -> Result<(reqwest::Body, String, u64), GwsError> {
let boundary = format!("gws_boundary_{:016x}", rand::random::<u64>());
let media_mime = media_mime.to_string();
let metadata_json = match metadata {
Some(m) => serde_json::to_string(m).map_err(|e| {
GwsError::Validation(format!("Failed to serialize upload metadata: {e}"))
})?,
None => "{}".to_string(),
};
let preamble = format!(
"--{boundary}\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n{metadata_json}\r\n\
--{boundary}\r\nContent-Type: {media_mime}\r\n\r\n"
);
let postamble = format!("\r\n--{boundary}--\r\n");
let content_length = preamble.len() as u64 + file_size + postamble.len() as u64;
let content_type = format!("multipart/related; boundary={boundary}");
let preamble_bytes: bytes::Bytes = preamble.into_bytes().into();
let postamble_bytes: bytes::Bytes = postamble.into_bytes().into();
let file_path_owned = file_path.to_owned();
let file_stream = futures_util::stream::once(async move {
tokio::fs::File::open(&file_path_owned).await.map_err(|e| {
std::io::Error::new(
e.kind(),
format!("failed to open upload file '{}': {}", file_path_owned, e),
)
})
})
.map_ok(tokio_util::io::ReaderStream::new)
.try_flatten();
let stream = futures_util::stream::once(async { Ok::<_, std::io::Error>(preamble_bytes) })
.chain(file_stream)
.chain(futures_util::stream::once(async {
Ok::<_, std::io::Error>(postamble_bytes)
}));
Ok((
reqwest::Body::wrap_stream(stream),
content_type,
content_length,
))
}
/// Builds a buffered multipart/related body for media upload requests.
///
/// This is the legacy implementation retained for unit tests that need
/// a fully materialized body to assert against.
///
/// Returns the body bytes and the Content-Type header value (with boundary).
#[cfg(test)]
fn build_multipart_body(
metadata: &Option<Value>,
file_bytes: &[u8],
@@ -1369,6 +1438,82 @@ mod tests {
);
}
#[tokio::test]
async fn test_build_multipart_stream_content_length() {
let dir = tempfile::tempdir().unwrap();
let file_path = dir.path().join("small.txt");
let file_content = b"Hello stream";
std::fs::write(&file_path, file_content).unwrap();
let metadata = Some(json!({ "name": "small.txt" }));
let file_size = file_content.len() as u64;
let (_body, content_type, declared_len) = build_multipart_stream(
&metadata,
file_path.to_str().unwrap(),
file_size,
"text/plain",
)
.unwrap();
assert!(content_type.starts_with("multipart/related; boundary="));
let boundary = content_type.split("boundary=").nth(1).unwrap();
// Manually compute expected content length:
// preamble = "--{boundary}\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n{json}\r\n--{boundary}\r\nContent-Type: text/plain\r\n\r\n"
// postamble = "\r\n--{boundary}--\r\n"
let metadata_json = serde_json::to_string(&metadata.unwrap()).unwrap();
let preamble = format!(
"--{boundary}\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n{metadata_json}\r\n\
--{boundary}\r\nContent-Type: text/plain\r\n\r\n"
);
let postamble = format!("\r\n--{boundary}--\r\n");
let expected = preamble.len() as u64 + file_size + postamble.len() as u64;
assert_eq!(
declared_len, expected,
"declared Content-Length must match expected preamble + file + postamble"
);
}
#[tokio::test]
async fn test_build_multipart_stream_large_file() {
let dir = tempfile::tempdir().unwrap();
let file_path = dir.path().join("large.bin");
// 256 KB — larger than the default 64 KB ReaderStream chunk size
let data = vec![0xABu8; 256 * 1024];
std::fs::write(&file_path, &data).unwrap();
let metadata = None;
let file_size = data.len() as u64;
let (_body, _content_type, declared_len) = build_multipart_stream(
&metadata,
file_path.to_str().unwrap(),
file_size,
"application/octet-stream",
)
.unwrap();
// Content-Length must account for the empty-metadata preamble + large file + postamble
assert!(
declared_len > file_size,
"Content-Length ({declared_len}) must be larger than file size ({file_size}) due to multipart framing"
);
// Verify exact arithmetic: preamble overhead + file_size + postamble
let boundary = _content_type.split("boundary=").nth(1).unwrap();
let preamble = format!(
"--{boundary}\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n{{}}\r\n\
--{boundary}\r\nContent-Type: application/octet-stream\r\n\r\n"
);
let postamble = format!("\r\n--{boundary}--\r\n");
let expected = preamble.len() as u64 + file_size + postamble.len() as u64;
assert_eq!(
declared_len, expected,
"Content-Length must match for multi-chunk files"
);
}
#[test]
fn test_build_url_basic() {
let doc = RestDescription {