From 29a476e190deaf979d4b858ff132192edfe15edd Mon Sep 17 00:00:00 2001
From: Mostafa Abdelraouf <mostafa.mohmmed@gmail.com>
Date: Fri, 30 Aug 2024 04:26:36 -0500
Subject: [PATCH] QueryRouter: route to primary when locks exists (select for
 update) (#782)

Authored-by: Javier Goday <jgoday@gmail.com>
---
 src/query_router.rs | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/query_router.rs b/src/query_router.rs
index 7acd6841..da67e96a 100644
--- a/src/query_router.rs
+++ b/src/query_router.rs
@@ -427,8 +427,12 @@ impl QueryRouter {
                         None => (),
                     };
 
-                    // If we already visited a write statement, we should be going to the primary.
-                    if !visited_write_statement {
+                    let has_locks = !query.locks.is_empty();
+
+                    if has_locks {
+                        self.active_role = Some(Role::Primary);
+                    } else if !visited_write_statement {
+                        // If we already visited a write statement, we should be going to the primary.
                         self.active_role = match self.primary_reads_enabled() {
                             false => Some(Role::Replica), // If primary should not be receiving reads, use a replica.
                             true => None,                 // Any server role is fine in this case.
@@ -1158,6 +1162,29 @@ mod test {
         }
     }
 
+    #[test]
+    fn test_select_for_update() {
+        QueryRouter::setup();
+        let mut qr = QueryRouter::new();
+        qr.pool_settings.query_parser_read_write_splitting = true;
+
+        let queries_in_primary_role = vec![
+            simple_query("BEGIN"), // Transaction start
+            simple_query("SELECT * FROM items WHERE id = 5 FOR UPDATE"),
+            simple_query("UPDATE items SET name = 'pumpkin' WHERE id = 5"),
+        ];
+
+        for query in queries_in_primary_role {
+            assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
+            assert_eq!(qr.role(), Some(Role::Primary));
+        }
+
+        // query without lock do not change role
+        let query = simple_query("SELECT * FROM items WHERE id = 5");
+        assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
+        assert_eq!(qr.role(), None);
+    }
+
     #[test]
     fn test_infer_primary_reads_enabled() {
         QueryRouter::setup();