Fix comments with SET (#492)

* Fix comments with SET

* Add test case

Author: levkk

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

@@ -121,8 +121,14 @@ impl QueryParser {
             }
         }
 
+        let shard = if let Shard::Direct(_) = self.shard {
+            self.shard.clone()
+        } else {
+            Shard::All
+        };
+
         Ok(Command::Query(
-            Route::write(Shard::All).set_read(context.read_only),
+            Route::write(shard).set_read(context.read_only),
         ))
     }
 }

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

@@ -545,3 +545,20 @@ fn test_dry_run_simple() {
     assert_eq!(stmt.stats.lock().multi, 0);
     assert_eq!(command.route().shard(), &Shard::Direct(0));
 }
+
+#[test]
+fn test_set_comments() {
+    let command = query_parser!(
+        QueryParser::default(),
+        Query::new("/* pgdog_sharding_key: 1234 */ SET statement_timeout TO 1"),
+        true
+    );
+    assert_eq!(command.route().shard(), &Shard::Direct(0));
+
+    let command = query_parser!(
+        QueryParser::default(),
+        Query::new("SET statement_timeout TO 1"),
+        true
+    );
+    assert_eq!(command.route().shard(), &Shard::All);
+}