datafusion_flight_sql_server/
state.rs

1use std::fmt::Display;
2
3use arrow_flight::{
4    error::FlightError,
5    sql::{self, Any, Command},
6};
7use prost::{bytes::Bytes, Message};
8
9pub type Result<T, E = FlightError> = std::result::Result<T, E>;
10
11#[derive(Debug, PartialEq, Clone)]
12pub struct CommandTicket {
13    pub command: sql::Command,
14}
15
16impl CommandTicket {
17    pub fn new(cmd: sql::Command) -> Self {
18        Self { command: cmd }
19    }
20
21    pub fn try_decode(msg: Bytes) -> Result<Self> {
22        let msg = CommandTicketMessage::decode(msg).map_err(decode_error_flight_error)?;
23
24        Self::try_decode_command(msg.command)
25    }
26
27    pub fn try_decode_command(cmd: Bytes) -> Result<Self> {
28        let content_msg = Any::decode(cmd).map_err(decode_error_flight_error)?;
29        let command = Command::try_from(content_msg).map_err(FlightError::Arrow)?;
30
31        Ok(Self { command })
32    }
33
34    pub fn try_encode(self) -> Result<Bytes> {
35        let content_msg = self.command.into_any().encode_to_vec();
36
37        let msg = CommandTicketMessage {
38            command: content_msg.into(),
39        };
40
41        Ok(msg.encode_to_vec().into())
42    }
43}
44
45#[derive(Clone, PartialEq, Message)]
46struct CommandTicketMessage {
47    #[prost(bytes = "bytes", tag = "2")]
48    command: Bytes,
49}
50
51fn decode_error_flight_error(err: prost::DecodeError) -> FlightError {
52    FlightError::DecodeError(format!("{err:?}"))
53}
54
55/// Represents a query handle for use in prepared statements.
56/// All state required to run the prepared statement is passed
57/// back and forth to the client, so any service instance can run it
58#[derive(Debug, Clone)]
59pub struct QueryHandle {
60    /// The raw SQL query text
61    query: String,
62    parameters: Option<Bytes>,
63}
64
65impl QueryHandle {
66    pub fn new(query: String, parameters: Option<Bytes>) -> Self {
67        Self { query, parameters }
68    }
69
70    pub fn query(&self) -> &str {
71        self.query.as_ref()
72    }
73
74    pub fn parameters(&self) -> Option<&[u8]> {
75        self.parameters.as_deref()
76    }
77
78    pub fn set_parameters(&mut self, parameters: Option<Bytes>) {
79        self.parameters = parameters;
80    }
81
82    pub fn try_decode(msg: Bytes) -> Result<Self> {
83        let msg = QueryHandleMessage::decode(msg).map_err(decode_error_flight_error)?;
84
85        Ok(Self {
86            query: msg.query,
87            parameters: msg.parameters,
88        })
89    }
90
91    pub fn encode(self) -> Bytes {
92        let msg = QueryHandleMessage {
93            query: self.query,
94            parameters: self.parameters,
95        };
96
97        msg.encode_to_vec().into()
98    }
99}
100
101impl From<QueryHandle> for Bytes {
102    fn from(value: QueryHandle) -> Self {
103        value.encode()
104    }
105}
106
107impl Display for QueryHandle {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        write!(f, "Query({})", self.query)
110    }
111}
112
113#[derive(Clone, PartialEq, Message)]
114pub struct QueryHandleMessage {
115    /// The raw SQL query text
116    #[prost(string, tag = "1")]
117    query: String,
118    #[prost(bytes = "bytes", optional, tag = "2")]
119    parameters: Option<Bytes>,
120}
datafusion_flight_sql_server/state.rs

datafusion_flight_sql_server/
state.rs
