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

arrow_flight/
lib.rs
