1#![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_auto_cfg))]
45#![allow(rustdoc::invalid_html_tags)]
46#![warn(missing_docs)]
47#![allow(unused_crate_dependencies)]
49
50use arrow_ipc::{convert, writer, writer::EncodedData, writer::IpcWriteOptions};
51use arrow_schema::{ArrowError, Schema};
52
53use arrow_ipc::convert::try_schema_from_ipc_buffer;
54use base64::prelude::BASE64_STANDARD;
55use base64::Engine;
56use bytes::Bytes;
57use prost_types::Timestamp;
58use std::{fmt, ops::Deref};
59
60type ArrowResult<T> = std::result::Result<T, ArrowError>;
61
62#[allow(clippy::all)]
63mod gen {
64 #![allow(missing_docs)]
66 include!("arrow.flight.protocol.rs");
67}
68
69pub mod flight_descriptor {
71 use super::gen;
72 pub use gen::flight_descriptor::DescriptorType;
73}
74
75pub mod flight_service_client {
77 use super::gen;
78 pub use gen::flight_service_client::FlightServiceClient;
79}
80
81pub mod flight_service_server {
84 use super::gen;
85 pub use gen::flight_service_server::FlightService;
86 pub use gen::flight_service_server::FlightServiceServer;
87}
88
89pub mod client;
91pub use client::FlightClient;
92
93pub mod decode;
96
97pub mod encode;
100
101pub mod error;
103
104pub use gen::Action;
105pub use gen::ActionType;
106pub use gen::BasicAuth;
107pub use gen::CancelFlightInfoRequest;
108pub use gen::CancelFlightInfoResult;
109pub use gen::CancelStatus;
110pub use gen::Criteria;
111pub use gen::Empty;
112pub use gen::FlightData;
113pub use gen::FlightDescriptor;
114pub use gen::FlightEndpoint;
115pub use gen::FlightInfo;
116pub use gen::HandshakeRequest;
117pub use gen::HandshakeResponse;
118pub use gen::Location;
119pub use gen::PollInfo;
120pub use gen::PutResult;
121pub use gen::RenewFlightEndpointRequest;
122pub use gen::Result;
123pub use gen::SchemaResult;
124pub use gen::Ticket;
125
126mod trailers;
128
129pub mod utils;
130
131#[cfg(feature = "flight-sql")]
132pub mod sql;
133mod streams;
134
135use flight_descriptor::DescriptorType;
136
137pub struct SchemaAsIpc<'a> {
139 pub pair: (&'a Schema, &'a IpcWriteOptions),
141}
142
143#[derive(Debug)]
146pub struct IpcMessage(pub Bytes);
147
148fn flight_schema_as_encoded_data(arrow_schema: &Schema, options: &IpcWriteOptions) -> EncodedData {
151 let data_gen = writer::IpcDataGenerator::default();
152 let mut dict_tracker = writer::DictionaryTracker::new(false);
153 data_gen.schema_to_bytes_with_dictionary_tracker(arrow_schema, &mut dict_tracker, options)
154}
155
156fn flight_schema_as_flatbuffer(schema: &Schema, options: &IpcWriteOptions) -> IpcMessage {
157 let encoded_data = flight_schema_as_encoded_data(schema, options);
158 IpcMessage(encoded_data.ipc_message.into())
159}
160
161impl Deref for IpcMessage {
167 type Target = [u8];
168
169 fn deref(&self) -> &Self::Target {
170 &self.0
171 }
172}
173
174impl<'a> Deref for SchemaAsIpc<'a> {
175 type Target = (&'a Schema, &'a IpcWriteOptions);
176
177 fn deref(&self) -> &Self::Target {
178 &self.pair
179 }
180}
181
182fn limited_fmt(f: &mut fmt::Formatter<'_>, value: &[u8], limit: usize) -> fmt::Result {
186 if value.len() > limit {
187 write!(f, "{:?}", &value[..limit])
188 } else {
189 write!(f, "{:?}", &value)
190 }
191}
192
193impl fmt::Display for FlightData {
194 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
195 write!(f, "FlightData {{")?;
196 write!(f, " descriptor: ")?;
197 match &self.flight_descriptor {
198 Some(d) => write!(f, "{d}")?,
199 None => write!(f, "None")?,
200 };
201 write!(f, ", header: ")?;
202 limited_fmt(f, &self.data_header, 8)?;
203 write!(f, ", metadata: ")?;
204 limited_fmt(f, &self.app_metadata, 8)?;
205 write!(f, ", body: ")?;
206 limited_fmt(f, &self.data_body, 8)?;
207 write!(f, " }}")
208 }
209}
210
211impl fmt::Display for FlightDescriptor {
212 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213 write!(f, "FlightDescriptor {{")?;
214 write!(f, " type: ")?;
215 match self.r#type() {
216 DescriptorType::Cmd => {
217 write!(f, "cmd, value: ")?;
218 limited_fmt(f, &self.cmd, 8)?;
219 }
220 DescriptorType::Path => {
221 write!(f, "path: [")?;
222 let mut sep = "";
223 for element in &self.path {
224 write!(f, "{sep}{element}")?;
225 sep = ", ";
226 }
227 write!(f, "]")?;
228 }
229 DescriptorType::Unknown => {
230 write!(f, "unknown")?;
231 }
232 }
233 write!(f, " }}")
234 }
235}
236
237impl fmt::Display for FlightEndpoint {
238 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
239 write!(f, "FlightEndpoint {{")?;
240 write!(f, " ticket: ")?;
241 match &self.ticket {
242 Some(value) => write!(f, "{value}"),
243 None => write!(f, " None"),
244 }?;
245 write!(f, ", location: [")?;
246 let mut sep = "";
247 for location in &self.location {
248 write!(f, "{sep}{location}")?;
249 sep = ", ";
250 }
251 write!(f, "]")?;
252 write!(f, ", expiration_time:")?;
253 match &self.expiration_time {
254 Some(value) => write!(f, " {value}"),
255 None => write!(f, " None"),
256 }?;
257 write!(f, ", app_metadata: ")?;
258 limited_fmt(f, &self.app_metadata, 8)?;
259 write!(f, " }}")
260 }
261}
262
263impl fmt::Display for FlightInfo {
264 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
265 let ipc_message = IpcMessage(self.schema.clone());
266 let schema: Schema = ipc_message.try_into().map_err(|_err| fmt::Error)?;
267 write!(f, "FlightInfo {{")?;
268 write!(f, " schema: {schema}")?;
269 write!(f, ", descriptor:")?;
270 match &self.flight_descriptor {
271 Some(d) => write!(f, " {d}"),
272 None => write!(f, " None"),
273 }?;
274 write!(f, ", endpoint: [")?;
275 let mut sep = "";
276 for endpoint in &self.endpoint {
277 write!(f, "{sep}{endpoint}")?;
278 sep = ", ";
279 }
280 write!(f, "], total_records: {}", self.total_records)?;
281 write!(f, ", total_bytes: {}", self.total_bytes)?;
282 write!(f, ", ordered: {}", self.ordered)?;
283 write!(f, ", app_metadata: ")?;
284 limited_fmt(f, &self.app_metadata, 8)?;
285 write!(f, " }}")
286 }
287}
288
289impl fmt::Display for PollInfo {
290 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
291 write!(f, "PollInfo {{")?;
292 write!(f, " info:")?;
293 match &self.info {
294 Some(value) => write!(f, " {value}"),
295 None => write!(f, " None"),
296 }?;
297 write!(f, ", descriptor:")?;
298 match &self.flight_descriptor {
299 Some(d) => write!(f, " {d}"),
300 None => write!(f, " None"),
301 }?;
302 write!(f, ", progress:")?;
303 match &self.progress {
304 Some(value) => write!(f, " {value}"),
305 None => write!(f, " None"),
306 }?;
307 write!(f, ", expiration_time:")?;
308 match &self.expiration_time {
309 Some(value) => write!(f, " {value}"),
310 None => write!(f, " None"),
311 }?;
312 write!(f, " }}")
313 }
314}
315
316impl fmt::Display for CancelFlightInfoRequest {
317 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
318 write!(f, "CancelFlightInfoRequest {{")?;
319 write!(f, " info: ")?;
320 match &self.info {
321 Some(value) => write!(f, "{value}")?,
322 None => write!(f, "None")?,
323 };
324 write!(f, " }}")
325 }
326}
327
328impl fmt::Display for CancelFlightInfoResult {
329 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
330 write!(f, "CancelFlightInfoResult {{")?;
331 write!(f, " status: {}", self.status().as_str_name())?;
332 write!(f, " }}")
333 }
334}
335
336impl fmt::Display for RenewFlightEndpointRequest {
337 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
338 write!(f, "RenewFlightEndpointRequest {{")?;
339 write!(f, " endpoint: ")?;
340 match &self.endpoint {
341 Some(value) => write!(f, "{value}")?,
342 None => write!(f, "None")?,
343 };
344 write!(f, " }}")
345 }
346}
347
348impl fmt::Display for Location {
349 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350 write!(f, "Location {{")?;
351 write!(f, " uri: ")?;
352 write!(f, "{}", self.uri)
353 }
354}
355
356impl fmt::Display for Ticket {
357 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
358 write!(f, "Ticket {{")?;
359 write!(f, " ticket: ")?;
360 write!(f, "{}", BASE64_STANDARD.encode(&self.ticket))
361 }
362}
363
364impl From<EncodedData> for FlightData {
367 fn from(data: EncodedData) -> Self {
368 FlightData {
369 data_header: data.ipc_message.into(),
370 data_body: data.arrow_data.into(),
371 ..Default::default()
372 }
373 }
374}
375
376impl From<SchemaAsIpc<'_>> for FlightData {
377 fn from(schema_ipc: SchemaAsIpc) -> Self {
378 let IpcMessage(vals) = flight_schema_as_flatbuffer(schema_ipc.0, schema_ipc.1);
379 FlightData {
380 data_header: vals,
381 ..Default::default()
382 }
383 }
384}
385
386impl TryFrom<SchemaAsIpc<'_>> for SchemaResult {
387 type Error = ArrowError;
388
389 fn try_from(schema_ipc: SchemaAsIpc) -> ArrowResult<Self> {
390 let IpcMessage(vals) = schema_to_ipc_format(schema_ipc)?;
396 Ok(SchemaResult { schema: vals })
397 }
398}
399
400impl TryFrom<SchemaAsIpc<'_>> for IpcMessage {
401 type Error = ArrowError;
402
403 fn try_from(schema_ipc: SchemaAsIpc) -> ArrowResult<Self> {
404 schema_to_ipc_format(schema_ipc)
405 }
406}
407
408fn schema_to_ipc_format(schema_ipc: SchemaAsIpc) -> ArrowResult<IpcMessage> {
409 let pair = *schema_ipc;
410 let encoded_data = flight_schema_as_encoded_data(pair.0, pair.1);
411
412 let mut schema = vec![];
413 writer::write_message(&mut schema, encoded_data, pair.1)?;
414 Ok(IpcMessage(schema.into()))
415}
416
417impl TryFrom<&FlightData> for Schema {
418 type Error = ArrowError;
419 fn try_from(data: &FlightData) -> ArrowResult<Self> {
420 convert::try_schema_from_flatbuffer_bytes(&data.data_header[..]).map_err(|err| {
421 ArrowError::ParseError(format!(
422 "Unable to convert flight data to Arrow schema: {err}"
423 ))
424 })
425 }
426}
427
428impl TryFrom<FlightInfo> for Schema {
429 type Error = ArrowError;
430
431 fn try_from(value: FlightInfo) -> ArrowResult<Self> {
432 value.try_decode_schema()
433 }
434}
435
436impl TryFrom<IpcMessage> for Schema {
437 type Error = ArrowError;
438
439 fn try_from(value: IpcMessage) -> ArrowResult<Self> {
440 try_schema_from_ipc_buffer(&value)
441 }
442}
443
444impl TryFrom<&SchemaResult> for Schema {
445 type Error = ArrowError;
446 fn try_from(data: &SchemaResult) -> ArrowResult<Self> {
447 try_schema_from_ipc_buffer(&data.schema)
448 }
449}
450
451impl TryFrom<SchemaResult> for Schema {
452 type Error = ArrowError;
453 fn try_from(data: SchemaResult) -> ArrowResult<Self> {
454 (&data).try_into()
455 }
456}
457
458impl FlightData {
461 pub fn new() -> Self {
486 Default::default()
487 }
488
489 pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
491 self.flight_descriptor = Some(flight_descriptor);
492 self
493 }
494
495 pub fn with_data_header(mut self, data_header: impl Into<Bytes>) -> Self {
497 self.data_header = data_header.into();
498 self
499 }
500
501 pub fn with_data_body(mut self, data_body: impl Into<Bytes>) -> Self {
505 self.data_body = data_body.into();
506 self
507 }
508
509 pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
511 self.app_metadata = app_metadata.into();
512 self
513 }
514}
515
516impl FlightDescriptor {
517 pub fn new_cmd(cmd: impl Into<Bytes>) -> Self {
521 FlightDescriptor {
522 r#type: DescriptorType::Cmd.into(),
523 cmd: cmd.into(),
524 ..Default::default()
525 }
526 }
527
528 pub fn new_path(path: Vec<String>) -> Self {
532 FlightDescriptor {
533 r#type: DescriptorType::Path.into(),
534 path,
535 ..Default::default()
536 }
537 }
538}
539
540impl FlightInfo {
541 pub fn new() -> FlightInfo {
567 FlightInfo {
568 schema: Bytes::new(),
569 flight_descriptor: None,
570 endpoint: vec![],
571 ordered: false,
572 total_records: -1,
576 total_bytes: -1,
577 app_metadata: Bytes::new(),
578 }
579 }
580
581 pub fn try_decode_schema(self) -> ArrowResult<Schema> {
583 let msg = IpcMessage(self.schema);
584 msg.try_into()
585 }
586
587 pub fn try_with_schema(mut self, schema: &Schema) -> ArrowResult<Self> {
594 let options = IpcWriteOptions::default();
595 let IpcMessage(schema) = SchemaAsIpc::new(schema, &options).try_into()?;
596 self.schema = schema;
597 Ok(self)
598 }
599
600 pub fn with_endpoint(mut self, endpoint: FlightEndpoint) -> Self {
602 self.endpoint.push(endpoint);
603 self
604 }
605
606 pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
608 self.flight_descriptor = Some(flight_descriptor);
609 self
610 }
611
612 pub fn with_total_records(mut self, total_records: i64) -> Self {
614 self.total_records = total_records;
615 self
616 }
617
618 pub fn with_total_bytes(mut self, total_bytes: i64) -> Self {
620 self.total_bytes = total_bytes;
621 self
622 }
623
624 pub fn with_ordered(mut self, ordered: bool) -> Self {
628 self.ordered = ordered;
629 self
630 }
631
632 pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
634 self.app_metadata = app_metadata.into();
635 self
636 }
637}
638
639impl PollInfo {
640 pub fn new() -> Self {
657 Self {
658 info: None,
659 flight_descriptor: None,
660 progress: None,
661 expiration_time: None,
662 }
663 }
664
665 pub fn with_info(mut self, info: FlightInfo) -> Self {
667 self.info = Some(info);
668 self
669 }
670
671 pub fn with_descriptor(mut self, flight_descriptor: FlightDescriptor) -> Self {
674 self.flight_descriptor = Some(flight_descriptor);
675 self
676 }
677
678 pub fn try_with_progress(mut self, progress: f64) -> ArrowResult<Self> {
681 if !(0.0..=1.0).contains(&progress) {
682 return Err(ArrowError::InvalidArgumentError(format!(
683 "PollInfo progress must be in the range [0.0, 1.0], got {progress}"
684 )));
685 }
686 self.progress = Some(progress);
687 Ok(self)
688 }
689
690 pub fn with_expiration_time(mut self, expiration_time: Timestamp) -> Self {
692 self.expiration_time = Some(expiration_time);
693 self
694 }
695}
696
697impl<'a> SchemaAsIpc<'a> {
698 pub fn new(schema: &'a Schema, options: &'a IpcWriteOptions) -> Self {
700 SchemaAsIpc {
701 pair: (schema, options),
702 }
703 }
704}
705
706impl CancelFlightInfoRequest {
707 pub fn new(info: FlightInfo) -> Self {
710 Self { info: Some(info) }
711 }
712}
713
714impl CancelFlightInfoResult {
715 pub fn new(status: CancelStatus) -> Self {
717 Self {
718 status: status as i32,
719 }
720 }
721}
722
723impl RenewFlightEndpointRequest {
724 pub fn new(endpoint: FlightEndpoint) -> Self {
727 Self {
728 endpoint: Some(endpoint),
729 }
730 }
731}
732
733impl Action {
734 pub fn new(action_type: impl Into<String>, body: impl Into<Bytes>) -> Self {
736 Self {
737 r#type: action_type.into(),
738 body: body.into(),
739 }
740 }
741}
742
743impl Result {
744 pub fn new(body: impl Into<Bytes>) -> Self {
746 Self { body: body.into() }
747 }
748}
749
750impl Ticket {
751 pub fn new(ticket: impl Into<Bytes>) -> Self {
760 Self {
761 ticket: ticket.into(),
762 }
763 }
764}
765
766impl FlightEndpoint {
767 pub fn new() -> FlightEndpoint {
786 Default::default()
787 }
788
789 pub fn with_ticket(mut self, ticket: Ticket) -> Self {
791 self.ticket = Some(ticket);
792 self
793 }
794
795 pub fn with_location(mut self, uri: impl Into<String>) -> Self {
807 self.location.push(Location { uri: uri.into() });
808 self
809 }
810
811 pub fn with_expiration_time(mut self, expiration_time: Timestamp) -> Self {
813 self.expiration_time = Some(expiration_time);
814 self
815 }
816
817 pub fn with_app_metadata(mut self, app_metadata: impl Into<Bytes>) -> Self {
819 self.app_metadata = app_metadata.into();
820 self
821 }
822}
823
824#[cfg(test)]
825mod tests {
826 use super::*;
827 use arrow_ipc::MetadataVersion;
828 use arrow_schema::{DataType, Field, TimeUnit};
829
830 struct TestVector(Vec<u8>, usize);
831
832 impl fmt::Display for TestVector {
833 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
834 limited_fmt(f, &self.0, self.1)
835 }
836 }
837
838 #[test]
839 fn it_creates_flight_descriptor_command() {
840 let expected_cmd = "my_command".as_bytes();
841 let fd = FlightDescriptor::new_cmd(expected_cmd.to_vec());
842 assert_eq!(fd.r#type(), DescriptorType::Cmd);
843 assert_eq!(fd.cmd, expected_cmd.to_vec());
844 }
845
846 #[test]
847 fn it_accepts_equal_output() {
848 let input = TestVector(vec![91; 10], 10);
849
850 let actual = format!("{input}");
851 let expected = format!("{:?}", vec![91; 10]);
852 assert_eq!(actual, expected);
853 }
854
855 #[test]
856 fn it_accepts_short_output() {
857 let input = TestVector(vec![91; 6], 10);
858
859 let actual = format!("{input}");
860 let expected = format!("{:?}", vec![91; 6]);
861 assert_eq!(actual, expected);
862 }
863
864 #[test]
865 fn it_accepts_long_output() {
866 let input = TestVector(vec![91; 10], 9);
867
868 let actual = format!("{input}");
869 let expected = format!("{:?}", vec![91; 9]);
870 assert_eq!(actual, expected);
871 }
872
873 #[test]
874 fn ser_deser_schema_result() {
875 let schema = Schema::new(vec![
876 Field::new("c1", DataType::Utf8, false),
877 Field::new("c2", DataType::Float64, true),
878 Field::new("c3", DataType::UInt32, false),
879 Field::new("c4", DataType::Boolean, true),
880 Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
881 Field::new("c6", DataType::Time32(TimeUnit::Second), false),
882 ]);
883 let option = IpcWriteOptions::default();
886 let schema_ipc = SchemaAsIpc::new(&schema, &option);
887 let result: SchemaResult = schema_ipc.try_into().unwrap();
888 let des_schema: Schema = (&result).try_into().unwrap();
889 assert_eq!(schema, des_schema);
890
891 let option = IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap();
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
900 #[test]
901 fn test_dict_schema() {
902 let schema = Schema::new(vec![
904 Field::new(
905 "a",
906 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
907 false,
908 ),
909 Field::new(
910 "b",
911 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
912 false,
913 ),
914 ]);
915
916 let flight_info = FlightInfo::new().try_with_schema(&schema).unwrap();
917
918 let new_schema = Schema::try_from(flight_info).unwrap();
919 assert_eq!(schema, new_schema);
920 }
921}