search_path schema-based sharding (#621)

* search_path sharding

* Handle insert

* Finish up

* Fix explain

* remove debug

* move up

Author: levkk

pgdog/src/backend/server.rs

@@ -451,26 +451,41 @@ impl Server {
     ) -> Result<usize, Error> {
         // Sync application_name parameter
         // and update it in the stats.
-        let default_name = "PgDog";
-        let server_name = self
-            .client_params
-            .get_default("application_name", default_name);
-        let client_name = params.get_default("application_name", default_name);
+        let server_name = self.client_params.get_default("application_name", "PgDog");
+        let client_name = params.get_default("application_name", "PgDog");
+
         self.stats.link_client(client_name, server_name, id);
 
         // Clear any params previously tracked by SET.
         self.changed_params.clear();
 
         // Compare client and server params.
         if !params.identical(&self.client_params) {
+            // Construct client parameter SET queries.
             let tracked = params.tracked();
+            // Construct RESET queries to reset any current params
+            // to their default values.
             let mut queries = self.client_params.reset_queries();
+
+            // Combine both to create a new, fresh session state
+            // on this connection.
             queries.extend(tracked.set_queries());
+
+            // Set state on the connection only if
+            // there are any params to change.
             if !queries.is_empty() {
                 debug!("syncing {} params", queries.len());
+
                 self.execute_batch(&queries).await?;
+
+                // We can receive ParameterStatus messages here,
+                // but we should ignore them since we are managing the session state.
+                self.changed_params.clear();
             }
+
+            // Update params on this connection.
             self.client_params = tracked;
+
             Ok(queries.len())
         } else {
             Ok(0)

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

@@ -73,6 +73,7 @@ impl ExplainEntry {
     }
 }
 
+/// EXPLAIN recorder.
 #[derive(Debug, Default)]
 pub struct ExplainRecorder {
     entries: Vec<ExplainEntry>,
@@ -89,6 +90,12 @@ impl ExplainRecorder {
         self.entries.push(ExplainEntry::new(shard, description));
     }
 
+    pub fn clear(&mut self) {
+        self.entries.clear();
+        self.comment = None;
+        self.plugin = None;
+    }
+
     pub fn record_comment_override(&mut self, shard: Shard, role: Option<&str>) {
         let mut description = match shard {
             Shard::Direct(_) | Shard::Multi(_) | Shard::All => {

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

@@ -9,7 +9,17 @@ impl QueryParser {
         node: &Option<NodeEnum>,
         context: &mut QueryParserContext<'_>,
     ) -> Result<Command, Error> {
-        Self::shard_ddl(node, &context.sharding_schema)
+        let mut command = Self::shard_ddl(node, &context.sharding_schema)?;
+
+        if let Command::Query(ref mut route) = command {
+            if route.shard().all() {
+                if let Some(shard) = self.check_search_path_for_shard(context)? {
+                    route.set_shard_mut(shard);
+                }
+            }
+        }
+
+        Ok(command)
     }
 
     pub(super) fn shard_ddl(

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

@@ -11,6 +11,10 @@ impl QueryParser {
     ) -> Result<Command, Error> {
         let table = stmt.relation.as_ref().map(Table::from);
 
+        if let Some(shard) = self.check_search_path_for_shard(context)? {
+            return Ok(Command::Query(Route::write(shard)));
+        }
+
         if let Some(table) = table {
             // Schema-based sharding.
             if let Some(schema) = context.sharding_schema.schemas.get(table.schema()) {

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

@@ -28,7 +28,7 @@ mod ddl;
 mod delete;
 mod explain;
 mod plugins;
-pub mod schema_sharding;
+mod schema_sharding;
 mod select;
 mod set;
 mod shared;
@@ -382,6 +382,46 @@ impl QueryParser {
             // with the fingerprint.
             //
             if route.shard().all() {
+                // Check search_path for schema.
+                let search_path = context.router_context.params.get("search_path");
+
+                // Quick inline function to shard query
+                // based on schema in search_path.
+                fn shard_from_search_path(
+                    search_path: &str,
+                    context: &QueryParserContext<'_>,
+                    query_parser: &mut QueryParser,
+                    route: &mut Route,
+                ) {
+                    let schema = Schema::from(search_path);
+
+                    if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
+                        let shard: Shard = schema.shard().into();
+
+                        if let Some(recorder) = query_parser.recorder_mut() {
+                            recorder.record_entry(
+                                Some(shard.clone()),
+                                format!("matched schema {} in search_path", schema.name()),
+                            );
+                        }
+                        route.set_shard_mut(shard);
+                    }
+                }
+
+                match search_path {
+                    Some(ParameterValue::String(search_path)) => {
+                        shard_from_search_path(search_path, context, self, route);
+                    }
+
+                    Some(ParameterValue::Tuple(search_paths)) => {
+                        for schema in search_paths {
+                            shard_from_search_path(schema, context, self, route);
+                        }
+                    }
+
+                    None => (),
+                }
+
                 let databases = databases();
                 // Only fingerprint the query if some manual queries are configured.
                 // Otherwise, we're wasting time parsing SQL.
@@ -469,6 +509,10 @@ impl QueryParser {
             context.split_insert_mode(),
         )?;
 
+        if let Some(shard) = self.check_search_path_for_shard(context)? {
+            return Ok(Command::Query(Route::write(shard)));
+        }
+
         match routing {
             InsertRouting::Routed(shard) => {
                 if let Some(recorder) = self.recorder_mut() {

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

@@ -1 +1,64 @@
+use super::*;
 
+impl QueryParser {
+    pub(super) fn check_search_path_for_shard(
+        &mut self,
+        context: &QueryParserContext<'_>,
+    ) -> Result<Option<Shard>, Error> {
+        // Shortcut.
+        if context.sharding_schema.schemas.is_empty() {
+            return Ok(None);
+        }
+
+        // Check search_path for schema.
+        let search_path = context.router_context.params.get("search_path");
+
+        match search_path {
+            Some(ParameterValue::String(search_path)) => {
+                let schema = Schema::from(search_path.as_str());
+                if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
+                    let shard: Shard = schema.shard().into();
+                    if let Some(recorder) = self.recorder_mut() {
+                        recorder.record_entry(
+                            Some(shard.clone()),
+                            format!("matched schema \"{}\" in search_path", schema.name()),
+                        );
+                    }
+                    return Ok(Some(shard));
+                }
+            }
+
+            Some(ParameterValue::Tuple(search_paths)) => {
+                let mut candidates = vec![];
+
+                for (idx, schema) in search_paths.iter().enumerate() {
+                    let schema = Schema::from(schema.as_str());
+                    if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
+                        let shard: Shard = schema.shard().into();
+                        let catch_all = schema.is_default();
+                        candidates.push((shard, catch_all, idx));
+                    }
+                }
+
+                // false < true
+                // Catch-all schemas go first, more qualified ones go last.
+                candidates.sort_by_key(|cand| !cand.1);
+                if let Some(candidate) = candidates.pop() {
+                    if let Some(schema) = search_paths.get(candidate.2) {
+                        if let Some(recorder) = self.recorder_mut() {
+                            recorder.record_entry(
+                                Some(candidate.0.clone()),
+                                format!("matched schema mult \"{}\" in search_path", schema),
+                            );
+                        }
+                    }
+                    return Ok(Some(candidate.0));
+                }
+            }
+
+            None => (),
+        }
+
+        Ok(None)
+    }
+}

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

@@ -60,6 +60,10 @@ impl QueryParser {
             )?;
         }
 
+        if let Some(Shard::Direct(number)) = self.check_search_path_for_shard(context)? {
+            return Ok(Command::Query(Route::read(number).set_write(writes)));
+        }
+
         // Schema-based sharding.
         for table in cached_ast.tables() {
             if let Some(schema) = context.sharding_schema.schemas.get(table.schema()) {

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

@@ -20,6 +20,10 @@ impl QueryParser {
     ) -> Result<Command, Error> {
         let table = stmt.relation.as_ref().map(Table::from);
 
+        if let Some(shard) = self.check_search_path_for_shard(context)? {
+            return Ok(Command::Query(Route::write(shard)));
+        }
+
         if let Some(table) = table {
             // Schema-based sharding.
             if let Some(schema) = context.sharding_schema.schemas.get(table.schema()) {