1+ use std:: io:: { stdout, Write } ;
2+ use std:: path:: Path ;
3+ use std:: sync:: Arc ;
4+ use std:: time:: Instant ;
5+
6+ use anyhow:: { anyhow, Context , Result } ;
7+ use async_trait:: async_trait;
18use clap:: Parser ;
2- use libtest_mimic :: { run_tests , Arguments , Outcome , Test } ;
3- use std :: sync :: { Arc , Mutex } ;
9+ use console :: style ;
10+ use sqllogictest :: { Control , Record } ;
411
512#[ derive( Parser , Debug ) ]
613#[ clap( about, version, author) ]
@@ -35,66 +42,137 @@ struct Opt {
3542 pass : String ,
3643}
3744
38- fn main ( ) {
45+ #[ tokio:: main]
46+ async fn main ( ) -> Result < ( ) > {
3947 env_logger:: init ( ) ;
4048
4149 let opt = Opt :: parse ( ) ;
4250
4351 let files = glob:: glob ( & opt. files ) . expect ( "failed to read glob pattern" ) ;
44- let client = postgres:: Config :: new ( )
52+
53+ let ( client, connection) = tokio_postgres:: Config :: new ( )
4554 . host ( & opt. host )
4655 . port ( opt. port )
4756 . dbname ( & opt. db )
4857 . user ( & opt. user )
4958 . password ( & opt. pass )
50- . connect ( postgres:: NoTls )
51- . expect ( "failed to connect to postgres" ) ;
59+ . connect ( tokio_postgres:: NoTls )
60+ . await
61+ . context ( "failed to connect to postgres" ) ?;
62+
63+ tokio:: spawn ( async move {
64+ if let Err ( e) = connection. await {
65+ log:: error!( "Postgres connection error: {:?}" , e) ;
66+ }
67+ } ) ;
5268
5369 let pg = Postgres {
54- client : Arc :: new ( Mutex :: new ( client) ) ,
70+ client : Arc :: new ( client) ,
5571 engine_name : opt. engine ,
5672 } ;
57- let tests = files
58- . map ( |file| Test {
59- name : file. unwrap ( ) . to_str ( ) . unwrap ( ) . into ( ) ,
60- kind : String :: new ( ) ,
61- is_ignored : false ,
62- is_bench : false ,
63- data : pg. clone ( ) ,
64- } )
65- . collect ( ) ;
66-
67- // Parse command line arguments
68- let mut args = Arguments :: from_iter ( std:: env:: args ( ) . take ( 1 ) ) ;
69- args. num_threads = Some ( 1 ) ;
70-
71- // Run all tests and exit the application appropriatly (in this case, the
72- // test runner is a dummy runner which does nothing and says that all tests
73- // passed).
74- run_tests ( & args, tests, run_test) . exit ( ) ;
73+
74+ for file in files {
75+ let file = file?;
76+ if let Err ( e) = run_test_file ( pg. clone ( ) , & file) . await {
77+ println ! ( "{}\n \n {:?}" , style( "[FAILED]" ) . red( ) . bold( ) , e) ;
78+ println ! ( ) ;
79+ }
80+ }
81+
82+ Ok ( ( ) )
7583}
7684
77- fn run_test ( test : & Test < Postgres > ) -> Outcome {
78- let mut runner = sqllogictest:: Runner :: new ( test. data . clone ( ) ) ;
79- match runner. run_file ( & test. name ) {
80- Ok ( _) => Outcome :: Passed ,
81- Err ( err) => Outcome :: Failed {
82- msg : Some ( err. to_string ( ) ) ,
83- } ,
85+ async fn flush_stdout ( ) -> std:: io:: Result < ( ) > {
86+ tokio:: task:: block_in_place ( || stdout ( ) . flush ( ) )
87+ }
88+
89+ async fn run_test_file ( engine : Postgres , filename : impl AsRef < Path > ) -> Result < ( ) > {
90+ let filename = filename. as_ref ( ) ;
91+ let mut runner = sqllogictest:: Runner :: new ( engine) ;
92+ let records = tokio:: task:: block_in_place ( || {
93+ sqllogictest:: parse_file ( & filename) . map_err ( |e| anyhow ! ( "{:?}" , e) )
94+ } )
95+ . context ( "failed to parse sqllogictest file" ) ?;
96+
97+ let mut begin_times = vec ! [ ] ;
98+ let mut did_pop = false ;
99+
100+ print ! ( "{} .. " , filename. to_string_lossy( ) ) ;
101+ flush_stdout ( ) . await ?;
102+
103+ begin_times. push ( Instant :: now ( ) ) ;
104+
105+ let finish = |time_stack : & mut Vec < Instant > , did_pop : & mut bool , file : & str | {
106+ let begin_time = time_stack. pop ( ) . unwrap ( ) ;
107+
108+ if * did_pop {
109+ // start a new line if the result is not immediately after the item
110+ print ! (
111+ "\n {}{} {} {} in {} ms" ,
112+ " " . repeat( time_stack. len( ) ) ,
113+ file,
114+ style( "[END]" ) . blue( ) . bold( ) ,
115+ style( "[OK]" ) . green( ) . bold( ) ,
116+ begin_time. elapsed( ) . as_millis( )
117+ ) ;
118+ } else {
119+ // otherwise, append time to the previous line
120+ print ! (
121+ "{} in {} ms" ,
122+ style( "[OK]" ) . green( ) . bold( ) ,
123+ begin_time. elapsed( ) . as_millis( )
124+ ) ;
125+ }
126+
127+ * did_pop = true ;
128+ } ;
129+
130+ for record in records {
131+ match & record {
132+ Record :: Control ( Control :: BeginInclude ( file) ) => {
133+ begin_times. push ( Instant :: now ( ) ) ;
134+ if !did_pop {
135+ println ! ( "{}" , style( "[BEGIN]" ) . blue( ) . bold( ) ) ;
136+ } else {
137+ println ! ( ) ;
138+ }
139+ did_pop = false ;
140+ print ! ( "{}{} .. " , " " . repeat( begin_times. len( ) - 1 ) , file) ;
141+ flush_stdout ( ) . await ?;
142+ }
143+ Record :: Control ( Control :: EndInclude ( file) ) => {
144+ finish ( & mut begin_times, & mut did_pop, file) ;
145+ }
146+ _ => { }
147+ }
148+ runner
149+ . run_async ( record)
150+ . await
151+ . map_err ( |e| anyhow ! ( "{:?}" , e) )
152+ . context ( format ! (
153+ "failed to run `{}`" ,
154+ style( filename. to_string_lossy( ) ) . bold( )
155+ ) ) ?;
84156 }
157+
158+ finish ( & mut begin_times, & mut did_pop, & * filename. to_string_lossy ( ) ) ;
159+
160+ println ! ( ) ;
161+
162+ Ok ( ( ) )
85163}
86164
87165#[ derive( Clone ) ]
88166struct Postgres {
89- client : Arc < Mutex < postgres:: Client > > ,
90-
167+ client : Arc < tokio_postgres:: Client > ,
91168 engine_name : String ,
92169}
93170
94- impl sqllogictest:: DB for Postgres {
95- type Error = postgres:: error:: Error ;
171+ #[ async_trait]
172+ impl sqllogictest:: AsyncDB for Postgres {
173+ type Error = tokio_postgres:: error:: Error ;
96174
97- fn run ( & mut self , sql : & str ) -> Result < String , Self :: Error > {
175+ async fn run ( & mut self , sql : & str ) -> Result < String , Self :: Error > {
98176 use std:: fmt:: Write ;
99177
100178 let mut output = String :: new ( ) ;
@@ -104,10 +182,10 @@ impl sqllogictest::DB for Postgres {
104182 // and we have to follow the format given by the specific database (pg).
105183 // For example, postgres will output `t` as true and `f` as false,
106184 // thus we have to write `t`/`f` in the expected results.
107- let rows = self . client . lock ( ) . unwrap ( ) . simple_query ( sql) ?;
185+ let rows = self . client . simple_query ( sql) . await ?;
108186 for row in rows {
109187 match row {
110- postgres :: SimpleQueryMessage :: Row ( row) => {
188+ tokio_postgres :: SimpleQueryMessage :: Row ( row) => {
111189 for i in 0 ..row. len ( ) {
112190 if i != 0 {
113191 write ! ( output, " " ) . unwrap ( ) ;
@@ -118,7 +196,7 @@ impl sqllogictest::DB for Postgres {
118196 }
119197 }
120198 }
121- postgres :: SimpleQueryMessage :: CommandComplete ( _) => { }
199+ tokio_postgres :: SimpleQueryMessage :: CommandComplete ( _) => { }
122200 _ => unreachable ! ( ) ,
123201 }
124202 writeln ! ( output) . unwrap ( ) ;
0 commit comments