From 5efd77ed998d12ce4f4ae3e10f1965cb97cb5797 Mon Sep 17 00:00:00 2001
From: Julius de Bruijn <julius+github@nauk.io>
Date: Tue, 21 Feb 2023 17:37:39 +0100
Subject: [PATCH] Allow setting to return found rows when writing to the
 database

---
 src/conn/mod.rs | 52 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/opts/mod.rs | 20 +++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/src/conn/mod.rs b/src/conn/mod.rs
index 04c1e5f7..7aaeedf2 100644
--- a/src/conn/mod.rs
+++ b/src/conn/mod.rs
@@ -1144,6 +1144,58 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn should_return_found_rows_if_flag_is_set() -> super::Result<()> {
+        let opts = get_opts().writes_return_found_rows(true);
+        let mut conn = Conn::new(opts).await.unwrap();
+
+        "CREATE TEMPORARY TABLE mysql.found_rows (id INT PRIMARY KEY AUTO_INCREMENT, val INT)"
+            .ignore(&mut conn)
+            .await?;
+
+        "INSERT INTO mysql.found_rows (val) VALUES (1)"
+            .ignore(&mut conn)
+            .await?;
+
+        // Inserted one row, affected should be one.
+        assert_eq!(conn.affected_rows(), 1);
+
+        "UPDATE mysql.found_rows SET val = 1 WHERE val = 1"
+            .ignore(&mut conn)
+            .await?;
+
+        // The query doesn't affect any rows, but due to us wanting FOUND rows,
+        // this has to return one.
+        assert_eq!(conn.affected_rows(), 1);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn should_not_return_found_rows_if_flag_is_not_set() -> super::Result<()> {
+        let mut conn = Conn::new(get_opts()).await.unwrap();
+
+        "CREATE TEMPORARY TABLE mysql.found_rows (id INT PRIMARY KEY AUTO_INCREMENT, val INT)"
+            .ignore(&mut conn)
+            .await?;
+
+        "INSERT INTO mysql.found_rows (val) VALUES (1)"
+            .ignore(&mut conn)
+            .await?;
+
+        // Inserted one row, affected should be one.
+        assert_eq!(conn.affected_rows(), 1);
+
+        "UPDATE mysql.found_rows SET val = 1 WHERE val = 1"
+            .ignore(&mut conn)
+            .await?;
+
+        // The query doesn't affect any rows.
+        assert_eq!(conn.affected_rows(), 0);
+
+        Ok(())
+    }
+
     async fn read_binlog_streams_and_close_their_connections(
         pool: Option<&Pool>,
         binlog_server_ids: (u32, u32, u32),
diff --git a/src/opts/mod.rs b/src/opts/mod.rs
index a74a6fc0..10be818d 100644
--- a/src/opts/mod.rs
+++ b/src/opts/mod.rs
@@ -404,6 +404,10 @@ pub(crate) struct MysqlOpts {
     ///
     /// Available via `secure_auth` connection url parameter.
     secure_auth: bool,
+
+    /// Changes the behavior of the affected count returned for writes (UPDATE/INSERT etc).
+    /// It makes MySQL return the FOUND rows instead of the AFFECTED rows.
+    client_found_rows: bool,
 }
 
 /// Mysql connection options.
@@ -721,6 +725,11 @@ impl Opts {
         self.inner.mysql_opts.secure_auth
     }
 
+    /// If true, write queries return found rows and if false, affected rows.
+    pub fn writes_return_found_rows(&self) -> bool {
+        self.inner.mysql_opts.client_found_rows
+    }
+
     pub(crate) fn get_capabilities(&self) -> CapabilityFlags {
         let mut out = CapabilityFlags::CLIENT_PROTOCOL_41
             | CapabilityFlags::CLIENT_SECURE_CONNECTION
@@ -742,6 +751,9 @@ impl Opts {
         if self.inner.mysql_opts.compression.is_some() {
             out |= CapabilityFlags::CLIENT_COMPRESS;
         }
+        if self.writes_return_found_rows() {
+            out |= CapabilityFlags::CLIENT_FOUND_ROWS;
+        }
 
         out
     }
@@ -767,6 +779,7 @@ impl Default for MysqlOpts {
             max_allowed_packet: None,
             wait_timeout: None,
             secure_auth: true,
+            client_found_rows: false,
         }
     }
 }
@@ -1017,6 +1030,13 @@ impl OptsBuilder {
         self.opts.secure_auth = secure_auth;
         self
     }
+
+    /// Changes the behavior of the affected count returned for writes.
+    /// See [`Opts::writes_return_found_rows`].
+    pub fn writes_return_found_rows(mut self, client_found_rows: bool) -> Self {
+        self.opts.client_found_rows = client_found_rows;
+        self
+    }
 }
 
 impl From<OptsBuilder> for Opts {