arrow_flight/sql/
mod.rs

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

arrow_flight/sql/
mod.rs
