elefant_client/postgres_client/
statements.rs

1use crate::pool::ConnectionFactory;
2use crate::postgres_client::query::PreparedQueryResult;
3use crate::protocol::{BackendMessage, FrontendMessage, ValueFormat};
4use crate::{protocol, ElefantClientError, PostgresClient, QueryResult, ToSql};
5use std::borrow::Cow;
6use std::future::Future;
7use std::rc::Rc;
8use tracing::trace;
9
10pub struct PreparedQuery {
11    name: Option<String>,
12    client_id: u64,
13    parameter_description: protocol::ParameterDescription,
14    result: Rc<PreparedQueryResult>,
15}
16
17impl PreparedQuery {
18    pub(crate) fn new(
19        name: Option<String>,
20        client_id: u64,
21        parameter_description: protocol::ParameterDescription,
22        result: PreparedQueryResult,
23    ) -> Self {
24        PreparedQuery {
25            name,
26            client_id,
27            parameter_description,
28            result: Rc::new(result),
29        }
30    }
31
32    /// Execute this prepared statement with parameters, returning a binary mode result
33    pub async fn execute<'postgres_client, F: ConnectionFactory>(
34        &self,
35        client: &'postgres_client mut PostgresClient<F>,
36        parameters: &[&dyn ToSql],
37    ) -> Result<QueryResult<'postgres_client, F>, ElefantClientError> {
38        client.start_new_query().await?;
39        client.sync_required = true;
40
41        let mut parameter_values: Vec<Option<&[u8]>> = Vec::with_capacity(parameters.len());
42        client.write_buffer.clear();
43
44        let mut parameter_positions = Vec::with_capacity(parameters.len());
45
46        for param in parameters.iter() {
47            if param.is_null() {
48                parameter_positions.push(None);
49                continue;
50            }
51            let start_index = client.write_buffer.len();
52            param.to_sql_binary(&mut client.write_buffer).map_err(|e| {
53                ElefantClientError::IoError(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
54            })?;
55            let end_index = client.write_buffer.len();
56            parameter_positions.push(Some((start_index, end_index)));
57        }
58
59        for position in parameter_positions {
60            if let Some((start_index, end_index)) = position {
61                parameter_values.push(Some(&client.write_buffer[start_index..end_index]));
62            } else {
63                parameter_values.push(None);
64            }
65        }
66
67        let source_statement_name = self
68            .name
69            .as_ref()
70            .map(|n| Cow::Borrowed(n.as_str()))
71            .unwrap_or(Cow::Borrowed(""));
72
73        client
74            .connection
75            .write_frontend_message(&FrontendMessage::Bind(protocol::Bind {
76                source_statement_name,
77                destination_portal_name: Cow::Borrowed(""),
78                parameter_values,
79                result_column_formats: vec![ValueFormat::Binary],
80                parameter_formats: vec![ValueFormat::Binary],
81            }))
82            .await?;
83        client
84            .connection
85            .write_frontend_message(&FrontendMessage::Execute(protocol::Execute {
86                portal_name: Cow::Borrowed(""),
87                max_rows: 0,
88            }))
89            .await?;
90        client
91            .connection
92            .write_frontend_message(&FrontendMessage::Flush)
93            .await?;
94        client.connection.flush().await?;
95
96        let msg = client.read_next_backend_message().await?;
97
98        match msg {
99            BackendMessage::BindComplete => {
100                trace!("Bind complete");
101            }
102            BackendMessage::ErrorResponse(er) => {
103                return Err(ElefantClientError::PostgresError(format!("{er:?}")));
104            }
105            _ => {
106                return Err(ElefantClientError::UnexpectedBackendMessage(format!(
107                    "{msg:?}"
108                )));
109            }
110        }
111
112        Ok(QueryResult::new(client, Some(self.result.clone())))
113    }
114}
115
116trait Sealed {}
117
118#[allow(private_bounds)]
119pub trait Statement: Sealed {
120    fn prepare<F: ConnectionFactory>(
121        &self,
122        client: &mut PostgresClient<F>,
123    ) -> impl Future<Output = Result<PreparedQuery, ElefantClientError>>;
124}
125
126impl Sealed for PreparedQuery {}
127
128impl Statement for PreparedQuery {
129    async fn prepare<F: ConnectionFactory>(
130        &self,
131        _client: &mut PostgresClient<F>,
132    ) -> Result<PreparedQuery, ElefantClientError> {
133        // PreparedQuery is already prepared, so just clone it
134        Ok(PreparedQuery {
135            name: self.name.clone(),
136            client_id: self.client_id,
137            parameter_description: self.parameter_description.clone(),
138            result: self.result.clone(),
139        })
140    }
141}
142
143impl Sealed for str {}
144
145impl Statement for str {
146    async fn prepare<F: ConnectionFactory>(
147        &self,
148        client: &mut PostgresClient<F>,
149    ) -> Result<PreparedQuery, ElefantClientError> {
150        client.prepare_with_name(self, None).await
151    }
152}
153
154impl Sealed for String {}
155
156impl Statement for String {
157    async fn prepare<F: ConnectionFactory>(
158        &self,
159        client: &mut PostgresClient<F>,
160    ) -> Result<PreparedQuery, ElefantClientError> {
161        self.as_str().prepare(client).await
162    }
163}
elefant_client/postgres_client/statements.rs

elefant_client/postgres_client/
statements.rs
