From 606234559dae69ff988f04c3b74fb1dee8dd25c4 Mon Sep 17 00:00:00 2001 From: Mike Cronce Date: Fri, 16 Aug 2024 10:59:45 -0400 Subject: [PATCH 1/4] diesel/src/query_dsl/load_dsl.rs: + ExecuteInsertWithReturningId --- diesel/src/query_dsl/load_dsl.rs | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/diesel/src/query_dsl/load_dsl.rs b/diesel/src/query_dsl/load_dsl.rs index 7f2a8cb5e384..e2c3a8f1e69c 100644 --- a/diesel/src/query_dsl/load_dsl.rs +++ b/diesel/src/query_dsl/load_dsl.rs @@ -1,7 +1,9 @@ use self::private::LoadIter; use super::RunQueryDsl; use crate::backend::Backend; -use crate::connection::{Connection, DefaultLoadingMode, LoadConnection}; +use crate::connection::{ + Connection, ConnectionWithReturningId, DefaultLoadingMode, LoadConnection, +}; use crate::deserialize::FromSqlRow; use crate::expression::QueryMetadata; use crate::query_builder::{AsQuery, QueryFragment, QueryId}; @@ -91,6 +93,34 @@ where } } +/// The `execute_with_returning_id` method +/// +/// This is currently only implemented for [`MySQL`], because MySQL provides a +/// `last_insert_id` field as part of its response packet to `INSERT` queries, but +/// it's in a trait to avoid a backwards-incompatible change in the future if such +/// functionality becomes available in other back-ends. +/// +/// [`MySQL`]: crate::mysql::Mysql +pub trait ExecuteInsertWithReturningId< + Conn: ConnectionWithReturningId, + DB: Backend = ::Backend, +>: Sized +{ + /// Execute this command + fn execute_with_returning_id(query: Self, conn: &mut Conn) -> QueryResult; +} + +impl ExecuteInsertWithReturningId for T +where + Conn: ConnectionWithReturningId, + DB: Backend, + T: QueryFragment + QueryId, +{ + fn execute_with_returning_id(query: Self, conn: &mut Conn) -> QueryResult { + conn.execute_returning_id(&query) + } +} + // These types and traits are not part of the public API. // // * CompatibleType as we consider this as "sealed" trait. It shouldn't From a90716534b7a2ccbd924aa361cae9939b09355f4 Mon Sep 17 00:00:00 2001 From: Mike Cronce Date: Fri, 16 Aug 2024 11:08:29 -0400 Subject: [PATCH 2/4] diesel/src/connection/mod.rs: + ConnectionWithReturningId --- diesel/src/connection/mod.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/diesel/src/connection/mod.rs b/diesel/src/connection/mod.rs index 5740ca0f4c73..6fb481013024 100644 --- a/diesel/src/connection/mod.rs +++ b/diesel/src/connection/mod.rs @@ -486,6 +486,19 @@ pub trait BoxableConnection: SimpleConnection + std::any::Any { fn as_any_mut(&mut self) -> &mut dyn std::any::Any; } +/// An extension of the [`Connection`](trait.Connection.html) trait that provides the value +/// assigned to an inserted row with an auto-incremented ID column +pub trait ConnectionWithReturningId: Connection { + /// The type that the connection implementation returns for IDs + type ReturnedId; + + /// Execute a single INSERT statement given by a query and return the ID that the database + /// assigns + fn execute_returning_id(&mut self, source: &T) -> QueryResult + where + T: QueryFragment + QueryId; +} + impl BoxableConnection for C where C: Connection + std::any::Any, From edb451d1c97dc4327139ddc7d49dca0ef05d16f5 Mon Sep 17 00:00:00 2001 From: Mike Cronce Date: Fri, 16 Aug 2024 11:09:21 -0400 Subject: [PATCH 3/4] diesel/src/mysql/connection/stmt/mod.rs: + StatementUse::insert_id() --- diesel/src/mysql/connection/stmt/mod.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/diesel/src/mysql/connection/stmt/mod.rs b/diesel/src/mysql/connection/stmt/mod.rs index 986a344ed5af..638d303217ab 100644 --- a/diesel/src/mysql/connection/stmt/mod.rs +++ b/diesel/src/mysql/connection/stmt/mod.rs @@ -160,6 +160,11 @@ impl<'a> StatementUse<'a> { .map_err(|e| Error::DeserializationError(Box::new(e))) } + /// SAFETY: This should only be called for INSERT queries. + pub(in crate::mysql::connection) unsafe fn insert_id(&self) -> u64 { + ffi::mysql_stmt_insert_id(self.inner.stmt.as_ptr()) + } + /// This function should be called after `execute` only /// otherwise it's not guaranteed to return a valid result pub(in crate::mysql::connection) unsafe fn result_size(&mut self) -> QueryResult { From c6639eca8e77a4ebf22089847780f1cc9e1f1710 Mon Sep 17 00:00:00 2001 From: Mike Cronce Date: Fri, 16 Aug 2024 11:09:41 -0400 Subject: [PATCH 4/4] diesel/src/mysql/connection/mod.rs: impl ConnectionWithReturningId for MysqlConnection --- diesel/src/mysql/connection/mod.rs | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/diesel/src/mysql/connection/mod.rs b/diesel/src/mysql/connection/mod.rs index d43c21ad2642..c5e677558ba8 100644 --- a/diesel/src/mysql/connection/mod.rs +++ b/diesel/src/mysql/connection/mod.rs @@ -211,6 +211,34 @@ impl Connection for MysqlConnection { } } +impl ConnectionWithReturningId for MysqlConnection { + type ReturnedId = u64; + + fn execute_returning_id(&mut self, source: &T) -> QueryResult + where + T: QueryFragment + QueryId, + { + #[allow(unsafe_code)] // call to unsafe function + update_transaction_manager_status( + prepared_query( + &source, + &mut self.statement_cache, + &mut self.raw_connection, + &mut *self.instrumentation, + ) + .and_then(|stmt| { + // we have not called result yet, so calling `execute` is + // fine + let stmt_use = unsafe { stmt.execute() }?; + Ok(unsafe { stmt_use.insert_id() }) + }), + &mut self.transaction_state, + &mut self.instrumentation, + &crate::debug_query(source), + ) + } +} + #[inline(always)] fn update_transaction_manager_status( query_result: QueryResult,