elefant_client/postgres_client/
mod.rs

1mod copy;
2mod easy_client;
3mod establish;
4mod query;
5pub mod replication;
6mod statements;
7
8use crate::pool::{ConnectionFactory, PostgresPool};
9use crate::protocol::{
10    BackendMessage, CurrentTransactionStatus, FrontendMessage, PostgresConnection,
11};
12use crate::types::EnumTypeRegistry;
13use crate::{reborrow_until_polonius, ElefantClientError, PostgresConnectionSettings};
14use std::collections::HashMap;
15use std::sync::atomic::AtomicU64;
16use std::sync::Arc;
17use tracing::{debug, trace};
18
19pub use copy::{CopyReader, CopyWriter, OwnedCopyReader};
20pub use query::{PostgresDataRow, QueryResult, QueryResultSet, RowResultReader, SimpleQueryResult};
21pub use statements::*;
22
23pub struct PostgresClient<F: ConnectionFactory> {
24    pub(crate) connection: PostgresConnection<F::Connection>,
25    pub(crate) pool: Option<PostgresPool<F>>,
26    pub(crate) ready_for_query: bool,
27    write_buffer: Vec<u8>,
28    pub(crate) client_id: u64,
29    pub(crate) prepared_query_counter: u64,
30    sync_required: bool,
31    current_transaction_status: CurrentTransactionStatus,
32    pub(crate) enum_registry: Arc<EnumTypeRegistry>,
33    parameter_statuses: HashMap<String, String>,
34}
35
36impl<F: ConnectionFactory> PostgresClient<F> {
37    pub(crate) async fn start_new_query(&mut self) -> Result<(), ElefantClientError> {
38        if !self.ready_for_query {
39            if self.sync_required {
40                self.connection
41                    .write_frontend_message(&FrontendMessage::Sync)
42                    .await?;
43                self.connection.flush().await?;
44                self.sync_required = false;
45            }
46
47            loop {
48                match self.read_next_backend_message().await {
49                    Err(ElefantClientError::IoError(e)) => {
50                        return Err(ElefantClientError::IoError(e));
51                    }
52                    Err(e) => {
53                        debug!("Ignoring error while starting new query: {:?}", e);
54                    }
55                    Ok(msg) => match msg {
56                        BackendMessage::ReadyForQuery(_) => {
57                            break;
58                        }
59                        _ => {
60                            trace!("Ignoring message while starting new query: {:?}", msg);
61                        }
62                    },
63                }
64            }
65        }
66
67        self.ready_for_query = false;
68        Ok(())
69    }
70
71    pub async fn reset(&mut self) -> Result<(), ElefantClientError> {
72        if !self.ready_for_query {
73            if self.sync_required {
74                self.connection
75                    .write_frontend_message(&FrontendMessage::Sync)
76                    .await?;
77                self.connection.flush().await?;
78                self.sync_required = false;
79            }
80
81            loop {
82                match self.read_next_backend_message().await {
83                    Err(ElefantClientError::IoError(io_err)) => {
84                        return Err(ElefantClientError::IoError(io_err));
85                    }
86                    Err(e) => {
87                        debug!("Ignoring error while resetting elefant client: {:?}", e);
88                    }
89                    Ok(msg) => match msg {
90                        BackendMessage::ReadyForQuery(rfq) => {
91                            self.current_transaction_status = rfq.current_transaction_status;
92                            self.ready_for_query = true;
93                            break;
94                        }
95                        _ => {
96                            debug!("Ignoring message while resetting elefant client: {:?}", msg);
97                        }
98                    },
99                }
100            }
101        }
102
103        // If the connection was left in a transaction, roll it back to ensure a clean state.
104        if self.current_transaction_status == CurrentTransactionStatus::InTransaction
105            || self.current_transaction_status == CurrentTransactionStatus::InFailedTransaction
106        {
107            debug!("Rolling back lingering transaction during pool reset");
108            self.connection
109                .write_frontend_message(&FrontendMessage::Query(crate::protocol::Query {
110                    query: std::borrow::Cow::Borrowed("ROLLBACK;"),
111                }))
112                .await?;
113            self.connection.flush().await?;
114            self.ready_for_query = false;
115            loop {
116                if let BackendMessage::ReadyForQuery(rfq) = self.read_next_backend_message().await?
117                {
118                    self.current_transaction_status = rfq.current_transaction_status;
119                    self.ready_for_query = true;
120                    break;
121                }
122            }
123        }
124
125        Ok(())
126    }
127
128    /// Get a server parameter received via ParameterStatus messages.
129    ///
130    /// PostgreSQL sends these during connection startup and whenever a session
131    /// parameter changes. Common parameters include `server_version`,
132    /// `server_encoding`, `TimeZone`, etc.
133    pub fn get_parameter(&self, name: &str) -> Option<&str> {
134        self.parameter_statuses.get(name).map(|s| s.as_str())
135    }
136
137    /// Gracefully close this connection by sending a Terminate message to the backend
138    /// and a TLS close_notify (if applicable). Consumes the client so it cannot be used afterward.
139    pub async fn close(mut self) -> Result<(), ElefantClientError> {
140        self.connection
141            .write_frontend_message(&FrontendMessage::Terminate)
142            .await?;
143        self.connection.flush().await?;
144        // Best-effort TLS shutdown — the connection is already logically closed.
145        let _ = self.connection.shutdown().await;
146        Ok(())
147    }
148
149    pub(crate) async fn new(
150        connection: PostgresConnection<F::Connection>,
151        settings: &PostgresConnectionSettings,
152        enum_registry: Arc<EnumTypeRegistry>,
153        channel_binding_data: Option<Vec<u8>>,
154    ) -> Result<Self, ElefantClientError> {
155        let mut client = Self {
156            connection,
157            pool: None,
158            ready_for_query: false,
159            write_buffer: Vec::new(),
160            client_id: CLIENT_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
161            prepared_query_counter: 1,
162            sync_required: false,
163            current_transaction_status: CurrentTransactionStatus::Idle,
164            enum_registry,
165            parameter_statuses: HashMap::new(),
166        };
167
168        client.establish(settings, channel_binding_data).await?;
169        Ok(client)
170    }
171
172    /// Helper method for reading backend messages while ignoring and handling "async" messages.
173    pub(crate) async fn read_next_backend_message(
174        &mut self,
175    ) -> Result<BackendMessage<'_>, ElefantClientError> {
176        loop {
177            let connection: &mut PostgresConnection<F::Connection> =
178                reborrow_until_polonius!(&mut self.connection);
179            let msg = connection.read_backend_message().await?;
180            match msg {
181                BackendMessage::NoticeResponse(nr) => {
182                    debug!("Received notice response from postgres: {:?}", nr);
183                }
184                BackendMessage::ParameterStatus(ps) => {
185                    debug!("Received parameter status from postgres: {:?}", ps);
186                    self.parameter_statuses
187                        .insert(ps.name.into_owned(), ps.value.into_owned());
188                }
189                _ => {
190                    return Ok(msg);
191                }
192            }
193        }
194    }
195}
196
197static CLIENT_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
elefant_client/postgres_client/mod.rs

elefant_client/postgres_client/
mod.rs
