diesel/pg/connection/stmt/
mod.rs

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

diesel/pg/connection/stmt/
mod.rs
