fix(partition): avoid panic on missing route columns (#8645)

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-07-27 19:52:06 +08:00
committed by GitHub
parent 054d4046d5
commit b32a1de2fd
2 changed files with 178 additions and 7 deletions
+130 -5
View File
@@ -15,6 +15,9 @@
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use api::v1::helper::tag_column_schema;
use api::v1::value::ValueData;
use api::v1::{ColumnDataType, Row, Rows};
use common_base::hash::partition_expr_version;
use common_meta::cache::{TableRouteCacheRef, new_table_route_cache};
use common_meta::key::TableMetadataManager;
@@ -37,7 +40,11 @@ pub fn new_test_table_info(
table_name: &str,
_region_numbers: impl Iterator<Item = u32>,
) -> TableInfo {
let column_schemas = vec![
new_test_table_info_with_columns(table_id, table_name, test_column_schemas(true), vec![0])
}
fn test_column_schemas(include_b: bool) -> Vec<ColumnSchema> {
let mut column_schemas = vec![
ColumnSchema::new("a", ConcreteDataType::int32_datatype(), true),
ColumnSchema::new(
"ts",
@@ -45,8 +52,24 @@ pub fn new_test_table_info(
false,
)
.with_time_index(true),
ColumnSchema::new("b", ConcreteDataType::int32_datatype(), true),
];
if include_b {
column_schemas.push(ColumnSchema::new(
"b",
ConcreteDataType::int32_datatype(),
true,
));
}
column_schemas
}
fn new_test_table_info_with_columns(
table_id: u32,
table_name: &str,
column_schemas: Vec<ColumnSchema>,
partition_key_indices: Vec<usize>,
) -> TableInfo {
let next_column_id = column_schemas.len() as u32;
let schema = SchemaBuilder::try_from(column_schemas)
.unwrap()
.version(123)
@@ -55,10 +78,10 @@ pub fn new_test_table_info(
let meta = TableMetaBuilder::empty()
.schema(Arc::new(schema))
.primary_key_indices(vec![0])
.primary_key_indices(partition_key_indices.clone())
.engine("engine")
.next_column_id(3)
.partition_key_indices(vec![0])
.next_column_id(next_column_id)
.partition_key_indices(partition_key_indices)
.build()
.unwrap();
TableInfoBuilder::default()
@@ -70,6 +93,14 @@ pub fn new_test_table_info(
.unwrap()
}
fn new_physical_test_table_info(table_id: u32, table_name: &str) -> TableInfo {
new_test_table_info_with_columns(table_id, table_name, test_column_schemas(true), vec![0, 2])
}
fn new_logical_test_table_info(table_id: u32, table_name: &str) -> TableInfo {
new_test_table_info_with_columns(table_id, table_name, test_column_schemas(false), vec![0])
}
fn new_test_region_wal_options(regions: Vec<RegionNumber>) -> RegionWalOptions {
// TODO(niebayes): construct region wal options for test.
let _ = regions;
@@ -239,3 +270,97 @@ async fn test_partition_expr_version_cache() {
assert_ne!(None, *version_by_region.get(&2).unwrap());
assert_ne!(None, *version_by_region.get(&3).unwrap());
}
#[tokio::test]
async fn test_logical_split_rows_rejects_physical_only_partition_column() {
let kv_backend = Arc::new(common_meta::kv_backend::memory::MemoryKvBackend::new());
let table_metadata_manager = TableMetadataManager::new(kv_backend.clone());
let table_route_cache = test_new_table_route_cache(kv_backend.clone());
let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone());
let partition_manager =
PartitionRuleManager::new(kv_backend, table_route_cache, partition_info_cache);
let physical_table_info = new_physical_test_table_info(1024, "physical");
let logical_table_info = new_logical_test_table_info(1025, "logical");
assert_eq!(
vec!["a", "b"],
physical_table_info
.meta
.partition_column_names()
.map(|name| name.as_str())
.collect::<Vec<_>>()
);
assert_eq!(
vec!["a"],
logical_table_info
.meta
.partition_column_names()
.map(|name| name.as_str())
.collect::<Vec<_>>()
);
table_metadata_manager
.create_table_metadata(
physical_table_info,
TableRouteValue::physical(vec![RegionRoute {
region: Region {
id: 1.into(),
name: "r1".to_string(),
attrs: BTreeMap::new(),
partition_expr: PartitionExpr::new(
Operand::Expr(PartitionExpr::new(
Operand::Column("a".to_string()),
RestrictedOp::GtEq,
Operand::Value(datatypes::value::Value::Int32(10)),
)),
RestrictedOp::And,
Operand::Expr(PartitionExpr::new(
Operand::Column("b".to_string()),
RestrictedOp::Lt,
Operand::Value(datatypes::value::Value::Int32(50)),
)),
)
.as_json_str()
.unwrap(),
},
leader_peer: Some(Peer::new(1, "")),
follower_peers: vec![],
leader_state: None,
leader_down_since: None,
write_route_policy: None,
}]),
new_test_region_wal_options(vec![1]),
)
.await
.unwrap();
table_metadata_manager
.create_table_metadata(
logical_table_info.clone(),
TableRouteValue::logical(1024),
new_test_region_wal_options(vec![]),
)
.await
.unwrap();
let (partition_rule, _) = partition_manager
.find_table_partition_rule(&logical_table_info)
.await
.unwrap();
assert_eq!(&["a"], partition_rule.partition_columns());
let error = partition_manager
.split_rows(
&logical_table_info,
Rows {
schema: vec![tag_column_schema("a", ColumnDataType::Int32)],
rows: vec![Row {
values: vec![ValueData::I32Value(100).into()],
}],
},
)
.await
.unwrap_err();
assert!(matches!(
error,
partition::error::Error::UndefinedColumn { ref column, .. } if column == "b"
));
}
+48 -2
View File
@@ -121,12 +121,18 @@ impl MultiDimPartitionRule {
fn evaluate_expr(&self, expr: &PartitionExpr, values: &[Value]) -> Result<bool> {
match (expr.lhs.as_ref(), expr.rhs.as_ref()) {
(Operand::Column(name), Operand::Value(r)) => {
let index = self.name_to_index.get(name).unwrap();
let index = self
.name_to_index
.get(name)
.context(UndefinedColumnSnafu { column: name })?;
let l = &values[*index];
Self::perform_op(l, &expr.op, r)
}
(Operand::Value(l), Operand::Column(name)) => {
let index = self.name_to_index.get(name).unwrap();
let index = self
.name_to_index
.get(name)
.context(UndefinedColumnSnafu { column: name })?;
let r = &values[*index];
Self::perform_op(l, &expr.op, r)
}
@@ -416,6 +422,46 @@ mod tests {
assert_matches!(rule.find_region(&["zzzz".into()]), Ok(3));
}
#[test]
fn test_find_region_rejects_undeclared_column_on_lhs() {
let rule = MultiDimPartitionRule::try_new(
vec!["host".to_string()],
vec![1],
vec![PartitionExpr::new(
Operand::Column("rack".to_string()),
RestrictedOp::Lt,
Operand::Value(Value::String("n".into())),
)],
false,
)
.unwrap();
assert_matches!(
rule.find_region(&[Value::String("z".into())]),
Err(Error::UndefinedColumn { column, .. }) if column == "rack"
);
}
#[test]
fn test_find_region_rejects_undeclared_column_on_rhs() {
let rule = MultiDimPartitionRule::try_new(
vec!["host".to_string()],
vec![1],
vec![PartitionExpr::new(
Operand::Value(Value::String("n".into())),
RestrictedOp::Gt,
Operand::Column("rack".to_string()),
)],
false,
)
.unwrap();
assert_matches!(
rule.find_region(&[Value::String("z".into())]),
Err(Error::UndefinedColumn { column, .. }) if column == "rack"
);
}
#[test]
fn invalid_expr_case_1() {
// PARTITION ON COLUMNS (b) (