@@ -5,16 +5,36 @@ use crate::http::dbt::schemas::{
55} ;
66use crate :: http:: session:: DFSessionId ;
77use crate :: state:: AppState ;
8+ use arrow:: ipc:: writer:: { IpcWriteOptions , StreamWriter } ;
9+ use arrow:: ipc:: MetadataVersion ;
10+ use arrow:: json:: writer:: JsonArray ;
11+ use arrow:: json:: WriterBuilder ;
12+ use arrow:: record_batch:: RecordBatch ;
813use axum:: body:: Bytes ;
914use axum:: extract:: { Query , State } ;
1015use axum:: http:: HeaderMap ;
1116use axum:: Json ;
17+ use base64;
18+ use base64:: engine:: general_purpose:: STANDARD as engine_base64;
19+ use base64:: prelude:: * ;
1220use flate2:: read:: GzDecoder ;
1321use regex:: Regex ;
1422use snafu:: ResultExt ;
1523use std:: io:: Read ;
24+ use tracing:: debug;
1625use uuid:: Uuid ;
1726
27+ // TODO: move out as a configurable parameter
28+ const SERIALIZATION_FORMAT : & str = "json" ; // or "arrow"
29+
30+ // https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding
31+ // Buffer Alignment and Padding: Implementations are recommended to allocate memory
32+ // on aligned addresses (multiple of 8- or 64-bytes) and pad (overallocate) to a
33+ // length that is a multiple of 8 or 64 bytes. When serializing Arrow data for interprocess
34+ // communication, these alignment and padding requirements are enforced.
35+ // For more info see issue #115
36+ const ARROW_IPC_ALIGNMENT : usize = 8 ;
37+
1838#[ tracing:: instrument( level = "debug" , skip( state) , err, ret( level = tracing:: Level :: TRACE ) ) ]
1939pub async fn login (
2040 State ( state) : State < AppState > ,
@@ -32,8 +52,6 @@ pub async fn login(
3252 let _body_json: LoginRequestBody =
3353 serde_json:: from_str ( & s) . context ( dbt_error:: LoginRequestParseSnafu ) ?;
3454
35- //println!("Received login request: {:?}", query);
36- //println!("Body data parameters: {:?}", body_json);
3755 let token = Uuid :: new_v4 ( ) . to_string ( ) ;
3856
3957 let warehouses = state
@@ -56,6 +74,37 @@ pub async fn login(
5674 } ) )
5775}
5876
77+ fn records_to_arrow_string ( recs : & Vec < RecordBatch > ) -> Result < String , DbtError > {
78+ let mut buf = Vec :: new ( ) ;
79+ let options = IpcWriteOptions :: try_new ( ARROW_IPC_ALIGNMENT , false , MetadataVersion :: V5 )
80+ . context ( dbt_error:: ArrowSnafu ) ?;
81+ if !recs. is_empty ( ) {
82+ let mut writer =
83+ StreamWriter :: try_new_with_options ( & mut buf, recs[ 0 ] . schema_ref ( ) , options)
84+ . context ( dbt_error:: ArrowSnafu ) ?;
85+ for rec in recs {
86+ writer. write ( rec) . context ( dbt_error:: ArrowSnafu ) ?;
87+ }
88+ writer. finish ( ) . context ( dbt_error:: ArrowSnafu ) ?;
89+ drop ( writer) ;
90+ } ;
91+ Ok ( engine_base64. encode ( buf) )
92+ }
93+
94+ fn records_to_json_string ( recs : & [ RecordBatch ] ) -> Result < String , DbtError > {
95+ let buf = Vec :: new ( ) ;
96+ let write_builder = WriterBuilder :: new ( ) . with_explicit_nulls ( true ) ;
97+ let mut writer = write_builder. build :: < _ , JsonArray > ( buf) ;
98+ let record_refs: Vec < & RecordBatch > = recs. iter ( ) . collect ( ) ;
99+ writer
100+ . write_batches ( & record_refs)
101+ . context ( dbt_error:: ArrowSnafu ) ?;
102+ writer. finish ( ) . context ( dbt_error:: ArrowSnafu ) ?;
103+
104+ // Get the underlying buffer back,
105+ String :: from_utf8 ( writer. into_inner ( ) ) . context ( dbt_error:: Utf8Snafu )
106+ }
107+
59108#[ tracing:: instrument( level = "debug" , skip( state) , err, ret( level = tracing:: Level :: TRACE ) ) ]
60109pub async fn query (
61110 DFSessionId ( session_id) : DFSessionId ,
@@ -86,30 +135,46 @@ pub async fn query(
86135
87136 // let _ = log_query(&body_json.sql_text).await;
88137
89- let ( result , columns) = state
138+ let ( records , columns) = state
90139 . control_svc
91- . query_dbt ( & session_id, & body_json. sql_text )
140+ . query ( & session_id, & body_json. sql_text )
92141 . await
93142 . context ( dbt_error:: ControlServiceSnafu ) ?;
94143
95- // query_result_format now is JSON since arrow IPC has flatbuffers bug
96- // https://github.com/apache/arrow-rs/pull/6426
97- Ok ( Json ( JsonResponse {
144+ debug ! (
145+ "serialized json: {}" ,
146+ records_to_json_string( & records) ?. as_str( )
147+ ) ;
148+
149+ let json_resp = Json ( JsonResponse {
98150 data : Option :: from ( ResponseData {
99151 row_type : columns. into_iter ( ) . map ( Into :: into) . collect ( ) ,
100- // row_set_base_64: Option::from(result.clone()),
101- row_set_base_64 : None ,
102- #[ allow( clippy:: unwrap_used) ]
103- row_set : ResponseData :: rows_to_vec ( result. as_str ( ) ) ?,
152+ query_result_format : Option :: from ( String :: from ( SERIALIZATION_FORMAT ) ) ,
153+ row_set : if SERIALIZATION_FORMAT == "json" {
154+ Option :: from ( ResponseData :: rows_to_vec (
155+ records_to_json_string ( & records) ?. as_str ( ) ,
156+ ) ?)
157+ } else {
158+ None
159+ } ,
160+ row_set_base_64 : if SERIALIZATION_FORMAT == "arrow" {
161+ Option :: from ( records_to_arrow_string ( & records) ?)
162+ } else {
163+ None
164+ } ,
104165 total : Some ( 1 ) ,
105- query_result_format : Option :: from ( "json" . to_string ( ) ) ,
106166 error_code : None ,
107167 sql_state : Option :: from ( "ok" . to_string ( ) ) ,
108168 } ) ,
109169 success : true ,
110170 message : Option :: from ( "successfully executed" . to_string ( ) ) ,
111171 code : Some ( format ! ( "{:06}" , 200 ) ) ,
112- } ) )
172+ } ) ;
173+ debug ! (
174+ "query {:?}, response: {:?}, records: {:?}" ,
175+ body_json. sql_text, json_resp, records
176+ ) ;
177+ Ok ( json_resp)
113178}
114179
115180pub async fn abort ( ) -> DbtResult < Json < serde_json:: value:: Value > > {
@@ -127,15 +192,3 @@ pub fn extract_token(headers: &HeaderMap) -> Option<String> {
127192 } )
128193 } )
129194}
130-
131- /*async fn log_query(query: &str) -> Result<(), std::io::Error> {
132- let mut file = OpenOptions::new()
133- .create(true)
134- .append(true)
135- .open("queries.log")
136- .await
137- .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
138- file.write_all(query.as_bytes()).await?;
139- file.write_all(b"\n").await?;
140- Ok(())
141- }*/
0 commit comments