diff --git a/crates/dbx-core/src/schema/kingbase.rs b/crates/dbx-core/src/schema/kingbase.rs index f52beac99c..b943244514 100644 --- a/crates/dbx-core/src/schema/kingbase.rs +++ b/crates/dbx-core/src/schema/kingbase.rs @@ -27,11 +27,31 @@ impl ExtensionCatalog { } } +pub(super) fn object_statistics_query_plan(schema: &str) -> Vec { + vec![ + ("catalog", object_statistics_sql(schema), true), + // Some V8 servers lack sys_total_relation_size(oid). Page statistics + // are an estimate, may be stale, and exclude indexes and TOAST storage. + ( + "estimated-pages", + object_statistics_size_sql( + schema, + "CAST(c.relpages AS BIGINT) * CAST(current_setting('block_size') AS BIGINT)", + ), + true, + ), + ] +} + pub(super) fn object_statistics_sql(schema: &str) -> String { + object_statistics_size_sql(schema, "sys_total_relation_size(c.oid)") +} + +fn object_statistics_size_sql(schema: &str, size_expression: &str) -> String { format!( "SELECT c.relname, n.nspname, \ CAST(CASE WHEN c.reltuples < 0 THEN 0 ELSE c.reltuples END AS BIGINT) AS estimated_rows, \ - CAST(sys_total_relation_size(c.oid) AS BIGINT) AS total_bytes \ + CAST({size_expression} AS BIGINT) AS total_bytes \ FROM sys_catalog.sys_class c \ JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = {} AND c.relkind IN ('r','m','f','p') \ diff --git a/crates/dbx-core/src/schema/mod.rs b/crates/dbx-core/src/schema/mod.rs index 22b362783e..3641c208a9 100644 --- a/crates/dbx-core/src/schema/mod.rs +++ b/crates/dbx-core/src/schema/mod.rs @@ -2150,9 +2150,7 @@ fn external_driver_statistics_query_plan( match dialect { ExternalDriverStatisticsDialect::Oracle => oracle_object_statistics_query_plan(schema), ExternalDriverStatisticsDialect::Dameng => dameng_object_statistics_query_plan(schema), - ExternalDriverStatisticsDialect::Kingbase => { - vec![("catalog", kingbase::object_statistics_sql(schema), true)] - } + ExternalDriverStatisticsDialect::Kingbase => kingbase::object_statistics_query_plan(schema), } } @@ -4216,7 +4214,8 @@ done assert!(dameng.iter().all(|(_, _, accept_empty)| *accept_empty)); let kingbase = external_driver_statistics_query_plan(ExternalDriverStatisticsDialect::Kingbase, "public"); - assert_eq!(kingbase.len(), 1); + assert_eq!(kingbase.len(), 2); + assert!(kingbase[1].1.contains("current_setting('block_size')")); assert!(kingbase[0].1.contains("sys_catalog.sys_class")); } @@ -6215,6 +6214,87 @@ done } } + #[tokio::test] + async fn kingbase_object_statistics_falls_back_when_size_function_is_missing() { + let mut queries = Vec::new(); + let stats = super::object_statistics_from_query_plan( + "Kingbase", + "core's", + super::kingbase::object_statistics_query_plan("core's"), + |sql| { + queries.push(sql.clone()); + async move { + if sql.contains("sys_total_relation_size") { + Err("kb: function sys_total_relation_size(oid) does not exist".to_string()) + } else { + Ok(statistics_query_result(vec![vec![ + serde_json::json!("ceshi"), + serde_json::json!("core's"), + serde_json::json!(100), + serde_json::json!(630784), + ]])) + } + } + }, + ) + .await + .unwrap(); + assert_eq!(queries.len(), 2); + assert!(queries[1].contains("CAST(c.relpages AS BIGINT) * CAST(current_setting('block_size') AS BIGINT)")); + assert!(queries[1].contains("n.nspname = 'core''s'")); + assert_eq!(stats[0].name, "ceshi"); + assert_eq!(stats[0].schema.as_deref(), Some("core's")); + assert_eq!(stats[0].total_bytes, Some(630784)); + assert_eq!(stats[0].estimated_rows, Some(100)); + } + + #[tokio::test] + async fn kingbase_object_statistics_keeps_successful_size_results() { + for bytes in [None, Some(0), Some(8192)] { + let mut calls = 0; + let stats = super::object_statistics_from_query_plan( + "Kingbase", + "public", + super::kingbase::object_statistics_query_plan("public"), + |_| { + calls += 1; + let rows = bytes + .map(|size| { + vec![vec![ + serde_json::json!("ceshi"), + serde_json::json!("public"), + serde_json::json!(0), + serde_json::json!(size), + ]] + }) + .unwrap_or_default(); + async move { Ok(statistics_query_result(rows)) } + }, + ) + .await + .unwrap(); + assert_eq!(calls, 1); + assert_eq!(stats.first().and_then(|stat| stat.total_bytes), bytes); + } + } + + #[tokio::test] + async fn kingbase_object_statistics_reports_error_when_both_queries_fail() { + let mut calls = 0; + let result = super::object_statistics_from_query_plan( + "Kingbase", + "public", + super::kingbase::object_statistics_query_plan("public"), + |_| { + calls += 1; + async { Err("permission denied".to_string()) } + }, + ) + .await; + assert_eq!(calls, 2); + assert_eq!(result.unwrap_err(), "permission denied"); + } + #[tokio::test] async fn dameng_object_statistics_falls_back_to_rows_when_size_lookup_fails() { let sqls = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); @@ -7009,13 +7089,18 @@ async fn list_object_statistics_once( .await; } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { - let sql = kingbase::object_statistics_sql(schema); - return agent_list_object_statistics( - client, - database, + let timeout_duration = agent_metadata_timeout(db_config.as_ref()); + return object_statistics_from_query_plan( + "Kingbase", schema, - sql, - agent_metadata_timeout(db_config.as_ref()), + kingbase::object_statistics_query_plan(schema), + |sql| { + let client = client.clone(); + async move { + let mut client = client.lock().await; + agent_object_statistics_query(&mut client, database, schema, &sql, timeout_duration).await + } + }, ) .await; }