fix(cross_shard_disabled): provide more specific error msg when dealing with unknown sharding key

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();
+    }
+}

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();

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?;

pgdog/src/frontend/router/parser/cache/ast.rs

@@ -30,6 +30,8 @@ pub struct Ast {
     pub comment_role: Option<Role>,
     /// Parser query engine used.
     pub query_parser_engine: QueryParserEngine,
+    /// Sharding Key.
+    pub comment_sharding_key: Option<String>,
     /// Inner sync.
     inner: Arc<AstInner>,
 }
@@ -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)),
         }
     }

pgdog/src/frontend/router/parser/cache/cache_impl.rs

@@ -110,7 +110,6 @@ impl Cache {
     ) -> Result<Ast, Error> {
         // 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;
 

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<ShardOrLookup>,
+    pub role: Option<Role>,
+    pub sharding_key: Option<String>,
+}
+
 pub(super) fn shard_role_from_comment(
     comment: &str,
     schema: &ShardingSchema,
-) -> Result<(Option<ShardOrLookup>, Option<Role>), Error> {
+) -> Result<Directive, Error> {
     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::<usize>()
@@ -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,
+    })
 }

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<Role>,
     pub shard: Option<ShardOrLookup>,
+    pub sharding_key: Option<String>,
 }
 
 /// 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,
     })
 }

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),
 }

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);

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(),