Fix dry run (#491)

Author: levkk

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

@@ -156,15 +156,15 @@ impl QueryParser {
         };
 
         // Parse hardcoded shard from a query comment.
-        if context.router_needed {
+        if context.router_needed || context.dry_run {
             self.shard = statement.shard.clone();
             if let Some(role) = statement.role {
                 self.write_override = role == Role::Primary;
             }
         }
 
         debug!("{}", context.query()?.query());
-        trace!("{:#?}", statement.ast());
+        trace!("{:#?}", statement);
 
         let rewrite = Rewrite::new(statement.ast());
         if rewrite.needs_rewrite() {

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

@@ -1,4 +1,7 @@
+use std::ops::Deref;
+
 use crate::{
+    config::{self, config},
     frontend::client::TransactionType,
     net::{
         messages::{parse::Parse, Parameter},
@@ -522,3 +525,23 @@ fn test_commit_prepared() {
     let stmt = pg_query::parse("COMMIT PREPARED 'test'").unwrap();
     println!("{:?}", stmt);
 }
+
+#[test]
+fn test_dry_run_simple() {
+    let mut config = config().deref().clone();
+    config.config.general.dry_run = true;
+    config::set(config).unwrap();
+
+    let cluster = Cluster::new_test_single_shard();
+    let command = query_parser!(
+        QueryParser::default(),
+        Query::new("/* pgdog_sharding_key: 1234 */ SELECT * FROM sharded"),
+        false,
+        cluster
+    );
+    let cache = Cache::queries();
+    let stmt = cache.values().next().unwrap();
+    assert_eq!(stmt.stats.lock().direct, 1);
+    assert_eq!(stmt.stats.lock().multi, 0);
+    assert_eq!(command.route().shard(), &Shard::Direct(0));
+}