polars_parquet_format/thrift/protocol/
mod.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//! Types used to send and receive primitives between a Thrift client and server.
19use std::convert::{From, TryFrom};
20use std::fmt;
21use std::fmt::{Display, Formatter};
22
23use crate::thrift::{ApplicationError, Error, ProtocolError, ProtocolErrorKind};
24
25mod compact;
26
27pub use compact::TCompactInputProtocol;
28
29mod compact_write;
30pub use compact_write::TCompactOutputProtocol;
31
32#[cfg(feature = "async")]
33#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
34mod stream;
35#[cfg(feature = "async")]
36#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
37pub use stream::{AsyncReadThrift, TInputStreamProtocol, TOutputStreamProtocol};
38
39#[cfg(feature = "async")]
40#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
41mod compact_stream;
42#[cfg(feature = "async")]
43#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
44pub use compact_stream::TCompactInputStreamProtocol;
45
46#[cfg(feature = "async")]
47#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
48mod compact_stream_write;
49#[cfg(feature = "async")]
50#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
51pub use compact_stream_write::TCompactOutputStreamProtocol;
52
53// Default maximum depth to which `TInputProtocol::skip` will skip a Thrift
54// field. A default is necessary because Thrift structs or collections may
55// contain nested structs and collections, which could result in indefinite
56// recursion.
57const MAXIMUM_SKIP_DEPTH: i8 = 64;
58
59pub trait ReadThrift: Sized {
60    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self>;
61}
62
63impl ReadThrift for String {
64    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self> {
65        i_prot.read_string()
66    }
67}
68
69impl ReadThrift for bool {
70    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self> {
71        i_prot.read_bool()
72    }
73}
74
75impl ReadThrift for u8 {
76    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self> {
77        i_prot.read_byte()
78    }
79}
80
81impl ReadThrift for i64 {
82    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self> {
83        i_prot.read_i64()
84    }
85}
86
87impl ReadThrift for Vec<u8> {
88    fn read_from_in_protocol<T: TInputProtocol>(i_prot: &mut T) -> crate::thrift::Result<Self> {
89        i_prot.read_bytes()
90    }
91}
92
93/// Converts a stream of bytes into Thrift identifiers, primitives,
94/// containers, or structs.
95///
96/// This trait does not deal with higher-level Thrift concepts like structs or
97/// exceptions - only with primitives and message or container boundaries. Once
98/// bytes are read they are deserialized and an identifier (for example
99/// `TMessageIdentifier`) or a primitive is returned.
100///
101/// All methods return a `thrift::Result`. If an `Err` is returned the protocol
102/// instance and its underlying transport should be terminated.
103pub trait TInputProtocol: Sized {
104    /// Read the beginning of a Thrift message.
105    fn read_message_begin(&mut self) -> crate::thrift::Result<TMessageIdentifier>;
106    /// Read the end of a Thrift message.
107    fn read_message_end(&mut self) -> crate::thrift::Result<()>;
108    /// Read the beginning of a Thrift struct.
109    fn read_struct_begin(&mut self) -> crate::thrift::Result<Option<TStructIdentifier>>;
110    /// Read the end of a Thrift struct.
111    fn read_struct_end(&mut self) -> crate::thrift::Result<()>;
112    /// Read the beginning of a Thrift struct field.
113    fn read_field_begin(&mut self) -> crate::thrift::Result<TFieldIdentifier>;
114    /// Read the end of a Thrift struct field.
115    fn read_field_end(&mut self) -> crate::thrift::Result<()>;
116    /// Read a bool.
117    fn read_bool(&mut self) -> crate::thrift::Result<bool>;
118    /// Read a fixed-length byte array.
119    fn read_bytes(&mut self) -> crate::thrift::Result<Vec<u8>>;
120    /// Read a word.
121    fn read_i8(&mut self) -> crate::thrift::Result<i8>;
122    /// Read a 16-bit signed integer.
123    fn read_i16(&mut self) -> crate::thrift::Result<i16>;
124    /// Read a 32-bit signed integer.
125    fn read_i32(&mut self) -> crate::thrift::Result<i32>;
126    /// Read a 64-bit signed integer.
127    fn read_i64(&mut self) -> crate::thrift::Result<i64>;
128    /// Read a 64-bit float.
129    fn read_double(&mut self) -> crate::thrift::Result<f64>;
130    /// Read a fixed-length string (not null terminated).
131    fn read_string(&mut self) -> crate::thrift::Result<String>;
132    /// Read the beginning of a list.
133    fn read_list_begin(&mut self) -> crate::thrift::Result<TListIdentifier>;
134    /// Read the end of a list.
135    fn read_list_end(&mut self) -> crate::thrift::Result<()>;
136    /// Read the beginning of a set.
137    fn read_set_begin(&mut self) -> crate::thrift::Result<TSetIdentifier>;
138    /// Read the end of a set.
139    fn read_set_end(&mut self) -> crate::thrift::Result<()>;
140    /// Read the beginning of a map.
141    fn read_map_begin(&mut self) -> crate::thrift::Result<TMapIdentifier>;
142    /// Read the end of a map.
143    fn read_map_end(&mut self) -> crate::thrift::Result<()>;
144    /// Skip a field with type `field_type` recursively until the default
145    /// maximum skip depth is reached.
146    fn skip(&mut self, field_type: TType) -> crate::thrift::Result<()> {
147        self.skip_till_depth(field_type, MAXIMUM_SKIP_DEPTH)
148    }
149    /// Skip a field with type `field_type` recursively up to `depth` levels.
150    fn skip_till_depth(&mut self, field_type: TType, depth: i8) -> crate::thrift::Result<()> {
151        if depth == 0 {
152            return Err(crate::thrift::Error::Protocol(ProtocolError {
153                kind: ProtocolErrorKind::DepthLimit,
154                message: format!("cannot parse past {:?}", field_type),
155            }));
156        }
157
158        match field_type {
159            TType::Bool => self.read_bool().map(|_| ()),
160            TType::I08 => self.read_i8().map(|_| ()),
161            TType::I16 => self.read_i16().map(|_| ()),
162            TType::I32 => self.read_i32().map(|_| ()),
163            TType::I64 => self.read_i64().map(|_| ()),
164            TType::Double => self.read_double().map(|_| ()),
165            TType::String => self.read_string().map(|_| ()),
166            TType::Struct => {
167                self.read_struct_begin()?;
168                loop {
169                    let field_ident = self.read_field_begin()?;
170                    if field_ident.field_type == TType::Stop {
171                        break;
172                    }
173                    self.skip_till_depth(field_ident.field_type, depth - 1)?;
174                }
175                self.read_struct_end()
176            }
177            TType::List => {
178                let list_ident = self.read_list_begin()?;
179                for _ in 0..list_ident.size {
180                    self.skip_till_depth(list_ident.element_type, depth - 1)?;
181                }
182                self.read_list_end()
183            }
184            TType::Set => {
185                let set_ident = self.read_set_begin()?;
186                for _ in 0..set_ident.size {
187                    self.skip_till_depth(set_ident.element_type, depth - 1)?;
188                }
189                self.read_set_end()
190            }
191            TType::Map => {
192                let map_ident = self.read_map_begin()?;
193                for _ in 0..map_ident.size {
194                    let key_type = map_ident
195                        .key_type
196                        .expect("non-zero sized map should contain key type");
197                    let val_type = map_ident
198                        .value_type
199                        .expect("non-zero sized map should contain value type");
200                    self.skip_till_depth(key_type, depth - 1)?;
201                    self.skip_till_depth(val_type, depth - 1)?;
202                }
203                self.read_map_end()
204            }
205            u => Err(crate::thrift::Error::Protocol(ProtocolError {
206                kind: ProtocolErrorKind::Unknown,
207                message: format!("cannot skip field type {:?}", &u),
208            })),
209        }
210    }
211
212    fn read_list<P: ReadThrift>(&mut self) -> crate::thrift::Result<Vec<P>> {
213        let list_ident = self.read_list_begin()?;
214        let mut val: Vec<P> = Vec::with_capacity(list_ident.size as usize);
215        for _ in 0..list_ident.size {
216            val.push(P::read_from_in_protocol(self)?);
217        }
218        self.read_list_end()?;
219        Ok(val)
220    }
221
222    // utility (DO NOT USE IN GENERATED CODE!!!!)
223    //
224
225    /// Read an unsigned byte.
226    ///
227    /// This method should **never** be used in generated code.
228    fn read_byte(&mut self) -> crate::thrift::Result<u8>;
229}
230
231/// Converts Thrift identifiers, primitives, containers or structs into a
232/// stream of bytes.
233///
234/// This trait does not deal with higher-level Thrift concepts like structs or
235/// exceptions - only with primitives and message or container boundaries.
236/// Write methods take an identifier (for example, `TMessageIdentifier`) or a
237/// primitive. Any or all of the fields in an identifier may be omitted when
238/// writing to the transport. Write methods may even be noops. All of this is
239/// transparent to the caller; as long as a matching `TInputProtocol`
240/// implementation is used, received messages will be decoded correctly.
241///
242/// All methods return a `thrift::Result`. If an `Err` is returned the protocol
243/// instance and its underlying transport should be terminated.
244pub trait TOutputProtocol {
245    /// Write the beginning of a Thrift message.
246    fn write_message_begin(
247        &mut self,
248        identifier: &TMessageIdentifier,
249    ) -> crate::thrift::Result<usize>;
250    /// Write the end of a Thrift message.
251    fn write_message_end(&mut self) -> crate::thrift::Result<usize>;
252    /// Write the beginning of a Thrift struct.
253    fn write_struct_begin(
254        &mut self,
255        identifier: &TStructIdentifier,
256    ) -> crate::thrift::Result<usize>;
257    /// Write the end of a Thrift struct.
258    fn write_struct_end(&mut self) -> crate::thrift::Result<usize>;
259    /// Write the beginning of a Thrift field.
260    fn write_field_begin(&mut self, identifier: &TFieldIdentifier) -> crate::thrift::Result<usize>;
261    /// Write the end of a Thrift field.
262    fn write_field_end(&mut self) -> crate::thrift::Result<usize>;
263    /// Write a STOP field indicating that all the fields in a struct have been
264    /// written.
265    fn write_field_stop(&mut self) -> crate::thrift::Result<usize>;
266    /// Write a bool.
267    fn write_bool(&mut self, b: bool) -> crate::thrift::Result<usize>;
268    /// Write a fixed-length byte array.
269    fn write_bytes(&mut self, b: &[u8]) -> crate::thrift::Result<usize>;
270    /// Write an 8-bit signed integer.
271    fn write_i8(&mut self, i: i8) -> crate::thrift::Result<usize>;
272    /// Write a 16-bit signed integer.
273    fn write_i16(&mut self, i: i16) -> crate::thrift::Result<usize>;
274    /// Write a 32-bit signed integer.
275    fn write_i32(&mut self, i: i32) -> crate::thrift::Result<usize>;
276    /// Write a 64-bit signed integer.
277    fn write_i64(&mut self, i: i64) -> crate::thrift::Result<usize>;
278    /// Write a 64-bit float.
279    fn write_double(&mut self, d: f64) -> crate::thrift::Result<usize>;
280    /// Write a fixed-length string.
281    fn write_string(&mut self, s: &str) -> crate::thrift::Result<usize>;
282    /// Write the beginning of a list.
283    fn write_list_begin(&mut self, identifier: &TListIdentifier) -> crate::thrift::Result<usize>;
284    /// Write the end of a list.
285    fn write_list_end(&mut self) -> crate::thrift::Result<usize>;
286    /// Write the beginning of a set.
287    fn write_set_begin(&mut self, identifier: &TSetIdentifier) -> crate::thrift::Result<usize>;
288    /// Write the end of a set.
289    fn write_set_end(&mut self) -> crate::thrift::Result<usize>;
290    /// Write the beginning of a map.
291    fn write_map_begin(&mut self, identifier: &TMapIdentifier) -> crate::thrift::Result<usize>;
292    /// Write the end of a map.
293    fn write_map_end(&mut self) -> crate::thrift::Result<usize>;
294    /// Flush buffered bytes to the underlying transport.
295    fn flush(&mut self) -> crate::thrift::Result<()>;
296
297    // utility (DO NOT USE IN GENERATED CODE!!!!)
298    //
299
300    /// Write an unsigned byte.
301    ///
302    /// This method should **never** be used in generated code.
303    fn write_byte(&mut self, b: u8) -> crate::thrift::Result<usize>; // FIXME: REMOVE
304}
305
306impl<P> TInputProtocol for Box<P>
307where
308    P: TInputProtocol,
309{
310    fn read_message_begin(&mut self) -> crate::thrift::Result<TMessageIdentifier> {
311        (**self).read_message_begin()
312    }
313
314    fn read_message_end(&mut self) -> crate::thrift::Result<()> {
315        (**self).read_message_end()
316    }
317
318    fn read_struct_begin(&mut self) -> crate::thrift::Result<Option<TStructIdentifier>> {
319        (**self).read_struct_begin()
320    }
321
322    fn read_struct_end(&mut self) -> crate::thrift::Result<()> {
323        (**self).read_struct_end()
324    }
325
326    fn read_field_begin(&mut self) -> crate::thrift::Result<TFieldIdentifier> {
327        (**self).read_field_begin()
328    }
329
330    fn read_field_end(&mut self) -> crate::thrift::Result<()> {
331        (**self).read_field_end()
332    }
333
334    fn read_bool(&mut self) -> crate::thrift::Result<bool> {
335        (**self).read_bool()
336    }
337
338    fn read_bytes(&mut self) -> crate::thrift::Result<Vec<u8>> {
339        (**self).read_bytes()
340    }
341
342    fn read_i8(&mut self) -> crate::thrift::Result<i8> {
343        (**self).read_i8()
344    }
345
346    fn read_i16(&mut self) -> crate::thrift::Result<i16> {
347        (**self).read_i16()
348    }
349
350    fn read_i32(&mut self) -> crate::thrift::Result<i32> {
351        (**self).read_i32()
352    }
353
354    fn read_i64(&mut self) -> crate::thrift::Result<i64> {
355        (**self).read_i64()
356    }
357
358    fn read_double(&mut self) -> crate::thrift::Result<f64> {
359        (**self).read_double()
360    }
361
362    fn read_string(&mut self) -> crate::thrift::Result<String> {
363        (**self).read_string()
364    }
365
366    fn read_list_begin(&mut self) -> crate::thrift::Result<TListIdentifier> {
367        (**self).read_list_begin()
368    }
369
370    fn read_list_end(&mut self) -> crate::thrift::Result<()> {
371        (**self).read_list_end()
372    }
373
374    fn read_set_begin(&mut self) -> crate::thrift::Result<TSetIdentifier> {
375        (**self).read_set_begin()
376    }
377
378    fn read_set_end(&mut self) -> crate::thrift::Result<()> {
379        (**self).read_set_end()
380    }
381
382    fn read_map_begin(&mut self) -> crate::thrift::Result<TMapIdentifier> {
383        (**self).read_map_begin()
384    }
385
386    fn read_map_end(&mut self) -> crate::thrift::Result<()> {
387        (**self).read_map_end()
388    }
389
390    fn read_byte(&mut self) -> crate::thrift::Result<u8> {
391        (**self).read_byte()
392    }
393}
394
395impl<P> TOutputProtocol for Box<P>
396where
397    P: TOutputProtocol + ?Sized,
398{
399    fn write_message_begin(
400        &mut self,
401        identifier: &TMessageIdentifier,
402    ) -> crate::thrift::Result<usize> {
403        (**self).write_message_begin(identifier)
404    }
405
406    fn write_message_end(&mut self) -> crate::thrift::Result<usize> {
407        (**self).write_message_end()
408    }
409
410    fn write_struct_begin(
411        &mut self,
412        identifier: &TStructIdentifier,
413    ) -> crate::thrift::Result<usize> {
414        (**self).write_struct_begin(identifier)
415    }
416
417    fn write_struct_end(&mut self) -> crate::thrift::Result<usize> {
418        (**self).write_struct_end()
419    }
420
421    fn write_field_begin(&mut self, identifier: &TFieldIdentifier) -> crate::thrift::Result<usize> {
422        (**self).write_field_begin(identifier)
423    }
424
425    fn write_field_end(&mut self) -> crate::thrift::Result<usize> {
426        (**self).write_field_end()
427    }
428
429    fn write_field_stop(&mut self) -> crate::thrift::Result<usize> {
430        (**self).write_field_stop()
431    }
432
433    fn write_bool(&mut self, b: bool) -> crate::thrift::Result<usize> {
434        (**self).write_bool(b)
435    }
436
437    fn write_bytes(&mut self, b: &[u8]) -> crate::thrift::Result<usize> {
438        (**self).write_bytes(b)
439    }
440
441    fn write_i8(&mut self, i: i8) -> crate::thrift::Result<usize> {
442        (**self).write_i8(i)
443    }
444
445    fn write_i16(&mut self, i: i16) -> crate::thrift::Result<usize> {
446        (**self).write_i16(i)
447    }
448
449    fn write_i32(&mut self, i: i32) -> crate::thrift::Result<usize> {
450        (**self).write_i32(i)
451    }
452
453    fn write_i64(&mut self, i: i64) -> crate::thrift::Result<usize> {
454        (**self).write_i64(i)
455    }
456
457    fn write_double(&mut self, d: f64) -> crate::thrift::Result<usize> {
458        (**self).write_double(d)
459    }
460
461    fn write_string(&mut self, s: &str) -> crate::thrift::Result<usize> {
462        (**self).write_string(s)
463    }
464
465    fn write_list_begin(&mut self, identifier: &TListIdentifier) -> crate::thrift::Result<usize> {
466        (**self).write_list_begin(identifier)
467    }
468
469    fn write_list_end(&mut self) -> crate::thrift::Result<usize> {
470        (**self).write_list_end()
471    }
472
473    fn write_set_begin(&mut self, identifier: &TSetIdentifier) -> crate::thrift::Result<usize> {
474        (**self).write_set_begin(identifier)
475    }
476
477    fn write_set_end(&mut self) -> crate::thrift::Result<usize> {
478        (**self).write_set_end()
479    }
480
481    fn write_map_begin(&mut self, identifier: &TMapIdentifier) -> crate::thrift::Result<usize> {
482        (**self).write_map_begin(identifier)
483    }
484
485    fn write_map_end(&mut self) -> crate::thrift::Result<usize> {
486        (**self).write_map_end()
487    }
488
489    fn flush(&mut self) -> crate::thrift::Result<()> {
490        (**self).flush()
491    }
492
493    fn write_byte(&mut self, b: u8) -> crate::thrift::Result<usize> {
494        (**self).write_byte(b)
495    }
496}
497
498/// Thrift message identifier.
499#[derive(Clone, Debug, Eq, PartialEq)]
500pub struct TMessageIdentifier {
501    /// Service call the message is associated with.
502    pub name: String,
503    /// Message type.
504    pub message_type: TMessageType,
505    /// Ordered sequence number identifying the message.
506    pub sequence_number: u32,
507}
508
509impl TMessageIdentifier {
510    /// Create a `TMessageIdentifier` for a Thrift service-call named `name`
511    /// with message type `message_type` and sequence number `sequence_number`.
512    pub fn new<S: Into<String>>(
513        name: S,
514        message_type: TMessageType,
515        sequence_number: u32,
516    ) -> TMessageIdentifier {
517        TMessageIdentifier {
518            name: name.into(),
519            message_type,
520            sequence_number,
521        }
522    }
523}
524
525/// Thrift struct identifier.
526#[derive(Clone, Debug, Eq, PartialEq)]
527pub struct TStructIdentifier {
528    /// Name of the encoded Thrift struct.
529    pub name: String,
530}
531
532impl TStructIdentifier {
533    /// Create a `TStructIdentifier` for a struct named `name`.
534    pub fn new<S: Into<String>>(name: S) -> TStructIdentifier {
535        TStructIdentifier { name: name.into() }
536    }
537}
538
539/// Thrift field identifier.
540#[derive(Clone, Debug, Eq, PartialEq)]
541pub struct TFieldIdentifier {
542    /// Name of the Thrift field.
543    ///
544    /// `None` if it's not sent over the wire.
545    pub name: Option<String>,
546    /// Field type.
547    ///
548    /// This may be a primitive, container, or a struct.
549    pub field_type: TType,
550    /// Thrift field id.
551    ///
552    /// `None` only if `field_type` is `TType::Stop`.
553    pub id: Option<i16>,
554}
555
556impl TFieldIdentifier {
557    /// Create a `TFieldIdentifier` for a field named `name` with type
558    /// `field_type` and field id `id`.
559    ///
560    /// `id` should be `None` if `field_type` is `TType::Stop`.
561    pub fn new<N, S, I>(name: N, field_type: TType, id: I) -> TFieldIdentifier
562    where
563        N: Into<Option<S>>,
564        S: Into<String>,
565        I: Into<Option<i16>>,
566    {
567        TFieldIdentifier {
568            name: name.into().map(|n| n.into()),
569            field_type,
570            id: id.into(),
571        }
572    }
573}
574
575/// Thrift list identifier.
576#[derive(Clone, Debug, Eq, PartialEq)]
577pub struct TListIdentifier {
578    /// Type of the elements in the list.
579    pub element_type: TType,
580    /// Number of elements in the list.
581    pub size: u32,
582}
583
584impl TListIdentifier {
585    /// Create a `TListIdentifier` for a list with `size` elements of type
586    /// `element_type`.
587    pub fn new(element_type: TType, size: u32) -> TListIdentifier {
588        TListIdentifier { element_type, size }
589    }
590}
591
592/// Thrift set identifier.
593#[derive(Clone, Debug, Eq, PartialEq)]
594pub struct TSetIdentifier {
595    /// Type of the elements in the set.
596    pub element_type: TType,
597    /// Number of elements in the set.
598    pub size: u32,
599}
600
601impl TSetIdentifier {
602    /// Create a `TSetIdentifier` for a set with `size` elements of type
603    /// `element_type`.
604    pub fn new(element_type: TType, size: u32) -> TSetIdentifier {
605        TSetIdentifier { element_type, size }
606    }
607}
608
609/// Thrift map identifier.
610#[derive(Clone, Debug, Eq, PartialEq)]
611pub struct TMapIdentifier {
612    /// Map key type.
613    pub key_type: Option<TType>,
614    /// Map value type.
615    pub value_type: Option<TType>,
616    /// Number of entries in the map.
617    pub size: u32,
618}
619
620impl TMapIdentifier {
621    /// Create a `TMapIdentifier` for a map with `size` entries of type
622    /// `key_type -> value_type`.
623    pub fn new<K, V>(key_type: K, value_type: V, size: u32) -> TMapIdentifier
624    where
625        K: Into<Option<TType>>,
626        V: Into<Option<TType>>,
627    {
628        TMapIdentifier {
629            key_type: key_type.into(),
630            value_type: value_type.into(),
631            size,
632        }
633    }
634}
635
636/// Thrift message types.
637#[derive(Clone, Copy, Debug, Eq, PartialEq)]
638pub enum TMessageType {
639    /// Service-call request.
640    Call,
641    /// Service-call response.
642    Reply,
643    /// Unexpected error in the remote service.
644    Exception,
645    /// One-way service-call request (no response is expected).
646    OneWay,
647}
648
649impl Display for TMessageType {
650    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
651        match *self {
652            TMessageType::Call => write!(f, "Call"),
653            TMessageType::Reply => write!(f, "Reply"),
654            TMessageType::Exception => write!(f, "Exception"),
655            TMessageType::OneWay => write!(f, "OneWay"),
656        }
657    }
658}
659
660impl From<TMessageType> for u8 {
661    fn from(message_type: TMessageType) -> Self {
662        match message_type {
663            TMessageType::Call => 0x01,
664            TMessageType::Reply => 0x02,
665            TMessageType::Exception => 0x03,
666            TMessageType::OneWay => 0x04,
667        }
668    }
669}
670
671impl TryFrom<u8> for TMessageType {
672    type Error = Error;
673    fn try_from(b: u8) -> Result<Self, Self::Error> {
674        match b {
675            0x01 => Ok(TMessageType::Call),
676            0x02 => Ok(TMessageType::Reply),
677            0x03 => Ok(TMessageType::Exception),
678            0x04 => Ok(TMessageType::OneWay),
679            unkn => Err(Error::Protocol(ProtocolError {
680                kind: ProtocolErrorKind::InvalidData,
681                message: format!("cannot convert {} to TMessageType", unkn),
682            })),
683        }
684    }
685}
686
687/// Thrift struct-field types.
688#[derive(Clone, Copy, Debug, Eq, PartialEq)]
689pub enum TType {
690    /// Indicates that there are no more serialized fields in this Thrift struct.
691    Stop,
692    /// Void (`()`) field.
693    Void,
694    /// Boolean.
695    Bool,
696    /// Signed 8-bit int.
697    I08,
698    /// Double-precision number.
699    Double,
700    /// Signed 16-bit int.
701    I16,
702    /// Signed 32-bit int.
703    I32,
704    /// Signed 64-bit int.
705    I64,
706    /// UTF-8 string.
707    String,
708    /// UTF-7 string. *Unsupported*.
709    Utf7,
710    /// Thrift struct.
711    Struct,
712    /// Map.
713    Map,
714    /// Set.
715    Set,
716    /// List.
717    List,
718    /// UTF-8 string.
719    Utf8,
720    /// UTF-16 string. *Unsupported*.
721    Utf16,
722}
723
724impl Display for TType {
725    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
726        match *self {
727            TType::Stop => write!(f, "STOP"),
728            TType::Void => write!(f, "void"),
729            TType::Bool => write!(f, "bool"),
730            TType::I08 => write!(f, "i08"),
731            TType::Double => write!(f, "double"),
732            TType::I16 => write!(f, "i16"),
733            TType::I32 => write!(f, "i32"),
734            TType::I64 => write!(f, "i64"),
735            TType::String => write!(f, "string"),
736            TType::Utf7 => write!(f, "UTF7"),
737            TType::Struct => write!(f, "struct"),
738            TType::Map => write!(f, "map"),
739            TType::Set => write!(f, "set"),
740            TType::List => write!(f, "list"),
741            TType::Utf8 => write!(f, "UTF8"),
742            TType::Utf16 => write!(f, "UTF16"),
743        }
744    }
745}
746
747/// Compare the expected message sequence number `expected` with the received
748/// message sequence number `actual`.
749///
750/// Return `()` if `actual == expected`, `Err` otherwise.
751pub fn verify_expected_sequence_number(expected: i32, actual: i32) -> crate::thrift::Result<()> {
752    if expected == actual {
753        Ok(())
754    } else {
755        Err(crate::thrift::Error::Application(
756            crate::thrift::ApplicationError {
757                kind: crate::thrift::ApplicationErrorKind::BadSequenceId,
758                message: format!("expected {} got {}", expected, actual),
759            },
760        ))
761    }
762}
763
764/// Compare the expected service-call name `expected` with the received
765/// service-call name `actual`.
766///
767/// Return `()` if `actual == expected`, `Err` otherwise.
768pub fn verify_expected_service_call(expected: &str, actual: &str) -> crate::thrift::Result<()> {
769    if expected == actual {
770        Ok(())
771    } else {
772        Err(crate::thrift::Error::Application(ApplicationError {
773            kind: crate::thrift::ApplicationErrorKind::WrongMethodName,
774            message: format!("expected {} got {}", expected, actual),
775        }))
776    }
777}
778
779/// Compare the expected message type `expected` with the received message type
780/// `actual`.
781///
782/// Return `()` if `actual == expected`, `Err` otherwise.
783pub fn verify_expected_message_type(
784    expected: TMessageType,
785    actual: TMessageType,
786) -> crate::thrift::Result<()> {
787    if expected == actual {
788        Ok(())
789    } else {
790        Err(crate::thrift::Error::Application(ApplicationError {
791            kind: crate::thrift::ApplicationErrorKind::InvalidMessageType,
792            message: format!("expected {} got {}", expected, actual),
793        }))
794    }
795}
796
797/// Check if a required Thrift struct field exists.
798///
799/// Return `()` if it does, `Err` otherwise.
800pub fn verify_required_field_exists<T>(
801    field_name: &str,
802    field: &Option<T>,
803) -> crate::thrift::Result<()> {
804    match *field {
805        Some(_) => Ok(()),
806        None => Err(Error::Protocol(ProtocolError {
807            kind: ProtocolErrorKind::Unknown,
808            message: format!("missing required field {}", field_name),
809        })),
810    }
811}
812
813/// Extract the field id from a Thrift field identifier.
814///
815/// `field_ident` must *not* have `TFieldIdentifier.field_type` of type `TType::Stop`.
816///
817/// Return `TFieldIdentifier.id` if an id exists, `Err` otherwise.
818pub fn field_id(field_ident: &TFieldIdentifier) -> crate::thrift::Result<i16> {
819    field_ident.id.ok_or_else(|| {
820        crate::thrift::Error::Protocol(crate::thrift::ProtocolError {
821            kind: crate::thrift::ProtocolErrorKind::Unknown,
822            message: format!("missing field in in {:?}", field_ident),
823        })
824    })
825}
polars_parquet_format/thrift/protocol/mod.rs

polars_parquet_format/thrift/protocol/
mod.rs
