11use crate :: api:: ApiClient ;
2- use indicatif:: { ProgressBar , ProgressStyle } ;
32use serde:: { Deserialize , Serialize } ;
43use serde_json:: json;
5- use std:: path:: Path ;
64
75#[ derive( Deserialize , Serialize ) ]
86struct Dataset {
@@ -72,187 +70,22 @@ struct UpdateResponse {
7270 updated_at : String ,
7371}
7472
75- struct FileType {
76- content_type : & ' static str ,
77- format : & ' static str ,
78- }
79-
80- fn detect_from_bytes ( bytes : & [ u8 ] ) -> FileType {
81- if bytes. starts_with ( b"PAR1" ) {
82- return FileType {
83- content_type : "application/octet-stream" ,
84- format : "parquet" ,
85- } ;
86- }
87- let first = bytes. iter ( ) . find ( |& & b| !b. is_ascii_whitespace ( ) ) . copied ( ) ;
88- if matches ! ( first, Some ( b'{' ) | Some ( b'[' ) ) {
89- return FileType {
90- content_type : "application/json" ,
91- format : "json" ,
92- } ;
93- }
94- FileType {
95- content_type : "text/csv" ,
96- format : "csv" ,
97- }
98- }
99-
100- fn detect_from_path ( path : & str ) -> Option < FileType > {
101- match Path :: new ( path) . extension ( ) . and_then ( |e| e. to_str ( ) ) {
102- Some ( "csv" ) => Some ( FileType {
103- content_type : "text/csv" ,
104- format : "csv" ,
105- } ) ,
106- Some ( "json" ) => Some ( FileType {
107- content_type : "application/json" ,
108- format : "json" ,
109- } ) ,
110- Some ( "parquet" ) => Some ( FileType {
111- content_type : "application/octet-stream" ,
112- format : "parquet" ,
113- } ) ,
114- _ => None ,
115- }
116- }
117-
118- /// Try to resolve the filename of the file redirected into stdin.
119- /// Works for `cmd < file.csv` but not for pipes (`cat file.csv | cmd`).
120- fn stdin_redirect_filename ( ) -> Option < String > {
121- #[ cfg( target_os = "linux" ) ]
122- {
123- std:: fs:: read_link ( "/proc/self/fd/0" )
124- . ok ( )
125- . and_then ( |p| p. file_stem ( ) . map ( |s| s. to_string_lossy ( ) . into_owned ( ) ) )
126- }
127- #[ cfg( target_os = "macos" ) ]
128- {
129- use nix:: fcntl:: { FcntlArg , fcntl} ;
130- use std:: os:: unix:: io:: AsRawFd ;
131- let fd = std:: io:: stdin ( ) . as_raw_fd ( ) ;
132- let mut path = std:: path:: PathBuf :: new ( ) ;
133- match fcntl ( fd, FcntlArg :: F_GETPATH ( & mut path) ) {
134- Ok ( _) => path. file_stem ( ) . map ( |s| s. to_string_lossy ( ) . into_owned ( ) ) ,
135- Err ( _) => None ,
136- }
137- }
138- #[ cfg( not( any( target_os = "linux" , target_os = "macos" ) ) ) ]
139- {
140- None
141- }
142- }
143-
144- fn make_progress_bar ( total : u64 ) -> ProgressBar {
145- let pb = ProgressBar :: new ( total) ;
146- pb. set_style (
147- ProgressStyle :: with_template (
148- "{spinner:.green} [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})" ,
149- )
150- . unwrap ( )
151- . progress_chars ( "=>-" ) ,
152- ) ;
153- pb
154- }
155-
156- fn do_upload < R : std:: io:: Read + Send + ' static > (
157- api : & ApiClient ,
158- content_type : & str ,
159- reader : R ,
160- pb : ProgressBar ,
161- content_length : Option < u64 > ,
162- ) -> String {
163- let ( status, resp_body) = api. post_body ( "/files" , content_type, reader, content_length) ;
164-
165- pb. finish_and_clear ( ) ;
166-
167- if !status. is_success ( ) {
168- use crossterm:: style:: Stylize ;
169- eprintln ! ( "{}" , crate :: util:: api_error( resp_body) . red( ) ) ;
170- std:: process:: exit ( 1 ) ;
171- }
172-
173- let body: serde_json:: Value = match serde_json:: from_str ( & resp_body) {
174- Ok ( v) => v,
175- Err ( e) => {
176- eprintln ! ( "error parsing upload response: {e}" ) ;
177- std:: process:: exit ( 1 ) ;
178- }
179- } ;
180-
181- match body[ "id" ] . as_str ( ) {
182- Some ( id) => id. to_string ( ) ,
183- None => {
184- eprintln ! ( "error: upload response missing id" ) ;
185- std:: process:: exit ( 1 ) ;
186- }
187- }
188- }
189-
190- // Returns (upload_id, format)
191- fn upload_from_file ( api : & ApiClient , path : & str ) -> ( String , & ' static str ) {
192- let mut f = match std:: fs:: File :: open ( path) {
193- Ok ( f) => f,
194- Err ( e) => {
195- eprintln ! ( "error opening file '{path}': {e}" ) ;
196- std:: process:: exit ( 1 ) ;
197- }
198- } ;
199-
200- let ft = detect_from_path ( path) . unwrap_or_else ( || {
201- use std:: io:: { Read , Seek } ;
202- let mut probe = [ 0u8 ; 512 ] ;
203- let n = f. read ( & mut probe) . unwrap_or ( 0 ) ;
204- let _ = f. seek ( std:: io:: SeekFrom :: Start ( 0 ) ) ;
205- detect_from_bytes ( & probe[ ..n] )
206- } ) ;
207-
208- let file_size = f. metadata ( ) . map ( |m| m. len ( ) ) . unwrap_or ( 0 ) ;
209- let pb = make_progress_bar ( file_size) ;
210- let reader = pb. wrap_read ( f) ;
211-
212- let id = do_upload ( api, ft. content_type , reader, pb, Some ( file_size) ) ;
213- ( id, ft. format )
214- }
215-
216- // Returns (upload_id, format)
217- fn upload_from_stdin ( api : & ApiClient ) -> ( String , & ' static str ) {
218- use std:: io:: Read ;
219- let mut probe = [ 0u8 ; 512 ] ;
220- let n = std:: io:: stdin ( ) . read ( & mut probe) . unwrap_or ( 0 ) ;
221- let ft = detect_from_bytes ( & probe[ ..n] ) ;
222-
223- let reader = std:: io:: Cursor :: new ( probe[ ..n] . to_vec ( ) ) . chain ( std:: io:: stdin ( ) ) ;
224-
225- let pb = ProgressBar :: new_spinner ( ) ;
226- pb. set_style (
227- ProgressStyle :: with_template ( "{spinner:.green} {bytes} uploaded ({elapsed})" ) . unwrap ( ) ,
228- ) ;
229- pb. enable_steady_tick ( std:: time:: Duration :: from_millis ( 80 ) ) ;
230- let reader = pb. wrap_read ( reader) ;
231-
232- let id = do_upload ( api, ft. content_type , reader, pb, None ) ;
233- ( id, ft. format )
234- }
235-
23673fn create_dataset (
23774 api : & ApiClient ,
238- label : & str ,
239- table_name : Option < & str > ,
75+ description : Option < & str > ,
76+ name : & str ,
24077 source : serde_json:: Value ,
241- on_failure : Option < Box < dyn FnOnce ( ) > > ,
24278) {
243- let mut body = json ! ( { "label " : label , "source" : source } ) ;
244- if let Some ( tn ) = table_name {
245- body[ "table_name " ] = json ! ( tn ) ;
79+ let mut body = json ! ( { "table_name " : name , "source" : source } ) ;
80+ if let Some ( desc ) = description {
81+ body[ "label " ] = json ! ( desc ) ;
24682 }
24783
24884 let ( status, resp_body) = api. post_raw ( "/datasets" , & body) ;
24985
25086 if !status. is_success ( ) {
25187 use crossterm:: style:: Stylize ;
25288 eprintln ! ( "{}" , crate :: util:: api_error( resp_body) . red( ) ) ;
253- if let Some ( f) = on_failure {
254- f ( ) ;
255- }
25689 std:: process:: exit ( 1 ) ;
25790 }
25891
@@ -274,144 +107,19 @@ fn create_dataset(
274107 ) ;
275108}
276109
277- pub fn create_from_upload (
278- workspace_id : & str ,
279- label : Option < & str > ,
280- table_name : Option < & str > ,
281- file : Option < & str > ,
282- upload_id : Option < & str > ,
283- source_format : & str ,
284- ) {
285- let api = ApiClient :: new ( Some ( workspace_id) ) ;
286-
287- let label_derived;
288- let label: & str = match label {
289- Some ( l) => l,
290- None => match file {
291- Some ( path) => {
292- label_derived = Path :: new ( path)
293- . file_stem ( )
294- . and_then ( |s| s. to_str ( ) )
295- . unwrap_or ( "dataset" )
296- . to_string ( ) ;
297- & label_derived
298- }
299- None => {
300- if upload_id. is_some ( ) {
301- eprintln ! ( "error: no label provided. Use --label to name the dataset." ) ;
302- std:: process:: exit ( 1 ) ;
303- }
304- match stdin_redirect_filename ( ) {
305- Some ( name) => {
306- label_derived = name;
307- & label_derived
308- }
309- None => {
310- eprintln ! ( "error: no label provided. Use --label to name the dataset." ) ;
311- std:: process:: exit ( 1 ) ;
312- }
313- }
314- }
315- } ,
316- } ;
317-
318- let ( upload_id, format, upload_id_was_uploaded) : ( String , & str , bool ) = if let Some ( id) =
319- upload_id
320- {
321- ( id. to_string ( ) , source_format, false )
322- } else {
323- let ( id, fmt) = match file {
324- Some ( path) => upload_from_file ( & api, path) ,
325- None => {
326- use std:: io:: IsTerminal ;
327- if std:: io:: stdin ( ) . is_terminal ( ) {
328- eprintln ! (
329- "error: no input data. Use --file <path>, --upload-id <id>, or pipe data via stdin."
330- ) ;
331- std:: process:: exit ( 1 ) ;
332- }
333- upload_from_stdin ( & api)
334- }
335- } ;
336- ( id, fmt, true )
337- } ;
338-
339- let source = json ! ( { "upload_id" : upload_id, "format" : format } ) ;
340-
341- let on_failure: Option < Box < dyn FnOnce ( ) > > = if upload_id_was_uploaded {
342- let uid = upload_id. clone ( ) ;
343- Some ( Box :: new ( move || {
344- use crossterm:: style:: Stylize ;
345- eprintln ! (
346- "{}" ,
347- format!(
348- "Resume dataset creation without re-uploading by passing --upload-id {uid}"
349- )
350- . yellow( )
351- ) ;
352- } ) )
353- } else {
354- None
355- } ;
356-
357- create_dataset ( & api, label, table_name, source, on_failure) ;
358- }
359-
360- pub fn create_from_url (
361- workspace_id : & str ,
362- url : & str ,
363- label : Option < & str > ,
364- table_name : Option < & str > ,
365- ) {
366- let label = match label {
367- Some ( l) => l,
368- None => {
369- eprintln ! ( "error: --label is required when using --url" ) ;
370- std:: process:: exit ( 1 ) ;
371- }
372- } ;
373- let api = ApiClient :: new ( Some ( workspace_id) ) ;
374- create_dataset ( & api, label, table_name, json ! ( { "url" : url } ) , None ) ;
375- }
376-
377- pub fn create_from_query (
378- workspace_id : & str ,
379- sql : & str ,
380- label : Option < & str > ,
381- table_name : Option < & str > ,
382- ) {
383- let label = match label {
384- Some ( l) => l,
385- None => {
386- eprintln ! ( "error: --label is required when using --sql" ) ;
387- std:: process:: exit ( 1 ) ;
388- }
389- } ;
110+ pub fn create_from_query ( workspace_id : & str , sql : & str , description : Option < & str > , name : & str ) {
390111 let api = ApiClient :: new ( Some ( workspace_id) ) ;
391- create_dataset ( & api, label , table_name , json ! ( { "sql" : sql } ) , None ) ;
112+ create_dataset ( & api, description , name , json ! ( { "sql" : sql } ) ) ;
392113}
393114
394115pub fn create_from_saved_query (
395116 workspace_id : & str ,
396117 query_id : & str ,
397- label : Option < & str > ,
398- table_name : Option < & str > ,
118+ description : Option < & str > ,
119+ name : & str ,
399120) {
400- let label = match label {
401- Some ( l) => l,
402- None => {
403- eprintln ! ( "error: --label is required when using --query-id" ) ;
404- std:: process:: exit ( 1 ) ;
405- }
406- } ;
407121 let api = ApiClient :: new ( Some ( workspace_id) ) ;
408- create_dataset (
409- & api,
410- label,
411- table_name,
412- json ! ( { "saved_query_id" : query_id } ) ,
413- None ,
414- ) ;
122+ create_dataset ( & api, description, name, json ! ( { "saved_query_id" : query_id } ) ) ;
415123}
416124
417125pub fn list ( workspace_id : & str , limit : Option < u32 > , offset : Option < u32 > , format : & str ) {
0 commit comments