mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-08-18 12:08:22 +00:00
fix(partition): avoid panic on missing route columns (#8645)
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
@@ -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"
|
||||
));
|
||||
}
|
||||
|
||||
@@ -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) (
|
||||
|
||||
Reference in New Issue
Block a user