libpq/
lib.rs

1use std::{
2    ffi::{CString, NulError},
3    fmt::Display,
4    io::{Read, Seek},
5    ops::ControlFlow,
6    os::raw::{c_char, c_void},
7    ptr::{null, null_mut},
8};
9
10use std::fmt::Debug;
11
12use tempfile::Builder;
13
14include!("bindings.rs");
15
16pub struct PgSocket {
17    socket: i32,
18}
19
20pub enum PgSocketPollResult {
21    Timeout,
22    Error(String),
23}
24
25impl Display for PgSocketPollResult {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            PgSocketPollResult::Timeout => write!(f, "Timeout"),
29            PgSocketPollResult::Error(s) => write!(f, "Error: {}", s),
30        }
31    }
32}
33
34impl PgSocket {
35    pub fn poll(
36        &self,
37        read: bool,
38        write: bool,
39        timeout: Option<f64>,
40    ) -> Result<(), PgSocketPollResult> {
41        unsafe {
42            let timeout_ms = match timeout {
43                Some(t) => PQgetCurrentTimeUSec() + (t * 1000000.0) as i64,
44                None => -1,
45            };
46
47            match PQsocketPoll(self.socket, read.into(), write.into(), timeout_ms) {
48                a if a > 0 => Ok(()),
49                0 => Err(PgSocketPollResult::Timeout),
50                _ => Err(PgSocketPollResult::Error(
51                    std::io::Error::last_os_error().to_string(),
52                )),
53            }
54        }
55    }
56}
57pub struct PgConn {
58    conn: *mut PGconn,
59}
60
61unsafe impl Send for PgConn {}
62
63unsafe impl Sync for PgConn {}
64
65pub struct PgResult {
66    res: *mut PGresult,
67}
68
69pub struct PgNotify {
70    notify: *mut PGnotify,
71}
72
73impl PgNotify {
74    pub fn relname(&self) -> String {
75        unsafe {
76            let s = (*self.notify).relname;
77            std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
78        }
79    }
80
81    pub fn be_pid(&self) -> i32 {
82        unsafe { (*self.notify).be_pid }
83    }
84
85    pub fn extra(&self) -> String {
86        unsafe {
87            let s = (*self.notify).extra;
88            std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
89        }
90    }
91}
92
93impl Drop for PgConn {
94    fn drop(&mut self) {
95        unsafe {
96            PQfinish(self.conn);
97        }
98    }
99}
100
101impl Drop for PgNotify {
102    fn drop(&mut self) {
103        unsafe {
104            PQfreemem(self.notify as *mut c_void);
105        }
106    }
107}
108
109impl Drop for PgResult {
110    fn drop(&mut self) {
111        unsafe {
112            PQclear(self.res);
113        }
114    }
115}
116
117impl PgConn {
118    /// Connect to the database using environment variables.
119    ///
120    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-envars.html).
121    pub fn connect_db_env_vars() -> Result<PgConn, NulError> {
122        Self::connect_db("")
123    }
124
125    pub fn connect_db(s: &str) -> Result<PgConn, NulError> {
126        unsafe {
127            let conninfo = std::ffi::CString::new(s)?;
128            let conn = PQconnectdb(conninfo.as_ptr());
129            Ok(PgConn { conn })
130        }
131    }
132
133    pub fn status(&self) -> ConnStatusType {
134        unsafe { PQstatus(self.conn) }
135    }
136
137    pub fn exec(&self, query: &str) -> Result<PgResult, NulError> {
138        unsafe {
139            let c_query = std::ffi::CString::new(query)?;
140            let res = PQexec(self.conn, c_query.as_ptr());
141            Ok(PgResult { res })
142        }
143    }
144
145    pub fn exec_file(&self, file_path: &str) -> Result<PgResult, NulError> {
146        let content = std::fs::read_to_string(file_path).expect("Failed to read file.");
147        self.exec(&content)
148    }
149
150    /// Submit a command along with parameters, and wait for the result. Parameters are
151    /// passed as text; a `None` value is sent as SQL `NULL`. `param_types` may be shorter
152    /// than `param_values` (or empty) to let the server infer unspecified types.
153    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQEXECPARAMS).
154    pub fn exec_params(
155        &self,
156        query: &str,
157        param_values: &[Option<&str>],
158    ) -> Result<PgResult, NulError> {
159        unsafe {
160            let c_query = CString::new(query)?;
161
162            let c_params = param_values
163                .iter()
164                .map(|v| v.map(CString::new).transpose())
165                .collect::<Result<Vec<_>, NulError>>()?;
166
167            let ptrs: Vec<*const c_char> = c_params
168                .iter()
169                .map(|v| match v {
170                    Some(c) => c.as_ptr(),
171                    None => null(),
172                })
173                .collect();
174
175            let res = PQexecParams(
176                self.conn,
177                c_query.as_ptr(),
178                ptrs.len() as i32,
179                null(),
180                ptrs.as_ptr(),
181                null(),
182                null(),
183                0,
184            );
185
186            Ok(PgResult { res })
187        }
188    }
189
190    /// Submit a request to create a prepared statement, and wait for completion.
191    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQPREPARE).
192    pub fn prepare(&self, stmt_name: &str, query: &str) -> Result<PgResult, NulError> {
193        unsafe {
194            let c_stmt_name = CString::new(stmt_name)?;
195            let c_query = CString::new(query)?;
196
197            let res = PQprepare(self.conn, c_stmt_name.as_ptr(), c_query.as_ptr(), 0, null());
198
199            Ok(PgResult { res })
200        }
201    }
202
203    /// Send a request to execute a prepared statement with given parameters, and wait for
204    /// the result. Parameters are passed as text; a `None` value is sent as SQL `NULL`.
205    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-EXEC-PQEXECPREPARED).
206    pub fn exec_prepared(
207        &self,
208        stmt_name: &str,
209        param_values: &[Option<&str>],
210    ) -> Result<PgResult, NulError> {
211        unsafe {
212            let c_stmt_name = CString::new(stmt_name)?;
213
214            let c_params = param_values
215                .iter()
216                .map(|v| v.map(CString::new).transpose())
217                .collect::<Result<Vec<_>, NulError>>()?;
218
219            let ptrs: Vec<*const c_char> = c_params
220                .iter()
221                .map(|v| match v {
222                    Some(c) => c.as_ptr(),
223                    None => null(),
224                })
225                .collect();
226
227            let res = PQexecPrepared(
228                self.conn,
229                c_stmt_name.as_ptr(),
230                ptrs.len() as i32,
231                ptrs.as_ptr(),
232                null(),
233                null(),
234                0,
235            );
236
237            Ok(PgResult { res })
238        }
239    }
240
241    /// Submit a request to obtain information about the specified prepared statement, and
242    /// wait for completion.
243    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQDESCRIBEPREPARED).
244    pub fn describe_prepared(&self, stmt_name: &str) -> Result<PgResult, NulError> {
245        unsafe {
246            let c_stmt_name = CString::new(stmt_name)?;
247            let res = PQdescribePrepared(self.conn, c_stmt_name.as_ptr());
248            Ok(PgResult { res })
249        }
250    }
251
252    /// Submit a request to obtain information about the specified portal, and wait for
253    /// completion.
254    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQDESCRIBEPORTAL).
255    pub fn describe_portal(&self, portal_name: &str) -> Result<PgResult, NulError> {
256        unsafe {
257            let c_portal_name = CString::new(portal_name)?;
258            let res = PQdescribePortal(self.conn, c_portal_name.as_ptr());
259            Ok(PgResult { res })
260        }
261    }
262
263    /// Submit a request to close the specified prepared statement, and wait for completion.
264    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQCLOSEPREPARED).
265    pub fn close_prepared(&self, stmt_name: &str) -> Result<PgResult, NulError> {
266        unsafe {
267            let c_stmt_name = CString::new(stmt_name)?;
268            let res = PQclosePrepared(self.conn, c_stmt_name.as_ptr());
269            Ok(PgResult { res })
270        }
271    }
272
273    pub fn trace(&mut self, file: &str) {
274        unsafe {
275            let c_file = std::ffi::CString::new(file).unwrap();
276            let mode = std::ffi::CString::new("w").unwrap();
277            let fp = fopen(c_file.as_ptr(), mode.as_ptr());
278            PQtrace(self.conn, fp);
279            assert_eq!(fflush(fp), 0);
280        }
281    }
282
283    pub fn untrace(&mut self) {
284        unsafe {
285            PQuntrace(self.conn);
286        }
287    }
288
289    pub fn socket(&self) -> PgSocket {
290        unsafe {
291            PgSocket {
292                socket: PQsocket(self.conn),
293            }
294        }
295    }
296
297    pub fn consume_input(&mut self) -> Result<(), String> {
298        unsafe {
299            if PQconsumeInput(self.conn) == 0 {
300                Err(self.error_message())
301            } else {
302                Ok(())
303            }
304        }
305    }
306
307    pub fn notifies(&mut self) -> Option<PgNotify> {
308        unsafe {
309            let notify = PQnotifies(self.conn);
310            if notify.is_null() {
311                None
312            } else {
313                Some(PgNotify { notify })
314            }
315        }
316    }
317
318    pub fn error_message(&self) -> String {
319        unsafe {
320            let s = PQerrorMessage(self.conn);
321            if s.is_null() {
322                "".to_string()
323            } else {
324                std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
325            }
326        }
327    }
328
329    pub fn notify(&mut self, channel: &str, payload: Option<&str>) -> Result<PgResult, NulError> {
330        let query = match payload {
331            Some(p) => format!("NOTIFY {}, '{}';", channel, p),
332            None => format!("NOTIFY {};", channel),
333        };
334        self.exec(&query)
335    }
336
337    pub fn listen(&mut self, channel: &str) -> Result<PgResult, NulError> {
338        let query = format!("LISTEN {};", channel);
339        self.exec(&query)
340    }
341
342    ///
343    /// A callback function to receive notices from the server.
344    /// https://stackoverflow.com/questions/24191249/working-with-c-void-in-an-ffi
345    /// https://adventures.michaelfbryan.com/posts/rust-closures-in-ffi/
346    extern "C" fn ffi_notice_processor<F>(arg: *mut c_void, data: *const c_char)
347    where
348        F: FnMut(String),
349    {
350        unsafe {
351            let s = std::ffi::CStr::from_ptr(data)
352                .to_string_lossy()
353                .into_owned();
354
355            let f = &mut *(arg as *mut F);
356
357            f(s);
358        }
359    }
360
361    pub fn set_notice_processor<F>(&mut self, proc: F) -> Box<F>
362    where
363        F: FnMut(String),
364    {
365        unsafe {
366            let mut b = Box::new(proc);
367            let a = b.as_mut() as *mut F as *mut c_void;
368            PQsetNoticeProcessor(self.conn, Some(Self::ffi_notice_processor::<F>), a);
369            b
370        }
371    }
372
373    extern "C" fn ffi_notice_receiver<F>(arg: *mut c_void, data: *const PGresult)
374    where
375        F: FnMut(PgResult),
376    {
377        unsafe {
378            let s = PgResult {
379                res: data as *mut PGresult,
380            };
381
382            let f = &mut *(arg as *mut F);
383
384            f(s);
385        }
386    }
387
388    /// Sets a notice receiver function to receive notices from the server.
389    /// Notices are sent to the receiver after command execution is completed.
390    /// https://www.postgresql.org/docs/current/libpq-notice-processing.html
391    pub fn set_notice_receiver<F>(&mut self, proc: F) -> Box<F>
392    where
393        F: FnMut(PgResult),
394    {
395        unsafe {
396            let mut b = Box::new(proc);
397            let a = b.as_mut() as *mut F as *mut c_void;
398            PQsetNoticeReceiver(self.conn, Some(Self::ffi_notice_receiver::<F>), a);
399            b
400        }
401    }
402
403    pub fn listen_loop<F, T>(&mut self, timeout_sec: Option<f64>, proc: F) -> Vec<T>
404    where
405        F: Fn(usize, PgNotify) -> ControlFlow<(), Option<T>>,
406    {
407        let mut recvs = Vec::new();
408
409        let mut count = 0;
410
411        loop {
412            match self.socket().poll(true, false, timeout_sec) {
413                Ok(()) => {
414                    self.consume_input().expect("Failed to consume input.");
415
416                    while let Some(notify) = self.notifies() {
417                        match proc(count, notify) {
418                            ControlFlow::Continue(Some(p)) => recvs.push(p),
419                            ControlFlow::Break(()) => {
420                                break;
421                            }
422                            _ => {} // Do nothing
423                        }
424                        self.consume_input().expect("Failed to consume input.");
425                        count += 1;
426                    }
427                }
428                Err(_e) => break,
429            }
430        }
431
432        recvs
433    }
434}
435
436impl PgResult {
437    pub fn status(&self) -> ExecStatusType {
438        unsafe { PQresultStatus(self.res) }
439    }
440
441    pub fn cmd_status(&mut self) -> String {
442        unsafe {
443            let s = PQcmdStatus(self.res);
444            std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
445        }
446    }
447
448    pub fn error_message(&self) -> String {
449        unsafe {
450            let s = PQresultErrorMessage(self.res);
451            std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
452        }
453    }
454
455    pub fn error_field(&self, field_code: u8) -> Option<String> {
456        unsafe {
457            let s = PQresultErrorField(self.res, field_code.into());
458            if s.is_null() {
459                None
460            } else {
461                Some(std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned())
462            }
463        }
464    }
465
466    pub fn get_value<T>(&self, row: i32, col: i32) -> Option<T>
467    where
468        T: std::str::FromStr,
469    {
470        unsafe {
471            let s = PQgetvalue(self.res, row, col);
472            if s.is_null() {
473                None
474            } else {
475                let s = std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned();
476                match s.parse::<T>() {
477                    Ok(v) => Some(v),
478                    Err(_) => None,
479                }
480            }
481        }
482    }
483
484    pub fn ntuples(&self) -> i32 {
485        unsafe { PQntuples(self.res) }
486    }
487
488    /// Print the result to a file.
489    /// See the [official doc](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQPRINT
490    pub fn print(
491        &self,
492        filename: &str,
493        header: bool,
494        align: bool,
495        fieldsep: &str,
496        standard: bool,
497        html3: bool,
498        expanded: bool,
499        pager: bool,
500    ) {
501        unsafe {
502            let sep = CString::new(fieldsep).unwrap();
503
504            let printopt = PQprintOpt {
505                header: header.into(),
506                align: align.into(),
507                fieldSep: sep.as_ptr() as *mut c_char,
508                tableOpt: null_mut(),
509                caption: null_mut(),
510                standard: standard.into(),
511                html3: html3.into(),
512                expanded: expanded.into(),
513                pager: pager.into(),
514                fieldName: null_mut(),
515            };
516
517            let fp = fopen(
518                CString::new(filename).unwrap().as_ptr(),
519                CString::new("w").unwrap().as_ptr(),
520            );
521
522            PQprint(fp, self.res, &printopt);
523
524            assert_eq!(fflush(fp), 0);
525            assert_eq!(fclose(fp), 0);
526        }
527    }
528
529    /// Get the value at the specified row and column.
530    /// See also [PQgetvalue](https://www.postgresql.org/docs/current/libpq-exec.html#LIBPQ-PQGETVALUE).
531    pub fn get_value_raw(&self, row: i32, col: i32) -> String {
532        unsafe {
533            let s = PQgetvalue(self.res, row, col);
534            if s.is_null() {
535                "".to_string()
536            } else {
537                std::ffi::CStr::from_ptr(s).to_string_lossy().into_owned()
538            }
539        }
540    }
541}
542
543impl Display for PgResult {
544    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
545        let mut temp_file = Builder::new()
546            .prefix("pg-res-")
547            .suffix(".json")
548            .tempfile()
549            .unwrap();
550
551        let temp_path = temp_file.path().to_path_buf();
552
553        self.print(
554            temp_path.as_path().to_str().unwrap(),
555            true,
556            true,
557            "|",
558            true,
559            false,
560            false,
561            false,
562        );
563
564        let mut s = String::new();
565        temp_file
566            .seek(std::io::SeekFrom::Start(0))
567            .expect("Failed to seek to start of temp file.");
568        temp_file
569            .read_to_string(&mut s)
570            .expect("Failed to read temp file.");
571
572        write!(f, "{}", s)
573    }
574}
libpq/lib.rs

libpq/
lib.rs
