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

arrow_flight/
lib.rs
