fix: handle SET after query inside transaction correctly with schema sharding (#684)

Example:

```
BEGIN;
SELECT 1;
SET LOCAL work_mem TO '128MB';
COMMIT;
```

Author: levkk

integration/python/test_sqlalchemy.py

@@ -372,12 +372,9 @@ impl Server {
                             }
                         }
                         Err(err) => {
-                            match err {
-                                Error::ProtocolOutOfSync => {
-                                    // conservatively, we do not know for sure if this is recoverable
-                                    self.stats.state(State::Error);
-                                }
-                                _ => {}
+                            if let Error::ProtocolOutOfSync = err {
+                                // conservatively, we do not know for sure if this is recoverable
+                                self.stats.state(State::Error);
                             }
                             error!(
                                 "{:?} got: {}, extended buffer: {:?}, state: {}",

pgdog/src/backend/server.rs

@@ -239,7 +239,9 @@ impl QueryEngine {
                     .await?
             }
             Command::Unlisten(channel) => self.unlisten(context, &channel.clone()).await?,
-            Command::Set { name, value, local } => {
+            Command::Set {
+                name, value, local, ..
+            } => {
                 // FIXME: parameters set in between statements inside a transaction won't
                 // be recorded in the client parameters.
                 if self.backend.connected() {

pgdog/src/frontend/client/query_engine/mod.rs

@@ -27,6 +27,7 @@ pub enum Command {
         name: String,
         value: ParameterValue,
         local: bool,
+        route: Route,
     },
     PreparedStatement(Prepare),
     InternalField {
@@ -61,6 +62,7 @@ impl Command {
         match self {
             Self::Query(route) => route,
             Self::ShardKeyRewrite(plan) => plan.route(),
+            Self::Set { route, .. } => route,
             _ => &DEFAULT_ROUTE,
         }
     }

pgdog/src/frontend/router/parser/command.rs

@@ -128,11 +128,7 @@ impl QueryParser {
             }
         }
 
-        debug!(
-            "query router decision: {:#?} (shard: {:#?})",
-            command,
-            context.shards_calculator.peek(),
-        );
+        debug!("query router decision: {:#?}", command);
 
         self.attach_explain(&mut command);
 

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

@@ -20,6 +20,7 @@ impl QueryParser {
                 name: stmt.name.to_string(),
                 value,
                 local: stmt.is_local,
+                route: Route::write(context.shards_calculator.shard()),
             });
         }
 

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

@@ -595,7 +595,9 @@ fn test_transaction() {
         cluster.clone()
     );
     match route {
-        Command::Set { name, value, local } => {
+        Command::Set {
+            name, value, local, ..
+        } => {
             assert_eq!(name, "application_name");
             assert_eq!(value.as_str().unwrap(), "test");
             assert!(!cluster.read_only());

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

@@ -1,4 +1,5 @@
-use crate::frontend::router::parser::Shard;
+use crate::frontend::router::parser::{route::ShardSource, Shard};
+use crate::net::parameter::ParameterValue;
 
 use super::setup::{QueryParserTest, *};
 
@@ -273,3 +274,19 @@ fn test_schema_sharding_priority_on_delete() {
 
     assert_eq!(command.route().shard(), &Shard::Direct(0));
 }
+
+// --- SET commands with schema sharding via search_path ---
+
+#[test]
+fn test_set_routes_to_shard_from_search_path() {
+    let mut test =
+        QueryParserTest::new().with_param("search_path", ParameterValue::String("shard_0".into()));
+
+    let command = test.execute(vec![Query::new("SET statement_timeout TO 1000").into()]);
+
+    assert_eq!(command.route().shard(), &Shard::Direct(0));
+    assert_eq!(
+        command.route().shard_with_priority().source(),
+        &ShardSource::SearchPath("shard_0".into())
+    );
+}

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

@@ -12,7 +12,8 @@ fn test_set_comment() {
     .into()]);
 
     assert!(
-        matches!(command, Command::Set { name, value, local } if name == "statement_timeout" && !local && value == ParameterValue::String("1".into())),
-        "expected Command::Set"
+        matches!(command.clone(), Command::Set { name, value, local, route } if name == "statement_timeout" && !local && value == ParameterValue::String("1".into()) && route.shard().is_direct()),
+        "expected Command::Set, got {:#?}",
+        command,
     );
 }

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

@@ -79,13 +79,22 @@ fn test_set_application_name_in_transaction() {
 
     let command = test.execute(vec![Query::new("SET application_name TO 'test'").into()]);
 
-    match command {
-        Command::Set { name, value, local } => {
+    match command.clone() {
+        Command::Set {
+            name,
+            value,
+            local,
+            route,
+        } => {
             assert_eq!(name, "application_name");
             assert_eq!(value.as_str().unwrap(), "test");
             assert!(!local);
             assert!(!test.cluster().read_only());
+            assert!(route.is_write());
+            assert!(route.shard().is_all());
         }
         _ => panic!("expected Set, got {command:?}"),
     }
+
+    assert!(command.route().shard().is_all());
 }