From e87d42799f4c2fb52d6dd558be5e29f10a52c0c9 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 14 Aug 2026 12:58:50 -0400 Subject: [PATCH 1/6] fix(cross_shard_disabled): provide more specific error msg when dealing with unknown sharding key --- .../tests/integration/cross_shard_disabled.rs | 75 ++++++++++++++++++- .../src/frontend/client/query_engine/query.rs | 15 +++- pgdog/src/net/messages/error_response.rs | 15 ++++ 3 files changed, 102 insertions(+), 3 deletions(-) diff --git a/integration/rust/tests/integration/cross_shard_disabled.rs b/integration/rust/tests/integration/cross_shard_disabled.rs index 87394fe08..b399d2734 100644 --- a/integration/rust/tests/integration/cross_shard_disabled.rs +++ b/integration/rust/tests/integration/cross_shard_disabled.rs @@ -1,7 +1,7 @@ use std::time::Duration; use crate::setup::{admin_sqlx, connections_sqlx}; -use sqlx::Executor; +use sqlx::{Connection, Executor}; use tokio::time::sleep; #[tokio::test] @@ -59,3 +59,76 @@ async fn test_cross_shard_disabled() { .await .unwrap(); } + +#[tokio::test] +async fn test_cross_shard_disabled_with_unknown_sharding_key() { + let admin = admin_sqlx().await; + + let mut conn = + sqlx::PgConnection::connect("postgres://pgdog:pgdog@127.0.0.1:6432/single_sharded_list") + .await + .unwrap(); + + // Get everything setup to test. + { + admin + .execute("SET cross_shard_disabled TO false") + .await + .unwrap(); + + conn.execute("DROP TABLE IF EXISTS test_unknown_sharding_key") + .await + .unwrap(); + + conn.execute( + "CREATE TABLE IF NOT EXISTS test_unknown_sharding_key(id BIGINT, value VARCHAR)", + ) + .await + .unwrap(); + } + + // Query has sharding key that is unknown using list-based sharding. + // 0-10 => shard 0, 11-20 => shard 1 in this instance + // 25 doesn't map to a shard. + // With a valid sharding key, SELECT * FROM statement would return the pertinent rows. + // However, since we have an unknown one, and cross-shard queries are denied, + // an error must be thrown (otherwise, it would be a cross-shard query). + { + admin + .execute("SET cross_shard_disabled TO true") + .await + .unwrap(); + + conn.execute("BEGIN").await.unwrap(); + + conn.execute(format!("SET pgdog.sharding_key TO '{}'", 25).as_str()) + .await + .unwrap(); + + let err = sqlx::query("SELECT * FROM test_unknown_sharding_key") + .fetch_one(&mut conn) + .await + .err() + .unwrap(); + assert!( + err.to_string() + .contains("error returned from database: unknown sharding key was specified") + ); + + conn.execute("COMMIT").await.unwrap(); + } + + // Reset back to normal. + { + admin + .execute("SET cross_shard_disabled TO false") + .await + .unwrap(); + + conn.execute("DROP TABLE test_unknown_sharding_key") + .await + .unwrap(); + + conn.close().await.unwrap(); + } +} diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index cb250d22e..597924a0c 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -363,7 +363,6 @@ impl QueryEngine { if !cross_shard_disabled { return Ok(true); } - let query_is_cross_shard = context.client_request.route().is_cross_shard(); // The query is direct-to-shard, we're good. @@ -389,7 +388,19 @@ impl QueryEngine { // should be cross-shard (e.g. BEGIN, COMMIT) but aren't really. if connected_shards == 0 || connected_shards > 1 { let query = context.client_request.query()?; - let error = ErrorResponse::cross_shard_disabled(query.as_ref().map(|q| q.query())); + + let error = if let Some(sharding_key) = context.params.get("pgdog.sharding_key") + && let Some(sharding_key_value) = sharding_key.as_str() + { + // The user specified a sharding key, which was determined to be unknown. + // This was done in list-based/range-based sharding. + // If not stopped, the query would be cross-shard (due to the unknown key). + // Rather than give them a generic 'cross shard disabled' error, + // tell them the sharding key is unknown. + ErrorResponse::unknown_sharding_key_in_cross_shard_disabled(sharding_key_value) + } else { + ErrorResponse::cross_shard_disabled(query.as_ref().map(|q| q.query())) + }; self.error_response(context, error).await?; diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs index b05b21fba..f0b19c94a 100644 --- a/pgdog/src/net/messages/error_response.rs +++ b/pgdog/src/net/messages/error_response.rs @@ -86,6 +86,21 @@ impl ErrorResponse { } } + // Cross-shard queries are disabled. + // User specified an unknown sharding key in list-based/range-based sharding, + // and, if not stopped, the query would be cross-shard. + pub fn unknown_sharding_key_in_cross_shard_disabled(sharding_key: &str) -> ErrorResponse { + ErrorResponse { + severity: "ERROR".into(), + code: "58000".into(), + message: "unknown sharding key was specified".into(), + detail: Some(format!("sharding key {} is unknown", sharding_key)), + context: None, + file: None, + routine: None, + } + } + pub fn set_shard_after_connect(name: &str) -> ErrorResponse { ErrorResponse { severity: "ERROR".into(), From af981705caa01290e3debfd0a922718a8a074fb5 Mon Sep 17 00:00:00 2001 From: John Kaczman Date: Fri, 14 Aug 2026 13:12:42 -0400 Subject: [PATCH 2/6] Update pgdog/src/net/messages/error_response.rs Co-authored-by: Lev Kokotov --- pgdog/src/net/messages/error_response.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs index f0b19c94a..e8b5400ef 100644 --- a/pgdog/src/net/messages/error_response.rs +++ b/pgdog/src/net/messages/error_response.rs @@ -94,7 +94,7 @@ impl ErrorResponse { severity: "ERROR".into(), code: "58000".into(), message: "unknown sharding key was specified".into(), - detail: Some(format!("sharding key {} is unknown", sharding_key)), + detail: Some(format!("sharding key '{}' is not mapped", sharding_key)), context: None, file: None, routine: None, From 6c8dbac4f51c893b1c9b9cb41981fc9ecc34b7d6 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 14 Aug 2026 17:12:34 -0400 Subject: [PATCH 3/6] fix: move check logic to parser, catch comment directive --- .../tests/integration/cross_shard_disabled.rs | 21 ++++++++++++--- .../src/frontend/client/query_engine/query.rs | 20 ++++---------- .../client/query_engine/route_query.rs | 11 ++++++++ pgdog/src/frontend/router/parser/cache/ast.rs | 5 ++++ .../router/parser/cache/cache_impl.rs | 11 +++++--- .../router/parser/comment/directive.rs | 17 +++++++++--- .../src/frontend/router/parser/comment/mod.rs | 11 +++++--- pgdog/src/frontend/router/parser/error.rs | 3 +++ pgdog/src/frontend/router/parser/query/mod.rs | 26 +++++++++++++++++++ pgdog/src/net/messages/error_response.rs | 6 ++--- 10 files changed, 98 insertions(+), 33 deletions(-) diff --git a/integration/rust/tests/integration/cross_shard_disabled.rs b/integration/rust/tests/integration/cross_shard_disabled.rs index b399d2734..627722dc5 100644 --- a/integration/rust/tests/integration/cross_shard_disabled.rs +++ b/integration/rust/tests/integration/cross_shard_disabled.rs @@ -110,12 +110,25 @@ async fn test_cross_shard_disabled_with_unknown_sharding_key() { .await .err() .unwrap(); - assert!( - err.to_string() - .contains("error returned from database: unknown sharding key was specified") + assert_eq!( + err.to_string(), + "error returned from database: unmapped sharding key was specified" ); - conn.execute("COMMIT").await.unwrap(); + // Clear SET parameter for prep for next test. + conn.execute("RESET pgdog.sharding_key").await.unwrap(); + + // Also verify it works for comment directives. + let err = + sqlx::query("/* pgdog_sharding_key: 25 */ SELECT * FROM test_unknown_sharding_key") + .fetch_one(&mut conn) + .await + .err() + .unwrap(); + assert_eq!( + err.to_string(), + "error returned from database: unmapped sharding key was specified" + ); } // Reset back to normal. diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 597924a0c..117adbd57 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -388,21 +388,11 @@ impl QueryEngine { // should be cross-shard (e.g. BEGIN, COMMIT) but aren't really. if connected_shards == 0 || connected_shards > 1 { let query = context.client_request.query()?; - - let error = if let Some(sharding_key) = context.params.get("pgdog.sharding_key") - && let Some(sharding_key_value) = sharding_key.as_str() - { - // The user specified a sharding key, which was determined to be unknown. - // This was done in list-based/range-based sharding. - // If not stopped, the query would be cross-shard (due to the unknown key). - // Rather than give them a generic 'cross shard disabled' error, - // tell them the sharding key is unknown. - ErrorResponse::unknown_sharding_key_in_cross_shard_disabled(sharding_key_value) - } else { - ErrorResponse::cross_shard_disabled(query.as_ref().map(|q| q.query())) - }; - - self.error_response(context, error).await?; + self.error_response( + context, + ErrorResponse::cross_shard_disabled(query.as_ref().map(|q| q.query())), + ) + .await?; if self.backend.connected() && self.backend.done() { self.backend.disconnect(); diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 0bb668bcd..b989deeca 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -180,6 +180,17 @@ impl QueryEngine { return Ok(false); } + Err(RouterError::Parser(ParserError::UnmappedShardKey(shard_key))) => { + self.error_response( + context, + ErrorResponse::unmapped_sharding_key_in_cross_shard_disabled( + shard_key.as_str(), + ), + ) + .await?; + + return Ok(false); + } Err(err) => { self.error_response(context, ErrorResponse::syntax(err.to_string().as_str())) .await?; diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index 471f5f531..c9675fffc 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -30,6 +30,8 @@ pub struct Ast { pub comment_role: Option, /// Parser query engine used. pub query_parser_engine: QueryParserEngine, + /// Sharding Key. + pub comment_sharding_key: Option, /// Inner sync. inner: Arc, } @@ -119,6 +121,7 @@ impl Ast { comment_shard: None, comment_role: None, query_parser_engine: schema.query_parser_engine, + comment_sharding_key: None, inner: Arc::new(AstInner { stats: Mutex::new(stats), ast, @@ -156,6 +159,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine, + comment_sharding_key: None, inner: Arc::new(AstInner::new(ast.into_inner())), }) } @@ -167,6 +171,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine: QueryParserEngine::default(), + comment_sharding_key: None, inner: Arc::new(AstInner::new(stmts)), } } diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 832a233c0..8155c105d 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -110,7 +110,6 @@ impl Cache { ) -> Result { // Separate query from comment, if one is present. let query_and_comment = parse_edge_comment(query.query(), &ctx.sharding_schema)?; - { let mut guard = self.inner.lock(); let ast = guard.queries.get_mut(query_and_comment.query).map(|entry| { @@ -120,7 +119,8 @@ impl Cache { if let Some(mut ast) = ast { guard.stats.hits += 1; ast.comment_role = query_and_comment.role; - ast.comment_shard = query_and_comment.shard.clone(); + ast.comment_shard = query_and_comment.shard; + ast.comment_sharding_key = query_and_comment.sharding_key; return Ok(ast); } @@ -136,7 +136,9 @@ impl Cache { prepared_statements, )?; entry.comment_role = query_and_comment.role; - entry.comment_shard = query_and_comment.shard.clone(); + entry.comment_shard = query_and_comment.shard; + entry.comment_sharding_key = query_and_comment.sharding_key; + let parse_time = entry.stats.lock().parse_time; let mut guard = self.inner.lock(); @@ -177,7 +179,8 @@ impl Cache { )?; entry.cached = false; entry.comment_role = query_and_comment.role; - entry.comment_shard = query_and_comment.shard.clone(); + entry.comment_shard = query_and_comment.shard; + entry.comment_sharding_key = query_and_comment.sharding_key; let parse_time = entry.stats.lock().parse_time; diff --git a/pgdog/src/frontend/router/parser/comment/directive.rs b/pgdog/src/frontend/router/parser/comment/directive.rs index eb6ffd695..75ebf9cc6 100644 --- a/pgdog/src/frontend/router/parser/comment/directive.rs +++ b/pgdog/src/frontend/router/parser/comment/directive.rs @@ -26,7 +26,7 @@ pub(super) fn get_matched_value<'a>(caps: &'a regex::Captures<'a>) -> Option<&'a pub(super) fn shard_role_from_comment( comment: &str, schema: &ShardingSchema, -) -> Result<(Option, Option), Error> { +) -> Result<(Option, Option, Option), Error> { let mut role = None; if let Some(cap) = ROLE.captures(comment) @@ -42,9 +42,17 @@ pub(super) fn shard_role_from_comment( && let Some(sharding_key) = get_matched_value(&cap) { if let Some(schema) = schema.schemas.get(Some(sharding_key.into())) { - return Ok((Some(ShardOrLookup::Shard(schema.shard().into())), role)); + return Ok(( + Some(ShardOrLookup::Shard(schema.shard().into())), + role, + Some(sharding_key.to_string()), + )); } - return Ok((Some(shard_for_bare_key(sharding_key, schema, None)?), role)); + return Ok(( + Some(shard_for_bare_key(sharding_key, schema, None)?), + role, + Some(sharding_key.to_string()), + )); } if let Some(cap) = SHARD.captures(comment) && let Some(shard) = cap.get(1) @@ -59,8 +67,9 @@ pub(super) fn shard_role_from_comment( .unwrap_or(Shard::All), )), role, + None, )); } - Ok((None, role)) + Ok((None, role, None)) } diff --git a/pgdog/src/frontend/router/parser/comment/mod.rs b/pgdog/src/frontend/router/parser/comment/mod.rs index b8ce5d442..630a30ad6 100644 --- a/pgdog/src/frontend/router/parser/comment/mod.rs +++ b/pgdog/src/frontend/router/parser/comment/mod.rs @@ -18,6 +18,7 @@ pub struct QueryAndComment<'a> { pub comment: String, pub role: Option, pub shard: Option, + pub sharding_key: Option, } /// Extract SQL C-style block comments from both the beginning and the end @@ -52,18 +53,21 @@ pub fn parse_edge_comment<'a>( // Leading wins per-field: extract from leading first, then fill in any // fields the leading didn't provide from trailing. - let (mut shard, mut role) = match leading { + let (mut shard, mut role, mut sharding_key) = match leading { Some(c) => directive::shard_role_from_comment(c, schema)?, - None => (None, None), + None => (None, None, None), }; if let Some(c) = trailing { - let (t_shard, t_role) = directive::shard_role_from_comment(c, schema)?; + let (t_shard, t_role, t_sharding_key) = directive::shard_role_from_comment(c, schema)?; if shard.is_none() { shard = t_shard; } if role.is_none() { role = t_role; } + if sharding_key.is_none() { + sharding_key = t_sharding_key; + } } Ok(QueryAndComment { @@ -75,6 +79,7 @@ pub fn parse_edge_comment<'a>( (None, Some(t)) => t.to_string(), (None, None) => String::new(), }, + sharding_key, shard, role, }) diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index 276a1237e..24ae38e37 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -108,4 +108,7 @@ pub enum Error { #[error("multi-statement queries cannot mix SET with other commands")] MultiStatementMixedSet, + + #[error("unmapped sharding key was specified")] + UnmappedShardKey(String), } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 5f480bdfc..fcb4da85b 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -144,6 +144,32 @@ impl QueryParser { _ => (), } + if let Command::Query(_) = command { + if context.router_context.cluster.cross_shard_disabled() { + if context.shards_calculator.shard().is_all() { + let statement = context + .router_context + .ast + .clone() + .ok_or(Error::EmptyQuery)?; + + // The user specified a sharding key, which is un-mapped (list-based/range-based). + // If not stopped, the query would be cross-shard. + // Rather than give them a generic 'cross shard disabled' error, + // tell them the sharding key is un-mapped. + if let Some(sharding_key) = + context.router_context.parameter_hints.pgdog_sharding_key + { + if let Some(sharding_key_value) = sharding_key.as_str() { + return Err(Error::UnmappedShardKey(sharding_key_value.to_string())); + } + } else if let Some(sharding_key) = &statement.comment_sharding_key { + return Err(Error::UnmappedShardKey(sharding_key.to_string())); + } + } + } + } + debug!("query router decision: {:#?}", command); self.attach_explain(&mut command); diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs index e8b5400ef..84d3251bb 100644 --- a/pgdog/src/net/messages/error_response.rs +++ b/pgdog/src/net/messages/error_response.rs @@ -87,13 +87,13 @@ impl ErrorResponse { } // Cross-shard queries are disabled. - // User specified an unknown sharding key in list-based/range-based sharding, + // User specified an unmapped sharding key in list-based/range-based sharding, // and, if not stopped, the query would be cross-shard. - pub fn unknown_sharding_key_in_cross_shard_disabled(sharding_key: &str) -> ErrorResponse { + pub fn unmapped_sharding_key_in_cross_shard_disabled(sharding_key: &str) -> ErrorResponse { ErrorResponse { severity: "ERROR".into(), code: "58000".into(), - message: "unknown sharding key was specified".into(), + message: "unmapped sharding key was specified".into(), detail: Some(format!("sharding key '{}' is not mapped", sharding_key)), context: None, file: None, From 2bb1b693edf794a1a9bb218bfa0f33a8c2f10a88 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 14 Aug 2026 17:27:54 -0400 Subject: [PATCH 4/6] fix: get rid of Option::None error prop; gate on Some --- pgdog/src/frontend/router/parser/query/mod.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index fcb4da85b..d8413db39 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -147,23 +147,22 @@ impl QueryParser { if let Command::Query(_) = command { if context.router_context.cluster.cross_shard_disabled() { if context.shards_calculator.shard().is_all() { - let statement = context - .router_context - .ast - .clone() - .ok_or(Error::EmptyQuery)?; - // The user specified a sharding key, which is un-mapped (list-based/range-based). // If not stopped, the query would be cross-shard. // Rather than give them a generic 'cross shard disabled' error, // tell them the sharding key is un-mapped. + if let Some(sharding_key) = context.router_context.parameter_hints.pgdog_sharding_key { + // SET sharding key if let Some(sharding_key_value) = sharding_key.as_str() { return Err(Error::UnmappedShardKey(sharding_key_value.to_string())); } - } else if let Some(sharding_key) = &statement.comment_sharding_key { + } else if let Some(statement) = context.router_context.ast + && let Some(sharding_key) = &statement.comment_sharding_key + { + // Comment directive sharding key return Err(Error::UnmappedShardKey(sharding_key.to_string())); } } From e137e274593ed08c80f3624e34c3a0ed32b4f34e Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 14 Aug 2026 17:35:06 -0400 Subject: [PATCH 5/6] fmt: collapse if stmts --- pgdog/src/frontend/router/parser/query/mod.rs | 41 +++++++++---------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index d8413db39..936140bd8 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -144,28 +144,25 @@ impl QueryParser { _ => (), } - if let Command::Query(_) = command { - if context.router_context.cluster.cross_shard_disabled() { - if context.shards_calculator.shard().is_all() { - // The user specified a sharding key, which is un-mapped (list-based/range-based). - // If not stopped, the query would be cross-shard. - // Rather than give them a generic 'cross shard disabled' error, - // tell them the sharding key is un-mapped. - - if let Some(sharding_key) = - context.router_context.parameter_hints.pgdog_sharding_key - { - // SET sharding key - if let Some(sharding_key_value) = sharding_key.as_str() { - return Err(Error::UnmappedShardKey(sharding_key_value.to_string())); - } - } else if let Some(statement) = context.router_context.ast - && let Some(sharding_key) = &statement.comment_sharding_key - { - // Comment directive sharding key - return Err(Error::UnmappedShardKey(sharding_key.to_string())); - } - } + if let Command::Query(_) = command + && context.router_context.cluster.cross_shard_disabled() + && context.shards_calculator.shard().is_all() + { + // The user specified a sharding key, which is un-mapped (list-based/range-based). + // If not stopped, the query would be cross-shard. + // Rather than give them a generic 'cross shard disabled' error, + // tell them the sharding key is un-mapped. + + if let Some(sharding_key) = context.router_context.parameter_hints.pgdog_sharding_key + && let Some(sharding_key_value) = sharding_key.as_str() + { + // SET sharding key + return Err(Error::UnmappedShardKey(sharding_key_value.to_string())); + } else if let Some(statement) = context.router_context.ast + && let Some(sharding_key) = &statement.comment_sharding_key + { + // Comment directive sharding key + return Err(Error::UnmappedShardKey(sharding_key.to_string())); } } From 13c955d9a5a20817c4248cd19b9f8a4f5683ef92 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 14 Aug 2026 21:19:33 -0400 Subject: [PATCH 6/6] fix: quick fix clippy warnings --- .../router/parser/comment/directive.rs | 38 ++++++++++++------- .../src/frontend/router/parser/comment/mod.rs | 29 ++++++++------ 2 files changed, 41 insertions(+), 26 deletions(-) diff --git a/pgdog/src/frontend/router/parser/comment/directive.rs b/pgdog/src/frontend/router/parser/comment/directive.rs index 75ebf9cc6..d2812e68a 100644 --- a/pgdog/src/frontend/router/parser/comment/directive.rs +++ b/pgdog/src/frontend/router/parser/comment/directive.rs @@ -23,10 +23,16 @@ pub(super) fn get_matched_value<'a>(caps: &'a regex::Captures<'a>) -> Option<&'a .map(|m| m.as_str()) } +pub struct Directive { + pub shard_or_lookup: Option, + pub role: Option, + pub sharding_key: Option, +} + pub(super) fn shard_role_from_comment( comment: &str, schema: &ShardingSchema, -) -> Result<(Option, Option, Option), Error> { +) -> Result { let mut role = None; if let Some(cap) = ROLE.captures(comment) @@ -42,23 +48,23 @@ pub(super) fn shard_role_from_comment( && let Some(sharding_key) = get_matched_value(&cap) { if let Some(schema) = schema.schemas.get(Some(sharding_key.into())) { - return Ok(( - Some(ShardOrLookup::Shard(schema.shard().into())), + return Ok(Directive { + shard_or_lookup: Some(ShardOrLookup::Shard(schema.shard().into())), role, - Some(sharding_key.to_string()), - )); + sharding_key: Some(sharding_key.to_string()), + }); } - return Ok(( - Some(shard_for_bare_key(sharding_key, schema, None)?), + return Ok(Directive { + shard_or_lookup: Some(shard_for_bare_key(sharding_key, schema, None)?), role, - Some(sharding_key.to_string()), - )); + sharding_key: Some(sharding_key.to_string()), + }); } if let Some(cap) = SHARD.captures(comment) && let Some(shard) = cap.get(1) { - return Ok(( - Some(ShardOrLookup::Shard( + return Ok(Directive { + shard_or_lookup: Some(ShardOrLookup::Shard( shard .as_str() .parse::() @@ -67,9 +73,13 @@ pub(super) fn shard_role_from_comment( .unwrap_or(Shard::All), )), role, - None, - )); + sharding_key: None, + }); } - Ok((None, role, None)) + Ok(Directive { + shard_or_lookup: None, + role, + sharding_key: None, + }) } diff --git a/pgdog/src/frontend/router/parser/comment/mod.rs b/pgdog/src/frontend/router/parser/comment/mod.rs index 630a30ad6..0f7c949fa 100644 --- a/pgdog/src/frontend/router/parser/comment/mod.rs +++ b/pgdog/src/frontend/router/parser/comment/mod.rs @@ -9,6 +9,7 @@ use crate::config::database::Role; use crate::frontend::router::sharding::ShardOrLookup; use super::Error; +use crate::frontend::router::parser::comment::directive::Directive; use strip::{leading_block_comment, trailing_block_comment}; #[derive(Default, Debug, Clone)] @@ -53,20 +54,24 @@ pub fn parse_edge_comment<'a>( // Leading wins per-field: extract from leading first, then fill in any // fields the leading didn't provide from trailing. - let (mut shard, mut role, mut sharding_key) = match leading { + let mut directive = match leading { Some(c) => directive::shard_role_from_comment(c, schema)?, - None => (None, None, None), + None => Directive { + shard_or_lookup: None, + role: None, + sharding_key: None, + }, }; if let Some(c) = trailing { - let (t_shard, t_role, t_sharding_key) = directive::shard_role_from_comment(c, schema)?; - if shard.is_none() { - shard = t_shard; + let t_directive = directive::shard_role_from_comment(c, schema)?; + if directive.shard_or_lookup.is_none() { + directive.shard_or_lookup = t_directive.shard_or_lookup; } - if role.is_none() { - role = t_role; + if directive.role.is_none() { + directive.role = t_directive.role; } - if sharding_key.is_none() { - sharding_key = t_sharding_key; + if directive.sharding_key.is_none() { + directive.sharding_key = t_directive.sharding_key; } } @@ -79,8 +84,8 @@ pub fn parse_edge_comment<'a>( (None, Some(t)) => t.to_string(), (None, None) => String::new(), }, - sharding_key, - shard, - role, + sharding_key: directive.sharding_key, + shard: directive.shard_or_lookup, + role: directive.role, }) }