arrow_flight/
lib.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//! A native Rust implementation of [Apache Arrow Flight](https://arrow.apache.org/docs/format/Flight.html)
19//! for exchanging [Arrow](https://arrow.apache.org) data between processes.
20//!
21//! Please see the [arrow-flight crates.io](https://crates.io/crates/arrow-flight)
22//! page for feature flags and more information.
23//!
24//! # Overview
25//!
26//! This crate contains:
27//!
28//! 1. Low level [prost] generated structs
29//!    for Flight gRPC protobuf messages, such as [`FlightData`], [`FlightInfo`],
30//!    [`Location`] and [`Ticket`].
31//!
32//! 2. Low level [tonic] generated [`flight_service_client`] and
33//!    [`flight_service_server`].
34//!
35//! 3. Support for [Flight SQL] in [`sql`]. Requires the
36//!    `flight-sql` feature of this crate to be activated.
37//!
38//! [Flight SQL]: https://arrow.apache.org/docs/format/FlightSql.html
39
40#![doc(
41    html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",
42    html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"
43)]
44#![cfg_attr(docsrs, feature(doc_cfg))]
45#![allow(rustdoc::invalid_html_tags)]
46#![warn(missing_docs)]
47// The unused_crate_dependencies lint does not work well for crates defining additional examples/bin targets
48#![allow(unused_crate_dependencies)]
49
50use arrow_ipc::{convert, writer, writer::EncodedData, writer::IpcWriteOptions};
51use arrow_schema::{ArrowError, Schema};
52
53use arrow_ipc::convert::try_schema_from_ipc_buffer;
54use base64::Engine;
55use base64::prelude::BASE64_STANDARD;
56use bytes::Bytes;
57use prost_types::Timestamp;
58use std::{fmt, ops::Deref};
59
60type ArrowResult<T> = std::result::Result<T, ArrowError>;
61
62// This code is generated so we don't want to fix any lint violations manually
63#[allow(clippy::allow_attributes)]
64#[allow(clippy::all)]
65mod r#gen {
66    #![allow(missing_docs)]
67    include!("arrow.flight.protocol.rs");
68}
69
70/// Defines a `Flight` for generation or retrieval.
71pub mod flight_descriptor {
72    use super::r#gen;
73    pub use r#gen::flight_descriptor::DescriptorType;
74}
75
76/// Low Level [tonic] [`FlightServiceClient`](gen::flight_service_client::FlightServiceClient).
77pub mod flight_service_client {
78    use super::r#gen;
79    pub use r#gen::flight_service_client::FlightServiceClient;
80}
81
82/// Low Level [tonic] [`FlightServiceServer`](gen::flight_service_server::FlightServiceServer)
83/// and [`FlightService`](gen::flight_service_server::FlightService).
84pub mod flight_service_server {
85    use super::r#gen;
86    pub use r#gen::flight_service_server::FlightService;
87    pub use r#gen::flight_service_server::FlightServiceServer;
88}
89
90/// Mid Level [`FlightClient`]
91pub mod client;
92pub use client::FlightClient;
93
94/// Decoder to create [`RecordBatch`](arrow_array::RecordBatch) streams from [`FlightData`] streams.
95/// See [`FlightRecordBatchStream`](decode::FlightRecordBatchStream).
96pub mod decode;
97
98/// Encoder to create [`FlightData`] streams from [`RecordBatch`](arrow_array::RecordBatch) streams.
99/// See [`FlightDataEncoderBuilder`](encode::FlightDataEncoderBuilder).
100pub mod encode;
101
102/// Common error types
103pub mod error;
104
105pub use r#gen::Action;
106pub use r#gen::ActionType;
107pub use r#gen::BasicAuth;
108pub use r#gen::CancelFlightInfoRequest;
109pub use r#gen::CancelFlightInfoResult;
110pub use r#gen::CancelStatus;
111pub use r#gen::Criteria;
112pub use r#gen::Empty;
113pub use r#gen::FlightData;
114pub use r#gen::FlightDescriptor;
115pub use r#gen::FlightEndpoint;
116pub use r#gen::FlightInfo;
117pub use r#gen::HandshakeRequest;
118pub use r#gen::HandshakeResponse;
119pub use r#gen::Location;
120pub use r#gen::PollInfo;
121pub use r#gen::PutResult;
122pub use r#gen::RenewFlightEndpointRequest;
123pub use r#gen::Result;
124pub use r#gen::SchemaResult;
125pub use r#gen::Ticket;
126
127/// Helper to extract HTTP/gRPC trailers from a tonic stream.
128mod trailers;
129
130pub mod utils;
131
132#[cfg(feature = "flight-sql")]
133pub mod sql;
134mod streams;
135
136use flight_descriptor::DescriptorType;
137
138/// SchemaAsIpc represents a pairing of a `Schema` with IpcWriteOptions
139pub struct SchemaAsIpc<'a> {
140    /// Data type representing a schema and its IPC write options
141    pub pair: (&'a Schema, &'a IpcWriteOptions),
142}
143
144/// IpcMessage represents a `Schema` in the format expected in
145/// `FlightInfo.schema`
146#[derive(Debug)]
147pub struct IpcMessage(pub Bytes);
148
149// Useful conversion functions
150
151fn flight_schema_as_encoded_data(arrow_schema: &Schema, options: &IpcWriteOptions) -> EncodedData {
152    let data_gen = writer::IpcDataGenerator::default();
153    let mut dict_tracker = writer::DictionaryTracker::new(false);
154    data_gen.schema_to_bytes_with_dictionary_tracker(arrow_schema, &mut dict_tracker, options)
155}
156
157fn flight_schema_as_flatbuffer(schema: &Schema, options: &IpcWriteOptions) -> IpcMessage {
158    let encoded_data = flight_schema_as_encoded_data(schema, options);
159    IpcMessage(encoded_data.ipc_message.into())
160}
161
162// Implement a bunch of useful traits for various conversions, displays,
163// etc...
164
165// Deref
166
167impl Deref for IpcMessage {
168    type Target = [u8];
169
170    fn deref(&self) -> &Self::Target {
171        &self.0
172    }
173}
174
175impl<'a> Deref for SchemaAsIpc<'a> {
176    type Target = (&'a Schema, &'a IpcWriteOptions);
177
178    fn deref(&self) -> &Self::Target {
179        &self.pair
180    }
181}
182
183// Display...
184
185/// Limits the output of value to limit...
186fn limited_fmt(f: &mut fmt::Formatter<'_>, value: &[u8], limit: usize) -> fmt::Result {
187    if value.len() > limit {
188        write!(f, "{:?}", &value[..limit])
189    } else {
190        write!(f, "{value:?}")
191    }
192}
193
194impl fmt::Display for FlightData {
195    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196        write!(f, "FlightData {{")?;
197        write!(f, " descriptor: ")?;
198        match &self.flight_descriptor {
199            Some(d) => write!(f, "{d}")?,
200            None => write!(f, "None")?,
201        }
202        write!(f, ", header: ")?;
203        limited_fmt(f, &self.data_header, 8)?;
204        write!(f, ", metadata: ")?;
205        limited_fmt(f, &self.app_metadata, 8)?;
206        write!(f, ", body: ")?;
207        limited_fmt(f, &self.data_body, 8)?;
208        write!(f, " }}")
209    }
210}
211
212impl fmt::Display for FlightDescriptor {
213    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
214        write!(f, "FlightDescriptor {{")?;
215        write!(f, " type: ")?;
216        match self.r#type() {
217            DescriptorType::Cmd => {
218                write!(f, "cmd, value: ")?;
219                limited_fmt(f, &self.cmd, 8)?;
220            }
221            DescriptorType::Path => {
222                write!(f, "path: [")?;
223                let mut sep = "";
224                for element in &self.path {
225                    write!(f, "{sep}{element}")?;
226                    sep = ", ";
227                }
228                write!(f, "]")?;
229            }
230            DescriptorType::Unknown => {
231                write!(f, "unknown")?;
232            }
233        }
234        write!(f, " }}")
235    }
236}
237
238impl fmt::Display for FlightEndpoint {
239    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
240        write!(f, "FlightEndpoint {{")?;
241        write!(f, " ticket: ")?;
242        match &self.ticket {
243            Some(value) => write!(f, "{value}"),
244            None => write!(f, " None"),
245        }?;
246        write!(f, ", location: [")?;
247        let mut sep = "";
248        for location in &self.location {
249            write!(f, "{sep}{location}")?;
250            sep = ", ";
251        }
252        write!(f, "]")?;
253        write!(f, ", expiration_time:")?;
254        match &self.expiration_time {
255            Some(value) => write!(f, " {value}"),
256            None => write!(f, " None"),
257        }?;
258        write!(f, ", app_metadata: ")?;
259        limited_fmt(f, &self.app_metadata, 8)?;
260        write!(f, " }}")
261    }
262}
263
264impl fmt::Display for FlightInfo {
265    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266        let ipc_message = IpcMessage(self.schema.clone());
267        let schema: Schema = ipc_message.try_into().map_err(|_err| fmt::Error)?;
268        write!(f, "FlightInfo {{")?;
269        write!(f, " schema: {schema}")?;
270        write!(f, ", descriptor:")?;
271        match &self.flight_descriptor {
272            Some(d) => write!(f, " {d}"),
273            None => write!(f, " None"),
274        }?;
275        write!(f, ", endpoint: [")?;
276        let mut sep = "";
277        for endpoint in &self.endpoint {
278            write!(f, "{sep}{endpoint}")?;
279            sep = ", ";
280        }
281        write!(f, "], total_records: {}", self.total_records)?;
282        write!(f, ", total_bytes: {}", self.total_bytes)?;
283        write!(f, ", ordered: {}", self.ordered)?;
284        write!(f, ", app_metadata: ")?;
285        limited_fmt(f, &self.app_metadata, 8)?;
286        write!(f, " }}")
287    }
288}
289
290impl fmt::Display for PollInfo {
291    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
292        write!(f, "PollInfo {{")?;
293        write!(f, " info:")?;
294        match &self.info {
295            Some(value) => write!(f, " {value}"),
296            None => write!(f, " None"),
297        }?;
298        write!(f, ", descriptor:")?;
299        match &self.flight_descriptor {
300            Some(d) => write!(f, " {d}"),
301            None => write!(f, " None"),
302        }?;
303        write!(f, ", progress:")?;
304        match &self.progress {
305            Some(value) => write!(f, " {value}"),
306            None => write!(f, " None"),
307        }?;
308        write!(f, ", expiration_time:")?;
309        match &self.expiration_time {
310            Some(value) => write!(f, " {value}"),
311            None => write!(f, " None"),
312        }?;
313        write!(f, " }}")
314    }
315}
316
317impl fmt::Display for CancelFlightInfoRequest {
318    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
319        write!(f, "CancelFlightInfoRequest {{")?;
320        write!(f, " info: ")?;
321        match &self.info {
322            Some(value) => write!(f, "{value}")?,
323            None => write!(f, "None")?,
324        }
325        write!(f, " }}")
326    }
327}
328
329impl fmt::Display for CancelFlightInfoResult {
330    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
331        write!(f, "CancelFlightInfoResult {{")?;
332        write!(f, " status: {}", self.status().as_str_name())?;
333        write!(f, " }}")
334    }
335}
336
337impl fmt::Display for RenewFlightEndpointRequest {
338    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
339        write!(f, "RenewFlightEndpointRequest {{")?;
340        write!(f, " endpoint: ")?;
341        match &self.endpoint {
342            Some(value) => write!(f, "{value}")?,
343            None => write!(f, "None")?,
344        }
345        write!(f, " }}")
346    }
347}
348
349impl fmt::Display for Location {
350    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
351        write!(f, "Location {{")?;
352        write!(f, " uri: ")?;
353        write!(f, "{}", self.uri)
354    }
355}
356
357impl fmt::Display for Ticket {
358    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
359        write!(f, "Ticket {{")?;
360        write!(f, " ticket: ")?;
361        write!(f, "{}", BASE64_STANDARD.encode(&self.ticket))
362    }
363}
364
365// From...
366
367impl From<EncodedData> for FlightData {
368    fn from(data: EncodedData) -> Self {
369        FlightData {
370            data_header: data.ipc_message.into(),
371            data_body: data.arrow_data.into(),
372            ..Default::default()
373        }
374    }
375}
376
377impl From<SchemaAsIpc<'_>> for FlightData {
378    fn from(schema_ipc: SchemaAsIpc) -> Self {
379        let IpcMessage(vals) = flight_schema_as_flatbuffer(schema_ipc.0, schema_ipc.1);
380        FlightData {
381            data_header: vals,
382            ..Default::default()
383        }
384    }
385}
386
387impl TryFrom<SchemaAsIpc<'_>> for SchemaResult {
388    type Error = ArrowError;
389
390    fn try_from(schema_ipc: SchemaAsIpc) -> ArrowResult<Self> {
391        // According to the definition from `Flight.proto`
392        // The schema of the dataset in its IPC form:
393        //   4 bytes - an optional IPC_CONTINUATION_TOKEN prefix
394        //   4 bytes - the byte length of the payload
395        //   a flatbuffer Message whose header is the Schema
396        let IpcMessage(vals) = schema_to_ipc_format(schema_ipc)?;
397        Ok(SchemaResult { schema: vals })
398    }
399}
400
401impl TryFrom<SchemaAsIpc<'_>> for IpcMessage {
402    type Error = ArrowError;
403
404    fn try_from(schema_ipc: SchemaAsIpc) -> ArrowResult<Self> {
405        schema_to_ipc_format(schema_ipc)
406    }
407}
408
409fn schema_to_ipc_format(schema_ipc: SchemaAsIpc) -> ArrowResult<IpcMessage> {
410    let pair = *schema_ipc;
411    let encoded_data = flight_schema_as_encoded_data(pair.0, pair.1);
412
413    let mut schema = vec![];
414    writer::write_message(&mut schema, encoded_data, pair.1)?;
415    Ok(IpcMessage(schema.into()))
416}
417
418impl TryFrom<&FlightData> for Schema {
419    type Error = ArrowError;
420    fn try_from(data: &FlightData) -> ArrowResult<Self> {
421        convert::try_schema_from_flatbuffer_bytes(&data.data_header[..]).map_err(|err| {
422            ArrowError::ParseError(format!(
423                "Unable to convert flight data to Arrow schema: {err}"
424            ))
425        })
426    }
427}
428
429impl TryFrom<FlightInfo> for Schema {
430    type Error = ArrowError;
431
432    fn try_from(value: FlightInfo) -> ArrowResult<Self> {
433        value.try_decode_schema()
434    }
435}
436
437impl TryFrom<IpcMessage> for Schema {
438    type Error = ArrowError;
439
440    fn try_from(value: IpcMessage) -> ArrowResult<Self> {
441        try_schema_from_ipc_buffer(&value)
442    }
443}
444
445impl TryFrom<&SchemaResult> for Schema {
446    type Error = ArrowError;
447    fn try_from(data: &SchemaResult) -> ArrowResult<Self> {
448        try_schema_from_ipc_buffer(&data.schema)
449    }
450}
451
452impl TryFrom<SchemaResult> for Schema {
453    type Error = ArrowError;
454    fn try_from(data: SchemaResult) -> ArrowResult<Self> {
455        (&data).try_into()
456    }
457}
458
459// FlightData, FlightDescriptor, etc..
460
461impl FlightData {
462    /// Create a new [`FlightData`].
463    ///
464    /// # See Also
465    ///
466    /// See [`FlightDataEncoderBuilder`] for a higher level API to
467    /// convert a stream of [`RecordBatch`]es to [`FlightData`]s
468    ///
469    /// # Example:
470    ///
471    /// ```
472    /// # use bytes::Bytes;
473    /// # use arrow_flight::{FlightData, FlightDescriptor};
474    /// # fn encode_data() -> Bytes { Bytes::new() } // dummy data
475    /// // Get encoded Arrow IPC data:
476    /// let data_body: Bytes = encode_data();
477    /// // Create the FlightData message
478    /// let flight_data = FlightData::new()
479    ///   .with_descriptor(FlightDescriptor::new_cmd("the command"))
480    ///   .with_app_metadata("My apps metadata")
481    ///   .with_data_body(data_body);
482    /// ```
483    ///
484    /// [`FlightDataEncoderBuilder`]: crate::encode::FlightDataEncoderBuilder
485    /// [`RecordBatch`]: arrow_array::RecordBatch
486    pub fn new() -> Self {
487        Default::default()
488    }
489
490    /// Add a [`FlightDescriptor`] describing the data
491    pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
492        self.flight_descriptor = Some(flight_descriptor);
493        self
494    }
495
496    /// Add a data header
497    pub fn with_data_header(mut self, data_header: impl Into<Bytes>) -> Self {
498        self.data_header = data_header.into();
499        self
500    }
501
502    /// Add a data body. See [`IpcDataGenerator`] to create this data.
503    ///
504    /// [`IpcDataGenerator`]: arrow_ipc::writer::IpcDataGenerator
505    pub fn with_data_body(mut self, data_body: impl Into<Bytes>) -> Self {
506        self.data_body = data_body.into();
507        self
508    }
509
510    /// Add optional application specific metadata to the message
511    pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
512        self.app_metadata = app_metadata.into();
513        self
514    }
515}
516
517impl FlightDescriptor {
518    /// Create a new opaque command [`CMD`] `FlightDescriptor` to generate a dataset.
519    ///
520    /// [`CMD`]: https://github.com/apache/arrow/blob/6bd31f37ae66bd35594b077cb2f830be57e08acd/format/Flight.proto#L224-L227
521    pub fn new_cmd(cmd: impl Into<Bytes>) -> Self {
522        FlightDescriptor {
523            r#type: DescriptorType::Cmd.into(),
524            cmd: cmd.into(),
525            ..Default::default()
526        }
527    }
528
529    /// Create a new named path [`PATH`] `FlightDescriptor` that identifies a dataset
530    ///
531    /// [`PATH`]: https://github.com/apache/arrow/blob/6bd31f37ae66bd35594b077cb2f830be57e08acd/format/Flight.proto#L217-L222
532    pub fn new_path(path: Vec<String>) -> Self {
533        FlightDescriptor {
534            r#type: DescriptorType::Path.into(),
535            path,
536            ..Default::default()
537        }
538    }
539}
540
541impl FlightInfo {
542    /// Create a new, empty `FlightInfo`, describing where to fetch flight data
543    ///
544    ///
545    /// # Example:
546    /// ```
547    /// # use arrow_flight::{FlightInfo, Ticket, FlightDescriptor, FlightEndpoint};
548    /// # use arrow_schema::{Schema, Field, DataType};
549    /// # fn get_schema() -> Schema {
550    /// #   Schema::new(vec![
551    /// #     Field::new("a", DataType::Utf8, false),
552    /// #   ])
553    /// # }
554    /// #
555    /// // Create a new FlightInfo
556    /// let flight_info = FlightInfo::new()
557    ///   // Encode the Arrow schema
558    ///   .try_with_schema(&get_schema())
559    ///   .expect("encoding failed")
560    ///   .with_endpoint(
561    ///      FlightEndpoint::new()
562    ///        .with_ticket(Ticket::new("ticket contents")
563    ///      )
564    ///    )
565    ///   .with_descriptor(FlightDescriptor::new_cmd("RUN QUERY"));
566    /// ```
567    pub fn new() -> FlightInfo {
568        FlightInfo {
569            schema: Bytes::new(),
570            flight_descriptor: None,
571            endpoint: vec![],
572            ordered: false,
573            // Flight says "Set these to -1 if unknown."
574            //
575            // https://github.com/apache/arrow-rs/blob/17ca4d51d0490f9c65f5adde144f677dbc8300e7/format/Flight.proto#L287-L289
576            total_records: -1,
577            total_bytes: -1,
578            app_metadata: Bytes::new(),
579        }
580    }
581
582    /// Try and convert the data in this  `FlightInfo` into a [`Schema`]
583    pub fn try_decode_schema(self) -> ArrowResult<Schema> {
584        let msg = IpcMessage(self.schema);
585        msg.try_into()
586    }
587
588    /// Specify the schema for the response.
589    ///
590    /// Note this takes the arrow [`Schema`] (not the IPC schema) and
591    /// encodes it using the default IPC options.
592    ///
593    /// Returns an error if `schema` can not be encoded into IPC form.
594    pub fn try_with_schema(mut self, schema: &Schema) -> ArrowResult<Self> {
595        let options = IpcWriteOptions::default();
596        let IpcMessage(schema) = SchemaAsIpc::new(schema, &options).try_into()?;
597        self.schema = schema;
598        Ok(self)
599    }
600
601    /// Add specific a endpoint for fetching the data
602    pub fn with_endpoint(mut self, endpoint: FlightEndpoint) -> Self {
603        self.endpoint.push(endpoint);
604        self
605    }
606
607    /// Add endpoints for fetching all data
608    pub fn with_endpoints(mut self, endpoints: Vec<FlightEndpoint>) -> Self {
609        self.endpoint = endpoints;
610        self
611    }
612
613    /// Add a [`FlightDescriptor`] describing what this data is
614    pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
615        self.flight_descriptor = Some(flight_descriptor);
616        self
617    }
618
619    /// Set the number of records in the result, if known
620    pub fn with_total_records(mut self, total_records: i64) -> Self {
621        self.total_records = total_records;
622        self
623    }
624
625    /// Set the number of bytes in the result, if known
626    pub fn with_total_bytes(mut self, total_bytes: i64) -> Self {
627        self.total_bytes = total_bytes;
628        self
629    }
630
631    /// Specify if the response is [ordered] across endpoints
632    ///
633    /// [ordered]: https://github.com/apache/arrow-rs/blob/17ca4d51d0490f9c65f5adde144f677dbc8300e7/format/Flight.proto#L269-L275
634    pub fn with_ordered(mut self, ordered: bool) -> Self {
635        self.ordered = ordered;
636        self
637    }
638
639    /// Add optional application specific metadata to the message
640    pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
641        self.app_metadata = app_metadata.into();
642        self
643    }
644}
645
646impl PollInfo {
647    /// Create a new, empty [`PollInfo`], providing information for a long-running query
648    ///
649    /// # Example:
650    /// ```
651    /// # use arrow_flight::{FlightInfo, PollInfo, FlightDescriptor};
652    /// # use prost_types::Timestamp;
653    /// // Create a new PollInfo
654    /// let poll_info = PollInfo::new()
655    ///   .with_info(FlightInfo::new())
656    ///   .with_descriptor(FlightDescriptor::new_cmd("RUN QUERY"))
657    ///   .try_with_progress(0.5)
658    ///   .expect("progress should've been valid")
659    ///   .with_expiration_time(
660    ///     "1970-01-01".parse().expect("invalid timestamp")
661    ///   );
662    /// ```
663    pub fn new() -> Self {
664        Self {
665            info: None,
666            flight_descriptor: None,
667            progress: None,
668            expiration_time: None,
669        }
670    }
671
672    /// Add the current available results for the poll call as a [`FlightInfo`]
673    pub fn with_info(mut self, info: FlightInfo) -> Self {
674        self.info = Some(info);
675        self
676    }
677
678    /// Add a [`FlightDescriptor`] that the client should use for the next poll call,
679    /// if the query is not yet complete
680    pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
681        self.flight_descriptor = Some(flight_descriptor);
682        self
683    }
684
685    /// Set the query progress if known. Must be in the range [0.0, 1.0] else this will
686    /// return an error
687    pub fn try_with_progress(mut self, progress: f64) -> ArrowResult<Self> {
688        if !(0.0..=1.0).contains(&progress) {
689            return Err(ArrowError::InvalidArgumentError(format!(
690                "PollInfo progress must be in the range [0.0, 1.0], got {progress}"
691            )));
692        }
693        self.progress = Some(progress);
694        Ok(self)
695    }
696
697    /// Specify expiration time for this request
698    pub fn with_expiration_time(mut self, expiration_time: Timestamp) -> Self {
699        self.expiration_time = Some(expiration_time);
700        self
701    }
702}
703
704impl<'a> SchemaAsIpc<'a> {
705    /// Create a new `SchemaAsIpc` from a `Schema` and `IpcWriteOptions`
706    pub fn new(schema: &'a Schema, options: &'a IpcWriteOptions) -> Self {
707        SchemaAsIpc {
708            pair: (schema, options),
709        }
710    }
711}
712
713impl CancelFlightInfoRequest {
714    /// Create a new [`CancelFlightInfoRequest`], providing the [`FlightInfo`]
715    /// of the query to cancel.
716    pub fn new(info: FlightInfo) -> Self {
717        Self { info: Some(info) }
718    }
719}
720
721impl CancelFlightInfoResult {
722    /// Create a new [`CancelFlightInfoResult`] from the provided [`CancelStatus`].
723    pub fn new(status: CancelStatus) -> Self {
724        Self {
725            status: status as i32,
726        }
727    }
728}
729
730impl RenewFlightEndpointRequest {
731    /// Create a new [`RenewFlightEndpointRequest`], providing the [`FlightEndpoint`]
732    /// for which is being requested an extension of its expiration.
733    pub fn new(endpoint: FlightEndpoint) -> Self {
734        Self {
735            endpoint: Some(endpoint),
736        }
737    }
738}
739
740impl Action {
741    /// Create a new Action with type and body
742    pub fn new(action_type: impl Into<String>, body: impl Into<Bytes>) -> Self {
743        Self {
744            r#type: action_type.into(),
745            body: body.into(),
746        }
747    }
748}
749
750impl Result {
751    /// Create a new Result with the specified body
752    pub fn new(body: impl Into<Bytes>) -> Self {
753        Self { body: body.into() }
754    }
755}
756
757impl Ticket {
758    /// Create a new `Ticket`
759    ///
760    /// # Example
761    ///
762    /// ```
763    /// # use arrow_flight::Ticket;
764    /// let ticket = Ticket::new("SELECT * from FOO");
765    /// ```
766    pub fn new(ticket: impl Into<Bytes>) -> Self {
767        Self {
768            ticket: ticket.into(),
769        }
770    }
771}
772
773impl FlightEndpoint {
774    /// Create a new, empty `FlightEndpoint` that represents a location
775    /// to retrieve Flight results.
776    ///
777    /// # Example
778    /// ```
779    /// # use arrow_flight::{FlightEndpoint, Ticket};
780    /// #
781    /// // Specify the client should fetch results from this server
782    /// let endpoint = FlightEndpoint::new()
783    ///   .with_ticket(Ticket::new("the ticket"));
784    ///
785    /// // Specify the client should fetch results from either
786    /// // `http://example.com` or `https://example.com`
787    /// let endpoint = FlightEndpoint::new()
788    ///   .with_ticket(Ticket::new("the ticket"))
789    ///   .with_location("http://example.com")
790    ///   .with_location("https://example.com");
791    /// ```
792    pub fn new() -> FlightEndpoint {
793        Default::default()
794    }
795
796    /// Set the [`Ticket`] used to retrieve data from the endpoint
797    pub fn with_ticket(mut self, ticket: Ticket) -> Self {
798        self.ticket = Some(ticket);
799        self
800    }
801
802    /// Add a location `uri` to this endpoint. Note each endpoint can
803    /// have multiple locations.
804    ///
805    /// If no `uri` is specified, the [Flight Spec] says:
806    ///
807    /// ```text
808    /// * If the list is empty, the expectation is that the ticket can only
809    /// * be redeemed on the current service where the ticket was
810    /// * generated.
811    /// ```
812    /// [Flight Spec]: https://github.com/apache/arrow-rs/blob/17ca4d51d0490f9c65f5adde144f677dbc8300e7/format/Flight.proto#L307C2-L312
813    pub fn with_location(mut self, uri: impl Into<String>) -> Self {
814        self.location.push(Location { uri: uri.into() });
815        self
816    }
817
818    /// Specify expiration time for this stream
819    pub fn with_expiration_time(mut self, expiration_time: Timestamp) -> Self {
820        self.expiration_time = Some(expiration_time);
821        self
822    }
823
824    /// Add optional application specific metadata to the message
825    pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
826        self.app_metadata = app_metadata.into();
827        self
828    }
829}
830
831#[cfg(test)]
832mod tests {
833    use super::*;
834    use arrow_ipc::MetadataVersion;
835    use arrow_schema::{DataType, Field, TimeUnit};
836
837    struct TestVector(Vec<u8>, usize);
838
839    impl fmt::Display for TestVector {
840        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
841            limited_fmt(f, &self.0, self.1)
842        }
843    }
844
845    #[test]
846    fn it_creates_flight_descriptor_command() {
847        let expected_cmd = b"my_command";
848        let fd = FlightDescriptor::new_cmd(expected_cmd.to_vec());
849        assert_eq!(fd.r#type(), DescriptorType::Cmd);
850        assert_eq!(fd.cmd, expected_cmd.to_vec());
851    }
852
853    #[test]
854    fn it_accepts_equal_output() {
855        let input = TestVector(vec![91; 10], 10);
856
857        let actual = format!("{input}");
858        let expected = format!("{:?}", vec![91; 10]);
859        assert_eq!(actual, expected);
860    }
861
862    #[test]
863    fn it_accepts_short_output() {
864        let input = TestVector(vec![91; 6], 10);
865
866        let actual = format!("{input}");
867        let expected = format!("{:?}", vec![91; 6]);
868        assert_eq!(actual, expected);
869    }
870
871    #[test]
872    fn it_accepts_long_output() {
873        let input = TestVector(vec![91; 10], 9);
874
875        let actual = format!("{input}");
876        let expected = format!("{:?}", vec![91; 9]);
877        assert_eq!(actual, expected);
878    }
879
880    #[test]
881    fn ser_deser_schema_result() {
882        let schema = Schema::new(vec![
883            Field::new("c1", DataType::Utf8, false),
884            Field::new("c2", DataType::Float64, true),
885            Field::new("c3", DataType::UInt32, false),
886            Field::new("c4", DataType::Boolean, true),
887            Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
888            Field::new("c6", DataType::Time32(TimeUnit::Second), false),
889        ]);
890        // V5 with write_legacy_ipc_format = false
891        // this will write the continuation marker
892        let option = IpcWriteOptions::default();
893        let schema_ipc = SchemaAsIpc::new(&schema, &option);
894        let result: SchemaResult = schema_ipc.try_into().unwrap();
895        let des_schema: Schema = (&result).try_into().unwrap();
896        assert_eq!(schema, des_schema);
897
898        // V4 with write_legacy_ipc_format = true
899        // this will not write the continuation marker
900        let option = IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap();
901        let schema_ipc = SchemaAsIpc::new(&schema, &option);
902        let result: SchemaResult = schema_ipc.try_into().unwrap();
903        let des_schema: Schema = (&result).try_into().unwrap();
904        assert_eq!(schema, des_schema);
905    }
906
907    #[test]
908    fn test_dict_schema() {
909        // Test for https://github.com/apache/arrow-rs/issues/7058
910        let schema = Schema::new(vec![
911            Field::new(
912                "a",
913                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
914                false,
915            ),
916            Field::new(
917                "b",
918                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
919                false,
920            ),
921        ]);
922
923        let flight_info = FlightInfo::new().try_with_schema(&schema).unwrap();
924
925        let new_schema = Schema::try_from(flight_info).unwrap();
926        assert_eq!(schema, new_schema);
927    }
928}
arrow_flight/lib.rs

arrow_flight/
lib.rs
