1use std::any::Any;
16
17use common_meta::key::table_route::PhysicalTableRouteValue;
18use common_procedure::{Context as ProcedureContext, Status};
19use common_telemetry::debug;
20use partition::collider::Collider;
21use partition::expr::PartitionExpr;
22use partition::subtask::{self, RepartitionSubtask};
23use serde::{Deserialize, Deserializer, Serialize};
24use snafu::{OptionExt, ResultExt, ensure};
25use tokio::time::Instant;
26use uuid::Uuid;
27
28use crate::error::{self, Result};
29use crate::procedure::repartition::allocate_region::AllocateRegion;
30use crate::procedure::repartition::plan::{AllocationPlanEntry, SourceRegionDescriptor};
31use crate::procedure::repartition::repartition_end::RepartitionEnd;
32use crate::procedure::repartition::update_partition_metadata::{
33 PartitionMetadataUpdate, UpdatePartitionMetadata,
34};
35use crate::procedure::repartition::{Context, State};
36
37#[derive(Debug, Clone, Serialize)]
38pub enum RepartitionFrom {
39 Partitioned {
40 exprs: Vec<PartitionExpr>,
41 target_partition_columns: Option<Vec<String>>,
47 },
48 Unpartitioned {
49 partition_columns: Vec<String>,
50 },
51}
52
53impl<'de> Deserialize<'de> for RepartitionFrom {
54 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
55 where
56 D: Deserializer<'de>,
57 {
58 #[derive(Deserialize)]
59 enum CurrentRepartitionFrom {
60 Partitioned {
61 exprs: Vec<PartitionExpr>,
62 #[serde(default)]
63 target_partition_columns: Option<Vec<String>>,
64 },
65 Unpartitioned {
66 partition_columns: Vec<String>,
67 },
68 }
69
70 #[derive(Deserialize)]
71 #[serde(untagged)]
72 enum RepartitionFromRepr {
73 Current(CurrentRepartitionFrom),
74 Legacy(Vec<PartitionExpr>),
75 }
76
77 match RepartitionFromRepr::deserialize(deserializer)? {
78 RepartitionFromRepr::Current(CurrentRepartitionFrom::Partitioned {
79 exprs,
80 target_partition_columns,
81 }) => Ok(Self::Partitioned {
82 exprs,
83 target_partition_columns,
84 }),
85 RepartitionFromRepr::Current(CurrentRepartitionFrom::Unpartitioned {
86 partition_columns,
87 }) => Ok(Self::Unpartitioned { partition_columns }),
88 RepartitionFromRepr::Legacy(exprs) => Ok(Self::Partitioned {
89 exprs,
90 target_partition_columns: None,
91 }),
92 }
93 }
94}
95
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct RepartitionStart {
98 #[serde(alias = "from_exprs")]
99 from: RepartitionFrom,
100 to_exprs: Vec<PartitionExpr>,
101}
102
103impl RepartitionStart {
104 pub fn new(from: RepartitionFrom, to_exprs: Vec<PartitionExpr>) -> Self {
105 Self { from, to_exprs }
106 }
107
108 pub(crate) fn submitted_intent(&self) -> RepartitionSubmittedIntent<'_> {
109 RepartitionSubmittedIntent {
110 from: &self.from,
111 target_partition_exprs: &self.to_exprs,
112 }
113 }
114}
115
116pub(crate) struct RepartitionSubmittedIntent<'a> {
117 from: &'a RepartitionFrom,
118 target_partition_exprs: &'a [PartitionExpr],
119}
120
121impl RepartitionSubmittedIntent<'_> {
122 pub(crate) fn source_type(&self) -> &'static str {
123 match self.from {
124 RepartitionFrom::Partitioned { .. } => "partitioned",
125 RepartitionFrom::Unpartitioned { .. } => "unpartitioned",
126 }
127 }
128
129 pub(crate) fn source_partition_exprs(&self) -> &[PartitionExpr] {
130 match self.from {
131 RepartitionFrom::Partitioned { exprs, .. } => exprs,
132 RepartitionFrom::Unpartitioned { .. } => &[],
133 }
134 }
135
136 pub(crate) fn target_partition_exprs(&self) -> &[PartitionExpr] {
137 self.target_partition_exprs
138 }
139
140 pub(crate) fn target_partition_columns(&self) -> Option<&[String]> {
141 match self.from {
142 RepartitionFrom::Partitioned {
143 target_partition_columns,
144 ..
145 } => target_partition_columns.as_deref(),
146 RepartitionFrom::Unpartitioned { partition_columns } => Some(partition_columns),
147 }
148 }
149}
150
151#[async_trait::async_trait]
152#[typetag::serde]
153impl State for RepartitionStart {
154 async fn next(
155 &mut self,
156 ctx: &mut Context,
157 _: &ProcedureContext,
158 ) -> Result<(Box<dyn State>, Status)> {
159 ensure!(
160 !self.to_exprs.is_empty(),
161 error::InvalidArgumentsSnafu {
162 err_msg: "Repartition expects non-empty target partition expressions".to_string(),
163 }
164 );
165
166 let timer = Instant::now();
167 let (physical_table_id, table_route) = ctx
168 .table_metadata_manager
169 .table_route_manager()
170 .get_physical_table_route(ctx.persistent_ctx.table_id)
171 .await
172 .context(error::TableMetadataManagerSnafu)?;
173 let table_id = ctx.persistent_ctx.table_id;
174 ensure!(
175 physical_table_id == table_id,
176 error::UnexpectedSnafu {
177 violated: format!(
178 "Repartition only works on the physical table, but got logical table: {}, physical table id: {}",
179 table_id, physical_table_id
180 ),
181 }
182 );
183
184 let from_exprs = self.prepare_from(ctx).await?;
185 let plans = Self::build_plan(&table_route, from_exprs, &self.to_exprs)?;
186 let plan_count = plans.len();
187 let total_source_regions: usize = plans.iter().map(|p| p.source_regions.len()).sum();
188 let total_target_regions: usize =
189 plans.iter().map(|p| p.target_partition_exprs.len()).sum();
190 common_telemetry::info!(
191 "Repartition start, table_id: {}, plans: {}, total_source_regions: {}, total_target_regions: {}",
192 table_id,
193 plan_count,
194 total_source_regions,
195 total_target_regions
196 );
197
198 ctx.update_build_plan_elapsed(timer.elapsed());
199
200 if plans.is_empty() {
201 return Ok((Box::new(RepartitionEnd), Status::done()));
202 }
203
204 if ctx.persistent_ctx.partition_metadata_update.is_some() {
205 Ok((
206 Box::new(UpdatePartitionMetadata::new(plans)),
207 Status::executing(true),
208 ))
209 } else {
210 Ok((
211 Box::new(AllocateRegion::new(plans)),
212 Status::executing(false),
213 ))
214 }
215 }
216
217 fn as_any(&self) -> &dyn Any {
218 self
219 }
220}
221
222impl RepartitionStart {
223 async fn prepare_from<'a>(&'a self, ctx: &mut Context) -> Result<&'a [PartitionExpr]> {
224 match &self.from {
225 RepartitionFrom::Partitioned {
226 exprs,
227 target_partition_columns,
228 } => {
229 Self::prepare_partitioned(ctx, target_partition_columns.as_deref()).await?;
230 Ok(exprs)
231 }
232 RepartitionFrom::Unpartitioned { partition_columns } => {
233 Self::prepare_unpartitioned(ctx, partition_columns).await?;
234 Ok(&[])
235 }
236 }
237 }
238
239 async fn prepare_unpartitioned(ctx: &mut Context, partition_columns: &[String]) -> Result<()> {
240 if ctx.persistent_ctx.partition_metadata_update.is_some() {
241 return Ok(());
242 }
243
244 ensure!(
245 !partition_columns.is_empty(),
246 error::InvalidArgumentsSnafu {
247 err_msg: "Unpartitioned repartition expects non-empty partition columns"
248 .to_string(),
249 }
250 );
251
252 let table_info_value = ctx.get_table_info_value().await?;
253 ensure!(
254 table_info_value
255 .table_info
256 .meta
257 .partition_key_indices
258 .is_empty(),
259 error::InvalidArgumentsSnafu {
260 err_msg: format!(
261 "Unpartitioned repartition expects an unpartitioned table, but table {} has partition key indices: {:?}",
262 ctx.persistent_ctx.table_id,
263 table_info_value.table_info.meta.partition_key_indices
264 ),
265 }
266 );
267
268 let schema = &table_info_value.table_info.meta.schema;
269 let partition_key_indices = partition_columns
270 .iter()
271 .map(|column_name| {
272 schema.column_index_by_name(column_name).with_context(|| {
273 error::InvalidArgumentsSnafu {
274 err_msg: format!(
275 "Partition column {} not found in table {}",
276 column_name, ctx.persistent_ctx.table_id
277 ),
278 }
279 })
280 })
281 .collect::<Result<Vec<_>>>()?;
282 ctx.persistent_ctx.partition_metadata_update = Some(
283 PartitionMetadataUpdate::from_unpartitioned(partition_key_indices),
284 );
285
286 Ok(())
287 }
288
289 async fn prepare_partitioned(
290 ctx: &mut Context,
291 target_partition_columns: Option<&[String]>,
292 ) -> Result<()> {
293 let Some(target_partition_columns) = target_partition_columns else {
294 return Ok(());
295 };
296 if ctx.persistent_ctx.partition_metadata_update.is_some() {
297 return Ok(());
298 }
299
300 ensure!(
301 !target_partition_columns.is_empty(),
302 error::InvalidArgumentsSnafu {
303 err_msg: "Partitioned source expects non-empty target partition columns"
304 .to_string(),
305 }
306 );
307
308 let table_info_value = ctx.get_table_info_value().await?;
309 let schema = &table_info_value.table_info.meta.schema;
310 let target_partition_key_indices = target_partition_columns
311 .iter()
312 .map(|column_name| {
313 schema.column_index_by_name(column_name).with_context(|| {
314 error::InvalidArgumentsSnafu {
315 err_msg: format!(
316 "Target partition column {} not found in table {}",
317 column_name, ctx.persistent_ctx.table_id
318 ),
319 }
320 })
321 })
322 .collect::<Result<Vec<_>>>()?;
323 ctx.persistent_ctx.partition_metadata_update =
324 Some(PartitionMetadataUpdate::from_partitioned(
325 table_info_value.table_info.meta.partition_key_indices,
326 target_partition_key_indices,
327 ));
328
329 Ok(())
330 }
331
332 pub(crate) fn build_plan(
333 physical_route: &PhysicalTableRouteValue,
334 from_exprs: &[PartitionExpr],
335 to_exprs: &[PartitionExpr],
336 ) -> Result<Vec<AllocationPlanEntry>> {
337 let subtasks = if from_exprs.is_empty() {
338 Self::default_source_subtasks(to_exprs)?
339 } else {
340 subtask::create_subtasks(from_exprs, to_exprs)
341 .context(error::RepartitionCreateSubtasksSnafu)?
342 };
343 if subtasks.is_empty() {
344 return Ok(vec![]);
345 }
346
347 let src_descriptors = Self::source_region_descriptors(from_exprs, physical_route)?;
348 Ok(Self::build_plan_entries(
349 subtasks,
350 &src_descriptors,
351 to_exprs,
352 ))
353 }
354
355 fn build_plan_entries(
356 subtasks: Vec<RepartitionSubtask>,
357 source_index: &[SourceRegionDescriptor],
358 target_exprs: &[PartitionExpr],
359 ) -> Vec<AllocationPlanEntry> {
360 subtasks
361 .into_iter()
362 .map(|subtask| {
363 let group_id = Uuid::new_v4();
364 let source_regions = subtask
365 .from_expr_indices
366 .iter()
367 .map(|&idx| source_index[idx].clone())
368 .collect::<Vec<_>>();
369
370 let target_partition_exprs = subtask
371 .to_expr_indices
372 .iter()
373 .map(|&idx| target_exprs[idx].clone())
374 .collect::<Vec<_>>();
375 AllocationPlanEntry {
376 group_id,
377 source_regions,
378 target_partition_exprs,
379 transition_map: subtask.transition_map,
380 }
381 })
382 .collect::<Vec<_>>()
383 }
384
385 fn default_source_subtasks(to_exprs: &[PartitionExpr]) -> Result<Vec<RepartitionSubtask>> {
386 ensure!(
387 !to_exprs.is_empty(),
388 error::UnexpectedSnafu {
389 violated: "Default source repartition expects non-empty target partition exprs",
390 }
391 );
392
393 Collider::new(to_exprs).context(error::RepartitionCreateSubtasksSnafu)?;
394
395 let to_expr_indices = (0..to_exprs.len()).collect::<Vec<_>>();
396 Ok(vec![RepartitionSubtask {
397 from_expr_indices: vec![0],
398 to_expr_indices: to_expr_indices.clone(),
399 transition_map: vec![to_expr_indices],
400 }])
401 }
402
403 fn source_region_descriptors(
404 from_exprs: &[PartitionExpr],
405 physical_route: &PhysicalTableRouteValue,
406 ) -> Result<Vec<SourceRegionDescriptor>> {
407 if from_exprs.is_empty() {
408 return Self::default_source_region_descriptors(physical_route);
409 }
410
411 let existing_regions = physical_route
412 .region_routes
413 .iter()
414 .map(|route| (route.region.id, route.region.partition_expr()))
415 .collect::<Vec<_>>();
416
417 let descriptors = from_exprs
418 .iter()
419 .map(|expr| {
420 let expr_json = expr
421 .as_json_str()
422 .context(error::SerializePartitionExprSnafu)?;
423
424 let matched_region_id = existing_regions
425 .iter()
426 .find_map(|(region_id, existing_expr)| {
427 (existing_expr == &expr_json).then_some(*region_id)
428 })
429 .with_context(|| error::RepartitionSourceExprMismatchSnafu { expr: &expr_json })
430 .inspect_err(|_| {
431 debug!("Failed to find matching region for partition expression: {}, existing regions: {:?}", expr_json, existing_regions);
432 })?;
433
434 Ok(SourceRegionDescriptor::partitioned(
435 matched_region_id,
436 expr.clone(),
437 ))
438 })
439 .collect::<Result<Vec<_>>>()?;
440
441 Ok(descriptors)
442 }
443
444 fn default_source_region_descriptors(
445 physical_route: &PhysicalTableRouteValue,
446 ) -> Result<Vec<SourceRegionDescriptor>> {
447 ensure!(
448 physical_route.region_routes.len() == 1,
449 error::UnexpectedSnafu {
450 violated: format!(
451 "Default source repartition expects exactly one source region, but got {}",
452 physical_route.region_routes.len()
453 ),
454 }
455 );
456 let source_region = &physical_route.region_routes[0].region;
457 ensure!(
458 source_region.partition_expr().is_empty(),
459 error::UnexpectedSnafu {
460 violated: format!(
461 "Default source repartition expects an empty partition expr, but got {}",
462 source_region.partition_expr()
463 ),
464 }
465 );
466
467 Ok(vec![SourceRegionDescriptor::Default {
468 region_id: source_region.id,
469 }])
470 }
471}
472
473#[cfg(test)]
474mod tests {
475 use std::sync::Arc;
476
477 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
478 use common_meta::key::table_route::PhysicalTableRouteValue;
479 use common_meta::peer::Peer;
480 use common_meta::rpc::router::{Region, RegionRoute};
481 use common_meta::test_util::MockDatanodeManager;
482 use datatypes::prelude::Value;
483 use partition::expr::{Operand, RestrictedOp};
484 use store_api::storage::RegionId;
485
486 use super::*;
487 use crate::procedure::repartition::test_util::{
488 TestingEnv, new_parent_context, range_expr, test_region_route, test_region_wal_options,
489 };
490
491 fn physical_route(region_routes: Vec<RegionRoute>) -> PhysicalTableRouteValue {
492 PhysicalTableRouteValue::new(region_routes)
493 }
494
495 async fn new_test_context(env: &TestingEnv, table_id: u32) -> Context {
496 env.create_physical_table_metadata_for_repartition(
497 table_id,
498 vec![test_region_route(RegionId::new(table_id, 1), "")],
499 test_region_wal_options(&[1]),
500 )
501 .await;
502 let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler));
503 new_parent_context(env, node_manager, table_id)
504 }
505
506 #[test]
507 fn test_build_plan_with_default_source_region() {
508 let table_id = 1024;
509 let physical_route =
510 physical_route(vec![test_region_route(RegionId::new(table_id, 1), "")]);
511 let to_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
512
513 let plans = RepartitionStart::build_plan(&physical_route, &[], &to_exprs).unwrap();
514
515 assert_eq!(plans.len(), 1);
516 let plan = &plans[0];
517 assert_eq!(
518 plan.source_regions,
519 vec![SourceRegionDescriptor::Default {
520 region_id: RegionId::new(table_id, 1)
521 }]
522 );
523 assert_eq!(plan.target_partition_exprs, to_exprs);
524 assert_eq!(plan.transition_map, vec![vec![0, 1]]);
525 }
526
527 #[test]
528 fn test_build_plan_with_default_source_rejects_non_empty_partition_expr() {
529 let table_id = 1024;
530 let physical_route = physical_route(vec![test_region_route(
531 RegionId::new(table_id, 1),
532 &range_expr("x", 0, 100).as_json_str().unwrap(),
533 )]);
534 let to_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
535
536 let err = RepartitionStart::build_plan(&physical_route, &[], &to_exprs).unwrap_err();
537
538 assert!(err.to_string().contains("empty partition expr"));
539 }
540
541 #[test]
542 fn test_build_plan_with_default_source_rejects_multiple_regions() {
543 let table_id = 1024;
544 let physical_route = physical_route(vec![
545 test_region_route(RegionId::new(table_id, 1), ""),
546 test_region_route(RegionId::new(table_id, 2), ""),
547 ]);
548 let to_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
549
550 let err = RepartitionStart::build_plan(&physical_route, &[], &to_exprs).unwrap_err();
551
552 assert!(err.to_string().contains("exactly one source region"));
553 }
554
555 #[test]
556 fn test_build_plan_with_default_source_rejects_empty_targets() {
557 let table_id = 1024;
558 let physical_route =
559 physical_route(vec![test_region_route(RegionId::new(table_id, 1), "")]);
560
561 let err = RepartitionStart::build_plan(&physical_route, &[], &[]).unwrap_err();
562
563 assert!(err.to_string().contains("non-empty target partition exprs"));
564 }
565
566 #[test]
567 fn test_build_plan_with_default_source_rejects_invalid_targets() {
568 let table_id = 1024;
569 let physical_route =
570 physical_route(vec![test_region_route(RegionId::new(table_id, 1), "")]);
571 let invalid_to_expr = PartitionExpr::new(
572 Operand::Value(Value::Int64(1)),
573 RestrictedOp::Eq,
574 Operand::Value(Value::Int64(2)),
575 );
576
577 let err =
578 RepartitionStart::build_plan(&physical_route, &[], &[invalid_to_expr]).unwrap_err();
579
580 assert!(
581 err.to_string()
582 .contains("Failed to create repartition subtasks")
583 );
584 }
585
586 #[test]
587 fn test_build_plan_keeps_partitioned_source_matching() {
588 let table_id = 1024;
589 let from_exprs = vec![range_expr("x", 0, 100)];
590 let to_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
591 let physical_route = physical_route(vec![RegionRoute {
592 region: Region {
593 id: RegionId::new(table_id, 1),
594 partition_expr: from_exprs[0].as_json_str().unwrap(),
595 ..Default::default()
596 },
597 leader_peer: Some(Peer::empty(1)),
598 ..Default::default()
599 }]);
600
601 let plans = RepartitionStart::build_plan(&physical_route, &from_exprs, &to_exprs).unwrap();
602
603 assert_eq!(plans.len(), 1);
604 assert_eq!(
605 plans[0].source_regions,
606 vec![SourceRegionDescriptor::partitioned(
607 RegionId::new(table_id, 1),
608 from_exprs[0].clone()
609 )]
610 );
611 }
612
613 #[test]
614 fn test_repartition_start_deserializes_legacy_from_exprs() {
615 let from_exprs = vec![range_expr("x", 0, 100)];
616 let to_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
617 let json = serde_json::json!({
618 "from_exprs": from_exprs,
619 "to_exprs": to_exprs,
620 })
621 .to_string();
622
623 let state: RepartitionStart = serde_json::from_str(&json).unwrap();
624
625 let RepartitionFrom::Partitioned {
626 exprs,
627 target_partition_columns,
628 } = state.from
629 else {
630 panic!("expected partition source");
631 };
632 assert_eq!(exprs, vec![range_expr("x", 0, 100)]);
633 assert!(target_partition_columns.is_none());
634 }
635
636 #[test]
637 fn test_repartition_start_deserializes_current_from() {
638 let state = RepartitionStart::new(
639 RepartitionFrom::Unpartitioned {
640 partition_columns: vec!["col1".to_string()],
641 },
642 vec![range_expr("col1", 0, 50)],
643 );
644 let json = serde_json::to_string(&state).unwrap();
645
646 let state: RepartitionStart = serde_json::from_str(&json).unwrap();
647
648 let RepartitionFrom::Unpartitioned { partition_columns } = state.from else {
649 panic!("expected unpartitioned source");
650 };
651 assert_eq!(partition_columns, vec!["col1"]);
652 }
653
654 #[tokio::test]
655 async fn test_partitioned_source_does_not_initialize_partition_metadata_update() {
656 let env = TestingEnv::new();
657 let table_id = 1024;
658 env.create_physical_table_metadata_for_repartition(
659 table_id,
660 vec![test_region_route(
661 RegionId::new(table_id, 1),
662 &range_expr("x", 0, 100).as_json_str().unwrap(),
663 )],
664 test_region_wal_options(&[1]),
665 )
666 .await;
667 let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler));
668 let mut ctx = new_parent_context(&env, node_manager, table_id);
669 let mut state = RepartitionStart::new(
670 RepartitionFrom::Partitioned {
671 exprs: vec![range_expr("x", 0, 100)],
672 target_partition_columns: None,
673 },
674 vec![range_expr("x", 0, 50), range_expr("x", 50, 100)],
675 );
676
677 let (next, status) = state
678 .next(&mut ctx, &TestingEnv::procedure_context())
679 .await
680 .unwrap();
681
682 assert!(!status.need_persist());
683 assert!(next.as_any().is::<AllocateRegion>());
684 assert!(ctx.persistent_ctx.partition_metadata_update.is_none());
685 }
686
687 #[tokio::test]
688 async fn test_partitioned_source_initializes_target_partition_metadata_update() {
689 let env = TestingEnv::new();
690 let table_id = 1024;
691 env.create_physical_table_metadata_for_repartition(
692 table_id,
693 vec![test_region_route(
694 RegionId::new(table_id, 1),
695 &range_expr("x", 0, 100).as_json_str().unwrap(),
696 )],
697 test_region_wal_options(&[1]),
698 )
699 .await;
700 let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler));
701 let mut ctx = new_parent_context(&env, node_manager, table_id);
702 let current = ctx.get_raw_table_info_value().await.unwrap();
703 let mut table_info = current.table_info.clone();
704 table_info.meta.partition_key_indices = vec![0];
705 ctx.update_table_info(¤t, current.update(table_info))
706 .await
707 .unwrap();
708 let mut state = RepartitionStart::new(
709 RepartitionFrom::Partitioned {
710 exprs: vec![range_expr("x", 0, 100)],
711 target_partition_columns: Some(vec!["col2".to_string(), "col1".to_string()]),
712 },
713 vec![range_expr("x", 0, 50), range_expr("x", 50, 100)],
714 );
715
716 let (next, status) = state
717 .next(&mut ctx, &TestingEnv::procedure_context())
718 .await
719 .unwrap();
720
721 assert!(status.need_persist());
722 assert!(next.as_any().is::<UpdatePartitionMetadata>());
723 let update = ctx
724 .persistent_ctx
725 .partition_metadata_update
726 .as_ref()
727 .unwrap();
728 assert_eq!(update.original_partition_key_indices, vec![0]);
729 assert_eq!(update.target_partition_key_indices, vec![2, 0]);
730 assert!(!update.expect_empty_partition_key_indices);
731 }
732
733 #[tokio::test]
734 async fn test_unpartitioned_source_initializes_partition_metadata_update() {
735 let env = TestingEnv::new();
736 let table_id = 1024;
737 let mut ctx = new_test_context(&env, table_id).await;
738 let mut state = RepartitionStart::new(
739 RepartitionFrom::Unpartitioned {
740 partition_columns: vec!["col2".to_string(), "col1".to_string()],
741 },
742 vec![range_expr("col2", 0, 50), range_expr("col2", 50, 100)],
743 );
744
745 let (next, status) = state
746 .next(&mut ctx, &TestingEnv::procedure_context())
747 .await
748 .unwrap();
749
750 assert!(status.need_persist());
751 assert!(next.as_any().is::<UpdatePartitionMetadata>());
752 assert_eq!(
753 ctx.persistent_ctx
754 .partition_metadata_update
755 .as_ref()
756 .unwrap()
757 .target_partition_key_indices,
758 vec![2, 0]
759 );
760 }
761
762 #[tokio::test]
763 async fn test_unpartitioned_source_rejects_existing_partition_metadata() {
764 let env = TestingEnv::new();
765 let table_id = 1024;
766 let mut ctx = new_test_context(&env, table_id).await;
767 let current = ctx.get_raw_table_info_value().await.unwrap();
768 let mut table_info = current.table_info.clone();
769 table_info.meta.partition_key_indices = vec![0];
770 ctx.update_table_info(¤t, current.update(table_info))
771 .await
772 .unwrap();
773 let mut state = RepartitionStart::new(
774 RepartitionFrom::Unpartitioned {
775 partition_columns: vec!["col1".to_string()],
776 },
777 vec![range_expr("col1", 0, 50)],
778 );
779
780 let err = state
781 .next(&mut ctx, &TestingEnv::procedure_context())
782 .await
783 .unwrap_err();
784
785 assert!(err.to_string().contains("expects an unpartitioned table"));
786 assert!(ctx.persistent_ctx.partition_metadata_update.is_none());
787 }
788
789 #[tokio::test]
790 async fn test_repartition_start_rejects_empty_target_partition_exprs() {
791 let env = TestingEnv::new();
792 let table_id = 1024;
793 let mut ctx = new_test_context(&env, table_id).await;
794 let mut state = RepartitionStart::new(
795 RepartitionFrom::Partitioned {
796 exprs: vec![],
797 target_partition_columns: None,
798 },
799 vec![],
800 );
801
802 let err = state
803 .next(&mut ctx, &TestingEnv::procedure_context())
804 .await
805 .unwrap_err();
806
807 assert!(
808 err.to_string()
809 .contains("non-empty target partition expressions")
810 );
811 }
812
813 #[tokio::test]
814 async fn test_unpartitioned_source_rejects_empty_target_partition_exprs() {
815 let env = TestingEnv::new();
816 let table_id = 1024;
817 let mut ctx = new_test_context(&env, table_id).await;
818 let mut state = RepartitionStart::new(
819 RepartitionFrom::Unpartitioned {
820 partition_columns: vec!["col1".to_string()],
821 },
822 vec![],
823 );
824
825 let err = state
826 .next(&mut ctx, &TestingEnv::procedure_context())
827 .await
828 .unwrap_err();
829
830 assert!(
831 err.to_string()
832 .contains("non-empty target partition expressions")
833 );
834 assert!(ctx.persistent_ctx.partition_metadata_update.is_none());
835 }
836
837 #[tokio::test]
838 async fn test_unpartitioned_source_rejects_empty_partition_columns() {
839 let env = TestingEnv::new();
840 let table_id = 1024;
841 let mut ctx = new_test_context(&env, table_id).await;
842 let mut state = RepartitionStart::new(
843 RepartitionFrom::Unpartitioned {
844 partition_columns: vec![],
845 },
846 vec![range_expr("col1", 0, 50)],
847 );
848
849 let err = state
850 .next(&mut ctx, &TestingEnv::procedure_context())
851 .await
852 .unwrap_err();
853
854 assert!(err.to_string().contains("non-empty partition columns"));
855 assert!(ctx.persistent_ctx.partition_metadata_update.is_none());
856 }
857
858 #[tokio::test]
859 async fn test_unpartitioned_source_rejects_missing_partition_column() {
860 let env = TestingEnv::new();
861 let table_id = 1024;
862 let mut ctx = new_test_context(&env, table_id).await;
863 let mut state = RepartitionStart::new(
864 RepartitionFrom::Unpartitioned {
865 partition_columns: vec!["missing_col".to_string()],
866 },
867 vec![range_expr("col1", 0, 50)],
868 );
869
870 let err = state
871 .next(&mut ctx, &TestingEnv::procedure_context())
872 .await
873 .unwrap_err();
874
875 assert!(
876 err.to_string()
877 .contains("Partition column missing_col not found")
878 );
879 assert!(ctx.persistent_ctx.partition_metadata_update.is_none());
880 }
881}