@@ -2,6 +2,7 @@ use crate::api::ApiClient;
22use rayon:: prelude:: * ;
33use serde:: { Deserialize , Serialize } ;
44use std:: collections:: HashMap ;
5+ use std:: ops:: ControlFlow ;
56
67#[ derive( Deserialize , Serialize ) ]
78struct Index {
@@ -52,16 +53,42 @@ struct ConnectionsBody {
5253 connections : Vec < ConnectionRef > ,
5354}
5455
55- fn connection_lookup ( api : & ApiClient ) -> HashMap < String , String > {
56- let body: ConnectionsBody = api. get ( "/connections" ) ;
56+ fn connection_label_to_id_map ( connections : & [ ConnectionRef ] ) -> HashMap < String , String > {
5757 let mut m = HashMap :: new ( ) ;
58- for c in body. connections {
59- m. insert ( c. id . clone ( ) , c. id . clone ( ) ) ;
60- m. insert ( c. name . clone ( ) , c. id ) ;
58+ for c in connections {
59+ m. insert ( c. name . clone ( ) , c. id . clone ( ) ) ;
6160 }
6261 m
6362}
6463
64+ fn connection_lookup ( api : & ApiClient ) -> HashMap < String , String > {
65+ let body: ConnectionsBody = api. get ( "/connections" ) ;
66+ connection_label_to_id_map ( & body. connections )
67+ }
68+
69+ /// How to continue after merging one `/information_schema` page.
70+ fn information_schema_followup (
71+ has_more : bool ,
72+ next_cursor : Option < String > ,
73+ ) -> ControlFlow < ( ) , String > {
74+ if !has_more {
75+ return ControlFlow :: Break ( ( ) ) ;
76+ }
77+ let Some ( c) = next_cursor else {
78+ return ControlFlow :: Break ( ( ) ) ;
79+ } ;
80+ ControlFlow :: Continue ( c)
81+ }
82+
83+ fn sort_info_tables ( tables : & mut [ InfoTable ] ) {
84+ tables. sort_by ( |a, b| {
85+ a. connection
86+ . cmp ( & b. connection )
87+ . then_with ( || a. schema . cmp ( & b. schema ) )
88+ . then_with ( || a. table . cmp ( & b. table ) )
89+ } ) ;
90+ }
91+
6592fn collect_tables (
6693 api : & ApiClient ,
6794 connection_id : Option < & str > ,
@@ -86,17 +113,12 @@ fn collect_tables(
86113 }
87114 let body: InfoListResponse = api. get_with_params ( "/information_schema" , & params) ;
88115 out. extend ( body. tables ) ;
89- if !body. has_more {
90- break ;
116+ match information_schema_followup ( body. has_more , body. next_cursor ) {
117+ ControlFlow :: Break ( ( ) ) => break ,
118+ ControlFlow :: Continue ( c) => cursor = Some ( c) ,
91119 }
92- cursor = body. next_cursor ;
93120 }
94- out. sort_by ( |a, b| {
95- a. connection
96- . cmp ( & b. connection )
97- . then_with ( || a. schema . cmp ( & b. schema ) )
98- . then_with ( || a. table . cmp ( & b. table ) )
99- } ) ;
121+ sort_info_tables ( & mut out) ;
100122 out
101123}
102124
@@ -266,3 +288,153 @@ pub fn create(
266288 println ! ( "{}" , "Index created." . green( ) ) ;
267289 }
268290}
291+
292+ #[ cfg( test) ]
293+ mod tests {
294+ use super :: * ;
295+
296+ #[ test]
297+ fn information_schema_followup_terminates_when_not_has_more ( ) {
298+ assert ! ( matches!(
299+ information_schema_followup( false , Some ( "c" . into( ) ) ) ,
300+ ControlFlow :: Break ( ( ) )
301+ ) ) ;
302+ }
303+
304+ #[ test]
305+ fn information_schema_followup_breaks_when_more_but_no_cursor ( ) {
306+ assert ! ( matches!(
307+ information_schema_followup( true , None ) ,
308+ ControlFlow :: Break ( ( ) )
309+ ) ) ;
310+ }
311+
312+ #[ test]
313+ fn information_schema_followup_continues_with_cursor ( ) {
314+ assert ! ( matches!(
315+ information_schema_followup( true , Some ( "next" . into( ) ) ) ,
316+ ControlFlow :: Continue ( ref s) if s == "next"
317+ ) ) ;
318+ }
319+
320+ #[ test]
321+ fn sort_info_tables_orders_by_connection_schema_table ( ) {
322+ let mut tables = vec ! [
323+ InfoTable {
324+ connection: "b" . into( ) ,
325+ schema: "s" . into( ) ,
326+ table: "t2" . into( ) ,
327+ } ,
328+ InfoTable {
329+ connection: "a" . into( ) ,
330+ schema: "s" . into( ) ,
331+ table: "t1" . into( ) ,
332+ } ,
333+ ] ;
334+ sort_info_tables ( & mut tables) ;
335+ assert_eq ! ( tables[ 0 ] . table, "t1" ) ;
336+ assert_eq ! ( tables[ 1 ] . table, "t2" ) ;
337+ }
338+
339+ #[ test]
340+ fn connection_label_to_id_map_maps_names_only ( ) {
341+ let connections = vec ! [
342+ ConnectionRef {
343+ id: "conn-id" . into( ) ,
344+ name: "Warehouse" . into( ) ,
345+ } ,
346+ ConnectionRef {
347+ id: "other" . into( ) ,
348+ name: "Lake" . into( ) ,
349+ } ,
350+ ] ;
351+ let m = connection_label_to_id_map ( & connections) ;
352+ assert_eq ! ( m. get( "Warehouse" ) . map( String :: as_str) , Some ( "conn-id" ) ) ;
353+ assert_eq ! ( m. get( "Lake" ) . map( String :: as_str) , Some ( "other" ) ) ;
354+ assert ! ( !m. contains_key( "conn-id" ) ) ;
355+ }
356+
357+ #[ test]
358+ fn collect_tables_single_page ( ) {
359+ let mut server = mockito:: Server :: new ( ) ;
360+ let mock = server
361+ . mock ( "GET" , "/information_schema" )
362+ . match_header ( "Authorization" , "Bearer k" )
363+ . match_header ( "X-Workspace-Id" , "ws1" )
364+ . with_status ( 200 )
365+ . with_body (
366+ r#"{"tables":[
367+ {"connection":"c1","schema":"public","table":"z"},
368+ {"connection":"c1","schema":"public","table":"a"}
369+ ],"has_more":false,"next_cursor":null}"# ,
370+ )
371+ . create ( ) ;
372+
373+ let api = ApiClient :: test_new ( & server. url ( ) , "k" , Some ( "ws1" ) ) ;
374+ let tables = collect_tables ( & api, None , None , None ) ;
375+ mock. assert ( ) ;
376+ assert_eq ! ( tables. len( ) , 2 ) ;
377+ assert_eq ! ( tables[ 0 ] . table, "a" ) ;
378+ assert_eq ! ( tables[ 1 ] . table, "z" ) ;
379+ }
380+
381+ #[ test]
382+ fn list_one_table_scan_returns_empty_on_404 ( ) {
383+ let mut server = mockito:: Server :: new ( ) ;
384+ let mock = server
385+ . mock (
386+ "GET" ,
387+ mockito:: Matcher :: Regex ( r"^/connections/.+/tables/.+/.+/indexes$" . into ( ) ) ,
388+ )
389+ . match_header ( "Authorization" , "Bearer k" )
390+ . with_status ( 404 )
391+ . create ( ) ;
392+
393+ let api = ApiClient :: test_new ( & server. url ( ) , "k" , Some ( "ws" ) ) ;
394+ let rows = list_one_table_scan ( & api, "cid" , "sch" , "tbl" ) ;
395+ mock. assert ( ) ;
396+ assert ! ( rows. is_empty( ) ) ;
397+ }
398+
399+ #[ test]
400+ fn list_one_table_returns_indexes ( ) {
401+ let mut server = mockito:: Server :: new ( ) ;
402+ let mock = server
403+ . mock ( "GET" , "/connections/cid/tables/sch/tbl/indexes" )
404+ . match_header ( "Authorization" , "Bearer k" )
405+ . with_status ( 200 )
406+ . with_body (
407+ r#"{"indexes":[{
408+ "index_name":"ix1",
409+ "index_type":"btree",
410+ "columns":["c1"],
411+ "metric":null,
412+ "status":"ready",
413+ "created_at":"2020-01-01T00:00:00Z",
414+ "updated_at":"2020-01-01T00:00:00Z"
415+ }]}"# ,
416+ )
417+ . create ( ) ;
418+
419+ let api = ApiClient :: test_new ( & server. url ( ) , "k" , None ) ;
420+ let rows = list_one_table ( & api, "cid" , "sch" , "tbl" ) ;
421+ mock. assert ( ) ;
422+ assert_eq ! ( rows. len( ) , 1 ) ;
423+ assert_eq ! ( rows[ 0 ] . index_name, "ix1" ) ;
424+ }
425+
426+ #[ test]
427+ fn list_one_table_scan_returns_indexes_on_200 ( ) {
428+ let mut server = mockito:: Server :: new ( ) ;
429+ let mock = server
430+ . mock ( "GET" , "/connections/x/tables/s/t/indexes" )
431+ . with_status ( 200 )
432+ . with_body ( r#"{"indexes":[]}"# )
433+ . create ( ) ;
434+
435+ let api = ApiClient :: test_new ( & server. url ( ) , "k" , None ) ;
436+ let rows = list_one_table_scan ( & api, "x" , "s" , "t" ) ;
437+ mock. assert ( ) ;
438+ assert ! ( rows. is_empty( ) ) ;
439+ }
440+ }
0 commit comments