1+ use std:: collections:: VecDeque ;
2+
13use color_eyre:: eyre:: { self , Result } ;
24use crossterm:: event:: { KeyEvent , MouseEventKind } ;
35use csv:: Writer ;
@@ -366,6 +368,11 @@ impl Component for Data<'_> {
366368 self . scrollable . transition_selection_mode ( Some ( SelectionMode :: Copied ) ) ;
367369 }
368370 } ,
371+ Input { key : Key :: Char ( 'Y' ) , .. } => {
372+ if let DataState :: HasResults ( rows) = & self . data_state {
373+ self . command_tx . clone ( ) . unwrap ( ) . send ( Action :: RequestYankAll ( rows. rows . len ( ) as i64 ) ) ?;
374+ }
375+ } ,
369376 Input { key : Key :: Esc , .. } => {
370377 self . scrollable . transition_selection_mode ( None ) ;
371378 } ,
@@ -389,6 +396,13 @@ impl Component for Data<'_> {
389396 }
390397 writer. flush ( ) ?;
391398 self . command_tx . clone ( ) . unwrap ( ) . send ( Action :: ExportDataFinished ) ?;
399+ } else if let Action :: YankAll = action {
400+ let DataState :: HasResults ( rows) = & self . data_state else {
401+ return Ok ( None ) ;
402+ } ;
403+ let table_for_yank = TableForYank :: new ( rows, app_state) . yank ( ) ;
404+ self . command_tx . clone ( ) . unwrap ( ) . send ( Action :: CopyData ( table_for_yank) ) ?;
405+ self . scrollable . transition_selection_mode ( Some ( SelectionMode :: Copied ) ) ;
392406 }
393407 Ok ( None )
394408 }
@@ -530,3 +544,138 @@ impl Component for Data<'_> {
530544 Ok ( ( ) )
531545 }
532546}
547+
548+ struct TableForYank {
549+ sql : Vec < String > ,
550+ table : Vec < VecDeque < String > > ,
551+ }
552+
553+ impl TableForYank {
554+ fn new ( rows : & Rows , app_state : & AppState ) -> Self {
555+ let sql = app_state. history . first ( ) . expect ( "expected the last SQL query in history" ) . query_lines . clone ( ) ;
556+
557+ let headers: & Vec < String > = & rows. headers . iter ( ) . map ( |h| h. name . clone ( ) ) . collect ( ) ;
558+ let rows = & rows. rows ;
559+
560+ let table = Self :: to_columns ( headers, rows) ;
561+
562+ Self { sql, table }
563+ }
564+
565+ fn yank ( & mut self ) -> String {
566+ let last_index = self . table . len ( ) - 1 ;
567+ self . table . iter_mut ( ) . enumerate ( ) . for_each ( |( index, col) | Self :: format_column ( col, index, last_index) ) ;
568+
569+ let mut buff = String :: new ( ) ;
570+
571+ for statement in & self . sql {
572+ buff. push_str ( statement) ;
573+ buff. push ( '\n' ) ;
574+ }
575+
576+ buff. push ( '\n' ) ;
577+
578+ while let Some ( col) = self . table . first ( ) {
579+ if col. is_empty ( ) {
580+ break ;
581+ }
582+
583+ for col in & mut self . table {
584+ if let Some ( cell) = col. pop_front ( ) {
585+ buff. push_str ( & cell) ;
586+ }
587+ }
588+ buff. push ( '\n' ) ;
589+ }
590+
591+ buff
592+ }
593+
594+ fn format_column ( col : & mut VecDeque < String > , index : usize , last_index : usize ) {
595+ let width = col. iter ( ) . map ( |s| s. len ( ) ) . max ( ) . unwrap_or ( 1 ) + 1 ;
596+
597+ let format_cell = |s : & str | {
598+ let prefix = if index == 0 { " " } else { "| " } ;
599+ let padding = if index == last_index { " " . repeat ( 0 ) } else { " " . repeat ( width. saturating_sub ( s. len ( ) ) ) } ;
600+ format ! ( "{prefix}{s}{padding}" )
601+ } ;
602+
603+ col. iter_mut ( ) . for_each ( |s| * s = format_cell ( s) ) ;
604+
605+ if let Some ( header) = col. pop_front ( ) {
606+ let div = if index == 0 { "-" . repeat ( width + 1 ) } else { format ! ( "+{}" , "-" . repeat( width + 1 ) ) } ;
607+ col. push_front ( div) ;
608+ col. push_front ( header) ;
609+ }
610+ }
611+
612+ fn to_columns ( headers : & [ String ] , rows : & [ Vec < String > ] ) -> Vec < VecDeque < String > > {
613+ headers
614+ . iter ( )
615+ . enumerate ( )
616+ . map ( |( i, h) | {
617+ let mut col: VecDeque < String > = VecDeque :: from ( [ h. clone ( ) ] ) ;
618+ rows. iter ( ) . filter_map ( |row| row. get ( i) ) . cloned ( ) . for_each ( |v| col. push_back ( v) ) ;
619+ col
620+ } )
621+ . collect ( )
622+ }
623+ }
624+
625+ #[ cfg( test) ]
626+ mod yank {
627+
628+ use std:: collections:: VecDeque ;
629+
630+ use crate :: components:: data:: TableForYank ;
631+
632+ #[ test]
633+ fn to_columns_is_works ( ) {
634+ let headers = vec ! [ "id" . to_string( ) , "name" . to_string( ) , "age" . to_string( ) ] ;
635+ let rows = vec ! [
636+ vec![ "id1" . to_string( ) , "name1" . to_string( ) , "age1" . to_string( ) ] ,
637+ vec![ "id2" . to_string( ) , "name2" . to_string( ) , "age2" . to_string( ) ] ,
638+ vec![ "id3" . to_string( ) , "name3" . to_string( ) , "age3" . to_string( ) ] ,
639+ ] ;
640+
641+ let result = TableForYank :: to_columns ( & headers, & rows) ;
642+
643+ let expected = vec ! [
644+ VecDeque :: from( [ "id" . to_string( ) , "id1" . to_string( ) , "id2" . to_string( ) , "id3" . to_string( ) ] ) ,
645+ VecDeque :: from( [ "name" . to_string( ) , "name1" . to_string( ) , "name2" . to_string( ) , "name3" . to_string( ) ] ) ,
646+ VecDeque :: from( [ "age" . to_string( ) , "age1" . to_string( ) , "age2" . to_string( ) , "age3" . to_string( ) ] ) ,
647+ ] ;
648+
649+ assert_eq ! ( expected, result)
650+ }
651+
652+ #[ test]
653+ fn yank_is_works ( ) {
654+ let headers = vec ! [ "id" . to_string( ) , "name" . to_string( ) , "age" . to_string( ) ] ;
655+ let rows = vec ! [
656+ vec![ "id1" . to_string( ) , "name1" . to_string( ) , "age1" . to_string( ) ] ,
657+ vec![ "id2" . to_string( ) , "name2" . to_string( ) , "age2" . to_string( ) ] ,
658+ vec![ "id3" . to_string( ) , "name3" . to_string( ) , "age3" . to_string( ) ] ,
659+ ] ;
660+
661+ let mut data_to_yank = TableForYank {
662+ sql : vec ! [ "select" . to_string( ) , "*" . to_string( ) , "from" . to_string( ) , "something" . to_string( ) ] ,
663+ table : TableForYank :: to_columns ( & headers, & rows) ,
664+ } ;
665+
666+ let result = data_to_yank. yank ( ) ;
667+
668+ let expected = "\
669+
670+ *
671+ from
672+ something
673+
674+ id | name | age
675+ -----+-------+------
676+ id1 | name1 | age1
677+ id2 | name2 | age2
678+ id3 | name3 | age3
679+ " ;
680+ }
681+ }
0 commit comments