fix: detect select for update in CTEs (#879)

fix #878

Author: levkk

pgdog/src/frontend/router/parser/query/select.rs

@@ -318,11 +318,28 @@ impl QueryParser {
             }
         }
 
-        Ok(if stmt.locking_clause.is_empty() {
-            FunctionBehavior::default()
-        } else {
-            FunctionBehavior::writes_only()
-        })
+        if !stmt.locking_clause.is_empty() {
+            return Ok(FunctionBehavior::writes_only());
+        }
+
+        // Recurse into CTEs so a locking clause or write-only function
+        // nested inside a WITH clause still routes to the primary.
+        if let Some(ref with_clause) = stmt.with_clause {
+            for cte in &with_clause.ctes {
+                if let Some(NodeEnum::CommonTableExpr(ref expr)) = cte.node {
+                    if let Some(ref query) = expr.ctequery {
+                        if let Some(NodeEnum::SelectStmt(ref inner)) = query.node {
+                            let behavior = Self::functions(inner)?;
+                            if behavior.writes {
+                                return Ok(behavior);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(FunctionBehavior::default())
     }
 
     /// Check for CTEs that could trigger this query to go to a primary.

pgdog/src/frontend/router/parser/query/test/mod.rs

@@ -292,6 +292,28 @@ fn test_select_for_update() {
     assert!(route.is_write());
 }
 
+#[test]
+fn test_select_for_update_in_cte() {
+    // FOR UPDATE buried inside a CTE should still route to the primary.
+    let route = query!(
+        "WITH locked AS (SELECT * FROM sharded WHERE id = 1 FOR UPDATE) SELECT * FROM locked"
+    );
+    assert!(route.is_write());
+
+    // Doubly-nested CTE: the locking clause is two WITH levels deep.
+    let route = query!(
+        "WITH outer_cte AS (\
+            WITH inner_cte AS (SELECT * FROM sharded WHERE id = 1 FOR UPDATE) \
+            SELECT * FROM inner_cte\
+         ) SELECT * FROM outer_cte"
+    );
+    assert!(route.is_write());
+
+    // Sanity check: a plain SELECT inside a CTE without locking is still a read.
+    let route = query!("WITH plain AS (SELECT * FROM sharded WHERE id = 1) SELECT * FROM plain");
+    assert!(route.is_read());
+}
+
 #[test]
 fn test_omni() {
     let mut omni_round_robin = HashSet::new();