qail_pg/driver/
types.rs

1//! Core types: ColumnInfo, PgRow, PgError, PgResult, QueryResult, ResultFormat,
2//! and wire-protocol message utilities.
3
4use bytes::Bytes;
5use std::collections::HashMap;
6use std::sync::Arc;
7
8/// Metadata about the columns returned by a query.
9///
10/// Maps column names to positional indices and stores OID / format
11/// information so that [`PgRow`] values can be decoded correctly.
12#[derive(Debug, Clone)]
13pub struct ColumnInfo {
14    /// Lookup table from column name to zero-based index.
15    pub name_to_index: HashMap<String, usize>,
16    /// PostgreSQL type OIDs, one per column.
17    pub oids: Vec<u32>,
18    /// Wire format codes (0 = text, 1 = binary), one per column.
19    pub formats: Vec<i16>,
20}
21
22impl ColumnInfo {
23    /// Build column metadata from the `RowDescription` field list
24    /// returned by the backend after a query.
25    pub fn from_fields(fields: &[crate::protocol::FieldDescription]) -> Self {
26        let mut name_to_index = HashMap::with_capacity(fields.len());
27        let mut oids = Vec::with_capacity(fields.len());
28        let mut formats = Vec::with_capacity(fields.len());
29
30        for (i, field) in fields.iter().enumerate() {
31            name_to_index.entry(field.name.clone()).or_insert(i);
32            oids.push(field.type_oid);
33            formats.push(field.format);
34        }
35
36        Self {
37            name_to_index,
38            oids,
39            formats,
40        }
41    }
42}
43
44/// PostgreSQL row with column data and metadata.
45pub struct PgRow {
46    /// Raw column values — `None` represents SQL `NULL`.
47    pub columns: Vec<Option<Vec<u8>>>,
48    /// Shared column metadata for decoding values by name or type.
49    pub column_info: Option<Arc<ColumnInfo>>,
50}
51
52/// PostgreSQL row backed by a single shared payload buffer.
53///
54/// This avoids per-cell byte copies by storing one `Bytes` payload plus
55/// column offsets into that payload.
56#[derive(Debug, Clone, Default)]
57pub struct PgBytesRow {
58    pub(crate) payload: Bytes,
59    pub(crate) spans: Vec<Option<(usize, usize)>>,
60    /// Shared column metadata for decoding values by name or type.
61    pub column_info: Option<Arc<ColumnInfo>>,
62}
63
64/// Error type for PostgreSQL driver operations.
65#[derive(Debug)]
66pub enum PgError {
67    /// TCP / TLS connection failure with the PostgreSQL server.
68    Connection(String),
69    /// Wire-protocol framing or decoding error.
70    Protocol(String),
71    /// Authentication failure (bad password, unsupported mechanism, etc.).
72    Auth(String),
73    /// Query execution error returned by the backend (e.g. constraint violation).
74    Query(String),
75    /// Structured server error with SQLSTATE and optional detail/hint fields.
76    QueryServer(PgServerError),
77    /// The query returned zero rows when at least one was expected.
78    NoRows,
79    /// I/O error (preserves inner error for chaining)
80    Io(std::io::Error),
81    /// Encoding error (parameter limit, etc.)
82    Encode(String),
83    /// Operation timed out (connection, acquire, query)
84    Timeout(String),
85    /// Pool exhausted — all connections are in use
86    PoolExhausted {
87        /// Maximum pool size that was reached.
88        max: usize,
89    },
90    /// Pool is closed and no longer accepting requests
91    PoolClosed,
92}
93
94/// Exact message emitted when the server answers the SSLRequest preface
95/// with anything other than `'S'`.
96///
97/// The `TlsMode::Prefer` fallback matches this message by full equality to
98/// decide a plaintext retry. TLS handshake and certificate-validation
99/// failures use different messages and must keep failing closed — never
100/// widen this to a substring match.
101pub(crate) const TLS_UNSUPPORTED_BY_SERVER: &str = "Server does not support TLS";
102
103impl PgError {
104    /// The SSLRequest preface was answered with a non-`'S'` byte: the server
105    /// does not offer TLS on this endpoint.
106    pub(crate) fn tls_unsupported_by_server() -> Self {
107        PgError::Connection(TLS_UNSUPPORTED_BY_SERVER.to_string())
108    }
109
110    /// True only for the exact [`Self::tls_unsupported_by_server`] sentinel.
111    /// A failed TLS handshake or certificate validation never matches.
112    pub(crate) fn is_tls_unsupported_by_server(&self) -> bool {
113        matches!(self, PgError::Connection(msg) if msg == TLS_UNSUPPORTED_BY_SERVER)
114    }
115}
116
117/// Structured PostgreSQL server error fields.
118#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct PgServerError {
120    /// Severity level (e.g. `ERROR`, `FATAL`, `WARNING`).
121    pub severity: String,
122    /// SQLSTATE error code (e.g. `23505`).
123    pub code: String,
124    /// Human-readable message.
125    pub message: String,
126    /// Optional detailed description.
127    pub detail: Option<String>,
128    /// Optional hint from server.
129    pub hint: Option<String>,
130}
131
132impl From<crate::protocol::ErrorFields> for PgServerError {
133    fn from(value: crate::protocol::ErrorFields) -> Self {
134        Self {
135            severity: value.severity,
136            code: value.code,
137            message: value.message,
138            detail: value.detail,
139            hint: value.hint,
140        }
141    }
142}
143
144impl std::fmt::Display for PgError {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        match self {
147            PgError::Connection(e) => write!(f, "Connection error: {}", e),
148            PgError::Protocol(e) => write!(f, "Protocol error: {}", e),
149            PgError::Auth(e) => write!(f, "Auth error: {}", e),
150            PgError::Query(e) => write!(f, "Query error: {}", e),
151            PgError::QueryServer(e) => write!(f, "Query error [{}]: {}", e.code, e.message),
152            PgError::NoRows => write!(f, "No rows returned"),
153            PgError::Io(e) => write!(f, "I/O error: {}", e),
154            PgError::Encode(e) => write!(f, "Encode error: {}", e),
155            PgError::Timeout(ctx) => write!(f, "Timeout: {}", ctx),
156            PgError::PoolExhausted { max } => write!(f, "Pool exhausted ({} max connections)", max),
157            PgError::PoolClosed => write!(f, "Connection pool is closed"),
158        }
159    }
160}
161
162impl std::error::Error for PgError {
163    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
164        match self {
165            PgError::Io(e) => Some(e),
166            _ => None,
167        }
168    }
169}
170
171impl From<std::io::Error> for PgError {
172    fn from(e: std::io::Error) -> Self {
173        PgError::Io(e)
174    }
175}
176
177impl From<crate::protocol::EncodeError> for PgError {
178    fn from(e: crate::protocol::EncodeError) -> Self {
179        PgError::Encode(e.to_string())
180    }
181}
182
183impl PgError {
184    /// Return structured server error fields when available.
185    pub fn server_error(&self) -> Option<&PgServerError> {
186        match self {
187            PgError::QueryServer(err) => Some(err),
188            _ => None,
189        }
190    }
191
192    /// Return SQLSTATE code when available.
193    pub fn sqlstate(&self) -> Option<&str> {
194        self.server_error().map(|e| e.code.as_str())
195    }
196
197    /// True when a cached prepared statement can be self-healed by clearing
198    /// local statement state and retrying once.
199    pub fn is_prepared_statement_retryable(&self) -> bool {
200        let Some(err) = self.server_error() else {
201            return false;
202        };
203
204        let code = err.code.as_str();
205        let message = err.message.to_ascii_lowercase();
206
207        // invalid_sql_statement_name
208        if code.eq_ignore_ascii_case("26000")
209            && message.contains("prepared statement")
210            && message.contains("does not exist")
211        {
212            return true;
213        }
214
215        // feature_not_supported + message heuristic used by PostgreSQL replans.
216        if code.eq_ignore_ascii_case("0A000") && message.contains("cached plan must be replanned") {
217            return true;
218        }
219
220        // Defensive message-only fallback for proxy/failover rewrites.
221        message.contains("cached plan must be replanned")
222    }
223
224    /// True when server reports the prepared statement name already exists.
225    ///
226    /// This typically means local cache eviction drifted from server state
227    /// (e.g. local entry dropped while backend statement still exists).
228    /// Callers can retry once without Parse after preserving local mapping.
229    pub fn is_prepared_statement_already_exists(&self) -> bool {
230        let Some(err) = self.server_error() else {
231            return false;
232        };
233        if !err.code.eq_ignore_ascii_case("42P05") {
234            return false;
235        }
236        let message = err.message.to_ascii_lowercase();
237        message.contains("prepared statement") && message.contains("already exists")
238    }
239
240    /// True when the error is a transient server condition that may succeed
241    /// on retry. Covers serialization failures, deadlocks, standby
242    /// unavailability, connection exceptions, and prepared-statement staleness.
243    ///
244    /// Callers should pair this with a bounded retry loop and backoff.
245    pub fn is_transient_server_error(&self) -> bool {
246        // Non-server errors that are inherently transient.
247        match self {
248            PgError::Timeout(_) => return true,
249            PgError::Io(io) => {
250                return matches!(
251                    io.kind(),
252                    std::io::ErrorKind::TimedOut
253                        | std::io::ErrorKind::ConnectionRefused
254                        | std::io::ErrorKind::ConnectionReset
255                        | std::io::ErrorKind::BrokenPipe
256                        | std::io::ErrorKind::Interrupted
257                );
258            }
259            PgError::Connection(_) => return true,
260            _ => {}
261        }
262
263        // Prepared-statement staleness is a subset of transient errors.
264        if self.is_prepared_statement_retryable() {
265            return true;
266        }
267
268        let Some(code) = self.sqlstate() else {
269            return false;
270        };
271
272        matches!(
273            code,
274            // serialization_failure — MVCC conflict, safe to retry
275            "40001"
276            // deadlock_detected — PG auto-aborts one participant
277            | "40P01"
278            // cannot_connect_now — hot-standby recovery in progress
279            | "57P03"
280            // admin_shutdown / crash_shutdown — server restarting
281            | "57P01"
282            | "57P02"
283        ) || code.starts_with("08") // connection_exception class
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::{ColumnInfo, PgError, TLS_UNSUPPORTED_BY_SERVER};
290    use crate::protocol::FieldDescription;
291
292    #[test]
293    fn tls_sentinel_matches_only_the_exact_message() {
294        assert!(PgError::tls_unsupported_by_server().is_tls_unsupported_by_server());
295        // Substring-containing messages must NOT match — a widening back to
296        // .contains() would silently downgrade cert failures to plaintext.
297        let prefixed = PgError::Connection(format!("connect failed: {TLS_UNSUPPORTED_BY_SERVER}"));
298        let suffixed = PgError::Connection(format!("{TLS_UNSUPPORTED_BY_SERVER}: retrying"));
299        let handshake = PgError::Connection("TLS handshake failed: bad cert".to_string());
300        assert!(!prefixed.is_tls_unsupported_by_server());
301        assert!(!suffixed.is_tls_unsupported_by_server());
302        assert!(!handshake.is_tls_unsupported_by_server());
303        assert!(
304            !PgError::Protocol(TLS_UNSUPPORTED_BY_SERVER.to_string())
305                .is_tls_unsupported_by_server()
306        );
307    }
308
309    fn field(name: &str, type_oid: u32) -> FieldDescription {
310        FieldDescription {
311            name: name.to_string(),
312            table_oid: 0,
313            column_attr: 0,
314            type_oid,
315            type_size: -1,
316            type_modifier: -1,
317            format: 0,
318        }
319    }
320
321    #[test]
322    fn column_info_preserves_first_duplicate_column_name() {
323        let info = ColumnInfo::from_fields(&[field("id", 23), field("id", 25)]);
324
325        assert_eq!(info.name_to_index.get("id").copied(), Some(0));
326        assert_eq!(info.oids, vec![23, 25]);
327    }
328}
329
330/// Result type for PostgreSQL operations.
331pub type PgResult<T> = Result<T, PgError>;
332
333#[inline]
334pub(crate) fn is_ignorable_session_message(msg: &crate::protocol::BackendMessage) -> bool {
335    matches!(
336        msg,
337        crate::protocol::BackendMessage::NoticeResponse(_)
338            | crate::protocol::BackendMessage::ParameterStatus { .. }
339    )
340}
341
342#[inline]
343pub(crate) fn unexpected_backend_message(
344    phase: &str,
345    msg: &crate::protocol::BackendMessage,
346) -> PgError {
347    PgError::Protocol(format!(
348        "Unexpected backend message during {} phase: {:?}",
349        phase, msg
350    ))
351}
352
353#[inline]
354pub(crate) fn is_ignorable_session_msg_type(msg_type: u8) -> bool {
355    matches!(msg_type, b'N' | b'S')
356}
357
358#[inline]
359pub(crate) fn unexpected_backend_msg_type(phase: &str, msg_type: u8) -> PgError {
360    let printable = if msg_type.is_ascii_graphic() {
361        msg_type as char
362    } else {
363        '?'
364    };
365    PgError::Protocol(format!(
366        "Unexpected backend message type during {} phase: byte={} char={}",
367        phase, msg_type, printable
368    ))
369}
370
371/// Result of a query that returns rows (SELECT/GET).
372#[derive(Debug, Clone)]
373pub struct QueryResult {
374    /// Column names from RowDescription.
375    pub columns: Vec<String>,
376    /// Rows of text-decoded values (None = NULL).
377    pub rows: Vec<Vec<Option<String>>>,
378}
379
380/// PostgreSQL result-column wire format.
381///
382/// - `Text` (0): server sends textual column values.
383/// - `Binary` (1): server sends binary column values.
384#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
385pub enum ResultFormat {
386    /// Text format (`0`)
387    #[default]
388    Text,
389    /// Binary format (`1`)
390    Binary,
391}
392
393impl ResultFormat {
394    #[inline]
395    pub(crate) fn as_wire_code(self) -> i16 {
396        match self {
397            ResultFormat::Text => crate::protocol::PgEncoder::FORMAT_TEXT,
398            ResultFormat::Binary => crate::protocol::PgEncoder::FORMAT_BINARY,
399        }
400    }
401}
qail_pg/driver/types.rs

qail_pg/driver/
types.rs
