diff --git a/integration/rust/tests/integration/cross_shard_disabled.rs b/integration/rust/tests/integration/cross_shard_disabled.rs index 87394fe08..627722dc5 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,89 @@ 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_eq!( + err.to_string(), + "error returned from database: unmapped sharding key was specified" + ); + + // 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. + { + 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..117adbd57 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,9 +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 = 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..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), Error> { +) -> Result { let mut role = None; if let Some(cap) = ROLE.captures(comment) @@ -42,15 +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())), role)); + return Ok(Directive { + shard_or_lookup: Some(ShardOrLookup::Shard(schema.shard().into())), + role, + sharding_key: Some(sharding_key.to_string()), + }); } - return Ok((Some(shard_for_bare_key(sharding_key, schema, None)?), role)); + return Ok(Directive { + shard_or_lookup: Some(shard_for_bare_key(sharding_key, schema, None)?), + role, + 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::() @@ -59,8 +73,13 @@ pub(super) fn shard_role_from_comment( .unwrap_or(Shard::All), )), role, - )); + sharding_key: None, + }); } - Ok((None, role)) + 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 b8ce5d442..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)] @@ -18,6 +19,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,17 +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) = match leading { + let mut directive = match leading { Some(c) => directive::shard_role_from_comment(c, schema)?, - None => (None, None), + None => Directive { + shard_or_lookup: None, + role: None, + sharding_key: None, + }, }; if let Some(c) = trailing { - let (t_shard, t_role) = 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 directive.role.is_none() { + directive.role = t_directive.role; } - if role.is_none() { - role = t_role; + if directive.sharding_key.is_none() { + directive.sharding_key = t_directive.sharding_key; } } @@ -75,7 +84,8 @@ pub fn parse_edge_comment<'a>( (None, Some(t)) => t.to_string(), (None, None) => String::new(), }, - shard, - role, + sharding_key: directive.sharding_key, + shard: directive.shard_or_lookup, + role: directive.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..936140bd8 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -144,6 +144,28 @@ impl QueryParser { _ => (), } + 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())); + } + } + 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 b05b21fba..84d3251bb 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 unmapped sharding key in list-based/range-based sharding, + // and, if not stopped, the query would be cross-shard. + pub fn unmapped_sharding_key_in_cross_shard_disabled(sharding_key: &str) -> ErrorResponse { + ErrorResponse { + severity: "ERROR".into(), + code: "58000".into(), + message: "unmapped sharding key was specified".into(), + detail: Some(format!("sharding key '{}' is not mapped", sharding_key)), + context: None, + file: None, + routine: None, + } + } + pub fn set_shard_after_connect(name: &str) -> ErrorResponse { ErrorResponse { severity: "ERROR".into(),