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

arrow_flight/
lib.rs
