arrow_flight/sql/
mod.rs

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//! Support for execute SQL queries using [Apache Arrow] [Flight SQL].
19//!
20//! [Flight SQL] is built on top of Arrow Flight RPC framework, by
21//! defining specific messages, encoded using the protobuf format,
22//! sent in the[`FlightDescriptor::cmd`] field to [`FlightService`]
23//! endpoints such as[`get_flight_info`] and [`do_get`].
24//!
25//! This module contains:
26//! 1. [prost] generated structs for FlightSQL messages such as [`CommandStatementQuery`]
27//! 2. Helpers for encoding and decoding FlightSQL messages: [`Any`] and [`Command`]
28//! 3. A [`FlightSqlServiceClient`] for interacting with FlightSQL servers.
29//! 4. A [`FlightSqlService`] to help building FlightSQL servers from [`FlightService`].
30//! 5. Helpers to build responses for FlightSQL metadata APIs: [`metadata`]
31//!
32//! [Flight SQL]: https://arrow.apache.org/docs/format/FlightSql.html
33//! [Apache Arrow]: https://arrow.apache.org
34//! [`FlightDescriptor::cmd`]: crate::FlightDescriptor::cmd
35//! [`FlightService`]: crate::flight_service_server::FlightService
36//! [`get_flight_info`]: crate::flight_service_server::FlightService::get_flight_info
37//! [`do_get`]: crate::flight_service_server::FlightService::do_get
38//! [`FlightSqlServiceClient`]: client::FlightSqlServiceClient
39//! [`FlightSqlService`]: server::FlightSqlService
40//! [`metadata`]: crate::sql::metadata
41use arrow_schema::ArrowError;
42use bytes::Bytes;
43use prost::Message;
44
45#[allow(clippy::all)]
46mod r#gen {
47    // Since this file is auto-generated, we suppress all warnings
48    #![allow(missing_docs)]
49    include!("arrow.flight.protocol.sql.rs");
50}
51
52pub use r#gen::ActionBeginSavepointRequest;
53pub use r#gen::ActionBeginSavepointResult;
54pub use r#gen::ActionBeginTransactionRequest;
55pub use r#gen::ActionBeginTransactionResult;
56pub use r#gen::ActionCancelQueryRequest;
57pub use r#gen::ActionCancelQueryResult;
58pub use r#gen::ActionClosePreparedStatementRequest;
59pub use r#gen::ActionCreatePreparedStatementRequest;
60pub use r#gen::ActionCreatePreparedStatementResult;
61pub use r#gen::ActionCreatePreparedSubstraitPlanRequest;
62pub use r#gen::ActionEndSavepointRequest;
63pub use r#gen::ActionEndTransactionRequest;
64pub use r#gen::CommandGetCatalogs;
65pub use r#gen::CommandGetCrossReference;
66pub use r#gen::CommandGetDbSchemas;
67pub use r#gen::CommandGetExportedKeys;
68pub use r#gen::CommandGetImportedKeys;
69pub use r#gen::CommandGetPrimaryKeys;
70pub use r#gen::CommandGetSqlInfo;
71pub use r#gen::CommandGetTableTypes;
72pub use r#gen::CommandGetTables;
73pub use r#gen::CommandGetXdbcTypeInfo;
74pub use r#gen::CommandPreparedStatementQuery;
75pub use r#gen::CommandPreparedStatementUpdate;
76pub use r#gen::CommandStatementIngest;
77pub use r#gen::CommandStatementQuery;
78pub use r#gen::CommandStatementSubstraitPlan;
79pub use r#gen::CommandStatementUpdate;
80pub use r#gen::DoPutPreparedStatementResult;
81pub use r#gen::DoPutUpdateResult;
82pub use r#gen::Nullable;
83pub use r#gen::Searchable;
84pub use r#gen::SqlInfo;
85pub use r#gen::SqlNullOrdering;
86pub use r#gen::SqlOuterJoinsSupportLevel;
87pub use r#gen::SqlSupportedCaseSensitivity;
88pub use r#gen::SqlSupportedElementActions;
89pub use r#gen::SqlSupportedGroupBy;
90pub use r#gen::SqlSupportedPositionedCommands;
91pub use r#gen::SqlSupportedResultSetConcurrency;
92pub use r#gen::SqlSupportedResultSetType;
93pub use r#gen::SqlSupportedSubqueries;
94pub use r#gen::SqlSupportedTransaction;
95pub use r#gen::SqlSupportedTransactions;
96pub use r#gen::SqlSupportedUnions;
97pub use r#gen::SqlSupportsConvert;
98pub use r#gen::SqlTransactionIsolationLevel;
99pub use r#gen::SubstraitPlan;
100pub use r#gen::SupportedSqlGrammar;
101pub use r#gen::TicketStatementQuery;
102pub use r#gen::UpdateDeleteRules;
103pub use r#gen::XdbcDataType;
104pub use r#gen::XdbcDatetimeSubcode;
105pub use r#gen::action_end_transaction_request::EndTransaction;
106pub use r#gen::command_statement_ingest::TableDefinitionOptions;
107pub use r#gen::command_statement_ingest::table_definition_options::{
108    TableExistsOption, TableNotExistOption,
109};
110
111pub mod client;
112pub mod metadata;
113pub mod server;
114
115pub use crate::streams::FallibleRequestStream;
116
117/// ProstMessageExt are useful utility methods for prost::Message types
118pub trait ProstMessageExt: prost::Message + Default {
119    /// type_url for this Message
120    fn type_url() -> &'static str;
121
122    /// Convert this Message to [`Any`]
123    fn as_any(&self) -> Any;
124}
125
126macro_rules! prost_message_ext {
127    ($($name:tt,)*) => {
128        /// Helper to convert to/from protobuf [`Any`] message
129        /// to a specific FlightSQL command message.
130        ///
131        /// # Example
132        /// ```rust
133        /// # use arrow_flight::sql::{Any, CommandStatementQuery, Command};
134        /// let flightsql_message = CommandStatementQuery {
135        ///   query: "SELECT * FROM foo".to_string(),
136        ///   transaction_id: None,
137        /// };
138        ///
139        /// // Given a packed FlightSQL Any message
140        /// let any_message = Any::pack(&flightsql_message).unwrap();
141        ///
142        /// // decode it to Command:
143        /// match Command::try_from(any_message).unwrap() {
144        ///   Command::CommandStatementQuery(decoded) => {
145        ///    assert_eq!(flightsql_message, decoded);
146        ///   }
147        ///   _ => panic!("Unexpected decoded message"),
148        /// }
149        /// ```
150        #[derive(Clone, Debug, PartialEq)]
151        pub enum Command {
152            $(
153                #[doc = concat!(stringify!($name), "variant")]
154                $name($name),)*
155
156            /// Any message that is not any FlightSQL command.
157            Unknown(Any),
158        }
159
160        impl Command {
161            /// Convert the command to [`Any`].
162            pub fn into_any(self) -> Any {
163                match self {
164                    $(
165                    Self::$name(cmd) => cmd.as_any(),
166                    )*
167                    Self::Unknown(any) => any,
168                }
169            }
170
171            /// Get the URL for the command.
172            pub fn type_url(&self) -> &str {
173                match self {
174                    $(
175                    Self::$name(_) => <$name as ProstMessageExt>::type_url(),
176                    )*
177                    Self::Unknown(any) => any.type_url.as_str(),
178                }
179            }
180        }
181
182        impl TryFrom<Any> for Command {
183            type Error = ArrowError;
184
185            fn try_from(any: Any) -> Result<Self, Self::Error> {
186                match any.type_url.as_str() {
187                    $(
188                    concat!("type.googleapis.com/arrow.flight.protocol.sql.", stringify!($name))
189                        => {
190                            let m: $name = Message::decode(&*any.value).map_err(|err| {
191                                ArrowError::ParseError(format!("Unable to decode Any value: {err}"))
192                            })?;
193                            Ok(Self::$name(m))
194                        }
195                    )*
196                    _ => Ok(Self::Unknown(any)),
197                }
198            }
199        }
200
201        $(
202            impl ProstMessageExt for $name {
203                fn type_url() -> &'static str {
204                    concat!("type.googleapis.com/arrow.flight.protocol.sql.", stringify!($name))
205                }
206
207                fn as_any(&self) -> Any {
208                    Any {
209                        type_url: <$name>::type_url().to_string(),
210                        value: self.encode_to_vec().into(),
211                    }
212                }
213            }
214        )*
215    };
216}
217
218// Implement ProstMessageExt for all structs defined in FlightSql.proto
219prost_message_ext!(
220    ActionBeginSavepointRequest,
221    ActionBeginSavepointResult,
222    ActionBeginTransactionRequest,
223    ActionBeginTransactionResult,
224    ActionCancelQueryRequest,
225    ActionCancelQueryResult,
226    ActionClosePreparedStatementRequest,
227    ActionCreatePreparedStatementRequest,
228    ActionCreatePreparedStatementResult,
229    ActionCreatePreparedSubstraitPlanRequest,
230    ActionEndSavepointRequest,
231    ActionEndTransactionRequest,
232    CommandGetCatalogs,
233    CommandGetCrossReference,
234    CommandGetDbSchemas,
235    CommandGetExportedKeys,
236    CommandGetImportedKeys,
237    CommandGetPrimaryKeys,
238    CommandGetSqlInfo,
239    CommandGetTableTypes,
240    CommandGetTables,
241    CommandGetXdbcTypeInfo,
242    CommandPreparedStatementQuery,
243    CommandPreparedStatementUpdate,
244    CommandStatementIngest,
245    CommandStatementQuery,
246    CommandStatementSubstraitPlan,
247    CommandStatementUpdate,
248    DoPutPreparedStatementResult,
249    DoPutUpdateResult,
250    TicketStatementQuery,
251);
252
253/// An implementation of the protobuf [`Any`] message type
254///
255/// Encoded protobuf messages are not self-describing, nor contain any information
256/// on the schema of the encoded payload. Consequently to decode a protobuf a client
257/// must know the exact schema of the message.
258///
259/// This presents a problem for loosely typed APIs, where the exact message payloads
260/// are not enumerable, and therefore cannot be enumerated as variants in a [oneof].
261///
262/// One solution is [`Any`] where the encoded payload is paired with a `type_url`
263/// identifying the type of encoded message, and the resulting combination encoded.
264///
265/// Clients can then decode the outer [`Any`], inspect the `type_url` and if it is
266/// a type they recognise, proceed to decode the embedded message `value`
267///
268/// [`Any`]: https://developers.google.com/protocol-buffers/docs/proto3#any
269/// [oneof]: https://developers.google.com/protocol-buffers/docs/proto3#oneof
270#[derive(Clone, PartialEq, ::prost::Message)]
271pub struct Any {
272    /// A URL/resource name that uniquely identifies the type of the serialized
273    /// protocol buffer message. This string must contain at least
274    /// one "/" character. The last segment of the URL's path must represent
275    /// the fully qualified name of the type (as in
276    /// `path/google.protobuf.Duration`). The name should be in a canonical form
277    /// (e.g., leading "." is not accepted).
278    #[prost(string, tag = "1")]
279    pub type_url: String,
280    /// Must be a valid serialized protocol buffer of the above specified type.
281    #[prost(bytes = "bytes", tag = "2")]
282    pub value: Bytes,
283}
284
285impl Any {
286    /// Checks whether the message is of type `M`
287    pub fn is<M: ProstMessageExt>(&self) -> bool {
288        M::type_url() == self.type_url
289    }
290
291    /// Unpacks the contents of the message if it is of type `M`
292    pub fn unpack<M: ProstMessageExt>(&self) -> Result<Option<M>, ArrowError> {
293        if !self.is::<M>() {
294            return Ok(None);
295        }
296        let m = Message::decode(&*self.value)
297            .map_err(|err| ArrowError::ParseError(format!("Unable to decode Any value: {err}")))?;
298        Ok(Some(m))
299    }
300
301    /// Packs a message into an [`Any`] message
302    pub fn pack<M: ProstMessageExt>(message: &M) -> Result<Any, ArrowError> {
303        Ok(message.as_any())
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn test_type_url() {
313        assert_eq!(
314            TicketStatementQuery::type_url(),
315            "type.googleapis.com/arrow.flight.protocol.sql.TicketStatementQuery"
316        );
317        assert_eq!(
318            CommandStatementQuery::type_url(),
319            "type.googleapis.com/arrow.flight.protocol.sql.CommandStatementQuery"
320        );
321    }
322
323    #[test]
324    fn test_prost_any_pack_unpack() {
325        let query = CommandStatementQuery {
326            query: "select 1".to_string(),
327            transaction_id: None,
328        };
329        let any = Any::pack(&query).unwrap();
330        assert!(any.is::<CommandStatementQuery>());
331        let unpack_query: CommandStatementQuery = any.unpack().unwrap().unwrap();
332        assert_eq!(query, unpack_query);
333    }
334
335    #[test]
336    fn test_command() {
337        let query = CommandStatementQuery {
338            query: "select 1".to_string(),
339            transaction_id: None,
340        };
341        let any = Any::pack(&query).unwrap();
342        let cmd: Command = any.try_into().unwrap();
343
344        assert!(matches!(cmd, Command::CommandStatementQuery(_)));
345        assert_eq!(cmd.type_url(), CommandStatementQuery::type_url());
346
347        // Unknown variant
348
349        let any = Any {
350            type_url: "fake_url".to_string(),
351            value: Default::default(),
352        };
353
354        let cmd: Command = any.try_into().unwrap();
355        assert!(matches!(cmd, Command::Unknown(_)));
356        assert_eq!(cmd.type_url(), "fake_url");
357    }
358}
arrow_flight/sql/mod.rs

arrow_flight/sql/
mod.rs
