Allow setting to return found rows when writing to the database

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),

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 {