Skip to content

Commit 5bf459f

Browse files
eddietejedaclaude
andcommitted
feat(databases): add --url flag to tables load for remote parquet files
Allows loading a remote parquet file directly by URL without a local download step. The CLI fetches the URL via reqwest, streams it to POST /files, and shows a progress bar (bytes+eta if Content-Length is present, spinner otherwise). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 6d80ba6 commit 5bf459f

3 files changed

Lines changed: 81 additions & 28 deletions

File tree

src/command.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -612,11 +612,15 @@ pub enum DatabaseTablesCommands {
612612
schema: String,
613613

614614
/// Path to a local parquet file to upload and load
615-
#[arg(long, conflicts_with = "upload_id")]
615+
#[arg(long, conflicts_with_all = ["upload_id", "url"])]
616616
file: Option<String>,
617617

618+
/// URL of a remote parquet file to download and load
619+
#[arg(long, conflicts_with_all = ["file", "upload_id"])]
620+
url: Option<String>,
621+
618622
/// Use a previously staged upload ID from `POST /v1/files` instead of uploading
619-
#[arg(long)]
623+
#[arg(long, conflicts_with_all = ["file", "url"])]
620624
upload_id: Option<String>,
621625
},
622626

src/databases.rs

Lines changed: 73 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,32 @@ fn table_rows_for_database(db_name: &str, tables: Vec<InfoTable>) -> Vec<TableRo
177177
.collect()
178178
}
179179

180+
fn finish_upload(api: &ApiClient, reader: impl std::io::Read + Send + 'static, size: Option<u64>, pb: &ProgressBar) -> String {
181+
let (status, resp_body) = api.post_body("/files", "application/octet-stream", reader, size);
182+
pb.finish_and_clear();
183+
184+
if !status.is_success() {
185+
use crossterm::style::Stylize;
186+
eprintln!("{}", crate::util::api_error(resp_body).red());
187+
std::process::exit(1);
188+
}
189+
190+
let body: serde_json::Value = match serde_json::from_str(&resp_body) {
191+
Ok(v) => v,
192+
Err(e) => {
193+
eprintln!("error parsing upload response: {e}");
194+
std::process::exit(1);
195+
}
196+
};
197+
match body["id"].as_str() {
198+
Some(id) => id.to_string(),
199+
None => {
200+
eprintln!("error: upload response missing id");
201+
std::process::exit(1);
202+
}
203+
}
204+
}
205+
180206
fn upload_parquet_file(api: &ApiClient, path: &str) -> String {
181207
if !is_parquet_path(path) {
182208
eprintln!(
@@ -205,35 +231,54 @@ fn upload_parquet_file(api: &ApiClient, path: &str) -> String {
205231
.progress_chars("=>-"),
206232
);
207233
let reader = pb.wrap_read(f);
234+
finish_upload(api, reader, Some(file_size), &pb)
235+
}
208236

209-
let (status, resp_body) = api.post_body(
210-
"/files",
211-
"application/octet-stream",
212-
reader,
213-
Some(file_size),
214-
);
215-
pb.finish_and_clear();
216-
217-
if !status.is_success() {
218-
use crossterm::style::Stylize;
219-
eprintln!("{}", crate::util::api_error(resp_body).red());
237+
fn upload_parquet_url(api: &ApiClient, url: &str) -> String {
238+
if !is_parquet_path(url) {
239+
eprintln!(
240+
"error: managed table loads require a parquet URL ending in .parquet (got '{url}')."
241+
);
220242
std::process::exit(1);
221243
}
222244

223-
let body: serde_json::Value = match serde_json::from_str(&resp_body) {
224-
Ok(v) => v,
245+
let resp = match reqwest::blocking::get(url) {
246+
Ok(r) => r,
225247
Err(e) => {
226-
eprintln!("error parsing upload response: {e}");
248+
eprintln!("error fetching '{url}': {e}");
227249
std::process::exit(1);
228250
}
229251
};
230-
match body["id"].as_str() {
231-
Some(id) => id.to_string(),
252+
253+
if !resp.status().is_success() {
254+
eprintln!("error: remote server returned {} for '{url}'", resp.status());
255+
std::process::exit(1);
256+
}
257+
258+
let content_length = resp.content_length();
259+
let pb = match content_length {
260+
Some(len) => {
261+
let pb = ProgressBar::new(len);
262+
pb.set_style(
263+
ProgressStyle::with_template(
264+
"{spinner:.green} [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})",
265+
)
266+
.unwrap()
267+
.progress_chars("=>-"),
268+
);
269+
pb
270+
}
232271
None => {
233-
eprintln!("error: upload response missing id");
234-
std::process::exit(1);
272+
let pb = ProgressBar::new_spinner();
273+
pb.set_style(
274+
ProgressStyle::with_template("{spinner:.green} {bytes} downloaded ({elapsed})")
275+
.unwrap(),
276+
);
277+
pb
235278
}
236-
}
279+
};
280+
let reader = pb.wrap_read(resp);
281+
finish_upload(api, reader, content_length, &pb)
237282
}
238283

239284
fn collect_tables(api: &ApiClient, connection_id: &str, schema: Option<&str>) -> Vec<InfoTable> {
@@ -433,6 +478,7 @@ pub fn tables_load(
433478
table: &str,
434479
schema: Option<&str>,
435480
file: Option<&str>,
481+
url: Option<&str>,
436482
upload_id: Option<&str>,
437483
) {
438484
use crossterm::style::Stylize;
@@ -441,15 +487,16 @@ pub fn tables_load(
441487
let db = resolve_database(&api, database);
442488
let schema = schema_name(schema);
443489

444-
// clap rejects `--file` and `--upload-id` together; the `(Some, Some)` arm is unreachable.
445-
let upload_id = match (upload_id, file) {
446-
(Some(id), None) => id.to_string(),
447-
(None, Some(path)) => upload_parquet_file(&api, path),
448-
(None, None) => {
449-
eprintln!("error: --file <path> or --upload-id <id> is required");
490+
// clap enforces mutual exclusion; only one of these is ever Some.
491+
let upload_id = match (upload_id, file, url) {
492+
(Some(id), None, None) => id.to_string(),
493+
(None, Some(path), None) => upload_parquet_file(&api, path),
494+
(None, None, Some(u)) => upload_parquet_url(&api, u),
495+
(None, None, None) => {
496+
eprintln!("error: --file <path>, --url <url>, or --upload-id <id> is required");
450497
std::process::exit(1);
451498
}
452-
(Some(_), Some(_)) => unreachable!(),
499+
_ => unreachable!(),
453500
};
454501

455502
let path = managed_table_load_path(&db.id, schema, table);

src/main.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -412,13 +412,15 @@ fn main() {
412412
table,
413413
schema,
414414
file,
415+
url,
415416
upload_id,
416417
} => databases::tables_load(
417418
&workspace_id,
418419
&database,
419420
&table,
420421
Some(schema.as_str()),
421422
file.as_deref(),
423+
url.as_deref(),
422424
upload_id.as_deref(),
423425
),
424426
DatabaseTablesCommands::Delete {

0 commit comments

Comments
 (0)