diff --git a/Cargo.lock b/Cargo.lock index 19b9aaabca4..21648f12e98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -213,7 +213,7 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" [[package]] name = "api" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arrow-schema 58.3.0", "common-base", @@ -936,7 +936,7 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "auth" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -1574,7 +1574,7 @@ dependencies = [ [[package]] name = "cache" -version = "1.1.3" +version = "1.1.4" dependencies = [ "catalog", "common-error", @@ -1610,7 +1610,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow 58.3.0", @@ -1945,7 +1945,7 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "cli" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-stream", "async-trait", @@ -2004,7 +2004,7 @@ dependencies = [ [[package]] name = "client" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arc-swap", @@ -2036,7 +2036,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", - "substrait 1.1.3", + "substrait 1.1.4", "tokio", "tokio-stream", "tonic 0.14.2", @@ -2085,7 +2085,7 @@ dependencies = [ [[package]] name = "cmd" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -2235,7 +2235,7 @@ dependencies = [ [[package]] name = "common-base" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "anymap2", @@ -2255,14 +2255,14 @@ dependencies = [ [[package]] name = "common-catalog" -version = "1.1.3" +version = "1.1.4" dependencies = [ "const_format", ] [[package]] name = "common-config" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-base", "common-error", @@ -2286,7 +2286,7 @@ dependencies = [ [[package]] name = "common-datasource" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -2322,7 +2322,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "1.1.3" +version = "1.1.4" dependencies = [ "bigdecimal 0.4.8", "common-error", @@ -2335,7 +2335,7 @@ dependencies = [ [[package]] name = "common-error" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-macro", "http 1.3.1", @@ -2346,7 +2346,7 @@ dependencies = [ [[package]] name = "common-event-recorder" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -2369,7 +2369,7 @@ dependencies = [ [[package]] name = "common-frontend" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -2390,7 +2390,7 @@ dependencies = [ [[package]] name = "common-function" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -2453,7 +2453,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "common-runtime", @@ -2470,7 +2470,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow-flight", @@ -2505,7 +2505,7 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "common-base", @@ -2525,7 +2525,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "1.1.3" +version = "1.1.4" dependencies = [ "greptime-proto", "once_cell", @@ -2536,7 +2536,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "1.1.3" +version = "1.1.4" dependencies = [ "anyhow", "common-error", @@ -2552,7 +2552,7 @@ dependencies = [ [[package]] name = "common-memory-manager" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-error", "common-macro", @@ -2564,7 +2564,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "1.1.3" +version = "1.1.4" dependencies = [ "anymap2", "api", @@ -2630,12 +2630,13 @@ dependencies = [ "tonic 0.14.2", "tracing", "typetag", + "url", "uuid", ] [[package]] name = "common-options" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-grpc", "humantime-serde", @@ -2645,11 +2646,11 @@ dependencies = [ [[package]] name = "common-plugins" -version = "1.1.3" +version = "1.1.4" [[package]] name = "common-pprof" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-error", "common-macro", @@ -2660,7 +2661,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-stream", @@ -2689,7 +2690,7 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "common-procedure", @@ -2699,7 +2700,7 @@ dependencies = [ [[package]] name = "common-query" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -2729,7 +2730,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arc-swap", "common-base", @@ -2754,7 +2755,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "clap", @@ -2783,7 +2784,7 @@ dependencies = [ [[package]] name = "common-session" -version = "1.1.3" +version = "1.1.4" dependencies = [ "serde", "strum 0.27.1", @@ -2791,7 +2792,7 @@ dependencies = [ [[package]] name = "common-sql" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arrow-schema 58.3.0", "common-base", @@ -2811,7 +2812,7 @@ dependencies = [ [[package]] name = "common-stat" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-base", "common-runtime", @@ -2826,7 +2827,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "1.1.3" +version = "1.1.4" dependencies = [ "backtrace", "common-base", @@ -2855,7 +2856,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "1.1.3" +version = "1.1.4" dependencies = [ "client", "common-grpc", @@ -2868,7 +2869,7 @@ dependencies = [ [[package]] name = "common-time" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arrow 58.3.0", "chrono", @@ -2886,7 +2887,7 @@ dependencies = [ [[package]] name = "common-version" -version = "1.1.3" +version = "1.1.4" dependencies = [ "cargo-manifest", "const_format", @@ -2896,7 +2897,7 @@ dependencies = [ [[package]] name = "common-wal" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-base", "common-error", @@ -2919,7 +2920,7 @@ dependencies = [ [[package]] name = "common-workload" -version = "1.1.3" +version = "1.1.4" dependencies = [ "common-telemetry", "serde", @@ -4360,7 +4361,7 @@ dependencies = [ [[package]] name = "datanode" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow-flight", @@ -4434,7 +4435,7 @@ checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" [[package]] name = "datatypes" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arrow 58.3.0", "arrow-array 58.3.0", @@ -5171,7 +5172,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "file-engine" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -5302,7 +5303,7 @@ checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" [[package]] name = "flow" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow 58.3.0", @@ -5371,7 +5372,7 @@ dependencies = [ "sql", "store-api", "strum 0.27.1", - "substrait 1.1.3", + "substrait 1.1.4", "table", "tokio", "tokio-stream", @@ -5454,7 +5455,7 @@ checksum = "28dd6caf6059519a65843af8fe2a3ae298b14b80179855aeb4adc2c1934ee619" [[package]] name = "frontend" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arc-swap", @@ -6744,7 +6745,7 @@ dependencies = [ [[package]] name = "index" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "asynchronous-codec", @@ -7843,7 +7844,7 @@ checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] name = "log-query" -version = "1.1.3" +version = "1.1.4" dependencies = [ "chrono", "common-error", @@ -7855,7 +7856,7 @@ dependencies = [ [[package]] name = "log-store" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-stream", "async-trait", @@ -8174,7 +8175,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -8206,7 +8207,7 @@ dependencies = [ [[package]] name = "meta-srv" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -8306,7 +8307,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "aquamarine", @@ -8407,7 +8408,7 @@ dependencies = [ [[package]] name = "mito-codec" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "bytes", @@ -8432,7 +8433,7 @@ dependencies = [ [[package]] name = "mito2" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "aquamarine", @@ -9221,7 +9222,7 @@ dependencies = [ [[package]] name = "object-store" -version = "1.1.3" +version = "1.1.4" dependencies = [ "anyhow", "async-trait", @@ -9776,7 +9777,7 @@ dependencies = [ [[package]] name = "operator" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -9835,7 +9836,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.1.3", + "substrait 1.1.4", "table", "tokio", "tokio-util", @@ -10120,7 +10121,7 @@ checksum = "e3c406c9e2aa74554e662d2c2ee11cd3e73756988800be7e6f5eddb16fed4699" [[package]] name = "partition" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "async-trait", @@ -10514,7 +10515,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -10671,7 +10672,7 @@ dependencies = [ [[package]] name = "plugins" -version = "1.1.3" +version = "1.1.4" dependencies = [ "auth", "catalog", @@ -11011,7 +11012,7 @@ dependencies = [ [[package]] name = "promql" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "async-trait", @@ -11363,7 +11364,7 @@ dependencies = [ [[package]] name = "puffin" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-compression", "async-trait", @@ -11425,7 +11426,7 @@ dependencies = [ [[package]] name = "query" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -11495,7 +11496,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.1.3", + "substrait 1.1.4", "table", "tokio", "tokio-stream", @@ -13067,7 +13068,7 @@ dependencies = [ [[package]] name = "servers" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -13205,7 +13206,7 @@ dependencies = [ [[package]] name = "session" -version = "1.1.3" +version = "1.1.4" dependencies = [ "ahash 0.8.12", "api", @@ -13558,7 +13559,7 @@ dependencies = [ [[package]] name = "sql" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow-buffer 58.3.0", @@ -13619,7 +13620,7 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "clap", @@ -13899,7 +13900,7 @@ dependencies = [ [[package]] name = "standalone" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "catalog", @@ -13943,7 +13944,7 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "store-api" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "aquamarine", @@ -14136,7 +14137,7 @@ dependencies = [ [[package]] name = "substrait" -version = "1.1.3" +version = "1.1.4" dependencies = [ "async-trait", "bytes", @@ -14258,7 +14259,7 @@ dependencies = [ [[package]] name = "table" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arc-swap", @@ -14541,7 +14542,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tests-fuzz" -version = "1.1.3" +version = "1.1.4" dependencies = [ "arbitrary", "async-trait", @@ -14586,7 +14587,7 @@ dependencies = [ [[package]] name = "tests-integration" -version = "1.1.3" +version = "1.1.4" dependencies = [ "api", "arrow-flight", @@ -14666,7 +14667,7 @@ dependencies = [ "sqlx", "standalone", "store-api", - "substrait 1.1.3", + "substrait 1.1.4", "table", "tempfile", "time", diff --git a/Cargo.toml b/Cargo.toml index fd6a533ed11..1124fea5bee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,7 +75,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.1.3" +version = "1.1.4" edition = "2024" license = "Apache-2.0" diff --git a/src/common/grpc/src/channel_manager.rs b/src/common/grpc/src/channel_manager.rs index 1be0e22d1a2..01da70558ab 100644 --- a/src/common/grpc/src/channel_manager.rs +++ b/src/common/grpc/src/channel_manager.rs @@ -383,6 +383,21 @@ pub struct ChannelConfig { pub accept_compression: bool, } +/// Configures a generated tonic client with compression and message size limits. +#[macro_export] +macro_rules! configure_tonic_client { + ($client:expr, $channel_manager:expr $(,)?) => {{ + let channel_manager = &$channel_manager; + let config = channel_manager.config(); + $client + .accept_compressed(::tonic::codec::CompressionEncoding::Gzip) + .accept_compressed(::tonic::codec::CompressionEncoding::Zstd) + .send_compressed(::tonic::codec::CompressionEncoding::Zstd) + .max_decoding_message_size(config.max_recv_message_size.as_bytes() as usize) + .max_encoding_message_size(config.max_send_message_size.as_bytes() as usize) + }}; +} + impl Default for ChannelConfig { fn default() -> Self { Self { @@ -710,6 +725,75 @@ mod tests { ); } + #[derive(Default)] + struct FakeTonicClient { + accepted_gzip: bool, + accepted_zstd: bool, + sent_gzip: bool, + sent_zstd: bool, + max_decoding_message_size: usize, + max_encoding_message_size: usize, + } + + impl FakeTonicClient { + fn accept_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self { + match encoding { + ::tonic::codec::CompressionEncoding::Gzip => self.accepted_gzip = true, + ::tonic::codec::CompressionEncoding::Zstd => self.accepted_zstd = true, + _ => unreachable!(), + } + self + } + + fn send_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self { + match encoding { + ::tonic::codec::CompressionEncoding::Gzip => self.sent_gzip = true, + ::tonic::codec::CompressionEncoding::Zstd => self.sent_zstd = true, + _ => unreachable!(), + } + self + } + + fn max_decoding_message_size(mut self, size: usize) -> Self { + self.max_decoding_message_size = size; + self + } + + fn max_encoding_message_size(mut self, size: usize) -> Self { + self.max_encoding_message_size = size; + self + } + } + + #[test] + fn test_configure_tonic_client() { + let recv_message_size = ReadableSize::mb(2); + let send_message_size = ReadableSize::mb(3); + let channel_manager = ChannelManager::with_config( + ChannelConfig { + max_recv_message_size: recv_message_size, + max_send_message_size: send_message_size, + ..ChannelConfig::new() + }, + None, + ); + + let client = crate::configure_tonic_client!(FakeTonicClient::default(), channel_manager,); + + assert!(client.accepted_gzip); + assert!(client.accepted_zstd); + assert!(!client.sent_gzip); + assert!(client.sent_zstd); + assert_eq!( + recv_message_size.as_bytes() as usize, + client.max_decoding_message_size + ); + assert_eq!( + send_message_size.as_bytes() as usize, + client.max_encoding_message_size + ); + } + #[test] fn test_build_endpoint() { let config = ChannelConfig::new() diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 5134c40a804..dba0ec2d66a 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -92,6 +92,8 @@ tokio-postgres-rustls = { version = "0.12", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true +url.workspace = true +uuid.workspace = true [dev-dependencies] chrono.workspace = true diff --git a/src/common/meta/src/kv_backend/util.rs b/src/common/meta/src/kv_backend/util.rs index 1021d78a60b..efe85367a5f 100644 --- a/src/common/meta/src/kv_backend/util.rs +++ b/src/common/meta/src/kv_backend/util.rs @@ -12,46 +12,180 @@ // See the License for the specific language governing permissions and // limitations under the License. -/// Removes sensitive information like passwords from connection strings. +use url::Url; + +/// Placeholder substituted for any password found in a connection string. +const REDACTED: &str = "***"; + +/// Substituted for a whole URI-shaped connection string that no parser accepted. +const REDACTED_URI: &str = ""; + +/// Removes sensitive information (passwords) from a connection string before it +/// is logged. /// -/// This function sanitizes connection strings by removing credentials: -/// - For URL format (mysql://user:password@host:port/db): Removes everything before '@' -/// - For parameter format (host=localhost password=secret): Removes the password parameter -/// - For URL format without credentials (mysql://host:port/db): Removes the protocol prefix +/// `store_addrs` entries can be a PostgreSQL DSN (a `postgres(ql)://` URL or a +/// libpq keyword string), a MySQL/etcd URL, or an etcd host list. /// -/// # Arguments -/// -/// * `conn_str` - The connection string to sanitize -/// -/// # Returns -/// -/// A sanitized version of the connection string with sensitive information removed +/// The function fails closed: it never returns a URI-shaped input verbatim. In +/// order: +/// 1. PostgreSQL DSNs are parsed with `tokio_postgres::Config` — the backend's +/// own parser — and logged via its `Debug`, which redacts the password. This +/// matches the grammar exactly (multi-host URIs, `\`-escapes, any Unicode +/// whitespace, percent-encoded query keys, `&`/`;`/`://` inside values). +/// 2. Any other URL is redacted with the `url` crate. +/// 3. A URI-shaped input that neither parser accepted (e.g. a malformed +/// multi-host Postgres URI, with or without leading whitespace) is redacted +/// in full — credentials in an unparsable authority or query cannot be split +/// out reliably, so the whole string is replaced rather than reparsed. +/// 4. Anything else is treated as a keyword string and best-effort redacted. pub fn sanitize_connection_string(conn_str: &str) -> String { - // Case 1: URL format with credentials (mysql://user:password@host:port/db) - // Extract everything after the '@' symbol - if let Some(at_pos) = conn_str.find('@') { - return conn_str[at_pos + 1..].to_string(); + // Redact explicit keyword-style password assignments before parsing. A + // malformed string such as `host=x;password=secret` can otherwise parse as + // a host value, causing `Config`'s Debug output to retain the password text. + if !is_uri_like(conn_str) { + let redacted = redact_keyword_password(conn_str); + if redacted != conn_str { + return redacted; + } } - // Case 2: Parameter format with password (host=localhost password=secret dbname=mydb) - // Filter out any parameter that starts with "password=" - if conn_str.contains("password=") { - return conn_str - .split_whitespace() - .filter(|param| !param.starts_with("password=")) - .collect::>() - .join(" "); + #[cfg(feature = "pg_kvbackend")] + if let Ok(config) = conn_str.parse::() { + // `Config`'s Debug prints the password as `Some("REDACTED")`. + return format!("{config:?}"); } - // Case 3: URL format without credentials (mysql://host:port/db) - // Extract everything after the protocol prefix - if let Some(host_part) = conn_str.split("://").nth(1) { - return host_part.to_string(); + if let Ok(url) = Url::parse(conn_str) { + return redact_url(url); } - // Case 4: Already sanitized or unknown format - // Return as is - conn_str.to_string() + if is_uri_like(conn_str) { + return REDACTED_URI.to_string(); + } + + redact_keyword_password(conn_str) +} + +/// Redacts the userinfo password and any `password` query parameter of a parsed URL. +fn redact_url(mut url: Url) -> String { + if url.password().is_some() { + let _ = url.set_password(Some(REDACTED)); + } + if url.query_pairs().any(|(key, _)| key == "password") { + let redacted: Vec<(String, String)> = url + .query_pairs() + .map(|(key, value)| { + let value = if key == "password" { + REDACTED.to_string() + } else { + value.into_owned() + }; + (key.into_owned(), value) + }) + .collect(); + url.query_pairs_mut().clear().extend_pairs(redacted); + } + url.to_string() +} + +/// True if `s`, ignoring leading whitespace, begins with a valid URI scheme +/// followed by `://` (`scheme = ALPHA *( ALPHA / DIGIT / "+" / "-" / "." )`). +/// Leading whitespace (including Unicode, e.g. U+2003) does not stop a value from +/// being a URI, so it is trimmed first; a bare `contains("://")` would instead +/// misfire on a `://` inside a keyword value. +fn is_uri_like(s: &str) -> bool { + let s = s.trim_start(); + let Some(pos) = s.find("://") else { + return false; + }; + let mut chars = s[..pos].chars(); + matches!(chars.next(), Some(c) if c.is_ascii_alphabetic()) + && chars.all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.')) +} + +/// Best-effort redaction of a `password` value in a keyword string, used only +/// for inputs that neither `tokio_postgres::Config` nor [`Url`] accept (e.g. a +/// libpq keyword DSN when `pg_kvbackend` is disabled). Keys are separated by +/// Unicode whitespace; a value is single/double quoted or runs to the next +/// unescaped whitespace, with `\` escaping the next character. +fn redact_keyword_password(s: &str) -> String { + const KEY: &str = "password"; + let mut out = String::with_capacity(s.len()); + let mut i = 0; + while i < s.len() { + if is_keyword_at(s, i, KEY) { + let mut j = skip_ws(s, i + KEY.len()); + if s[j..].starts_with('=') { + j = skip_ws(s, j + 1); + out.push_str("password="); + out.push_str(REDACTED); + i = skip_value(s, j); + continue; + } + } + let ch = s[i..].chars().next().unwrap(); + out.push(ch); + i += ch.len_utf8(); + } + out +} + +fn skip_ws(s: &str, mut i: usize) -> usize { + while let Some(ch) = s[i..].chars().next() { + if ch.is_whitespace() { + i += ch.len_utf8(); + } else { + break; + } + } + i +} + +/// True if `keyword` starts at byte `i` preceded by a keyword boundary. In +/// addition to valid libpq whitespace separators, `;` and `&` are treated as +/// conservative boundaries so malformed DSNs cannot leak credentials. `=` is +/// deliberately excluded, so a value like `application_name=password=x` is not +/// mistaken for a password key. +fn is_keyword_at(s: &str, i: usize, keyword: &str) -> bool { + s[i..].starts_with(keyword) + && (i == 0 + || s[..i] + .chars() + .next_back() + .is_some_and(|c| c.is_whitespace() || matches!(c, ';' | '&'))) +} + +/// Returns the byte index just past a keyword value beginning at `start`. +fn skip_value(s: &str, start: usize) -> usize { + if start >= s.len() { + return start; + } + let quote = s[start..].chars().next().unwrap(); + let quoted = quote == '\'' || quote == '"'; + let mut i = if quoted { + start + quote.len_utf8() + } else { + start + }; + while i < s.len() { + let ch = s[i..].chars().next().unwrap(); + if ch == '\\' { + i += ch.len_utf8(); + if let Some(next) = s[i..].chars().next() { + i += next.len_utf8(); // skip the escaped character + } + continue; + } + if quoted { + if ch == quote { + return i + ch.len_utf8(); // include the closing quote + } + } else if ch.is_whitespace() { + return i; + } + i += ch.len_utf8(); + } + s.len() } #[cfg(test)] @@ -59,27 +193,143 @@ mod tests { use super::*; #[test] - fn test_sanitize_connection_string() { - // Test URL format with username/password - let conn_str = "mysql://user:password123@localhost:3306/db"; - assert_eq!(sanitize_connection_string(conn_str), "localhost:3306/db"); + fn test_sanitize_non_pg_url() { + let s = sanitize_connection_string("mysql://user:password123@localhost:3306/db"); + assert!(!s.contains("password123"), "{s}"); + assert!(s.contains("user:***@localhost:3306"), "{s}"); - // Test URL format without credentials - let conn_str = "mysql://localhost:3306/db"; - assert_eq!(sanitize_connection_string(conn_str), "localhost:3306/db"); + let s = sanitize_connection_string("mysql://localhost:3306/db"); + assert!(s.contains("localhost:3306"), "{s}"); - // Test parameter format with password - let conn_str = "host=localhost port=5432 user=postgres password=secret dbname=mydb"; - assert_eq!( - sanitize_connection_string(conn_str), - "host=localhost port=5432 user=postgres dbname=mydb" - ); + let s = sanitize_connection_string("mysql://user@localhost:3306/db?password=secret"); + assert!(!s.contains("secret"), "{s}"); - // Test parameter format without password - let conn_str = "host=localhost port=5432 user=postgres dbname=mydb"; - assert_eq!( - sanitize_connection_string(conn_str), - "host=localhost port=5432 user=postgres dbname=mydb" - ); + let s = sanitize_connection_string("http://etcd-host:2379"); + assert!(s.contains("etcd-host:2379"), "{s}"); + } + + // Fail-closed: a URI-shaped string that neither parser accepts is redacted in + // full. These multi-host URIs are rejected by both `url` (InvalidPort) and, + // when compiled, `tokio_postgres::Config` (unknown option), so the invariant + // holds regardless of the `pg_kvbackend` feature — including with leading + // (Unicode) whitespace and percent-encoded query keys. + #[test] + fn test_sanitize_fail_closed_uri() { + for dsn in [ + "postgresql://user:secret@host1:1234,host2:5678/db?unknown=x", + " postgresql://user:secret@host1:1234,host2:5678/db?unknown=x", + "\u{2003}postgresql://user:secret@host1:1234,host2:5678/db?unknown=x", + "postgresql://user@host1:1,host2:2/db?pass%77ord=qsecret&unknown=x", + ] { + let s = sanitize_connection_string(dsn); + // "secret" also matches "qsecret", so this covers the query case too. + assert!(!s.contains("secret"), "leaked for {dsn:?}: {s}"); + // Redacted whole, not reparsed: the authority is gone as well. + assert!(!s.contains("host1"), "not fully redacted for {dsn:?}: {s}"); + } + } + + #[test] + fn test_sanitize_malformed_keyword_fail_closed() { + for dsn in [ + "host=localhost;password=LEAK_CANARY", + "host=localhost&password=LEAK_CANARY", + ] { + let s = sanitize_connection_string(dsn); + assert!(!s.contains("LEAK_CANARY"), "leaked for {dsn:?}: {s}"); + } + + for dsn in [ + "notpassword=LEAK_CANARY", + "application_name=password=LEAK_CANARY", + ] { + let s = sanitize_connection_string(dsn); + assert!(s.contains("LEAK_CANARY"), "over-redacted for {dsn:?}: {s}"); + } + } + + // Keyword redaction must handle libpq quoting and Unicode whitespace without + // misfiring on `=`/`://` inside values. Exact output is asserted without the + // pg feature; feature-enabled coverage below asserts the no-leak invariant. + #[cfg(not(feature = "pg_kvbackend"))] + #[test] + fn test_sanitize_keyword() { + for (dsn, ok) in [ + ( + "host=localhost password=secret dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + "host=localhost password = secret dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + "host=localhost password='my secret' dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + "host=localhost password=\"my secret\" dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + r"host=localhost password='pa\'ss' dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + "host=localhost user=password dbname=x", + "host=localhost user=password dbname=x", + ), + // A value that merely contains `://` is not treated as a URI. + ( + "host=localhost password=secret application_name=svc://a", + "host=localhost password=*** application_name=svc://a", + ), + // `&` and a `\`-escaped space are part of an unquoted libpq value, so + // the whole value is consumed (no trailing suffix leaks). + ( + "host=localhost password=secret&suffix dbname=x", + "host=localhost password=*** dbname=x", + ), + ( + r"host=localhost password=secret\ suffix dbname=x", + "host=localhost password=*** dbname=x", + ), + // `;` is not a libpq delimiter, so it belongs to the value too. + ("password=secret;host=localhost", "password=***"), + ] { + assert_eq!(sanitize_connection_string(dsn), ok, "for {dsn:?}"); + } + + // Unicode whitespace as separator / around '=' must not leak. + for dsn in [ + "host=localhost\u{2003}password=secret", + "host=localhost password\u{2003}=\u{2003}secret", + ] { + assert!( + !sanitize_connection_string(dsn).contains("secret"), + "leaked for {dsn:?}" + ); + } + } + + // With the pg backend enabled, cover both the conservative keyword pre-pass + // and URI forms parsed by tokio-postgres. + #[cfg(feature = "pg_kvbackend")] + #[test] + fn test_sanitize_pg_dsn() { + for dsn in [ + "postgresql://user:secret@localhost:5432/db", + "postgresql://user:secret@host1:1234,host2,host3:5678/db", + "postgresql://user@host1:1,host2/db?pass%77ord=secret", + "host=localhost port=5432 user=postgres password=secret dbname=mydb", + "host=localhost\u{2003}password=secret", + "password=secret;host=localhost", + "host=localhost password=secret application_name=http://client", + ] { + assert!( + !sanitize_connection_string(dsn).contains("secret"), + "leaked for {dsn:?}" + ); + } } } diff --git a/src/common/time/src/date.rs b/src/common/time/src/date.rs index bda9105f559..38667abbcac 100644 --- a/src/common/time/src/date.rs +++ b/src/common/time/src/date.rs @@ -198,6 +198,10 @@ impl Date { pub fn negative(&self) -> Self { Self(-self.0) } + + pub fn checked_negative(&self) -> Option { + self.0.checked_neg().map(Self) + } } #[cfg(test)] diff --git a/src/common/time/src/duration.rs b/src/common/time/src/duration.rs index a1e582ca82c..b34ced336bd 100644 --- a/src/common/time/src/duration.rs +++ b/src/common/time/src/duration.rs @@ -97,6 +97,11 @@ impl Duration { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } /// Convert i64 to Duration Type. diff --git a/src/common/time/src/interval.rs b/src/common/time/src/interval.rs index 4cb8c45c55f..21bdc656d88 100644 --- a/src/common/time/src/interval.rs +++ b/src/common/time/src/interval.rs @@ -80,6 +80,10 @@ impl IntervalYearMonth { Self::new(-self.months) } + pub fn checked_negative(&self) -> Option { + self.months.checked_neg().map(Self::new) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -157,6 +161,13 @@ impl IntervalDayTime { Self::new(-self.days, -self.milliseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.days.checked_neg()?, + self.milliseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -278,6 +289,14 @@ impl IntervalMonthDayNano { Self::new(-self.months, -self.days, -self.nanoseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.months.checked_neg()?, + self.days.checked_neg()?, + self.nanoseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } diff --git a/src/common/time/src/time.rs b/src/common/time/src/time.rs index 00b123f6607..a715e2efec0 100644 --- a/src/common/time/src/time.rs +++ b/src/common/time/src/time.rs @@ -145,6 +145,11 @@ impl Time { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl From for Time { diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 120dc8e1e97..8ad38e280fe 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -484,6 +484,11 @@ impl Timestamp { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl Timestamp { diff --git a/src/datanode/src/region_server/registrations.rs b/src/datanode/src/region_server/registrations.rs index 92a1bb813c1..6d7ebc614d7 100644 --- a/src/datanode/src/region_server/registrations.rs +++ b/src/datanode/src/region_server/registrations.rs @@ -21,7 +21,7 @@ use common_query::request::{ DynFilterPayload, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg, InitialDynFilterRegs, }; -use common_telemetry::warn; +use common_telemetry::{debug, warn}; use dashmap::DashMap; use datafusion::arrow::datatypes::{Schema, SchemaRef}; use datafusion::execution::{SessionStateBuilder, TaskContext}; @@ -505,7 +505,7 @@ pub(super) fn apply_remote_dyn_filter_update( } let Some(query_regs) = regs_by_query.get_query(query_id) else { - warn!( + debug!( "Ignored remote dynamic filter update without query registration, query_id: {}, filter_id: {}", query_id, filter_id ); @@ -514,7 +514,7 @@ pub(super) fn apply_remote_dyn_filter_update( let mut query_regs = query_regs.lock().unwrap(); let Some(registered) = query_regs.get_mut(filter_id) else { - warn!( + debug!( "Ignored remote dynamic filter update without filter registration, query_id: {}, filter_id: {}", query_id, filter_id ); @@ -530,7 +530,7 @@ pub(super) fn unregister_remote_dyn_filter( filter_id: &RemoteDynFilterId, ) -> RemoteDynFilterUpdateOutcome { let Some(query_regs) = regs_by_query.get_query(query_id) else { - warn!( + debug!( "Ignored remote dynamic filter unregister without query registration, query_id: {}, filter_id: {}", query_id, filter_id ); @@ -540,7 +540,7 @@ pub(super) fn unregister_remote_dyn_filter( let (registered, should_remove_query) = { let mut locked = query_regs.lock().unwrap(); let Some(registered) = locked.remove(filter_id) else { - warn!( + debug!( "Ignored remote dynamic filter unregister without filter registration, query_id: {}, filter_id: {}", query_id, filter_id ); diff --git a/src/datanode/src/region_server/remote_dyn_filter.rs b/src/datanode/src/region_server/remote_dyn_filter.rs index f750cb896cd..cd3491e97aa 100644 --- a/src/datanode/src/region_server/remote_dyn_filter.rs +++ b/src/datanode/src/region_server/remote_dyn_filter.rs @@ -225,8 +225,7 @@ impl RegionServer { ) { if matches!( outcome, - RemoteDynFilterUpdateOutcome::MissingRegistration - | RemoteDynFilterUpdateOutcome::AlreadyComplete + RemoteDynFilterUpdateOutcome::AlreadyComplete | RemoteDynFilterUpdateOutcome::PayloadTooLarge | RemoteDynFilterUpdateOutcome::DecodeFailed ) { diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 5d703480d3f..9d50ac78af4 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -543,20 +543,20 @@ impl Value { None } } - Value::Int8(x) => Some(Value::Int8(-*x)), - Value::Int16(x) => Some(Value::Int16(-*x)), - Value::Int32(x) => Some(Value::Int32(-*x)), - Value::Int64(x) => Some(Value::Int64(-*x)), + Value::Int8(x) => x.checked_neg().map(Value::Int8), + Value::Int16(x) => x.checked_neg().map(Value::Int16), + Value::Int32(x) => x.checked_neg().map(Value::Int32), + Value::Int64(x) => x.checked_neg().map(Value::Int64), Value::Float32(x) => Some(Value::Float32(-*x)), Value::Float64(x) => Some(Value::Float64(-*x)), Value::Decimal128(x) => Some(Value::Decimal128(x.negative())), - Value::Date(x) => Some(Value::Date(x.negative())), - Value::Timestamp(x) => Some(Value::Timestamp(x.negative())), - Value::Time(x) => Some(Value::Time(x.negative())), - Value::Duration(x) => Some(Value::Duration(x.negative())), - Value::IntervalYearMonth(x) => Some(Value::IntervalYearMonth(x.negative())), - Value::IntervalDayTime(x) => Some(Value::IntervalDayTime(x.negative())), - Value::IntervalMonthDayNano(x) => Some(Value::IntervalMonthDayNano(x.negative())), + Value::Date(x) => x.checked_negative().map(Value::Date), + Value::Timestamp(x) => x.checked_negative().map(Value::Timestamp), + Value::Time(x) => x.checked_negative().map(Value::Time), + Value::Duration(x) => x.checked_negative().map(Value::Duration), + Value::IntervalYearMonth(x) => x.checked_negative().map(Value::IntervalYearMonth), + Value::IntervalDayTime(x) => x.checked_negative().map(Value::IntervalDayTime), + Value::IntervalMonthDayNano(x) => x.checked_negative().map(Value::IntervalMonthDayNano), Value::Binary(_) | Value::String(_) @@ -1744,6 +1744,49 @@ pub(crate) mod tests { use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::vectors::ListVectorBuilder; + #[test] + fn test_try_negative_overflow() { + // Negating a MIN value overflows, so try_negative returns None instead + // of panicking (consistent with the unsigned arms). + assert_eq!(Value::Int8(i8::MIN).try_negative(), None); + assert_eq!(Value::Int16(i16::MIN).try_negative(), None); + assert_eq!(Value::Int32(i32::MIN).try_negative(), None); + assert_eq!(Value::Int64(i64::MIN).try_negative(), None); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!(Value::Date(Date::new(i32::MIN)).try_negative(), None); + assert_eq!( + Value::Time(Time::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::Duration(Duration::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalYearMonth(IntervalYearMonth::new(i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalDayTime(IntervalDayTime::new(i32::MIN, i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalMonthDayNano(IntervalMonthDayNano::new(i32::MIN, i32::MIN, i64::MIN)) + .try_negative(), + None + ); + + // Non-MIN values still negate. + assert_eq!(Value::Int64(5).try_negative(), Some(Value::Int64(-5))); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(5)).try_negative(), + Some(Value::Timestamp(Timestamp::new_nanosecond(-5))) + ); + } + pub(crate) fn build_struct_type() -> StructType { StructType::new(Arc::new(vec![ StructField::new("id".to_string(), ConcreteDataType::int32_datatype(), false), diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index 46c01b81847..1843903a40a 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -76,6 +76,23 @@ pub(crate) mod table_source; use crate::FrontendInvoker; use crate::error::Error; +fn expire_after_secs_to_millis(expire_after_secs: i64) -> Result { + ensure!( + expire_after_secs >= 0, + InvalidQuerySnafu { + reason: format!("EXPIRE AFTER must be non-negative, got {expire_after_secs} seconds"), + } + ); + + expire_after_secs + .checked_mul(1_000) + .with_context(|| InvalidQuerySnafu { + reason: format!( + "EXPIRE AFTER value {expire_after_secs} seconds cannot be represented in milliseconds" + ), + }) +} + // `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder"; @@ -739,7 +756,7 @@ impl StreamingEngine { source_table_ids, create_if_not_exists, or_replace, - expire_after, + expire_after: expire_after_secs, eval_interval: _, comment, sql, @@ -747,6 +764,9 @@ impl StreamingEngine { query_ctx, .. } = args; + let expire_after = expire_after_secs + .map(expire_after_secs_to_millis) + .transpose()?; let mut node_ctx = self.node_context.write().await; // assign global id to source and sink table diff --git a/src/flow/src/adapter/refill.rs b/src/flow/src/adapter/refill.rs index 6d66505e89e..3672c0427e1 100644 --- a/src/flow/src/adapter/refill.rs +++ b/src/flow/src/adapter/refill.rs @@ -85,7 +85,10 @@ impl StreamingEngine { } } - let expire_after = info.expire_after(); + let expire_after = info + .expire_after() + .map(super::expire_after_secs_to_millis) + .transpose()?; // TODO(discord9): better way to get last point let now = self.tick_manager.tick(); let plan = self diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 9ed7e676c29..df6a58c46ce 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -61,3 +61,25 @@ pub fn new_test_table_info_with_name>( /// /// containing several default table info and schema fn mock_harness_flow_node_manager() {} + +#[test] +fn test_expire_after_secs_to_millis() { + assert_eq!(expire_after_secs_to_millis(300).unwrap(), 300_000); + assert_eq!(expire_after_secs_to_millis(0).unwrap(), 0); + + let max_secs = i64::MAX / 1_000; + assert_eq!( + expire_after_secs_to_millis(max_secs).unwrap(), + max_secs * 1_000 + ); +} + +#[test] +fn test_expire_after_secs_to_millis_invalid() { + let invalid_values = [i64::MAX / 1_000 + 1, -1]; + + for invalid_secs in invalid_values { + let error = expire_after_secs_to_millis(invalid_secs).unwrap_err(); + assert!(matches!(error, Error::InvalidQuery { .. })); + } +} diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index 6401adc07e4..b820a0c43f0 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -19,7 +19,6 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::time::Duration; use common_telemetry::debug; -use common_telemetry::tracing::warn; use common_time::Timestamp; use datatypes::value::Value; use session::context::QueryContextRef; @@ -609,7 +608,7 @@ impl DirtyTimeWindows { let last_time_window = self.windows.last_key_value(); if let Some(task_ctx) = task_ctx { - warn!( + debug!( "Flow id = {:?}, too many time windows: {}, only the first {} are taken for this query, the group by expression might be wrong. Time window expr={:?}, expire_after={:?}, first_time_window={:?}, last_time_window={:?}, the original query: {:?}", task_ctx.config.flow_id, self.windows.len(), @@ -621,7 +620,7 @@ impl DirtyTimeWindows { task_ctx.config.query ); } else { - warn!( + debug!( "Flow id = {:?}, too many time windows: {}, only the first {} are taken for this query, the group by expression might be wrong. first_time_window={:?}, last_time_window={:?}", flow_id, self.windows.len(), diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 089adbb0edd..a8c916fd8e0 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -695,7 +695,7 @@ impl BatchingTask { } } Err(err) => { - warn!( + debug!( "Failed to unparse full-snapshot SQL flow {} plan; \ falling back to InsertIntoPlan: {:?}", flow_id, err diff --git a/src/flow/src/batching_mode/task/ckpt.rs b/src/flow/src/batching_mode/task/ckpt.rs index 7966d314e13..c736a0fdbdf 100644 --- a/src/flow/src/batching_mode/task/ckpt.rs +++ b/src/flow/src/batching_mode/task/ckpt.rs @@ -279,6 +279,15 @@ impl BatchingTask { "Flow {flow_id} continued fenced repair, pending_windows={pending_windows}, watermarks={watermarks}" ); } + FlowCheckpointDecision::FallbackToFullSnapshot { + previous_mode: CheckpointMode::FullSnapshot, + reason: FlowQueryFallbackReason::IncrementalDisabled, + } => { + debug!( + "Flow {flow_id} remains in full snapshot mode, reason={}", + FlowQueryFallbackReason::IncrementalDisabled.as_label() + ); + } FlowCheckpointDecision::FallbackToFullSnapshot { previous_mode, reason, diff --git a/src/flow/src/engine.rs b/src/flow/src/engine.rs index 386610ad716..3a139a898a9 100644 --- a/src/flow/src/engine.rs +++ b/src/flow/src/engine.rs @@ -34,6 +34,7 @@ pub struct CreateFlowArgs { pub source_table_ids: Vec, pub create_if_not_exists: bool, pub or_replace: bool, + /// Duration in seconds. pub expire_after: Option, pub eval_interval: Option, pub comment: Option, diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index de41caf19eb..4fc6020fffb 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -891,9 +891,11 @@ impl MetaClient { #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; use api::v1::meta::{HeartbeatRequest, Peer}; - use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; + use common_base::readable_size::ReadableSize; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackendRef, TxnService}; use rand::Rng; use super::*; @@ -919,6 +921,27 @@ mod tests { } } + async fn new_with_grpc_message_sizes( + ns: impl Into, + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, + ) -> Self { + let (client, meta_ctx) = mocks::mock_client_with_memstore_and_grpc_message_sizes( + server_max_recv_message_size, + server_max_send_message_size, + client_max_recv_message_size, + client_max_send_message_size, + ) + .await; + Self { + ns: ns.into(), + client, + meta_ctx, + } + } + fn key(&self, name: &str) -> Vec { format!("{}-{}-{}", TEST_KEY_PREFIX, self.ns, name).into_bytes() } @@ -1331,20 +1354,77 @@ mod tests { } } - fn mock_decoder(_kv: KeyValue) -> MetaResult<()> { - Ok(()) + fn mock_decoder(kv: KeyValue) -> MetaResult> { + Ok(kv.value) } - #[tokio::test] - async fn test_cluster_client_adaptive_range() { - let tx = new_client("test_cluster_client").await; + struct RecordingKvBackend { + inner: KvBackendRef, + limits: Arc>>, + } + + #[async_trait::async_trait] + impl TxnService for RecordingKvBackend { + type Error = meta_error::Error; + } + + #[async_trait::async_trait] + impl KvBackend for RecordingKvBackend { + fn name(&self) -> &str { + "RecordingKvBackend" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn range(&self, req: RangeRequest) -> MetaResult { + self.limits.lock().unwrap().push(req.limit); + self.inner.range(req).await + } + + async fn put(&self, req: PutRequest) -> MetaResult { + self.inner.put(req).await + } + + async fn batch_put(&self, req: BatchPutRequest) -> MetaResult { + self.inner.batch_put(req).await + } + + async fn batch_get(&self, req: BatchGetRequest) -> MetaResult { + self.inner.batch_get(req).await + } + + async fn delete_range(&self, req: DeleteRangeRequest) -> MetaResult { + self.inner.delete_range(req).await + } + + async fn batch_delete(&self, req: BatchDeleteRequest) -> MetaResult { + self.inner.batch_delete(req).await + } + } + + async fn adaptive_range_with_message_sizes( + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, + ) -> (Vec>, Vec>, Vec) { + let tx = TestClient::new_with_grpc_message_sizes( + "test_cluster_client", + server_max_recv_message_size, + server_max_send_message_size, + client_max_recv_message_size, + client_max_send_message_size, + ) + .await; let in_memory = tx.in_memory().unwrap(); let cluster_client = tx.client.cluster_client().unwrap(); let mut rng = rand::rng(); - // Generates rough 10MB data, which is larger than the default grpc message size limit. - for i in 0..10 { - let data: Vec = (0..1024 * 1024).map(|_| rng.random::()).collect(); + let mut expected = Vec::new(); + for i in 0..4 { + let data: Vec = (0..256 * 1024).map(|_| rng.random::()).collect(); in_memory .put( PutRequest::new() @@ -1353,13 +1433,48 @@ mod tests { ) .await .unwrap(); + expected.push(data); } let req = RangeRequest::new().with_prefix(b"__prefix/"); + let limits = Arc::new(Mutex::new(Vec::new())); + let recording_backend = RecordingKvBackend { + inner: Arc::new(cluster_client), + limits: limits.clone(), + }; let stream = - PaginationStream::new(Arc::new(cluster_client), req, 10, mock_decoder).into_stream(); + PaginationStream::new(Arc::new(recording_backend), req, 4, mock_decoder).into_stream(); let res = stream.try_collect::>().await.unwrap(); - assert_eq!(10, res.len()); + let limits = limits.lock().unwrap().clone(); + (expected, res, limits) + } + + #[tokio::test] + async fn test_cluster_client_adaptive_range() { + let (expected, res, limits) = adaptive_range_with_message_sizes( + ReadableSize::mb(2), + ReadableSize::mb(16), + ReadableSize::mb(1), + ReadableSize::mb(2), + ) + .await; + + assert_eq!(expected, res); + assert_eq!(vec![4, 2, 2], limits); + } + + #[tokio::test] + async fn test_cluster_client_adaptive_range_server_limit() { + let (expected, res, limits) = adaptive_range_with_message_sizes( + ReadableSize::mb(2), + ReadableSize::mb(1), + ReadableSize::mb(16), + ReadableSize::mb(2), + ) + .await; + + assert_eq!(expected, res); + assert_eq!(vec![4, 2, 2], limits); } } diff --git a/src/meta-client/src/client/ask_leader.rs b/src/meta-client/src/client/ask_leader.rs index 9a41559e14b..ff665ef9d40 100644 --- a/src/meta-client/src/client/ask_leader.rs +++ b/src/meta-client/src/client/ask_leader.rs @@ -186,10 +186,13 @@ impl AskLeader { } fn create_asker(&self, addr: impl AsRef) -> Result> { - Ok(HeartbeatClient::new( - self.channel_manager - .get(addr) - .context(error::CreateChannelSnafu)?, + Ok(common_grpc::configure_tonic_client!( + HeartbeatClient::new( + self.channel_manager + .get(addr) + .context(error::CreateChannelSnafu)?, + ), + self.channel_manager, )) } diff --git a/src/meta-client/src/client/cluster.rs b/src/meta-client/src/client/cluster.rs index 976626b3a05..56d745284da 100644 --- a/src/meta-client/src/client/cluster.rs +++ b/src/meta-client/src/client/cluster.rs @@ -34,7 +34,6 @@ use common_telemetry::{error, info, warn}; use snafu::{ResultExt, ensure}; use tokio::sync::RwLock; use tonic::Status; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{LeaderProviderRef, util}; @@ -155,10 +154,10 @@ impl Inner { fn make_client(&self, addr: impl AsRef) -> Result> { let channel = self.channel_manager.get(addr).context(CreateChannelSnafu)?; - Ok(ClusterClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ClusterClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/config.rs b/src/meta-client/src/client/config.rs index a2bcb009041..bd81066681f 100644 --- a/src/meta-client/src/client/config.rs +++ b/src/meta-client/src/client/config.rs @@ -21,7 +21,6 @@ use common_meta::util; use common_telemetry::tracing_context::TracingContext; use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, LeaderProviderRef}; @@ -132,10 +131,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - Ok(ConfigClient::new(channel) - .accept_compressed(CompressionEncoding::Zstd) - .accept_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ConfigClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index 4cdc2cf0f6b..da0e7b3837a 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -27,7 +27,6 @@ use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::{RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; use tonic::Streaming; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, LeaderProviderRef}; @@ -280,10 +279,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - Ok(HeartbeatClient::new(channel) - .accept_compressed(CompressionEncoding::Zstd) - .accept_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + HeartbeatClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index 7658653e135..883ae936dd2 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -32,7 +32,6 @@ use common_telemetry::tracing_context::TracingContext; use common_telemetry::{error, info, warn}; use snafu::{ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use tonic::{Request, Status}; @@ -151,10 +150,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - Ok(ProcedureServiceClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ProcedureServiceClient::new(channel), + self.channel_manager, + )) } #[inline] @@ -576,10 +575,14 @@ mod tests { let server = tonic::transport::Server::builder() .add_service( - HeartbeatServer::new(heartbeat).accept_compressed(CompressionEncoding::Zstd), + HeartbeatServer::new(heartbeat) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd), ) .add_service( - ProcedureServiceServer::new(procedure).accept_compressed(CompressionEncoding::Zstd), + ProcedureServiceServer::new(procedure) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd), ) .serve_with_incoming(TcpListenerStream::new(listener)); let server_handle = tokio::spawn(server); diff --git a/src/meta-client/src/client/store.rs b/src/meta-client/src/client/store.rs index 2aee8d83231..3f9bb98862f 100644 --- a/src/meta-client/src/client/store.rs +++ b/src/meta-client/src/client/store.rs @@ -25,7 +25,6 @@ use common_grpc::channel_manager::ChannelManager; use common_telemetry::tracing_context::TracingContext; use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, load_balance as lb}; @@ -275,17 +274,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - let max_decoding_message_size = self - .channel_manager - .config() - .max_recv_message_size - .as_bytes() as usize; - - Ok(StoreClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd) - .max_decoding_message_size(max_decoding_message_size)) + Ok(common_grpc::configure_tonic_client!( + StoreClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 19577f672cc..aa708eec72f 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -29,7 +29,7 @@ pub enum Error { location: Location, }, - #[snafu(display("{}", msg))] + #[snafu(display("{}, code: {}, tonic code: {}", msg, code, tonic_code))] MetaServer { code: StatusCode, msg: String, @@ -163,7 +163,7 @@ impl Error { matches!( self, Error::MetaServer { - tonic_code: tonic::Code::OutOfRange, + tonic_code: tonic::Code::OutOfRange | tonic::Code::ResourceExhausted, .. } ) @@ -171,3 +171,32 @@ impl Error { } define_from_tonic_status!(Error, MetaServer); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_exceeded_size_limit_for_out_of_range() { + let err = Error::from(tonic::Status::new(tonic::Code::OutOfRange, "any message")); + + assert!(err.is_exceeded_size_limit()); + } + + #[test] + fn test_is_exceeded_size_limit_for_resource_exhausted() { + let err = Error::from(tonic::Status::new( + tonic::Code::ResourceExhausted, + "arbitrary message", + )); + + assert!(err.is_exceeded_size_limit()); + } + + #[test] + fn test_is_exceeded_size_limit_for_non_size_code() { + let err = Error::from(tonic::Status::new(tonic::Code::Internal, "message")); + + assert!(!err.is_exceeded_size_limit()); + } +} diff --git a/src/meta-client/src/mocks.rs b/src/meta-client/src/mocks.rs index 2c80091c587..6e1daa2efc0 100644 --- a/src/meta-client/src/mocks.rs +++ b/src/meta-client/src/mocks.rs @@ -12,8 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_grpc::channel_manager::ChannelManager; +use std::sync::Arc; + +use common_base::readable_size::ReadableSize; +use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; +use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; +use meta_srv::metasrv::MetasrvOptions; use meta_srv::mocks as server_mock; use meta_srv::mocks::MockInfo; @@ -41,6 +46,50 @@ pub async fn mock_client_with_memstore() -> (MetaClient, MockMetaContext) { ) } +pub async fn mock_client_with_memstore_and_grpc_message_sizes( + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, +) -> (MetaClient, MockMetaContext) { + let mut opts = MetasrvOptions::default(); + opts.grpc.server_addr = "127.0.0.1:3002".to_string(); + opts.grpc.max_recv_message_size = server_max_recv_message_size; + opts.grpc.max_send_message_size = server_max_send_message_size; + + let client_channel_config = ChannelConfig { + max_recv_message_size: client_max_recv_message_size, + max_send_message_size: client_max_send_message_size, + ..ChannelConfig::new() + }; + + let kv_backend = Arc::new(MemoryKvBackend::new()); + let in_memory = Arc::new(MemoryKvBackend::new()); + let MockInfo { + server_addr, + channel_manager, + kv_backend, + in_memory, + .. + } = server_mock::mock_with_client_channel_config( + opts, + kv_backend, + None, + None, + Some(in_memory), + client_channel_config, + ) + .await; + + ( + mock_client_by(server_addr, channel_manager).await, + MockMetaContext { + kv_backend, + in_memory, + }, + ) +} + #[allow(dead_code)] pub async fn mock_client_with_etcdstore(addr: &str) -> (MetaClient, MockMetaContext) { let MockInfo { diff --git a/src/meta-srv/src/bootstrap.rs b/src/meta-srv/src/bootstrap.rs index 05b29d73c62..e9607ae78d9 100644 --- a/src/meta-srv/src/bootstrap.rs +++ b/src/meta-srv/src/bootstrap.rs @@ -233,13 +233,15 @@ pub async fn bootstrap_metasrv_with_router( #[macro_export] macro_rules! add_compressed_service { - ($builder:expr, $server:expr) => { + ($builder:expr, $server:expr, $grpc_config:expr) => { $builder.add_service( $server .accept_compressed(CompressionEncoding::Gzip) .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd), + .send_compressed(CompressionEncoding::Zstd) + .max_decoding_message_size($grpc_config.max_recv_message_size) + .max_encoding_message_size($grpc_config.max_send_message_size), ) }; } @@ -251,11 +253,26 @@ pub fn router(metasrv: Arc) -> Router { // For quick network failures detection. .http2_keepalive_interval(Some(metasrv.options().grpc.http2_keep_alive_interval)) .http2_keepalive_timeout(Some(metasrv.options().grpc.http2_keep_alive_timeout)); - let router = add_compressed_service!(router, HeartbeatServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, StoreServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ClusterServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ProcedureServiceServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ConfigServer::from_arc(metasrv.clone())); + let grpc_config = metasrv.options().grpc.as_config(); + let router = add_compressed_service!( + router, + HeartbeatServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, StoreServer::from_arc(metasrv.clone()), grpc_config); + let router = add_compressed_service!( + router, + ClusterServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = add_compressed_service!( + router, + ProcedureServiceServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, ConfigServer::from_arc(metasrv.clone()), grpc_config); router.add_service(admin::make_admin_service(metasrv)) } diff --git a/src/meta-srv/src/cluster.rs b/src/meta-srv/src/cluster.rs index ef3ba077027..688dc3f3332 100644 --- a/src/meta-srv/src/cluster.rs +++ b/src/meta-srv/src/cluster.rs @@ -81,10 +81,7 @@ impl KvBackend for MetaPeerClient { let retry_interval_ms = self.retry_interval_ms; for _ in 0..max_retry_count { - match self - .remote_range(req.key.clone(), req.range_end.clone(), req.keys_only) - .await - { + match self.remote_range(req.clone()).await { Ok(res) => return Ok(res), Err(e) => { if need_retry(&e) { @@ -238,12 +235,7 @@ impl MetaPeerClient { to_stat_kv_map(res.kvs) } - async fn remote_range( - &self, - key: Vec, - range_end: Vec, - keys_only: bool, - ) -> Result { + async fn remote_range(&self, req: RangeRequest) -> Result { // Safety: when self.is_leader() == false, election must not empty. let election = self.election.as_ref().unwrap(); @@ -254,14 +246,11 @@ impl MetaPeerClient { .get(&leader_addr) .context(error::CreateChannelSnafu)?; - let request = tonic::Request::new(PbRangeRequest { - key, - range_end, - keys_only, - ..Default::default() - }); + let request = tonic::Request::new(PbRangeRequest::from(req)); - let response: PbRangeResponse = ClusterClient::new(channel) + let mut client = + common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager); + let response: PbRangeResponse = client .range(request) .await .context(error::RangeSnafu)? @@ -291,7 +280,9 @@ impl MetaPeerClient { ..Default::default() }); - let response: PbBatchGetResponse = ClusterClient::new(channel) + let mut client = + common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager); + let response: PbBatchGetResponse = client .batch_get(request) .await .context(error::BatchGetSnafu)? @@ -365,13 +356,168 @@ fn need_retry(error: &error::Error) -> bool { #[cfg(test)] mod tests { - use api::v1::meta::{Error, ErrorCode, ResponseHeader}; - use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; - use common_meta::rpc::KeyValue; + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; - use super::{Context, check_resp_header, to_stat_kv_map}; + use api::v1::meta::cluster_server::{Cluster, ClusterServer}; + use api::v1::meta::{ + BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, Error, + ErrorCode, MetasrvPeersRequest, MetasrvPeersResponse, RangeRequest as PbRangeRequest, + RangeResponse as PbRangeResponse, ResponseHeader, + }; + use common_grpc::channel_manager::ChannelManager; + use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; + use common_meta::election::{Election, LeaderChangeMessage, LeaderValue, MetasrvNodeInfo}; + use common_meta::kv_backend::KvBackend; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::KeyValue; + use common_meta::rpc::store::RangeRequest; + use hyper_util::rt::TokioIo; + use tonic::{Request, Response, Status}; + use tower::service_fn; + + use super::{Context, MetaPeerClientBuilder, check_resp_header, to_stat_kv_map}; use crate::error; + struct FollowerElection { + leader_addr: String, + } + + #[async_trait::async_trait] + impl Election for FollowerElection { + type Leader = LeaderValue; + + fn is_leader(&self) -> bool { + false + } + + fn in_leader_infancy(&self) -> bool { + false + } + + async fn register_candidate(&self, _: &MetasrvNodeInfo) -> common_meta::error::Result<()> { + Ok(()) + } + + async fn all_candidates(&self) -> common_meta::error::Result> { + Ok(vec![]) + } + + async fn campaign(&self) -> common_meta::error::Result<()> { + Ok(()) + } + + async fn leader(&self) -> common_meta::error::Result { + Ok(LeaderValue(self.leader_addr.clone())) + } + + async fn resign(&self) -> common_meta::error::Result<()> { + Ok(()) + } + + fn subscribe_leader_change(&self) -> tokio::sync::broadcast::Receiver { + let (_, receiver) = tokio::sync::broadcast::channel(1); + receiver + } + } + + struct RangeServer { + requested_limit: Arc, + } + + #[async_trait::async_trait] + impl Cluster for RangeServer { + async fn batch_get( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("batch_get is not used in this test")) + } + + async fn range( + &self, + request: Request, + ) -> std::result::Result, Status> { + self.requested_limit + .store(request.into_inner().limit, Ordering::Relaxed); + Ok(Response::new(PbRangeResponse { + header: Some(ResponseHeader::success()), + kvs: vec![api::v1::meta::KeyValue { + key: b"key".to_vec(), + value: b"value".to_vec(), + }], + more: true, + })) + } + + async fn metasrv_peers( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented( + "metasrv_peers is not used in this test", + )) + } + } + + #[tokio::test] + async fn test_follower_range_forwards_limit() { + let requested_limit = Arc::new(AtomicI64::new(0)); + let range_server = RangeServer { + requested_limit: requested_limit.clone(), + }; + let (client, server) = tokio::io::duplex(1024); + let _server_handle = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service( + ClusterServer::new(range_server) + .accept_compressed(tonic::codec::CompressionEncoding::Zstd) + .send_compressed(tonic::codec::CompressionEncoding::Zstd), + ) + .serve_with_incoming(futures::stream::iter([Ok::<_, std::io::Error>(server)])) + .await + }); + + let channel_manager = ChannelManager::new(); + let mut client = Some(client); + channel_manager + .reset_with_connector( + "leader:0", + service_fn(move |_| { + let client = client.take(); + async move { + client + .map(TokioIo::new) + .ok_or_else(|| std::io::Error::other("client already taken")) + } + }), + ) + .unwrap(); + + let follower = MetaPeerClientBuilder::default() + .election(Some(Arc::new(FollowerElection { + leader_addr: "leader:0".to_string(), + }))) + .in_memory(Arc::new(MemoryKvBackend::new())) + .channel_manager(channel_manager) + .max_retry_count(1) + .build() + .unwrap(); + + let response = follower + .range(RangeRequest { + key: b"key".to_vec(), + limit: 1, + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(1, requested_limit.load(Ordering::Relaxed)); + assert_eq!(1, response.kvs.len()); + assert!(response.more); + } + #[test] fn test_to_stat_kv_map() { let stat_key = DatanodeStatKey { node_id: 100 }; diff --git a/src/meta-srv/src/mocks.rs b/src/meta-srv/src/mocks.rs index 26b1687e582..1a6fa306743 100644 --- a/src/meta-srv/src/mocks.rs +++ b/src/meta-srv/src/mocks.rs @@ -87,12 +87,52 @@ pub async fn mock( selector: Option, datanode_clients: Option>, in_memory: Option, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + None, + ) + .await +} + +pub async fn mock_with_client_channel_config( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + client_channel_config: ChannelConfig, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + Some(client_channel_config), + ) + .await +} + +async fn mock_inner( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + client_channel_config: Option, ) -> MockInfo { let server_addr = opts.grpc.server_addr.clone(); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); table_metadata_manager.init().await.unwrap(); + let grpc_config = opts.grpc.as_config(); + let grpc_options = opts.grpc.clone(); let builder = MetasrvBuilder::new() .options(opts) .kv_backend(kv_backend.clone()); @@ -121,22 +161,43 @@ pub async fn mock( let _handle = tokio::spawn(async move { let mut router = tonic::transport::Server::builder(); - let router = add_compressed_service!(router, HeartbeatServer::from_arc(service.clone())); - let router = add_compressed_service!(router, StoreServer::from_arc(service.clone())); + let router = add_compressed_service!( + router, + HeartbeatServer::from_arc(service.clone()), + grpc_config + ); let router = - add_compressed_service!(router, ProcedureServiceServer::from_arc(service.clone())); - let router = add_compressed_service!(router, ClusterServer::from_arc(service.clone())); - let router = add_compressed_service!(router, ConfigServer::from_arc(service.clone())); + add_compressed_service!(router, StoreServer::from_arc(service.clone()), grpc_config); + let router = add_compressed_service!( + router, + ProcedureServiceServer::from_arc(service.clone()), + grpc_config + ); + let router = add_compressed_service!( + router, + ClusterServer::from_arc(service.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, ConfigServer::from_arc(service.clone()), grpc_config); router .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) .await }); - let config = ChannelConfig::new() - // Use an long timeout to prevent test failures due to slow operations (e.g., when testing with S3). - .timeout(Some(Duration::from_secs(60))) - .connect_timeout(Duration::from_secs(10)) - .tcp_nodelay(true); + // Keep the mock client's codec limits aligned with the server by default. + let config = client_channel_config.unwrap_or_else(|| { + let config = ChannelConfig::new() + // Use an long timeout to prevent test failures due to slow operations (e.g., when testing with S3). + .timeout(Some(Duration::from_secs(60))) + .connect_timeout(Duration::from_secs(10)) + .tcp_nodelay(true); + ChannelConfig { + max_recv_message_size: grpc_options.max_recv_message_size, + max_send_message_size: grpc_options.max_send_message_size, + ..config + } + }); let channel_manager = ChannelManager::with_config(config, None); // Move client to an option so we can _move_ the inner value diff --git a/src/mito2/src/cache/manifest_cache.rs b/src/mito2/src/cache/manifest_cache.rs index 08ab0fb14a2..22fd796facd 100644 --- a/src/mito2/src/cache/manifest_cache.rs +++ b/src/mito2/src/cache/manifest_cache.rs @@ -39,6 +39,8 @@ const MANIFEST_DIR: &str = "cache/object/manifest/"; /// Metric label for manifest files. const MANIFEST_TYPE: &str = "manifest"; +const CHUNK_SIZE: usize = 64 * 1024 * 1024; // 64MB + /// A manifest cache manages manifest files on local store and evicts files based /// on size. #[derive(Debug, Clone)] @@ -261,7 +263,12 @@ impl ManifestCache { pub(crate) async fn put_file(&self, key: String, data: Vec) { let cache_file_path = self.cache_file_path(&key); - if let Err(e) = self.local_store.write(&cache_file_path, data.clone()).await { + if let Err(e) = self + .local_store + .write_with(&cache_file_path, data.clone()) + .chunk(CHUNK_SIZE) + .await + { warn!(e; "Failed to write manifest to cache {}", cache_file_path); return; } diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 9c0c0520458..872beec81cb 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -344,6 +344,32 @@ impl CompactionScheduler { .unwrap_or(false) } + /// Removes the compaction status of the region if it has no active compaction. + /// + /// The worker calls this when it decides not to schedule the next compaction + /// after a compaction finishes (e.g. throttled by `min_compaction_interval`). + /// Otherwise a stale status without an active compaction would remain in the + /// map and all subsequent compaction requests would be swallowed as if the + /// region were still compacting. + /// + /// Returns whether the status is removed. + pub(crate) fn remove_inactive_status(&mut self, region_id: RegionId) -> bool { + let Some(status) = self.region_status.get(®ion_id) else { + return false; + }; + if status.active_compaction.is_some() { + return false; + } + // A status without an active compaction must have been drained by + // `on_compaction_finished`: no waiters, pending requests or DDLs can be + // enqueued while no compaction is running. + debug_assert!(status.waiters.is_empty()); + debug_assert!(status.pending_request.is_none()); + debug_assert!(status.pending_ddl_requests.is_empty()); + self.region_status.remove(®ion_id); + true + } + /// Schedules next compaction upon a finished compaction. /// Returns whether the compaction is scheduled. pub(crate) async fn schedule_next_compaction( @@ -2509,6 +2535,115 @@ mod tests { assert!(!scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_remove_inactive_status_keeps_running_compaction() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let region_id = builder.region_id(); + + // No status at all. + assert!(!scheduler.remove_inactive_status(region_id)); + + scheduler.region_status.insert( + region_id, + CompactionStatus::new(region_id, version_control, env.access_layer.clone()), + ); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .start_local_task(); + + // A status with an active compaction must not be removed. + assert!(!scheduler.remove_inactive_status(region_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_compaction_reschedulable_when_next_compaction_throttled() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + // Five overlapping L0 files are enough for the regular picker to create a task. + let version_control = Arc::new( + builder + .push_l0_file(0, end) + .push_l0_file(10, end) + .push_l0_file(50, end) + .push_l0_file(80, end) + .push_l0_file(90, end) + .build(), + ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap(); + assert!(scheduled); + assert_eq!(1, job_scheduler.num_jobs()); + + // The compaction finishes with no pending request/DDL. The status remains + // in the map with no active compaction; the worker then skips + // `schedule_next_compaction` because `min_compaction_interval` has not + // passed and must remove the stale status instead. + let pending_ddls = scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert!(pending_ddls.is_empty()); + assert!(!scheduler.is_compacting(region_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + + assert!(scheduler.remove_inactive_status(region_id)); + assert!(!scheduler.region_status.contains_key(®ion_id)); + + // The region can be scheduled for compaction again. + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap(); + assert!(scheduled); + assert_eq!(2, job_scheduler.num_jobs()); + } + #[tokio::test] async fn test_on_compaction_finished_next_schedule_error_cleans_status() { let env = SchedulerEnv::new() diff --git a/src/mito2/src/gc.rs b/src/mito2/src/gc.rs index 8b5aadf141f..a484b4b897c 100644 --- a/src/mito2/src/gc.rs +++ b/src/mito2/src/gc.rs @@ -53,6 +53,7 @@ use crate::metrics::{ use crate::region::{MitoRegionRef, RegionRoleState}; use crate::sst::file::{RegionFileId, RegionIndexId, delete_files, delete_indexes}; use crate::sst::location::{self}; +use crate::worker::DROPPING_MARKER_FILE; #[cfg(test)] mod worker_test; @@ -854,6 +855,10 @@ impl LocalGcWorker { }); for entry in entries { + if entry.name() == DROPPING_MARKER_FILE { + continue; + } + let (file_id, file_type) = match location::parse_file_id_type_from_path(entry.name()) { Ok((file_id, file_type)) => (file_id, file_type), Err(err) => { diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 1baa3917449..923d935bb17 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -16,6 +16,7 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; use api::v1::Rows; +use bytes::Bytes; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionEngine as _; use store_api::region_request::{RegionCompactRequest, RegionRequest}; @@ -26,10 +27,12 @@ use crate::engine::MitoEngine; use crate::engine::compaction_test::{delete_and_flush, put_and_flush}; use crate::gc::{GcConfig, LocalGcWorker}; use crate::manifest::action::RemovedFile; +use crate::metrics::GC_SKIPPED_UNPARSABLE_FILES; use crate::region::MitoRegionRef; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, }; +use crate::worker::DROPPING_MARKER_FILE; async fn create_gc_worker( mito_engine: &MitoEngine, @@ -61,6 +64,77 @@ async fn create_gc_worker( .unwrap() } +#[tokio::test] +async fn test_gc_worker_ignores_dropping_marker_file() { + init_default_ut_logging(); + + let mut env = TestEnv::new().await; + env.log_store = Some(env.create_log_store().await); + env.object_store_manager = Some(Arc::new(env.create_in_memory_object_store_manager())); + + let engine = env + .new_mito_engine(MitoConfig { + gc: GcConfig { + enable: true, + lingering_time: None, + ..Default::default() + }, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + let manifest = region.manifest_ctx.manifest().await; + let marker_path = object_store::util::join_path( + ®ion.access_layer.build_region_dir(region_id), + DROPPING_MARKER_FILE, + ); + region + .access_layer + .object_store() + .write(&marker_path, Bytes::new()) + .await + .unwrap(); + + let before = GC_SKIPPED_UNPARSABLE_FILES.get(); + let regions = BTreeMap::from([(region_id, Some(region.clone()))]); + let file_ref_manifest = FileRefsManifest { + file_refs: Default::default(), + manifest_version: [(region_id, manifest.manifest_version)].into(), + cross_region_refs: HashMap::new(), + }; + + let gc_worker = create_gc_worker(&engine, regions, &file_ref_manifest, true).await; + let report = gc_worker.run().await.unwrap(); + + assert!( + report + .deleted_files + .get(®ion_id) + .is_none_or(|files| files.is_empty()) + ); + assert!(report.need_retry_regions.is_empty()); + assert_eq!(GC_SKIPPED_UNPARSABLE_FILES.get(), before); +} + /// Test insert/flush then truncate can allow gc worker to delete files #[tokio::test] async fn test_gc_worker_basic_truncate() { diff --git a/src/mito2/src/manifest/storage/utils.rs b/src/mito2/src/manifest/storage/utils.rs index 05d7ccceb26..18fd351acec 100644 --- a/src/mito2/src/manifest/storage/utils.rs +++ b/src/mito2/src/manifest/storage/utils.rs @@ -19,6 +19,8 @@ use store_api::ManifestVersion; use crate::cache::manifest_cache::ManifestCache; use crate::error::{OpenDalSnafu, Result}; +const CHUNK_SIZE: usize = 64 * 1024 * 1024; // 64MB + /// Gets a manifest file from cache. /// Returns the file data if found in cache, None otherwise. pub(crate) async fn get_from_cache(cache: Option<&ManifestCache>, key: &str) -> Option> { @@ -57,7 +59,11 @@ pub(crate) async fn write_and_put_cache( }; // Write to object store - object_store.write(path, data).await.context(OpenDalSnafu)?; + object_store + .write_with(path, data) + .chunk(CHUNK_SIZE) // 64MB + .await + .context(OpenDalSnafu)?; // Put to cache if we cloned the data if let Some(data) = cache_data { diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index 6ada2d2765b..d7fdb1d7097 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -201,10 +201,13 @@ impl RegionWriteCtx { /// Sets error and marks all write operations are failed. pub(crate) fn set_error(&mut self, err: Arc) { - // Set error for all notifiers + // Set error for all notifiers. for notify in &mut self.notifiers { notify.err = Some(err.clone()); } + for notify in &mut self.bulk_notifiers { + notify.err = Some(err.clone()); + } // Fail the whole write operation. self.failed = true; @@ -354,3 +357,60 @@ impl RegionWriteCtx { .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_recordbatch::DfRecordBatch; + use datatypes::arrow::array::{ArrayRef, TimestampMillisecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use store_api::logstore::provider::Provider; + use tokio::sync::oneshot; + + use super::*; + use crate::error::UnexpectedSnafu; + use crate::memtable::bulk::part::BulkPart; + use crate::test_util::version_util::VersionControlBuilder; + + #[test] + fn test_set_error_marks_bulk_notifiers_failed() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = + RegionWriteCtx::new(region_id, &version_control, Provider::noop_provider(), None); + let (tx, rx) = oneshot::channel(); + + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None)); + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "wal failed".to_string(), + } + .build(), + )); + drop(ctx); + + let result = rx.blocking_recv().unwrap(); + assert!(result.is_err(), "bulk notifier should report WAL error"); + } + + fn new_bulk_part() -> BulkPart { + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + )])); + let arrays = vec![Arc::new(TimestampMillisecondArray::from(vec![1, 2])) as ArrayRef]; + let batch = DfRecordBatch::try_new(schema, arrays).unwrap(); + + BulkPart { + batch, + max_timestamp: 2, + min_timestamp: 1, + sequence: 0, + timestamp_index: 0, + raw_data: None, + } + } +} diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 3abdf0e349d..569fa61d14a 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -141,6 +141,13 @@ impl RegionWorkerLoop { { region.update_schedule_compaction_millis(); } + } else { + // The minimal compaction interval has not passed, so the next compaction + // is not scheduled now and the region has no running compaction. Removes + // its compaction status, otherwise the stale status would make the + // scheduler swallow all subsequent compaction requests and the region + // would never compact again. + self.compaction_scheduler.remove_inactive_status(region_id); } } diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index daa9d9485df..c877acdfa5d 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, accept, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{DataFusionError, internal_err}; +use datafusion_common::{DataFusionError, assert_eq_or_internal_err, internal_err}; use datafusion_physical_expr::{Distribution, EquivalenceProperties, Partitioning}; use futures::StreamExt; use serde::Serialize; @@ -125,8 +125,13 @@ impl ExecutionPlan for DistAnalyzeExec { self: Arc, mut children: Vec>, ) -> DfResult> { + assert_eq_or_internal_err!( + children.len(), + 1, + "DistAnalyzeExec requires exactly one child" + ); Ok(Arc::new(Self::new( - children.pop().unwrap(), + children.swap_remove(0), self.verbose, self.format, ))) @@ -303,7 +308,7 @@ impl JsonMetrics { let mut elapsed_compute = 0; let mut output_rows = 0; let mut other_metrics = HashMap::default(); - let (name, param) = raw_name.split_once(": ").unwrap_or_default(); + let (name, param) = raw_name.split_once(": ").unwrap_or((raw_name, "")); for (name, value) in plan_metrics.metrics.into_iter() { if name == "elapsed_compute" { @@ -334,3 +339,97 @@ impl Display for JsonMetrics { write!(f, "{}", serde_json::to_string(self).unwrap()) } } + +#[cfg(test)] +mod tests { + use datafusion::physical_plan::empty::EmptyExec; + + use super::*; + + fn empty_plan(name: &str) -> Arc { + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + name, + DataType::Utf8, + true, + )])))) + } + + #[test] + fn qbs_dist_analyze_rejects_zero_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + assert!(ExecutionPlan::with_new_children(analyze, vec![]).is_err()); + } + + #[test] + fn qbs_dist_analyze_rejects_multiple_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + let result = ExecutionPlan::with_new_children( + analyze, + vec![empty_plan("first"), empty_plan("second")], + ); + + if let Ok(plan) = result { + let retained = plan + .as_any() + .downcast_ref::() + .unwrap() + .input + .schema() + .field(0) + .name() + .clone(); + panic!("expected an arity error for multiple children, but retained `{retained}`"); + } + } + + #[test] + fn qbs_dist_analyze_accepts_exactly_one_child() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + let replacement = empty_plan("replacement"); + + let rebuilt = ExecutionPlan::with_new_children(analyze, vec![replacement]).unwrap(); + let rebuilt = rebuilt.as_any().downcast_ref::().unwrap(); + + assert_eq!(rebuilt.input.schema().field(0).name(), "replacement"); + } + + #[test] + fn qbs_analyze_json_preserves_plan_name_without_parameters() { + let (_, metrics) = JsonMetrics::from_plan_metrics(PlanMetrics { + plan: "EmptyExec".to_string(), + plan_name: "EmptyExec".to_string(), + level: 0, + metrics: vec![], + }); + + assert_eq!(metrics.name, "EmptyExec"); + assert!(metrics.param.is_empty()); + } + + #[test] + fn qbs_analyze_json_splits_plan_name_and_parameters() { + let (_, metrics) = JsonMetrics::from_plan_metrics(PlanMetrics { + plan: "FilterExec: predicate".to_string(), + plan_name: "FilterExec".to_string(), + level: 0, + metrics: vec![], + }); + + assert_eq!(metrics.name, "FilterExec"); + assert_eq!(metrics.param, "predicate"); + } +} diff --git a/src/query/src/plan.rs b/src/query/src/plan.rs index d3eafd84c0a..b4e5e664b69 100644 --- a/src/query/src/plan.rs +++ b/src/query/src/plan.rs @@ -16,7 +16,7 @@ use std::collections::HashSet; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_expr::{Expr, LogicalPlan}; use session::context::QueryContextRef; pub use table::metadata::TableType; @@ -116,7 +116,7 @@ pub fn extract_and_rewrite_full_table_names( query_ctx: QueryContextRef, ) -> Result<(HashSet, LogicalPlan)> { let mut extractor = TableNamesExtractAndRewriter::new(query_ctx); - let plan = plan.rewrite(&mut extractor)?; + let plan = plan.rewrite_with_subqueries(&mut extractor)?; Ok((extractor.table_names, plan.data)) } @@ -142,23 +142,26 @@ pub(crate) mod tests { use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use common_catalog::consts::DEFAULT_CATALOG_NAME; use datafusion::logical_expr::builder::LogicalTableSource; - use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; + use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit, scalar_subquery}; use session::context::QueryContextBuilder; use super::*; - fn mock_plan() -> LogicalPlan { + fn mock_table_source() -> Arc { let schema = Schema::new(vec![ Field::new("id", DataType::Int32, true), Field::new("name", DataType::Utf8, true), Field::new("ts", DataType::Timestamp(TimeUnit::Millisecond, None), true), ]); - let table_source = LogicalTableSource::new(SchemaRef::new(schema)); + Arc::new(LogicalTableSource::new(SchemaRef::new(schema))) + } + + fn mock_plan() -> LogicalPlan { + let table_source = mock_table_source(); let projection = None; - let builder = - LogicalPlanBuilder::scan("devices", Arc::new(table_source), projection).unwrap(); + let builder = LogicalPlanBuilder::scan("devices", table_source, projection).unwrap(); builder .filter(col("id").gt(lit(500))) @@ -167,6 +170,55 @@ pub(crate) mod tests { .unwrap() } + fn scalar_subquery_plan(table_name: TableReference) -> LogicalPlan { + let subquery = LogicalPlanBuilder::scan(table_name, mock_table_source(), None) + .unwrap() + .project(vec![col("id")]) + .unwrap() + .build() + .unwrap(); + + LogicalPlanBuilder::empty(false) + .project(vec![scalar_subquery(Arc::new(subquery))]) + .unwrap() + .build() + .unwrap() + } + + fn assert_dependencies(actual: &HashSet, expected: &[(&str, &str, &str)]) { + let expected = expected + .iter() + .map(|(catalog, schema, table)| TableName::new(*catalog, *schema, *table)) + .collect::>(); + assert_eq!(&expected, actual); + } + + fn assert_nested_scalar_subquery_table_name(plan: &LogicalPlan, expected: TableReference) { + let LogicalPlan::Projection(projection) = plan else { + panic!("expected scalar-subquery projection, got {plan:?}"); + }; + let [Expr::ScalarSubquery(subquery)] = projection.expr.as_slice() else { + panic!( + "expected one scalar-subquery expression, got {:?}", + projection.expr + ); + }; + let LogicalPlan::Projection(projection) = subquery.subquery.as_ref() else { + panic!( + "expected scalar-subquery projection, got {:?}", + subquery.subquery + ); + }; + let LogicalPlan::TableScan(scan) = projection.input.as_ref() else { + panic!( + "expected scalar-subquery table scan, got {:?}", + projection.input + ); + }; + + assert_eq!(expected, scan.table_name); + } + #[test] fn test_extract_full_table_names() { let ctx = QueryContextBuilder::default() @@ -176,16 +228,84 @@ pub(crate) mod tests { let (table_names, plan) = extract_and_rewrite_full_table_names(mock_plan(), Arc::new(ctx)).unwrap(); - assert_eq!(1, table_names.len()); - assert!(table_names.contains(&TableName::new( - DEFAULT_CATALOG_NAME.to_string(), - "test".to_string(), - "devices".to_string() - ))); + assert_dependencies(&table_names, &[(DEFAULT_CATALOG_NAME, "test", "devices")]); assert_eq!( "Filter: devices.id > Int32(500)\n TableScan: greptime.test.devices", plan.to_string() ); } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_bare_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::bare("lookup")), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_catalog", "qp031_current_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_catalog", "qp031_current_schema", "lookup"), + ); + } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_partial_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::partial("qp031_lookup_schema", "lookup")), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_catalog", "qp031_lookup_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_catalog", "qp031_lookup_schema", "lookup"), + ); + } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_full_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::full( + "qp031_external_catalog", + "qp031_external_schema", + "lookup", + )), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_external_catalog", "qp031_external_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_external_catalog", "qp031_external_schema", "lookup"), + ); + } } diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index 2ff2c4481cb..c39ca86224b 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt::Write as FmtWrite; use std::io; use std::time::Duration; @@ -36,6 +37,7 @@ use datatypes::types::jsonb_to_string; use futures::StreamExt; use opensrv_mysql::{ Column, ColumnFlags, ColumnType, ErrorKind, OkResponse, QueryResultWriter, RowWriter, + ToMysqlValue, }; use session::SessionRef; use session::context::QueryContextRef; @@ -165,6 +167,21 @@ pub fn handle_err(e: impl ErrorExt, query_ctx: QueryContextRef) -> (ErrorKind, S struct MysqlResultWriter; +struct PrecisionTimestamp<'a> { + formatted: &'a str, + datetime: chrono::NaiveDateTime, +} + +impl<'a> ToMysqlValue for PrecisionTimestamp<'a> { + fn to_mysql_text(&self, w: &mut W) -> io::Result<()> { + self.formatted.to_mysql_text(w) + } + + fn to_mysql_bin(&self, w: &mut W, c: &Column) -> io::Result<()> { + self.datetime.to_mysql_bin(w, c) + } +} + impl MysqlResultWriter { async fn write_affected_rows<'a, W: AsyncWrite + Unpin>( w: QueryResultWriter<'a, W>, @@ -190,6 +207,9 @@ impl MysqlResultWriter { ) -> Result<()> { let schema = record_batch.schema.clone(); let record_batch = record_batch.into_df_record_batch(); + // Reusable buffer for timestamp formatting — avoids one heap allocation per + // timestamp value per row. Cleared and refilled in the Timestamp arm below. + let mut format_buf = String::new(); for i in 0..record_batch.num_rows() { for (j, column) in record_batch.columns().iter().enumerate() { if column.is_null(i) { @@ -265,8 +285,27 @@ impl MysqlResultWriter { } DataType::Timestamp(_, _) => { let v = datatypes::arrow_array::timestamp_array_value(column, i); - let v = v.to_chrono_datetime_with_timezone(Some(&query_context.timezone())); - row_writer.write_col(v)?; + // Reuse `format_buf` to avoid a per-row heap allocation. + // This mirrors what `Timestamp::as_formatted_string` does + // internally, but writes directly into our pre-allocated buffer. + format_buf.clear(); + if let Some(datetime) = + v.to_chrono_datetime_with_timezone(Some(&query_context.timezone())) + { + let _ = write!( + &mut format_buf, + "{}", + datetime.format("%Y-%m-%d %H:%M:%S%.f") + ); + row_writer.write_col(PrecisionTimestamp { + formatted: &format_buf, + datetime, + })?; + } else { + let _ = + write!(&mut format_buf, "[Timestamp{}: {}]", v.unit(), v.value()); + row_writer.write_col(&*format_buf)?; + } } DataType::Interval(interval_unit) => match interval_unit { IntervalUnit::YearMonth => { @@ -333,10 +372,7 @@ impl MysqlResultWriter { } } -pub(crate) fn create_mysql_column( - data_type: &ConcreteDataType, - column_name: &str, -) -> Result { +pub fn create_mysql_column(data_type: &ConcreteDataType, column_name: &str) -> Result { let column_type = match data_type { ConcreteDataType::Null(_) => Ok(ColumnType::MYSQL_TYPE_NULL), ConcreteDataType::Boolean(_) | ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => { diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 8d55fb3a6cc..99dc6ced912 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -1055,17 +1055,22 @@ pub(super) fn parameters_to_scalar_values( TimestampType::Nanosecond(_) => { let values = data .into_iter() - .filter_map(|ts| { - ts.and_then(|ts| { - ts.and_utc().timestamp_nanos_opt().map(|nanos| { + .map(|ts| match ts { + None => { + Ok(ScalarValue::TimestampNanosecond(None, None)) + } + Some(ts) => ts + .and_utc() + .timestamp_nanos_opt() + .map(|nanos| { ScalarValue::TimestampNanosecond( Some(nanos), None, ) }) - }) + .ok_or_else(|| numeric_out_of_range_error(ts)), }) - .collect::>(); + .collect::>>()?; ScalarValue::List(ScalarValue::new_list( &values, &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), @@ -2051,6 +2056,53 @@ mod test { } } + fn qbs_timestamp_nanosecond_array_plan() -> LogicalPlan { + build_plan_with_params(vec![( + "$1", + DataType::List(Arc::new(Field::new( + "item", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ))), + )]) + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_preserves_null_slot() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s( + r#"{"2024-01-01 00:00:00.000001",NULL,"2024-01-01 00:00:00.000003"}"#, + )], + ); + + let values = + parameters_to_scalar_values(&qbs_timestamp_nanosecond_array_plan(), &portal).unwrap(); + let expected = ScalarValue::List(ScalarValue::new_list( + &[ + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_001_000), None), + ScalarValue::TimestampNanosecond(None, None), + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_003_000), None), + ], + &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )); + assert_eq!(expected, values[0]); + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_out_of_range_rejected() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s(r#"{"3000-01-01 00:00:00"}"#)], + ); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &qbs_timestamp_nanosecond_array_plan(), + &portal, + )); + } + #[test] fn test_numeric_uint64_scalar_negative_rejected() { let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); diff --git a/src/servers/tests/mysql/mysql_writer_test.rs b/src/servers/tests/mysql/mysql_writer_test.rs index 8737fdf5841..473a6300263 100644 --- a/src/servers/tests/mysql/mysql_writer_test.rs +++ b/src/servers/tests/mysql/mysql_writer_test.rs @@ -14,11 +14,23 @@ use std::sync::Arc; -use datatypes::schema::Schema; -use servers::mysql::writer::create_mysql_column_def; +use common_time::timestamp::TimeUnit; +use common_time::{Timestamp, Timezone}; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::types::{ + IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, +}; +use opensrv_mysql::{ColumnFlags, ColumnType}; +use servers::mysql::writer::{create_mysql_column, create_mysql_column_def}; use crate::mysql::{TestingData, all_datatype_testing_data}; +// ── Schema-to-column-def tests ──────────────────────────────────────────────── + +/// Round-trip: every column schema in `all_datatype_testing_data()` must produce +/// a column definition whose type and name match the expected values. #[test] fn test_create_mysql_column_def() { let TestingData { @@ -28,12 +40,548 @@ fn test_create_mysql_column_def() { } = all_datatype_testing_data(); let schema = Arc::new(Schema::new(column_schemas.clone())); let columns_def = create_mysql_column_def(&schema).unwrap(); - assert_eq!(column_schemas.len(), columns_def.len()); + assert_eq!( + column_schemas.len(), + columns_def.len(), + "column count must match" + ); for (i, column_def) in columns_def.iter().enumerate() { - let column_schema = &column_schemas[i]; - assert_eq!(column_schema.name, column_def.column); - let expected_coltype = mysql_columns_def[i]; - assert_eq!(column_def.coltype, expected_coltype); + assert_eq!( + column_schemas[i].name, column_def.column, + "column name mismatch at index {i}" + ); + assert_eq!( + mysql_columns_def[i], column_def.coltype, + "column type mismatch at index {i} ({})", + column_schemas[i].name + ); + } +} + +// ── Individual create_mysql_column type-mapping tests ──────────────────────── + +/// Verify every concrete data type maps to the correct MySQL column type. +#[test] +fn test_create_mysql_column_type_mapping() { + let cases: &[(&str, ConcreteDataType, ColumnType)] = &[ + ( + "null_col", + ConcreteDataType::null_datatype(), + ColumnType::MYSQL_TYPE_NULL, + ), + ( + "bool_col", + ConcreteDataType::boolean_datatype(), + ColumnType::MYSQL_TYPE_TINY, + ), + ( + "int8_col", + ConcreteDataType::int8_datatype(), + ColumnType::MYSQL_TYPE_TINY, + ), + ( + "int16_col", + ConcreteDataType::int16_datatype(), + ColumnType::MYSQL_TYPE_SHORT, + ), + ( + "int32_col", + ConcreteDataType::int32_datatype(), + ColumnType::MYSQL_TYPE_LONG, + ), + ( + "int64_col", + ConcreteDataType::int64_datatype(), + ColumnType::MYSQL_TYPE_LONGLONG, + ), + ( + "uint8_col", + ConcreteDataType::uint8_datatype(), + ColumnType::MYSQL_TYPE_TINY, + ), + ( + "uint16_col", + ConcreteDataType::uint16_datatype(), + ColumnType::MYSQL_TYPE_SHORT, + ), + ( + "uint32_col", + ConcreteDataType::uint32_datatype(), + ColumnType::MYSQL_TYPE_LONG, + ), + ( + "uint64_col", + ConcreteDataType::uint64_datatype(), + ColumnType::MYSQL_TYPE_LONGLONG, + ), + ( + "float32_col", + ConcreteDataType::float32_datatype(), + ColumnType::MYSQL_TYPE_FLOAT, + ), + ( + "float64_col", + ConcreteDataType::float64_datatype(), + ColumnType::MYSQL_TYPE_DOUBLE, + ), + ( + "string_col", + ConcreteDataType::string_datatype(), + ColumnType::MYSQL_TYPE_VARCHAR, + ), + ( + "binary_col", + ConcreteDataType::binary_datatype(), + ColumnType::MYSQL_TYPE_VARCHAR, + ), + ( + "ts_s_col", + ConcreteDataType::timestamp_second_datatype(), + ColumnType::MYSQL_TYPE_TIMESTAMP, + ), + ( + "ts_ms_col", + ConcreteDataType::timestamp_millisecond_datatype(), + ColumnType::MYSQL_TYPE_TIMESTAMP, + ), + ( + "ts_us_col", + ConcreteDataType::timestamp_microsecond_datatype(), + ColumnType::MYSQL_TYPE_TIMESTAMP, + ), + ( + "ts_ns_col", + ConcreteDataType::timestamp_nanosecond_datatype(), + ColumnType::MYSQL_TYPE_TIMESTAMP, + ), + ( + "date_col", + ConcreteDataType::date_datatype(), + ColumnType::MYSQL_TYPE_DATE, + ), + ( + "time_s_col", + ConcreteDataType::time_second_datatype(), + ColumnType::MYSQL_TYPE_TIME, + ), + ( + "time_ms_col", + ConcreteDataType::time_millisecond_datatype(), + ColumnType::MYSQL_TYPE_TIME, + ), + ( + "time_us_col", + ConcreteDataType::time_microsecond_datatype(), + ColumnType::MYSQL_TYPE_TIME, + ), + ( + "time_ns_col", + ConcreteDataType::time_nanosecond_datatype(), + ColumnType::MYSQL_TYPE_TIME, + ), + ( + "interval_ym", + ConcreteDataType::interval_year_month_datatype(), + ColumnType::MYSQL_TYPE_VARCHAR, + ), + ( + "interval_dt", + ConcreteDataType::interval_day_time_datatype(), + ColumnType::MYSQL_TYPE_VARCHAR, + ), + ( + "interval_mdn", + ConcreteDataType::interval_month_day_nano_datatype(), + ColumnType::MYSQL_TYPE_VARCHAR, + ), + ( + "decimal_col", + ConcreteDataType::decimal128_datatype(10, 2), + ColumnType::MYSQL_TYPE_DECIMAL, + ), + ]; + + for (col_name, data_type, expected_type) in cases { + let col = create_mysql_column(data_type, col_name) + .unwrap_or_else(|_| panic!("create_mysql_column failed for {col_name}")); + assert_eq!( + *expected_type, col.coltype, + "type mismatch for column '{col_name}'" + ); + assert_eq!(*col_name, col.column, "name mismatch for '{col_name}'"); + // collen == 0 means "use MySQL default length" + assert_eq!(0, col.collen, "collen should be 0 for '{col_name}'"); + // table should be empty (not populated yet) + assert!( + col.table.is_empty(), + "table should be empty for '{col_name}'" + ); + } +} + +// ── Column flags tests ──────────────────────────────────────────────────────── + +/// Unsigned integer types must carry `UNSIGNED_FLAG`; all other types must not. +#[test] +fn test_create_mysql_column_unsigned_flags() { + // Types that must have UNSIGNED_FLAG + let unsigned_types: &[(&str, ConcreteDataType)] = &[ + ("u8", ConcreteDataType::uint8_datatype()), + ("u16", ConcreteDataType::uint16_datatype()), + ("u32", ConcreteDataType::uint32_datatype()), + ("u64", ConcreteDataType::uint64_datatype()), + ]; + for (name, dt) in unsigned_types { + let col = create_mysql_column(dt, name).unwrap(); + assert!( + col.colflags.contains(ColumnFlags::UNSIGNED_FLAG), + "UNSIGNED_FLAG must be set for {name}" + ); + } + + // Types that must NOT have UNSIGNED_FLAG + let signed_types: &[(&str, ConcreteDataType)] = &[ + ("i8", ConcreteDataType::int8_datatype()), + ("i16", ConcreteDataType::int16_datatype()), + ("i32", ConcreteDataType::int32_datatype()), + ("i64", ConcreteDataType::int64_datatype()), + ("f32", ConcreteDataType::float32_datatype()), + ("f64", ConcreteDataType::float64_datatype()), + ("str", ConcreteDataType::string_datatype()), + ("ts", ConcreteDataType::timestamp_millisecond_datatype()), + ("bool", ConcreteDataType::boolean_datatype()), + ("date", ConcreteDataType::date_datatype()), + ]; + for (name, dt) in signed_types { + let col = create_mysql_column(dt, name).unwrap(); + assert!( + !col.colflags.contains(ColumnFlags::UNSIGNED_FLAG), + "UNSIGNED_FLAG must NOT be set for {name}" + ); + } +} + +// ── Timestamp precision tests (issue #8227) ─────────────────────────────────── + +/// Timestamps displayed through the MySQL writer must respect their storage unit. +/// +/// Before the fix, `to_chrono_datetime_with_timezone` → `NaiveDateTime::Display` +/// always emitted 6 fractional-second digits, so: +/// - `TIMESTAMP(3)` showed `.195000` instead of `.195` +/// - `TIMESTAMP(9)` showed `.195123` instead of `.195123456` +/// +#[test] +fn test_timestamp_display_precision_utc() { + // Use UTC timezone for deterministic, timezone-independent output. + let tz = Timezone::from_tz_string("UTC").unwrap(); + + // Base instant: 2026-06-02 03:50:00 UTC (unix epoch seconds = 1780372200) + // Sub-second component: 195_123_456 ns = 195.123456 ms = 0.195123456 s + let base_s: i64 = 1_780_372_200; + + // TIMESTAMP(0) — second precision: no fractional part + let ts = Timestamp::new(base_s, TimeUnit::Second); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00", + "TIMESTAMP(0) must have no fractional seconds" + ); + + // TIMESTAMP(3) — millisecond precision: trailing zeros must be stripped + let ts = Timestamp::new(base_s * 1_000 + 195, TimeUnit::Millisecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.195", + "TIMESTAMP(3) must show exactly 3 fractional digits (no trailing zeros)" + ); + + // TIMESTAMP(6) — microsecond precision + let ts = Timestamp::new(base_s * 1_000_000 + 195_123, TimeUnit::Microsecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.195123", + "TIMESTAMP(6) must show exactly 6 fractional digits" + ); + + // TIMESTAMP(9) — nanosecond precision: full 9-digit fidelity + let ts = Timestamp::new(base_s * 1_000_000_000 + 195_123_456, TimeUnit::Nanosecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.195123456", + "TIMESTAMP(9) must show all 9 fractional digits" + ); +} + +/// When the subsecond part is exactly zero, no fractional portion is emitted +/// for any precision level (the `%.f` specifier omits the dot entirely). +#[test] +fn test_timestamp_display_precision_no_subseconds() { + let tz = Timezone::from_tz_string("UTC").unwrap(); + let base_s: i64 = 1_780_372_200; // 2026-06-02 03:50:00 UTC + + for (unit, value) in [ + (TimeUnit::Second, base_s), + (TimeUnit::Millisecond, base_s * 1_000), + (TimeUnit::Microsecond, base_s * 1_000_000), + (TimeUnit::Nanosecond, base_s * 1_000_000_000), + ] { + let ts = Timestamp::new(value, unit); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00", + "zero subseconds must produce no fractional part (unit = {unit:?})" + ); + } +} + +/// Subsecond digits are rendered in groups of 3 (chrono's `%.f`): a value +/// representable in milliseconds shows 3 digits, microseconds 6, nanoseconds 9. +/// e.g. 100 ms → `.100` (not `.100000`), 10 ms → `.010`. +#[test] +fn test_timestamp_display_precision_trailing_zeros_stripped() { + let tz = Timezone::from_tz_string("UTC").unwrap(); + let base_s: i64 = 1_780_372_200; + + // 100 ms → ".100" + let ts = Timestamp::new(base_s * 1_000 + 100, TimeUnit::Millisecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.100", + "100 ms must display as '.100'" + ); + + // 100_000 µs → ".100" (multiple of 1ms, so 3 digits) + let ts = Timestamp::new(base_s * 1_000_000 + 100_000, TimeUnit::Microsecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.100", + "100_000 µs must display as '.100'" + ); + + // 100_000_000 ns → ".100" (multiple of 1ms, so 3 digits) + let ts = Timestamp::new(base_s * 1_000_000_000 + 100_000_000, TimeUnit::Nanosecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.100", + "100_000_000 ns must display as '.100'" + ); + + // 10 ms → ".010" + let ts = Timestamp::new(base_s * 1_000 + 10, TimeUnit::Millisecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.010", + "10 ms must display as '.010'" + ); + + // 1 ms → ".001" + let ts = Timestamp::new(base_s * 1_000 + 1, TimeUnit::Millisecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "2026-06-02 03:50:00.001", + "1 ms must display as '.001'" + ); +} + +/// Verify the Unix epoch (timestamp == 0) formats correctly in UTC. +#[test] +fn test_timestamp_display_unix_epoch() { + let tz = Timezone::from_tz_string("UTC").unwrap(); + + let ts_s = Timestamp::new(0, TimeUnit::Second); + assert_eq!( + ts_s.to_timezone_aware_string(Some(&tz)), + "1970-01-01 00:00:00" + ); + + let ts_ms = Timestamp::new(0, TimeUnit::Millisecond); + assert_eq!( + ts_ms.to_timezone_aware_string(Some(&tz)), + "1970-01-01 00:00:00" + ); + + let ts_us = Timestamp::new(0, TimeUnit::Microsecond); + assert_eq!( + ts_us.to_timezone_aware_string(Some(&tz)), + "1970-01-01 00:00:00" + ); + + let ts_ns = Timestamp::new(0, TimeUnit::Nanosecond); + assert_eq!( + ts_ns.to_timezone_aware_string(Some(&tz)), + "1970-01-01 00:00:00" + ); +} + +/// Negative timestamps (before the Unix epoch) must format correctly. +#[test] +fn test_timestamp_display_negative_value() { + let tz = Timezone::from_tz_string("UTC").unwrap(); + + // -1 second = 1969-12-31 23:59:59 + let ts = Timestamp::new(-1, TimeUnit::Second); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "1969-12-31 23:59:59" + ); + + // -1 millisecond = 1969-12-31 23:59:59.999 + let ts = Timestamp::new(-1, TimeUnit::Millisecond); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz)), + "1969-12-31 23:59:59.999" + ); +} + +/// Verify that timezone offset shifts the display time correctly. +#[test] +fn test_timestamp_display_timezone_offset() { + // 2026-06-02 03:50:00 UTC = 2026-06-02 09:20:00 in +05:30 (India) + let base_s: i64 = 1_780_372_200; + + let tz_utc = Timezone::from_tz_string("UTC").unwrap(); + let tz_india = Timezone::from_tz_string("+05:30").unwrap(); + let tz_neg = Timezone::from_tz_string("-08:00").unwrap(); + + let ts = Timestamp::new(base_s * 1_000 + 195, TimeUnit::Millisecond); + + assert_eq!( + ts.to_timezone_aware_string(Some(&tz_utc)), + "2026-06-02 03:50:00.195" + ); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz_india)), + "2026-06-02 09:20:00.195" + ); + assert_eq!( + ts.to_timezone_aware_string(Some(&tz_neg)), + "2026-06-01 19:50:00.195" + ); +} + +// ── Column schema convenience constructors ──────────────────────────────────── + +/// All four timestamp types should map to `MYSQL_TYPE_TIMESTAMP`. +#[test] +fn test_create_mysql_column_all_timestamp_precisions() { + let timestamp_types: &[(&str, ConcreteDataType)] = &[ + ( + "ts_s", + ConcreteDataType::Timestamp(datatypes::types::TimestampType::Second( + TimestampSecondType, + )), + ), + ( + "ts_ms", + ConcreteDataType::Timestamp(datatypes::types::TimestampType::Millisecond( + TimestampMillisecondType, + )), + ), + ( + "ts_us", + ConcreteDataType::Timestamp(datatypes::types::TimestampType::Microsecond( + TimestampMicrosecondType, + )), + ), + ( + "ts_ns", + ConcreteDataType::Timestamp(datatypes::types::TimestampType::Nanosecond( + TimestampNanosecondType, + )), + ), + ]; + + for (name, dt) in timestamp_types { + let col = create_mysql_column(dt, name).unwrap(); + assert_eq!( + ColumnType::MYSQL_TYPE_TIMESTAMP, + col.coltype, + "timestamp type '{name}' must map to MYSQL_TYPE_TIMESTAMP" + ); + assert!( + !col.colflags.contains(ColumnFlags::UNSIGNED_FLAG), + "timestamps must not be unsigned" + ); + } +} + +/// All three interval types should map to `MYSQL_TYPE_VARCHAR`. +#[test] +fn test_create_mysql_column_all_interval_types() { + let interval_types: &[(&str, ConcreteDataType)] = &[ + ( + "iv_ym", + ConcreteDataType::Interval(datatypes::types::IntervalType::YearMonth( + IntervalYearMonthType, + )), + ), + ( + "iv_dt", + ConcreteDataType::Interval(datatypes::types::IntervalType::DayTime( + IntervalDayTimeType, + )), + ), + ( + "iv_mdn", + ConcreteDataType::Interval(datatypes::types::IntervalType::MonthDayNano( + IntervalMonthDayNanoType, + )), + ), + ]; + + for (name, dt) in interval_types { + let col = create_mysql_column(dt, name).unwrap(); + assert_eq!( + ColumnType::MYSQL_TYPE_VARCHAR, + col.coltype, + "interval type '{name}' must map to MYSQL_TYPE_VARCHAR" + ); + } +} + +/// create_mysql_column_def must preserve the ordering and names of all columns. +#[test] +fn test_create_mysql_column_def_preserves_order() { + let schemas: Vec = vec![ + ColumnSchema::new("alpha", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("beta", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("gamma", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new( + "delta", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let schema = Arc::new(Schema::new(schemas.clone())); + let cols = create_mysql_column_def(&schema).unwrap(); + + assert_eq!(4, cols.len()); + assert_eq!("alpha", cols[0].column); + assert_eq!(ColumnType::MYSQL_TYPE_LONG, cols[0].coltype); + assert_eq!("beta", cols[1].column); + assert_eq!(ColumnType::MYSQL_TYPE_VARCHAR, cols[1].coltype); + assert_eq!("gamma", cols[2].column); + assert_eq!(ColumnType::MYSQL_TYPE_DOUBLE, cols[2].coltype); + assert_eq!("delta", cols[3].column); + assert_eq!(ColumnType::MYSQL_TYPE_TIMESTAMP, cols[3].coltype); +} + +/// An empty schema must produce an empty column-def list (no panic). +#[test] +fn test_create_mysql_column_def_empty_schema() { + let schema = Arc::new(Schema::new(vec![])); + let cols = create_mysql_column_def(&schema).unwrap(); + assert!(cols.is_empty(), "empty schema must produce no columns"); +} + +/// Decimal128 must map to MYSQL_TYPE_DECIMAL regardless of precision/scale. +#[test] +fn test_create_mysql_column_decimal_variants() { + for (p, s) in [(10u8, 2i8), (38, 10), (5, 0)] { + let dt = ConcreteDataType::decimal128_datatype(p, s); + let col = create_mysql_column(&dt, "dec").unwrap(); + assert_eq!(ColumnType::MYSQL_TYPE_DECIMAL, col.coltype); + assert!(!col.colflags.contains(ColumnFlags::UNSIGNED_FLAG)); } } diff --git a/tests/cases/standalone/common/insert/mysql_insert.result b/tests/cases/standalone/common/insert/mysql_insert.result index 12176b2e887..3e94070f6e3 100644 --- a/tests/cases/standalone/common/insert/mysql_insert.result +++ b/tests/cases/standalone/common/insert/mysql_insert.result @@ -25,14 +25,14 @@ Failed to execute query, err: MySqlError { ERROR 1210 (HY000): (InvalidArguments -- SQLNESS PROTOCOL MYSQL SELECT * FROM integers ORDER BY ts; -+-----+----------------------------+ -| val | ts | -+-----+----------------------------+ -| 11 | 1970-01-01 00:00:00.001000 | -| 12 | 1970-01-01 00:00:00.002000 | -| 13 | 1970-01-01 00:00:00.003000 | -| 14 | 1970-01-01 00:00:00.004000 | -+-----+----------------------------+ ++-----+-------------------------+ +| val | ts | ++-----+-------------------------+ +| 11 | 1970-01-01 00:00:00.001 | +| 12 | 1970-01-01 00:00:00.002 | +| 13 | 1970-01-01 00:00:00.003 | +| 14 | 1970-01-01 00:00:00.004 | ++-----+-------------------------+ -- SQLNESS PROTOCOL MYSQL DROP TABLE integers; diff --git a/tests/cases/standalone/common/prepare/mysql_prepare.result b/tests/cases/standalone/common/prepare/mysql_prepare.result index 5ef242a891b..cd525d8cbec 100644 --- a/tests/cases/standalone/common/prepare/mysql_prepare.result +++ b/tests/cases/standalone/common/prepare/mysql_prepare.result @@ -113,11 +113,11 @@ affected_rows: 0 -- SQLNESS PROTOCOL MYSQL EXECUTE stmt USING 'happy', 42, 0; -+--------+-----------+----------------------------+ -| domain | is_expire | ts | -+--------+-----------+----------------------------+ -| happy | 0 | 2025-03-18 12:55:51.758000 | -+--------+-----------+----------------------------+ ++--------+-----------+-------------------------+ +| domain | is_expire | ts | ++--------+-----------+-------------------------+ +| happy | 0 | 2025-03-18 12:55:51.758 | ++--------+-----------+-------------------------+ -- SQLNESS PROTOCOL MYSQL DEALLOCATE stmt; diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result new file mode 100644 index 00000000000..45b930336f5 --- /dev/null +++ b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result @@ -0,0 +1,69 @@ +-- description: Timestamp display precision over the MySQL protocol should +-- respect the column's storage unit (issue #8227): +-- - TIMESTAMP(3) must not pad to 6 fractional digits +-- - TIMESTAMP(9) must not truncate to 6 fractional digits +-- SQLNESS PROTOCOL MYSQL +CREATE TABLE ts_precision_display ( + ts_s TIMESTAMP(0), + ts_ms TIMESTAMP(3), + ts_us TIMESTAMP(6), + ts_ns TIMESTAMP(9), + v INT, + TIME INDEX (ts_ms), + PRIMARY KEY (v) +); + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:00', + '2026-06-02 03:50:00.195', + '2026-06-02 03:50:00.195123', + '2026-06-02 03:50:00.195123456', + 1 +); + +affected_rows: 1 + +-- subsecond digits render in groups of 3: values representable in ms show +-- 3 digits regardless of storage unit (.100, not .100000) +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:01', + '2026-06-02 03:50:01.100', + '2026-06-02 03:50:01.100000', + '2026-06-02 03:50:01.100000000', + 2 +); + +affected_rows: 1 + +-- zero subseconds render with no fractional part for every precision +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + 3 +); + +affected_rows: 1 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM ts_precision_display ORDER BY v; + ++---------------------+-------------------------+----------------------------+-------------------------------+---+ +| ts_s | ts_ms | ts_us | ts_ns | v | ++---------------------+-------------------------+----------------------------+-------------------------------+---+ +| 2026-06-02 03:50:00 | 2026-06-02 03:50:00.195 | 2026-06-02 03:50:00.195123 | 2026-06-02 03:50:00.195123456 | 1 | +| 2026-06-02 03:50:01 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2 | +| 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 3 | ++---------------------+-------------------------+----------------------------+-------------------------------+---+ + +-- SQLNESS PROTOCOL MYSQL +DROP TABLE ts_precision_display; + +affected_rows: 0 + diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql new file mode 100644 index 00000000000..8eb82ed711d --- /dev/null +++ b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql @@ -0,0 +1,51 @@ +-- description: Timestamp display precision over the MySQL protocol should +-- respect the column's storage unit (issue #8227): +-- - TIMESTAMP(3) must not pad to 6 fractional digits +-- - TIMESTAMP(9) must not truncate to 6 fractional digits + +-- SQLNESS PROTOCOL MYSQL +CREATE TABLE ts_precision_display ( + ts_s TIMESTAMP(0), + ts_ms TIMESTAMP(3), + ts_us TIMESTAMP(6), + ts_ns TIMESTAMP(9), + v INT, + TIME INDEX (ts_ms), + PRIMARY KEY (v) +); + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:00', + '2026-06-02 03:50:00.195', + '2026-06-02 03:50:00.195123', + '2026-06-02 03:50:00.195123456', + 1 +); + +-- subsecond digits render in groups of 3: values representable in ms show +-- 3 digits regardless of storage unit (.100, not .100000) +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:01', + '2026-06-02 03:50:01.100', + '2026-06-02 03:50:01.100000', + '2026-06-02 03:50:01.100000000', + 2 +); + +-- zero subseconds render with no fractional part for every precision +-- SQLNESS PROTOCOL MYSQL +INSERT INTO ts_precision_display VALUES ( + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + '2026-06-02 03:50:02', + 3 +); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM ts_precision_display ORDER BY v; + +-- SQLNESS PROTOCOL MYSQL +DROP TABLE ts_precision_display; diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 6eb69919952..1753d59da91 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -192,7 +192,7 @@ TQL ANALYZE FORMAT JSON (0, 10, '5s') test; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_| {"name":"","param":"","output_rows":0,"REDACTED +| 0_| 0_| {"name":"CooperativeExec","param":"","output_rows":0,"REDACTED | 1_| 0_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED | 1_| 1_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED |_|_| Total rows: 0_| @@ -213,7 +213,7 @@ TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_| {"name":"","param":"","output_rows":0,"REDACTED +| 0_| 0_| {"name":"CooperativeExec","param":"","output_rows":0,"REDACTED | 1_| 0_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED | 1_| 1_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED |_|_| Total rows: 0_| diff --git a/tests/runner/src/formatter.rs b/tests/runner/src/formatter.rs index a379ad3244a..0ae3b5ab207 100644 --- a/tests/runner/src/formatter.rs +++ b/tests/runner/src/formatter.rs @@ -40,6 +40,7 @@ impl Display for ErrorFormatter { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let status_code = self.0.status_code(); let root_cause = self.0.output_msg(); + let root_cause = normalize_meta_client_error(&root_cause, status_code); write!( f, "Error: {}({status_code}), {root_cause}", @@ -48,6 +49,17 @@ impl Display for ErrorFormatter { } } +fn normalize_meta_client_error( + root_cause: &str, + status_code: common_error::status_code::StatusCode, +) -> &str { + let details_prefix = format!(", code: {status_code}, tonic code: "); + match root_cause.rsplit_once(&details_prefix) { + Some((message, tonic_code)) if !tonic_code.is_empty() => message, + _ => root_cause, + } +} + /// A formatter for [`Output`]. pub struct OutputFormatter(Output); @@ -214,3 +226,52 @@ pub fn build_recordbatches_from_mysql_rows(rows: &[MySqlRow]) -> RecordBatches { RecordBatches::try_from_columns(schema, columns) .expect("Failed to construct recordbatches from columns. Please check the schema.") } + +#[cfg(test)] +mod tests { + use common_error::ext::PlainError; + use common_error::status_code::StatusCode; + + use super::*; + + #[test] + fn test_normalize_meta_client_error() { + let error = PlainError::new( + "Invalid options, code: InvalidArguments, tonic code: Client specified an invalid argument" + .to_string(), + StatusCode::InvalidArguments, + ); + + assert_eq!( + "Error: 1004(InvalidArguments), Invalid options", + ErrorFormatter::from(error).to_string() + ); + } + + #[test] + fn test_preserve_regular_error() { + let error = PlainError::new( + "Invalid options without transport details".to_string(), + StatusCode::InvalidArguments, + ); + + assert_eq!( + "Error: 1004(InvalidArguments), Invalid options without transport details", + ErrorFormatter::from(error).to_string() + ); + } + + #[test] + fn test_preserve_error_with_mismatched_code() { + let error = PlainError::new( + "Invalid options, code: Unsupported, tonic code: Operation is not supported" + .to_string(), + StatusCode::InvalidArguments, + ); + + assert_eq!( + "Error: 1004(InvalidArguments), Invalid options, code: Unsupported, tonic code: Operation is not supported", + ErrorFormatter::from(error).to_string() + ); + } +}