diesel/pg/connection/stmt/
mod.rs

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(|_: std::num::TryFromIntError| {
45                crate::result::Error::SerializationError(
46                    "A bind parameter's serialized size is bigger than fits on an i32".into(),
47                )
48            })?;
49        let param_count: libc::c_int =
50            params_pointer
51                .len()
52                .try_into()
53                .map_err(|_: std::num::TryFromIntError| {
54                    crate::result::Error::SerializationError(
55                        "There are more than i32::MAX bind parameters".into(),
56                    )
57                })?;
58
59        match &self.kind {
60            StatementKind::Named { name } => {
61                unsafe {
62                    // execute the previously prepared statement
63                    // in autocommit mode, this will be a new transaction
64                    raw_connection.send_query_prepared(
65                        name.as_ptr(),
66                        param_count,
67                        params_pointer.as_ptr(),
68                        param_lengths.as_ptr(),
69                        self.param_formats.as_ptr(),
70                        1,
71                    )
72                }?
73            }
74            StatementKind::Unnamed { sql, param_types } => unsafe {
75                // execute the unnamed prepared statement using send_query_params
76                // which internally calls PQsendQueryParams, making sure the
77                // prepare and execute happens in a single transaction. This
78                // makes sure these are handled by PgBouncer.
79                // See https://github.com/diesel-rs/diesel/pull/4539
80                raw_connection.send_query_params(
81                    sql.as_ptr(),
82                    param_count,
83                    param_types.as_ptr(),
84                    params_pointer.as_ptr(),
85                    param_lengths.as_ptr(),
86                    self.param_formats.as_ptr(),
87                    1,
88                )
89            }?,
90        };
91
92        if row_by_row {
93            raw_connection.enable_row_by_row_mode()?;
94        }
95        Ok(raw_connection.get_next_result()?.expect("Is never none"))
96    }
97
98    pub(super) fn prepare(
99        raw_connection: &mut RawConnection,
100        sql: &str,
101        is_cached: PrepareForCache,
102        param_types: &[PgTypeMetadata],
103    ) -> QueryResult<Self> {
104        let sql_cstr = CString::new(sql)?;
105        let param_types_vec = param_types
106            .iter()
107            .map(|x| x.oid())
108            .collect::<Result<Vec<_>, _>>()
109            .map_err(|e| crate::result::Error::SerializationError(Box::new(e)))?;
110
111        match is_cached {
112            PrepareForCache::Yes { counter } => {
113                // For named/cached statements, prepare as usual using a prepare phase and then
114                // an execute phase
115                let name_cstr = CString::new(::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("__diesel_stmt_{0}", counter))
    })format!("__diesel_stmt_{counter}"))?;
116                let internal_result =
117                    unsafe {
118                        let param_count: libc::c_int = param_types.len().try_into().map_err(
119                            |_: std::num::TryFromIntError| {
120                                crate::result::Error::SerializationError(
121                                    "There are more than i32::MAX bind parameters".into(),
122                                )
123                            },
124                        )?;
125                        raw_connection.prepare(
126                            name_cstr.as_ptr(),
127                            sql_cstr.as_ptr(),
128                            param_count,
129                            param_types_vec.as_ptr(),
130                        )
131                    };
132                PgResult::new(internal_result?, raw_connection)?;
133
134                Ok(Statement {
135                    kind: StatementKind::Named { name: name_cstr },
136                    param_formats: ::alloc::vec::from_elem(1, param_types.len())vec![1; param_types.len()],
137                })
138            }
139            PrepareForCache::No => {
140                // For unnamed statements, we'll return a Statement object without
141                // actually preparing it. This allows us to use send_query_params
142                // later in the execute call. This is needed to better interface
143                // with PgBouncer which cannot handle unnamed prepared statements
144                // when those are prepared and executed in separate transactions.
145                // See https://github.com/diesel-rs/diesel/pull/4539
146                Ok(Statement {
147                    kind: StatementKind::Unnamed {
148                        sql: sql_cstr,
149                        param_types: param_types_vec,
150                    },
151                    param_formats: ::alloc::vec::from_elem(1, param_types.len())vec![1; param_types.len()],
152                })
153            }
154        }
155    }
156}
diesel/pg/connection/stmt/mod.rs

diesel/pg/connection/stmt/
mod.rs
