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

arrow_flight/sql/
mod.rs
