Skip to content

Commit 7da031d

Browse files
wangfenjinwangfenjin
committed
WIP: impl flight sql client
Change-Id: I4cfc1412c05933c9c903d678a42462852d99fd28
1 parent 3e933ea commit 7da031d

5 files changed

Lines changed: 400 additions & 9 deletions

File tree

‎arrow-flight/src/sql/client.rs‎

Lines changed: 373 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,373 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use std::{
19+
cell::{RefCell, RefMut},
20+
pin::Pin,
21+
sync::Arc,
22+
};
23+
24+
use arrow::{
25+
datatypes::Schema,
26+
error::{ArrowError, Result},
27+
ipc::RecordBatch,
28+
};
29+
use futures::{stream, Stream};
30+
use prost::Message;
31+
use prost_types::Any;
32+
use tonic::{
33+
codegen::{Body, StdError},
34+
Streaming,
35+
};
36+
37+
use crate::{
38+
flight_descriptor, flight_service_client::FlightServiceClient,
39+
flight_service_server::FlightService, Action, ActionType, FlightData,
40+
FlightDescriptor, FlightInfo, IpcMessage, Ticket,
41+
};
42+
43+
use super::{
44+
ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
45+
ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference,
46+
CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
47+
CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
48+
CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
49+
DoPutUpdateResult, ProstAnyExt, ProstMessageExt, SqlInfo, TicketStatementQuery,
50+
ACTION_TYPE_CLOSE_PREPARED_STATEMENT, ACTION_TYPE_CREATE_PREPARED_STATEMENT,
51+
};
52+
53+
/// A FlightSQLServiceClient is an endpoint for retrieving or storing Arrow data
54+
/// by FlightSQL protocol.
55+
#[derive(Debug, Clone)]
56+
pub struct FlightSqlServiceClient<T> {
57+
inner: RefCell<FlightServiceClient<T>>,
58+
}
59+
60+
impl<T> FlightSqlServiceClient<T>
61+
where
62+
T: tonic::client::GrpcService<tonic::body::BoxBody>,
63+
T::ResponseBody: Body + Send + 'static,
64+
T::Error: Into<StdError>,
65+
<T::ResponseBody as Body>::Error: Into<StdError> + Send,
66+
{
67+
/// create FlightSqlServiceClient using FlightServiceClient
68+
pub fn new(client: RefCell<FlightServiceClient<T>>) -> Self {
69+
FlightSqlServiceClient { inner: client }
70+
}
71+
72+
/// borrow mut FlightServiceClient
73+
fn mut_client(&self) -> RefMut<'_, FlightServiceClient<T>> {
74+
self.inner.borrow_mut()
75+
}
76+
77+
async fn get_flight_info_for_command<M: ProstMessageExt>(
78+
&mut self,
79+
cmd: M,
80+
) -> Result<FlightInfo> {
81+
let request =
82+
tonic::Request::new(FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec()));
83+
Ok(self
84+
.mut_client()
85+
.get_flight_info(request)
86+
.await?
87+
.into_inner())
88+
}
89+
90+
/// Execute a query on the server.
91+
pub async fn execute(&mut self, query: String) -> Result<FlightInfo> {
92+
let cmd = CommandStatementQuery { query };
93+
self.get_flight_info_for_command(cmd).await
94+
}
95+
96+
/// Execute a update query on the server.
97+
pub async fn execute_update(&mut self, query: String) -> Result<i64> {
98+
let cmd = CommandStatementUpdate { query };
99+
let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
100+
let mut result = self
101+
.mut_client()
102+
.do_put(stream::iter(vec![FlightData {
103+
flight_descriptor: Some(descriptor),
104+
..Default::default()
105+
}]))
106+
.await?
107+
.into_inner();
108+
let result = result.message().await?.unwrap();
109+
let any: prost_types::Any = prost::Message::decode(&*result.app_metadata)
110+
.map_err(decode_error_to_arrow_error)?;
111+
let result: DoPutUpdateResult = any.unpack()?.unwrap();
112+
Ok(result.record_count)
113+
}
114+
115+
/// Request a list of catalogs.
116+
pub async fn get_catalogs(&mut self) -> Result<FlightInfo> {
117+
self.get_flight_info_for_command(CommandGetCatalogs {})
118+
.await
119+
}
120+
121+
/// Request a list of database schemas.
122+
pub async fn get_db_schemas(
123+
&mut self,
124+
request: CommandGetDbSchemas,
125+
) -> Result<FlightInfo> {
126+
self.get_flight_info_for_command(request).await
127+
}
128+
129+
/// Given a flight ticket and schema, request to be sent the
130+
/// stream. Returns record batch stream reader
131+
pub async fn do_get(
132+
&mut self,
133+
ticket: TicketStatementQuery,
134+
) -> Result<Streaming<FlightData>> {
135+
Ok(self
136+
.mut_client()
137+
.do_get(tonic::Request::new(Ticket {
138+
ticket: ticket.statement_handle,
139+
}))
140+
.await?
141+
.into_inner())
142+
}
143+
144+
/// Request a list of tables.
145+
pub async fn get_tables(&mut self, request: CommandGetTables) -> Result<FlightInfo> {
146+
self.get_flight_info_for_command(request).await
147+
}
148+
149+
/// Request the primary keys for a table.
150+
pub async fn get_primary_keys(
151+
&mut self,
152+
request: CommandGetPrimaryKeys,
153+
) -> Result<FlightInfo> {
154+
self.get_flight_info_for_command(request).await
155+
}
156+
157+
/// Retrieves a description about the foreign key columns that reference the
158+
/// primary key columns of the given table.
159+
pub async fn get_exported_keys(
160+
&mut self,
161+
request: CommandGetExportedKeys,
162+
) -> Result<FlightInfo> {
163+
self.get_flight_info_for_command(request).await
164+
}
165+
166+
/// Retrieves the foreign key columns for the given table.
167+
pub async fn get_imported_keys(
168+
&mut self,
169+
request: CommandGetImportedKeys,
170+
) -> Result<FlightInfo> {
171+
self.get_flight_info_for_command(request).await
172+
}
173+
174+
/// Retrieves a description of the foreign key columns in the given foreign key
175+
/// table that reference the primary key or the columns representing a unique
176+
/// constraint of the parent table (could be the same or a different table).
177+
pub async fn get_cross_reference(
178+
&mut self,
179+
request: CommandGetCrossReference,
180+
) -> Result<FlightInfo> {
181+
self.get_flight_info_for_command(request).await
182+
}
183+
184+
/// Request a list of table types.
185+
pub async fn get_table_types(&mut self) -> Result<FlightInfo> {
186+
self.get_flight_info_for_command(CommandGetTableTypes {})
187+
.await
188+
}
189+
190+
/// Request a list of SQL information.
191+
pub async fn get_sql_info(&mut self, sql_infos: Vec<SqlInfo>) -> Result<FlightInfo> {
192+
let request = CommandGetSqlInfo {
193+
info: sql_infos.iter().map(|sql_info| *sql_info as u32).collect(),
194+
};
195+
self.get_flight_info_for_command(request).await
196+
}
197+
198+
/// Create a prepared statement object.
199+
pub async fn prepare(&mut self, query: String) -> Result<PreparedStatement<'_, T>> {
200+
let cmd = ActionCreatePreparedStatementRequest { query };
201+
let action = Action {
202+
r#type: ACTION_TYPE_CREATE_PREPARED_STATEMENT.to_string(),
203+
body: cmd.as_any().encode_to_vec(),
204+
};
205+
let mut result = self
206+
.mut_client()
207+
.do_action(tonic::Request::new(action))
208+
.await?
209+
.into_inner();
210+
let result = result.message().await?.unwrap();
211+
let any: prost_types::Any =
212+
prost::Message::decode(&*result.body).map_err(decode_error_to_arrow_error)?;
213+
let prepared_result: ActionCreatePreparedStatementResult = any.unpack()?.unwrap();
214+
let dataset_schema =
215+
Schema::try_from(IpcMessage(prepared_result.dataset_schema))?;
216+
let parameter_schema =
217+
Schema::try_from(IpcMessage(prepared_result.parameter_schema))?;
218+
Ok(PreparedStatement::new(
219+
&self.inner,
220+
prepared_result.prepared_statement_handle,
221+
dataset_schema,
222+
parameter_schema,
223+
))
224+
}
225+
226+
/// Explicitly shut down and clean up the client.
227+
pub async fn close(&mut self) -> Result<()> {
228+
Ok(())
229+
}
230+
}
231+
232+
/// A PreparedStatement
233+
#[derive(Debug, Clone)]
234+
pub struct PreparedStatement<'a, T> {
235+
inner: &'a RefCell<FlightServiceClient<T>>,
236+
is_closed: bool,
237+
parameter_binding: Option<RecordBatch<'a>>,
238+
handle: Vec<u8>,
239+
dataset_schema: Schema,
240+
parameter_schema: Schema,
241+
}
242+
243+
impl<'a, T> PreparedStatement<'a, T>
244+
where
245+
T: tonic::client::GrpcService<tonic::body::BoxBody>,
246+
T::ResponseBody: Body + Send + 'static,
247+
T::Error: Into<StdError>,
248+
<T::ResponseBody as Body>::Error: Into<StdError> + Send,
249+
{
250+
pub(crate) fn new(
251+
client: &'a RefCell<FlightServiceClient<T>>,
252+
handle: Vec<u8>,
253+
dataset_schema: Schema,
254+
parameter_schema: Schema,
255+
) -> Self {
256+
PreparedStatement {
257+
inner: client,
258+
is_closed: false,
259+
parameter_binding: None,
260+
handle,
261+
dataset_schema,
262+
parameter_schema,
263+
}
264+
}
265+
/// Executes the prepared statement query on the server.
266+
pub async fn execute(&mut self) -> Result<FlightInfo> {
267+
if self.is_closed() {
268+
return Err(ArrowError::TonicRequestError(
269+
"Statement already closed.".to_string(),
270+
));
271+
}
272+
let cmd = CommandPreparedStatementQuery {
273+
prepared_statement_handle: self.handle.clone(),
274+
};
275+
let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
276+
let mut result = self
277+
.mut_client()
278+
.do_put(stream::iter(vec![FlightData {
279+
flight_descriptor: Some(descriptor),
280+
..Default::default()
281+
}]))
282+
.await?
283+
.into_inner();
284+
let result = result.message().await?.unwrap();
285+
let any: prost_types::Any = prost::Message::decode(&*result.app_metadata)
286+
.map_err(decode_error_to_arrow_error)?;
287+
Err(ArrowError::NotYetImplemented(
288+
"Not yet implemented".to_string(),
289+
))
290+
}
291+
292+
/// Executes the prepared statement update query on the server.
293+
pub async fn execute_update(&self) -> Result<i64> {
294+
if self.is_closed() {
295+
return Err(ArrowError::TonicRequestError(
296+
"Statement already closed.".to_string(),
297+
));
298+
}
299+
let cmd = CommandPreparedStatementQuery {
300+
prepared_statement_handle: self.handle.clone(),
301+
};
302+
let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
303+
let mut result = self
304+
.mut_client()
305+
.do_put(stream::iter(vec![FlightData {
306+
flight_descriptor: Some(descriptor),
307+
..Default::default()
308+
}]))
309+
.await?
310+
.into_inner();
311+
let result = result.message().await?.unwrap();
312+
let any: prost_types::Any = prost::Message::decode(&*result.app_metadata)
313+
.map_err(decode_error_to_arrow_error)?;
314+
let result: DoPutUpdateResult = any.unpack()?.unwrap();
315+
Ok(result.record_count)
316+
}
317+
318+
/// Retrieve the parameter schema from the query.
319+
pub async fn parameter_schema(&self) -> Result<&Schema> {
320+
Ok(&self.parameter_schema)
321+
}
322+
323+
/// Retrieve the ResultSet schema from the query.
324+
pub async fn dataset_schema(&self) -> Result<&Schema> {
325+
Ok(&self.dataset_schema)
326+
}
327+
328+
/// Set a RecordBatch that contains the parameters that will be bind.
329+
pub async fn set_parameters(
330+
&mut self,
331+
parameter_binding: RecordBatch<'a>,
332+
) -> Result<()> {
333+
self.parameter_binding = Some(parameter_binding);
334+
Ok(())
335+
}
336+
337+
/// Close the prepared statement, so that this PreparedStatement can not used
338+
/// anymore and server can free up any resources.
339+
pub async fn close(&mut self) -> Result<()> {
340+
if self.is_closed() {
341+
return Err(ArrowError::TonicRequestError(
342+
"Statement already closed.".to_string(),
343+
));
344+
}
345+
let cmd = ActionClosePreparedStatementRequest {
346+
prepared_statement_handle: self.handle.clone(),
347+
};
348+
let action = Action {
349+
r#type: ACTION_TYPE_CLOSE_PREPARED_STATEMENT.to_string(),
350+
body: cmd.as_any().encode_to_vec(),
351+
};
352+
let _ = self
353+
.mut_client()
354+
.do_action(tonic::Request::new(action))
355+
.await?;
356+
self.is_closed = true;
357+
Ok(())
358+
}
359+
360+
/// Check if the prepared statement is closed.
361+
pub fn is_closed(&self) -> bool {
362+
self.is_closed
363+
}
364+
365+
/// borrow mut FlightServiceClient
366+
fn mut_client(&self) -> RefMut<'_, FlightServiceClient<T>> {
367+
self.inner.borrow_mut()
368+
}
369+
}
370+
371+
fn decode_error_to_arrow_error(err: prost::DecodeError) -> ArrowError {
372+
ArrowError::IoError(err.to_string())
373+
}

‎arrow-flight/src/sql/mod.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,10 +58,13 @@ pub use gen::SupportedSqlGrammar;
5858
pub use gen::TicketStatementQuery;
5959
pub use gen::UpdateDeleteRules;
6060

61+
use crate::FlightDescriptor;
62+
63+
pub mod client;
6164
pub mod server;
6265

6366
/// ProstMessageExt are useful utility methods for prost::Message types
64-
pub trait ProstMessageExt: prost::Message + Default {
67+
pub trait ProstMessageExt: prost::Message + Default + Sized {
6568
/// type_url for this Message
6669
fn type_url() -> &'static str;
6770

@@ -110,6 +113,9 @@ prost_message_ext!(
110113
TicketStatementQuery,
111114
);
112115

116+
pub static ACTION_TYPE_CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement";
117+
pub static ACTION_TYPE_CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement";
118+
113119
/// ProstAnyExt are useful utility methods for prost_types::Any
114120
/// The API design is inspired by [rust-protobuf](https://github.com/stepancheg/rust-protobuf/blob/master/protobuf/src/well_known_types_util/any.rs)
115121
pub trait ProstAnyExt {

0 commit comments

Comments
 (0)