diesel/pg/connection/stmt/
mod.rs

Help
1#![allow(unsafe_code)] // ffi code
2extern crate pq_sys;
3
4use std::ffi::CString;
5use std::os::raw as libc;
6use std::ptr;
7
8use super::result::PgResult;
9use super::statement_cache::PrepareForCache;
10use crate::pg::PgTypeMetadata;
11use crate::result::QueryResult;
12
13use super::raw::RawConnection;
14
15enum StatementKind {
16    Unnamed { sql: CString, param_types: Vec<u32> },
17    Named { name: CString },
18}
19
20pub(crate) struct Statement {
21    kind: StatementKind,
22    param_formats: Vec<libc::c_int>,
23}
24
25impl Statement {
26    pub(super) fn execute(
27        &self,
28        raw_connection: &mut RawConnection,
29        param_data: &[Option<Vec<u8>>],
30        row_by_row: bool,
31    ) -> QueryResult<PgResult> {
32        let params_pointer = param_data
33            .iter()
34            .map(|data| {
35                data.as_ref()
36                    .map(|d| d.as_ptr() as *const libc::c_char)
37                    .unwrap_or(ptr::null())
38            })
39            .collect::<Vec<_>>();
40        let param_lengths = param_data
41            .iter()
42            .map(|data| data.as_ref().map(|d| d.len().try_into()).unwrap_or(Ok(0)))
43            .collect::<Result<Vec<_>, _>>()
44            .map_err(|e| crate::result::Error::SerializationError(Box::new(e)))?;
45        let param_count: libc::c_int = params_pointer
46            .len()
47            .try_into()
48            .map_err(|e| crate::result::Error::SerializationError(Box::new(e)))?;
49
50        match &self.kind {
51            StatementKind::Named { name } => {
52                unsafe {
53                    // execute the previously prepared statement
54                    // in autocommit mode, this will be a new transaction
55                    raw_connection.send_query_prepared(
56                        name.as_ptr(),
57                        param_count,
58                        params_pointer.as_ptr(),
59                        param_lengths.as_ptr(),
60                        self.param_formats.as_ptr(),
61                        1,
62                    )
63                }?
64            }
65            StatementKind::Unnamed { sql, param_types } => unsafe {
66                // execute the unnamed prepared statement using send_query_params
67                // which internally calls PQsendQueryParams, making sure the
68                // prepare and execute happens in a single transaction. This
69                // makes sure these are handled by PgBouncer.
70                // See https://github.com/diesel-rs/diesel/pull/4539
71                raw_connection.send_query_params(
72                    sql.as_ptr(),
73                    param_count,
74                    param_types.as_ptr(),
75                    params_pointer.as_ptr(),
76                    param_lengths.as_ptr(),
77                    self.param_formats.as_ptr(),
78                    1,
79                )
80            }?,
81        };
82
83        if row_by_row {
84            raw_connection.enable_row_by_row_mode()?;
85        }
86        Ok(raw_connection.get_next_result()?.expect("Is never none"))
87    }
88
89    pub(super) fn prepare(
90        raw_connection: &mut RawConnection,
91        sql: &str,
92        is_cached: PrepareForCache,
93        param_types: &[PgTypeMetadata],
94    ) -> QueryResult<Self> {
95        let sql_cstr = CString::new(sql)?;
96        let param_types_vec = param_types
97            .iter()
98            .map(|x| x.oid())
99            .collect::<Result<Vec<_>, _>>()
100            .map_err(|e| crate::result::Error::SerializationError(Box::new(e)))?;
101
102        match is_cached {
103            PrepareForCache::Yes { counter } => {
104                // For named/cached statements, prepare as usual using a prepare phase and then
105                // an execute phase
106                let name_cstr = CString::new(format!("__diesel_stmt_{counter}"))?;
107                let internal_result = unsafe {
108                    let param_count: libc::c_int = param_types
109                        .len()
110                        .try_into()
111                        .map_err(|e| crate::result::Error::SerializationError(Box::new(e)))?;
112                    raw_connection.prepare(
113                        name_cstr.as_ptr(),
114                        sql_cstr.as_ptr(),
115                        param_count,
116                        param_types_vec.as_ptr(),
117                    )
118                };
119                PgResult::new(internal_result?, raw_connection)?;
120
121                Ok(Statement {
122                    kind: StatementKind::Named { name: name_cstr },
123                    param_formats: vec![1; param_types.len()],
124                })
125            }
126            PrepareForCache::No => {
127                // For unnamed statements, we'll return a Statement object without
128                // actually preparing it. This allows us to use send_query_params
129                // later in the execute call. This is needed to better interface
130                // with PgBouncer which cannot handle unnamed prepared statements
131                // when those are prepared and executed in separate transactions.
132                // See https://github.com/diesel-rs/diesel/pull/4539
133                Ok(Statement {
134                    kind: StatementKind::Unnamed {
135                        sql: sql_cstr,
136                        param_types: param_types_vec,
137                    },
138                    param_formats: vec![1; param_types.len()],
139                })
140            }
141        }
142    }
143}
diesel/pg/connection/stmt/mod.rs

diesel/pg/connection/stmt/
mod.rs
