diesel/pg/connection/
raw.rs

1#![allow(clippy::too_many_arguments)]
2#![allow(unsafe_code)] // ffi code
3
4extern crate pq_sys;
5
6use self::pq_sys::*;
7use alloc::ffi::CString;
8use core::ffi as libc;
9use core::ffi::CStr;
10use core::ptr::NonNull;
11use core::{ptr, str};
12
13use crate::result::*;
14
15use super::result::PgResult;
16use crate::pg::PgNotification;
17
18#[allow(missing_debug_implementations, missing_copy_implementations)]
19pub(super) struct RawConnection {
20    pub(super) internal_connection: NonNull<PGconn>,
21}
22
23impl RawConnection {
24    pub(super) fn establish(database_url: &str) -> ConnectionResult<Self> {
25        let connection_string = CString::new(database_url)?;
26        let connection_ptr = unsafe { PQconnectdb(connection_string.as_ptr()) };
27        let connection_status = unsafe { PQstatus(connection_ptr) };
28
29        match connection_status {
30            ConnStatusType::CONNECTION_OK => {
31                let connection_ptr = unsafe { NonNull::new_unchecked(connection_ptr) };
32                Ok(RawConnection {
33                    internal_connection: connection_ptr,
34                })
35            }
36            _ => {
37                let message = last_error_message(connection_ptr);
38
39                if !connection_ptr.is_null() {
40                    // Note that even if the server connection attempt fails (as indicated by PQstatus),
41                    // the application should call PQfinish to free the memory used by the PGconn object.
42                    // https://www.postgresql.org/docs/current/libpq-connect.html
43                    unsafe { PQfinish(connection_ptr) }
44                }
45
46                Err(ConnectionError::BadConnection(message))
47            }
48        }
49    }
50
51    pub(super) fn last_error_message(&self) -> String {
52        last_error_message(self.internal_connection.as_ptr())
53    }
54
55    pub(super) fn set_notice_processor(&self, notice_processor: NoticeProcessor) {
56        unsafe {
57            PQsetNoticeProcessor(
58                self.internal_connection.as_ptr(),
59                Some(notice_processor),
60                ptr::null_mut(),
61            );
62        }
63    }
64
65    pub(super) unsafe fn exec(&self, query: *const libc::c_char) -> QueryResult<RawResult> {
66        let result_ptr = unsafe { PQexec(self.internal_connection.as_ptr(), query) };
67        RawResult::new(result_ptr, self)
68    }
69
70    /// Sends a query and parameters to the server without using the prepare/bind cycle.
71    ///
72    /// This method uses PQsendQueryParams which combines the prepare and bind steps
73    /// and is more compatible with connection poolers like PgBouncer.
74    pub(super) unsafe fn send_query_params(
75        &self,
76        query: *const libc::c_char,
77        param_count: libc::c_int,
78        param_types: *const Oid,
79        param_values: *const *const libc::c_char,
80        param_lengths: *const libc::c_int,
81        param_formats: *const libc::c_int,
82        result_format: libc::c_int,
83    ) -> QueryResult<()> {
84        let res = unsafe {
85            PQsendQueryParams(
86                self.internal_connection.as_ptr(),
87                query,
88                param_count,
89                param_types,
90                param_values,
91                param_lengths,
92                param_formats,
93                result_format,
94            )
95        };
96        if res == 1 {
97            Ok(())
98        } else {
99            Err(Error::DatabaseError(
100                DatabaseErrorKind::UnableToSendCommand,
101                Box::new(self.last_error_message()),
102            ))
103        }
104    }
105
106    pub(super) unsafe fn send_query_prepared(
107        &self,
108        stmt_name: *const libc::c_char,
109        param_count: libc::c_int,
110        param_values: *const *const libc::c_char,
111        param_lengths: *const libc::c_int,
112        param_formats: *const libc::c_int,
113        result_format: libc::c_int,
114    ) -> QueryResult<()> {
115        let res = unsafe {
116            PQsendQueryPrepared(
117                self.internal_connection.as_ptr(),
118                stmt_name,
119                param_count,
120                param_values,
121                param_lengths,
122                param_formats,
123                result_format,
124            )
125        };
126        if res == 1 {
127            Ok(())
128        } else {
129            Err(Error::DatabaseError(
130                DatabaseErrorKind::UnableToSendCommand,
131                Box::new(self.last_error_message()),
132            ))
133        }
134    }
135
136    pub(super) unsafe fn prepare(
137        &self,
138        stmt_name: *const libc::c_char,
139        query: *const libc::c_char,
140        param_count: libc::c_int,
141        param_types: *const Oid,
142    ) -> QueryResult<RawResult> {
143        let ptr = unsafe {
144            PQprepare(
145                self.internal_connection.as_ptr(),
146                stmt_name,
147                query,
148                param_count,
149                param_types,
150            )
151        };
152        RawResult::new(ptr, self)
153    }
154
155    /// This is reasonably inexpensive as it just accesses variables internal to the connection
156    /// that are kept up to date by the `ReadyForQuery` messages from the PG server
157    pub(super) fn transaction_status(&self) -> PgTransactionStatus {
158        unsafe { PQtransactionStatus(self.internal_connection.as_ptr()) }.into()
159    }
160
161    /// The server version in `PQserverVersion` encoding, zero once the connection is bad.
162    pub(super) fn server_version(&self) -> libc::c_int {
163        // SAFETY: `internal_connection` is the live `PGconn` owned by `self`, and
164        // `PQserverVersion` only reads the version cached during connection startup.
165        unsafe { PQserverVersion(self.internal_connection.as_ptr()) }
166    }
167
168    pub(super) fn get_status(&self) -> ConnStatusType {
169        unsafe { PQstatus(self.internal_connection.as_ptr()) }
170    }
171
172    pub(crate) fn get_next_result(&self) -> Result<Option<PgResult>, Error> {
173        let res = unsafe { PQgetResult(self.internal_connection.as_ptr()) };
174        if res.is_null() {
175            Ok(None)
176        } else {
177            let raw = RawResult::new(res, self)?;
178            Ok(Some(PgResult::new(raw, self)?))
179        }
180    }
181
182    pub(crate) fn enable_row_by_row_mode(&self) -> QueryResult<()> {
183        let res = unsafe { PQsetSingleRowMode(self.internal_connection.as_ptr()) };
184        if res == 1 {
185            Ok(())
186        } else {
187            Err(Error::DatabaseError(
188                DatabaseErrorKind::Unknown,
189                Box::new(self.last_error_message()),
190            ))
191        }
192    }
193
194    pub(super) fn put_copy_data(&mut self, buf: &[u8]) -> QueryResult<()> {
195        for c in buf.chunks(i32::MAX as usize) {
196            let res = unsafe {
197                pq_sys::PQputCopyData(
198                    self.internal_connection.as_ptr(),
199                    c.as_ptr() as *const libc::c_char,
200                    c.len()
201                        .try_into()
202                        .map_err(|e| Error::SerializationError(Box::new(e)))?,
203                )
204            };
205            if res != 1 {
206                return Err(Error::DatabaseError(
207                    DatabaseErrorKind::Unknown,
208                    Box::new(self.last_error_message()),
209                ));
210            }
211        }
212        Ok(())
213    }
214
215    pub(crate) fn finish_copy_from(&self, err: Option<String>) -> QueryResult<()> {
216        let error = err.map(CString::new).map(|r| {
217            r.unwrap_or_else(|_| {
218                CString::new("Error message contains a \\0 byte")
219                    .expect("Does not contain a null byte")
220            })
221        });
222        let error = error
223            .as_ref()
224            .map(|l| l.as_ptr())
225            .unwrap_or(core::ptr::null());
226        let ret = unsafe { pq_sys::PQputCopyEnd(self.internal_connection.as_ptr(), error) };
227        if ret == 1 {
228            Ok(())
229        } else {
230            Err(Error::DatabaseError(
231                DatabaseErrorKind::Unknown,
232                Box::new(self.last_error_message()),
233            ))
234        }
235    }
236
237    pub(super) fn pq_notifies(&self) -> Result<Option<PgNotification>, Error> {
238        let conn = self.internal_connection;
239        let ret = unsafe { PQconsumeInput(conn.as_ptr()) };
240        if ret == 0 {
241            return Err(Error::DatabaseError(
242                DatabaseErrorKind::Unknown,
243                Box::new(self.last_error_message()),
244            ));
245        }
246
247        let pgnotify = unsafe { PQnotifies(conn.as_ptr()) };
248        if pgnotify.is_null() {
249            Ok(None)
250        } else {
251            // we use a drop guard here to
252            // make sure that we always free
253            // the provided pointer, even if we
254            // somehow return an error below
255            struct Guard<'a> {
256                value: &'a mut pgNotify,
257            }
258
259            impl Drop for Guard<'_> {
260                fn drop(&mut self) {
261                    unsafe {
262                        // SAFETY: We know that this value is not null here
263                        PQfreemem(self.value as *mut pgNotify as *mut core::ffi::c_void)
264                    };
265                }
266            }
267
268            let pgnotify = unsafe {
269                // SAFETY: We checked for null values above
270                Guard {
271                    value: &mut *pgnotify,
272                }
273            };
274            if pgnotify.value.relname.is_null() {
275                return Err(Error::DeserializationError(
276                    "Received an unexpected null value for `relname` from the notification".into(),
277                ));
278            }
279            if pgnotify.value.extra.is_null() {
280                return Err(Error::DeserializationError(
281                    "Received an unexpected null value for `extra` from the notification".into(),
282                ));
283            }
284
285            let channel = unsafe {
286                // SAFETY: We checked for null values above
287                CStr::from_ptr(pgnotify.value.relname)
288            }
289            .to_str()
290            .map_err(|e| Error::DeserializationError(e.into()))?
291            .to_string();
292            let payload = unsafe {
293                // SAFETY: We checked for null values above
294                CStr::from_ptr(pgnotify.value.extra)
295            }
296            .to_str()
297            .map_err(|e| Error::DeserializationError(e.into()))?
298            .to_string();
299            let ret = PgNotification {
300                process_id: pgnotify.value.be_pid,
301                channel,
302                payload,
303            };
304            Ok(Some(ret))
305        }
306    }
307
308    #[allow(unsafe_code, reason = "FFI calls")]
309    pub(super) fn consume_results(&mut self) -> QueryResult<()> {
310        if self.get_status() != pq_sys::ConnStatusType::CONNECTION_OK {
311            return Err(crate::result::Error::DeserializationError(
312                "Database connection lost".into(),
313            ));
314        }
315        // https://www.postgresql.org/docs/current/libpq-copy.html#LIBPQ-COPY-RECEIVE
316        //
317        // After PQgetCopyData returns -1, call PQgetResult to obtain the final result status of the COPY command.
318        // One can wait for this result to be available in the usual way. Then return to normal operation.
319        if let Some(result) = self.get_next_result()? {
320            match result.result_status() {
321                pq_sys::ExecStatusType::PGRES_COPY_OUT => {
322                    let mut ptr: *mut libc::c_char = core::ptr::null_mut();
323                    let mut ret = Ok(());
324                    loop {
325                        if self.get_status() != pq_sys::ConnStatusType::CONNECTION_OK {
326                            return Err(crate::result::Error::DeserializationError(
327                                "Database connection lost".into(),
328                            ));
329                        }
330                        let len = unsafe {
331                            pq_sys::PQgetCopyData(self.internal_connection.as_ptr(), &mut ptr, 0)
332                        };
333                        match len {
334                            len if len >= 0 => {
335                                if !ptr.is_null() {
336                                    unsafe {
337                                        pq_sys::PQfreemem(ptr as *mut libc::c_void);
338                                    }
339                                    ptr = core::ptr::null_mut();
340                                }
341                            }
342                            -1 => {
343                                break;
344                            }
345                            _ => {
346                                if ret.is_ok() {
347                                    let error = self.last_error_message();
348                                    ret = Err(crate::result::Error::DeserializationError(
349                                        error.into(),
350                                    ));
351                                }
352                            }
353                        }
354                    }
355                    ret?;
356                    self.consume_results()
357                }
358                pq_sys::ExecStatusType::PGRES_COPY_IN => {
359                    self.finish_copy_from(Some(String::from("Unexpected error")))?;
360                    self.consume_results()
361                }
362                _ => {
363                    let mut ret = Ok(());
364                    loop {
365                        if self.get_status() != pq_sys::ConnStatusType::CONNECTION_OK {
366                            return Err(crate::result::Error::DeserializationError(
367                                "Database connection lost".into(),
368                            ));
369                        }
370                        match self.get_next_result() {
371                            Ok(None) => break,
372                            Err(e) if ret.is_ok() => ret = Err(e),
373                            Ok(_) | Err(_) => {}
374                        }
375                    }
376                    ret
377                }
378            }
379        } else {
380            Ok(())
381        }
382    }
383}
384
385/// Represents the current in-transaction status of the connection
386#[derive(#[automatically_derived]impl ::core::fmt::Debug for PgTransactionStatus {
    #[inline]
    #[rustc_trivial_field_reads]
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        ::core::fmt::Formatter::debug_c_like_enum_write_str(f,
            "Idle Active InTransaction InError Unknown",
            match self {
                PgTransactionStatus::Idle => 0usize,
                PgTransactionStatus::Active => 5usize,
                PgTransactionStatus::InTransaction => 12usize,
                PgTransactionStatus::InError => 26usize,
                PgTransactionStatus::Unknown => 34usize,
            })
    }
}Debug, #[automatically_derived]impl ::core::marker::StructuralPartialEq for
    PgTransactionStatus {
}
#[automatically_derived]impl ::core::cmp::PartialEq for PgTransactionStatus {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        ::core::intrinsics::discriminant_value(self) ==
            ::core::intrinsics::discriminant_value(other)
    }
}PartialEq, #[automatically_derived]impl ::core::cmp::Eq for PgTransactionStatus { }Eq, #[automatically_derived]#[doc(hidden)]unsafe impl ::core::clone::TrivialClone
    for PgTransactionStatus {
}
#[automatically_derived]impl ::core::clone::Clone for PgTransactionStatus {
    #[inline]
    #[rustc_trivial_field_reads]
    fn clone(&self) -> Self { *self }
}Clone, #[automatically_derived]impl ::core::marker::Copy for PgTransactionStatus { }Copy)]
387pub(super) enum PgTransactionStatus {
388    /// Currently idle
389    Idle,
390    /// A command is in progress (sent to the server but not yet completed)
391    Active,
392    /// Idle, in a valid transaction block
393    InTransaction,
394    /// Idle, in a failed transaction block
395    InError,
396    /// Bad connection
397    Unknown,
398}
399
400impl From<PGTransactionStatusType> for PgTransactionStatus {
401    fn from(trans_status_type: PGTransactionStatusType) -> Self {
402        match trans_status_type {
403            PGTransactionStatusType::PQTRANS_IDLE => PgTransactionStatus::Idle,
404            PGTransactionStatusType::PQTRANS_ACTIVE => PgTransactionStatus::Active,
405            PGTransactionStatusType::PQTRANS_INTRANS => PgTransactionStatus::InTransaction,
406            PGTransactionStatusType::PQTRANS_INERROR => PgTransactionStatus::InError,
407            PGTransactionStatusType::PQTRANS_UNKNOWN => PgTransactionStatus::Unknown,
408        }
409    }
410}
411
412pub(super) type NoticeProcessor =
413    extern "C" fn(arg: *mut libc::c_void, message: *const libc::c_char);
414
415impl Drop for RawConnection {
416    fn drop(&mut self) {
417        unsafe { PQfinish(self.internal_connection.as_ptr()) };
418    }
419}
420
421fn last_error_message(conn: *const PGconn) -> String {
422    unsafe {
423        let error_ptr = PQerrorMessage(conn);
424        let bytes = CStr::from_ptr(error_ptr).to_bytes();
425        String::from_utf8_lossy(bytes).to_string()
426    }
427}
428
429/// Internal wrapper around a `*mut PGresult` which is known to be not-null, and
430/// have no aliases.  This wrapper is to ensure that it's always properly
431/// dropped.
432///
433/// If `Unique` is ever stabilized, we should use it here.
434#[allow(missing_debug_implementations)]
435pub(super) struct RawResult(NonNull<PGresult>);
436
437// SAFETY:
438// https://www.postgresql.org/docs/current/libpq-threading.html
439//
440// PGresult objects are normally read-only after creation,
441// and so can be passed around freely between threads. However,
442// if you use any of the PGresult-modifying functions described in
443// Section 31.12 or Section 31.14, it's up to you to avoid concurrent operations on the same PGresult, too.
444//
445// The type doesn't expose the raw pointer
446// and we don't call such Pgresult-modifying below
447unsafe impl Send for RawResult {}
448unsafe impl Sync for RawResult {}
449
450// Make sure to only call non-modifying functions here
451impl RawResult {
452    #[allow(clippy::new_ret_no_self)]
453    fn new(ptr: *mut PGresult, conn: &RawConnection) -> QueryResult<Self> {
454        NonNull::new(ptr).map(RawResult).ok_or_else(|| {
455            Error::DatabaseError(
456                DatabaseErrorKind::UnableToSendCommand,
457                Box::new(conn.last_error_message()),
458            )
459        })
460    }
461
462    pub(super) fn error_message(&self) -> &str {
463        let ptr = unsafe { PQresultErrorMessage(self.0.as_ptr()) };
464        let cstr = unsafe { CStr::from_ptr(ptr) };
465        cstr.to_str().unwrap_or_default()
466    }
467
468    pub(super) fn get_result_field(&self, field: ResultField) -> Option<&str> {
469        let ptr = unsafe { PQresultErrorField(self.0.as_ptr(), field as libc::c_int) };
470        if ptr.is_null() {
471            return None;
472        }
473
474        let c_str = unsafe { CStr::from_ptr(ptr) };
475        c_str.to_str().ok()
476    }
477
478    pub(super) fn result_status(&self) -> pq_sys::ExecStatusType {
479        unsafe {
480            // SAFETY:
481            // We have a unique not null pointer here
482            PQresultStatus(self.0.as_ptr())
483        }
484    }
485
486    pub(super) fn column_count(&self) -> libc::c_int {
487        unsafe {
488            // SAFETY:
489            // We have a unique not null pointer here
490            PQnfields(self.0.as_ptr())
491        }
492    }
493
494    pub(super) fn row_count(&self) -> libc::c_int {
495        unsafe {
496            // SAFETY:
497            // We have a unique not null pointer here
498            PQntuples(self.0.as_ptr())
499        }
500    }
501
502    pub(super) fn rows_affected(&self) -> &CStr {
503        unsafe {
504            // SAFETY:
505            // We have a unique not null pointer here
506            let count_char_ptr = PQcmdTuples(self.0.as_ptr());
507            CStr::from_ptr(count_char_ptr)
508        }
509    }
510
511    pub(super) fn get_bytes(&self, row_idx: i32, col_idx: i32) -> &[u8] {
512        unsafe {
513            // SAFETY:
514            // We have a unique not null pointer here
515            let value_ptr = PQgetvalue(self.0.as_ptr(), row_idx, col_idx) as *const u8;
516            let num_bytes = PQgetlength(self.0.as_ptr(), row_idx, col_idx);
517            if value_ptr.is_null() {
518                &[]
519            } else {
520                // SAFETY:
521                // * we rely on correct information from libpq here, if the provided length and pointer are invalid
522                core::slice::from_raw_parts(
523                    value_ptr,
524                    num_bytes
525                        .try_into()
526                        .expect("Diesel expects at least a 32 bit operating system"),
527                )
528            }
529        }
530    }
531
532    pub(super) fn is_null(&self, row_idx: libc::c_int, col_idx: libc::c_int) -> libc::c_int {
533        unsafe {
534            // SAFETY:
535            // We have a unique not null pointer here
536            PQgetisnull(self.0.as_ptr(), row_idx, col_idx)
537        }
538    }
539
540    pub(super) fn column_type(&self, col_idx: libc::c_int) -> pq_sys::Oid {
541        unsafe {
542            // SAFETY:
543            // We have a unique not null pointer here
544            PQftype(self.0.as_ptr(), col_idx)
545        }
546    }
547
548    pub(super) fn column_name(&self, col_idx: libc::c_int) -> Option<&CStr> {
549        unsafe {
550            // https://www.postgresql.org/docs/13/libpq-exec.html#LIBPQ-PQFNAME
551            // states that the returned ptr is valid till the underlying result is freed
552            // That means we can couple the lifetime to self
553            let ptr = PQfname(self.0.as_ptr(), col_idx);
554            if ptr.is_null() {
555                None
556            } else {
557                Some(CStr::from_ptr(ptr))
558            }
559        }
560    }
561}
562
563impl Drop for RawResult {
564    fn drop(&mut self) {
565        unsafe { PQclear(self.0.as_ptr()) }
566    }
567}
568
569/// Represents valid options to
570/// [`PQresultErrorField`](https://www.postgresql.org/docs/current/static/libpq-exec.html#LIBPQ-PQRESULTERRORFIELD)
571/// Their values are defined as C preprocessor macros, and therefore are not exported by libpq-sys.
572/// Their values can be found in `postgres_ext.h`
573#[repr(i32)]
574pub(super) enum ResultField {
575    SqlState = 'C' as i32,
576    MessagePrimary = 'M' as i32,
577    MessageDetail = 'D' as i32,
578    MessageHint = 'H' as i32,
579    TableName = 't' as i32,
580    ColumnName = 'c' as i32,
581    ConstraintName = 'n' as i32,
582    StatementPosition = 'P' as i32,
583}
diesel/pg/connection/raw.rs

diesel/pg/connection/
raw.rs
