sqlx_mysql/connection/
executor.rs

1use super::MySqlStream;
2use crate::connection::stream::Waiting;
3use crate::error::Error;
4use crate::executor::{Execute, Executor};
5use crate::ext::ustr::UStr;
6use crate::io::MySqlBufExt;
7use crate::logger::QueryLogger;
8use crate::protocol::response::Status;
9use crate::protocol::statement::{
10    BinaryRow, Execute as StatementExecute, Prepare, PrepareOk, StmtClose,
11};
12use crate::protocol::text::{ColumnDefinition, Query, TextRow};
13use crate::statement::{MySqlStatement, MySqlStatementMetadata};
14use crate::HashMap;
15use crate::{
16    MySql, MySqlArguments, MySqlColumn, MySqlConnection, MySqlQueryResult, MySqlRow, MySqlTypeInfo,
17    MySqlValueFormat,
18};
19use either::Either;
20use futures_core::future::BoxFuture;
21use futures_core::stream::BoxStream;
22use futures_core::Stream;
23use futures_util::TryStreamExt;
24use sqlx_core::column::{ColumnOrigin, TableColumn};
25use sqlx_core::sql_str::SqlStr;
26use std::{pin::pin, sync::Arc};
27
28impl MySqlConnection {
29    async fn prepare_statement(
30        &mut self,
31        sql: &str,
32    ) -> Result<(u32, MySqlStatementMetadata), Error> {
33        // https://dev.mysql.com/doc/internals/en/com-stmt-prepare.html
34        // https://dev.mysql.com/doc/internals/en/com-stmt-prepare-response.html#packet-COM_STMT_PREPARE_OK
35
36        self.inner
37            .stream
38            .send_packet(Prepare { query: sql })
39            .await?;
40
41        let ok: PrepareOk = self.inner.stream.recv().await?;
42
43        // the parameter definitions are very unreliable so we skip over them
44        // as we have little use
45
46        if ok.params > 0 {
47            for _ in 0..ok.params {
48                let _def: ColumnDefinition = self.inner.stream.recv().await?;
49            }
50
51            self.inner.stream.maybe_recv_eof().await?;
52        }
53
54        // the column definitions are berefit the type information from the
55        // to-be-bound parameters; we will receive the output column definitions
56        // once more on execute so we wait for that
57
58        let mut columns = Vec::new();
59
60        let column_names = if ok.columns > 0 {
61            recv_result_metadata(&mut self.inner.stream, ok.columns as usize, &mut columns).await?
62        } else {
63            Default::default()
64        };
65
66        let id = ok.statement_id;
67        let metadata = MySqlStatementMetadata {
68            parameters: ok.params as usize,
69            columns: Arc::new(columns),
70            column_names: Arc::new(column_names),
71        };
72
73        Ok((id, metadata))
74    }
75
76    async fn get_or_prepare_statement(
77        &mut self,
78        sql: &str,
79    ) -> Result<(u32, MySqlStatementMetadata), Error> {
80        if let Some(statement) = self.inner.cache_statement.get_mut(sql) {
81            // <MySqlStatementMetadata> is internally reference-counted
82            return Ok((*statement).clone());
83        }
84
85        let (id, metadata) = self.prepare_statement(sql).await?;
86
87        // in case of the cache being full, close the least recently used statement
88        if let Some((id, _)) = self
89            .inner
90            .cache_statement
91            .insert(sql, (id, metadata.clone()))
92        {
93            self.inner
94                .stream
95                .send_packet(StmtClose { statement: id })
96                .await?;
97        }
98
99        Ok((id, metadata))
100    }
101
102    #[allow(clippy::needless_lifetimes)]
103    pub(crate) async fn run<'e, 'c: 'e, 'q: 'e>(
104        &'c mut self,
105        sql: SqlStr,
106        arguments: Option<MySqlArguments>,
107        persistent: bool,
108    ) -> Result<impl Stream<Item = Result<Either<MySqlQueryResult, MySqlRow>, Error>> + 'e, Error>
109    {
110        let mut logger = QueryLogger::new(sql, self.inner.log_settings.clone());
111
112        self.inner.stream.wait_until_ready().await?;
113        self.inner.stream.waiting.push_back(Waiting::Result);
114
115        Ok(try_stream! {
116        let sql = logger.sql().as_str();
117
118            // make a slot for the shared column data
119            // as long as a reference to a row is not held past one iteration, this enables us
120            // to re-use this memory freely between result sets
121            let mut columns = Arc::new(Vec::new());
122
123            let format = if let Some(arguments) = arguments {
124                if persistent && self.inner.cache_statement.is_enabled() {
125                    let (id, metadata) = self
126                        .get_or_prepare_statement(sql)
127                        .await?;
128
129                    if arguments.types.len() != metadata.parameters {
130                        return Err(
131                            err_protocol!(
132                                "prepared statement expected {} parameters but {} parameters were provided",
133                                metadata.parameters,
134                                arguments.types.len()
135                            )
136                        );
137                    }
138
139                    // https://dev.mysql.com/doc/internals/en/com-stmt-execute.html
140                    self.inner.stream
141                        .send_packet(StatementExecute {
142                            statement: id,
143                            arguments: &arguments,
144                        })
145                        .await?;
146
147                    MySqlValueFormat::Binary
148                } else {
149                    let (id, metadata) = self
150                        .prepare_statement(sql)
151                        .await?;
152
153                    if arguments.types.len() != metadata.parameters {
154                        return Err(
155                            err_protocol!(
156                                "prepared statement expected {} parameters but {} parameters were provided",
157                                metadata.parameters,
158                                arguments.types.len()
159                            )
160                        );
161                    }
162
163                    // https://dev.mysql.com/doc/internals/en/com-stmt-execute.html
164                    self.inner.stream
165                        .send_packet(StatementExecute {
166                            statement: id,
167                            arguments: &arguments,
168                        })
169                        .await?;
170
171                    self.inner.stream.send_packet(StmtClose { statement: id }).await?;
172
173                    MySqlValueFormat::Binary
174                }
175            } else {
176                // https://dev.mysql.com/doc/internals/en/com-query.html
177                self.inner.stream.send_packet(Query(sql)).await?;
178
179                MySqlValueFormat::Text
180            };
181
182            loop {
183                // query response is a meta-packet which may be one of:
184                //  Ok, Err, ResultSet, or (unhandled) LocalInfileRequest
185                let mut packet = self.inner.stream.recv_packet().await?;
186
187                if packet[0] == 0x00 || packet[0] == 0xff {
188                    // first packet in a query response is OK or ERR
189                    // this indicates either a successful query with no rows at all or a failed query
190                    let ok = packet.ok()?;
191
192                    self.inner.status_flags = ok.status;
193
194                    let rows_affected = ok.affected_rows;
195                    logger.increase_rows_affected(rows_affected);
196                    let done = MySqlQueryResult {
197                        rows_affected,
198                        last_insert_id: ok.last_insert_id,
199                    };
200
201                    r#yield!(Either::Left(done));
202
203                    if ok.status.contains(Status::SERVER_MORE_RESULTS_EXISTS) {
204                        // more result sets exist, continue to the next one
205                        continue;
206                    }
207
208                    self.inner.stream.waiting.pop_front();
209                    return Ok(());
210                }
211
212                // otherwise, this first packet is the start of the result-set metadata,
213                *self.inner.stream.waiting.front_mut().unwrap() = Waiting::Row;
214
215                let num_columns = packet.get_uint_lenenc()?; // column count
216                let num_columns = usize::try_from(num_columns)
217                    .map_err(|_| err_protocol!("column count overflows usize: {num_columns}"))?;
218
219                // Always reload column names, even for prepared statements (the schema
220                // may change between PREPARE and EXECUTE).
221                let column_names = Arc::new(recv_result_metadata(&mut self.inner.stream, num_columns, Arc::make_mut(&mut columns)).await?);
222
223                // finally, there will be none or many result-rows
224                loop {
225                    let packet = self.inner.stream.recv_packet().await?;
226
227                    if packet[0] == 0xfe {
228                        let (rows_affected, last_insert_id, status) = if packet.len() < 9 {
229                            // EOF packet
230                            let eof = packet.eof(self.inner.stream.capabilities)?;
231                            (0, 0, eof.status)
232                        } else {
233                            // OK packet
234                            let ok = packet.ok()?;
235                            (ok.affected_rows, ok.last_insert_id, ok.status)
236                        };
237
238                        self.inner.status_flags = status;
239                        r#yield!(Either::Left(MySqlQueryResult {
240                            rows_affected,
241                            last_insert_id,
242                        }));
243
244                        if status.contains(Status::SERVER_MORE_RESULTS_EXISTS) {
245                            *self.inner.stream.waiting.front_mut().unwrap() = Waiting::Result;
246                            break;
247                        }
248                        self.inner.stream.waiting.pop_front();
249                        return Ok(());
250                    }
251
252                    let row = match format {
253                        MySqlValueFormat::Binary => packet.decode_with::<BinaryRow, _>(&columns)?.0,
254                        MySqlValueFormat::Text => packet.decode_with::<TextRow, _>(&columns)?.0,
255                    };
256
257                    let v = Either::Right(MySqlRow {
258                        row,
259                        format,
260                        columns: Arc::clone(&columns),
261                        column_names: Arc::clone(&column_names),
262                    });
263
264                    logger.increment_rows_returned();
265
266                    r#yield!(v);
267                }
268            }
269        })
270    }
271}
272
273impl<'c> Executor<'c> for &'c mut MySqlConnection {
274    type Database = MySql;
275
276    fn fetch_many<'e, 'q, E>(
277        self,
278        mut query: E,
279    ) -> BoxStream<'e, Result<Either<MySqlQueryResult, MySqlRow>, Error>>
280    where
281        'c: 'e,
282        E: Execute<'q, Self::Database>,
283        'q: 'e,
284        E: 'q,
285    {
286        let arguments = query.take_arguments().map_err(Error::Encode);
287        let persistent = query.persistent();
288
289        Box::pin(try_stream! {
290        let sql = query.sql();
291            let arguments = arguments?;
292            let mut s = pin!(self.run(sql, arguments, persistent).await?);
293
294            while let Some(v) = s.try_next().await? {
295                r#yield!(v);
296            }
297
298            Ok(())
299        })
300    }
301
302    fn fetch_optional<'e, 'q, E>(self, query: E) -> BoxFuture<'e, Result<Option<MySqlRow>, Error>>
303    where
304        'c: 'e,
305        E: Execute<'q, Self::Database>,
306        'q: 'e,
307        E: 'q,
308    {
309        let mut s = self.fetch_many(query);
310
311        Box::pin(async move {
312            while let Some(v) = s.try_next().await? {
313                if let Either::Right(r) = v {
314                    return Ok(Some(r));
315                }
316            }
317
318            Ok(None)
319        })
320    }
321
322    fn prepare_with<'e>(
323        self,
324        sql: SqlStr,
325        _parameters: &'e [MySqlTypeInfo],
326    ) -> BoxFuture<'e, Result<MySqlStatement, Error>>
327    where
328        'c: 'e,
329    {
330        Box::pin(async move {
331            self.inner.stream.wait_until_ready().await?;
332
333            let metadata = if self.inner.cache_statement.is_enabled() {
334                self.get_or_prepare_statement(sql.as_str()).await?.1
335            } else {
336                let (id, metadata) = self.prepare_statement(sql.as_str()).await?;
337
338                self.inner
339                    .stream
340                    .send_packet(StmtClose { statement: id })
341                    .await?;
342
343                metadata
344            };
345
346            Ok(MySqlStatement {
347                sql,
348                // metadata has internal Arcs for expensive data structures
349                metadata: metadata.clone(),
350            })
351        })
352    }
353
354    #[doc(hidden)]
355    #[cfg(feature = "offline")]
356    fn describe<'e>(
357        self,
358        sql: SqlStr,
359    ) -> BoxFuture<'e, Result<crate::describe::Describe<MySql>, Error>>
360    where
361        'c: 'e,
362    {
363        Box::pin(async move {
364            self.inner.stream.wait_until_ready().await?;
365
366            let (id, metadata) = self.prepare_statement(sql.as_str()).await?;
367
368            self.inner
369                .stream
370                .send_packet(StmtClose { statement: id })
371                .await?;
372
373            let columns = (*metadata.columns).clone();
374
375            let nullable = columns
376                .iter()
377                .map(|col| {
378                    col.flags
379                        .map(|flags| !flags.contains(crate::protocol::text::ColumnFlags::NOT_NULL))
380                })
381                .collect();
382
383            Ok(crate::describe::Describe {
384                parameters: Some(Either::Right(metadata.parameters)),
385                columns,
386                nullable,
387            })
388        })
389    }
390}
391
392fn recv_next_result_column(def: &ColumnDefinition, ordinal: usize) -> Result<MySqlColumn, Error> {
393    // if the alias is empty, use the alias
394    // only then use the name
395    let column_name = def.name()?;
396
397    let name = match (def.name()?, def.alias()?) {
398        (_, alias) if !alias.is_empty() => UStr::new(alias),
399        (name, _) => UStr::new(name),
400    };
401
402    let table = def.table()?;
403
404    let origin = if table.is_empty() {
405        ColumnOrigin::Expression
406    } else {
407        let schema = def.schema()?;
408
409        ColumnOrigin::Table(TableColumn {
410            table: if !schema.is_empty() {
411                format!("{schema}.{table}").into()
412            } else {
413                table.into()
414            },
415            name: column_name.into(),
416        })
417    };
418
419    let type_info = MySqlTypeInfo::from_column(def);
420
421    Ok(MySqlColumn {
422        name,
423        type_info,
424        ordinal,
425        flags: Some(def.flags),
426        origin,
427    })
428}
429
430async fn recv_result_metadata(
431    stream: &mut MySqlStream,
432    num_columns: usize,
433    columns: &mut Vec<MySqlColumn>,
434) -> Result<HashMap<UStr, usize>, Error> {
435    // the result-set metadata is primarily a listing of each output
436    // column in the result-set
437
438    let mut column_names = HashMap::with_capacity(num_columns);
439
440    columns.clear();
441    columns.reserve(num_columns);
442
443    for ordinal in 0..num_columns {
444        let def: ColumnDefinition = stream.recv().await?;
445
446        let column = recv_next_result_column(&def, ordinal)?;
447
448        column_names.insert(column.name.clone(), ordinal);
449        columns.push(column);
450    }
451
452    stream.maybe_recv_eof().await?;
453
454    Ok(column_names)
455}
sqlx_mysql/connection/executor.rs

sqlx_mysql/connection/
executor.rs
