common_meta/key/
topic_region.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::fmt::{self, Display};
17
18use common_wal::options::WalOptions;
19use serde::{Deserialize, Serialize};
20use snafu::OptionExt;
21use store_api::storage::RegionId;
22use table::metadata::TableId;
23
24use crate::error::{Error, InvalidMetadataSnafu, Result};
25use crate::key::{MetadataKey, MetadataValue, TOPIC_REGION_PATTERN, TOPIC_REGION_PREFIX};
26use crate::kv_backend::KvBackendRef;
27use crate::kv_backend::txn::{Txn, TxnOp};
28use crate::rpc::KeyValue;
29use crate::rpc::store::{
30    BatchDeleteRequest, BatchGetRequest, BatchPutRequest, PutRequest, RangeRequest,
31};
32use crate::wal_provider::RegionWalOptions;
33
34// The TopicRegionKey is a key for the topic-region mapping in the kvbackend.
35// The layout of the key is `__topic_region/{topic_name}/{region_id}`.
36#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37pub struct TopicRegionKey<'a> {
38    pub region_id: RegionId,
39    pub topic: &'a str,
40}
41
42/// Represents additional information for a region when using a shared WAL.
43#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
44pub struct TopicRegionValue {
45    #[serde(default, skip_serializing_if = "Option::is_none")]
46    pub checkpoint: Option<ReplayCheckpoint>,
47}
48
49#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq, PartialOrd, Ord)]
50pub struct ReplayCheckpoint {
51    #[serde(default)]
52    pub entry_id: u64,
53    #[serde(default, skip_serializing_if = "Option::is_none")]
54    pub metadata_entry_id: Option<u64>,
55}
56
57impl<'a> TopicRegionKey<'a> {
58    pub fn new(region_id: RegionId, topic: &'a str) -> Self {
59        Self { region_id, topic }
60    }
61
62    pub fn range_topic_key(topic: &str) -> String {
63        format!("{}/{}/", TOPIC_REGION_PREFIX, topic)
64    }
65}
66
67impl<'a> MetadataKey<'a, TopicRegionKey<'a>> for TopicRegionKey<'a> {
68    fn to_bytes(&self) -> Vec<u8> {
69        self.to_string().into_bytes()
70    }
71
72    fn from_bytes(bytes: &'a [u8]) -> Result<TopicRegionKey<'a>> {
73        let key = std::str::from_utf8(bytes).map_err(|e| {
74            InvalidMetadataSnafu {
75                err_msg: format!(
76                    "TopicRegionKey '{}' is not a valid UTF8 string: {e}",
77                    String::from_utf8_lossy(bytes)
78                ),
79            }
80            .build()
81        })?;
82        TopicRegionKey::try_from(key)
83    }
84}
85
86impl Display for TopicRegionKey<'_> {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        write!(
89            f,
90            "{}{}",
91            Self::range_topic_key(self.topic),
92            self.region_id.as_u64()
93        )
94    }
95}
96
97impl<'a> TryFrom<&'a str> for TopicRegionKey<'a> {
98    type Error = Error;
99
100    /// Value is of format `{prefix}/{topic}/{region_id}`
101    fn try_from(value: &'a str) -> Result<TopicRegionKey<'a>> {
102        let captures = TOPIC_REGION_PATTERN
103            .captures(value)
104            .context(InvalidMetadataSnafu {
105                err_msg: format!("Invalid TopicRegionKey: {}", value),
106            })?;
107        let topic = captures.get(1).map(|m| m.as_str()).unwrap();
108        let region_id = captures[2].parse::<u64>().map_err(|_| {
109            InvalidMetadataSnafu {
110                err_msg: format!("Invalid region id in TopicRegionKey: {}", value),
111            }
112            .build()
113        })?;
114        Ok(TopicRegionKey {
115            region_id: RegionId::from_u64(region_id),
116            topic,
117        })
118    }
119}
120
121impl ReplayCheckpoint {
122    /// Creates a new [`ReplayCheckpoint`] with the given entry id and metadata entry id.
123    pub fn new(entry_id: u64, metadata_entry_id: Option<u64>) -> Self {
124        Self {
125            entry_id,
126            metadata_entry_id,
127        }
128    }
129
130    /// Merges the checkpoint with the topic pruned entry id.
131    pub fn merge_with_topic_pruned_entry_id(
132        checkpoint: Option<Self>,
133        pruned_entry_id: Option<u64>,
134        is_metric_engine: bool,
135    ) -> Option<Self> {
136        match (checkpoint, pruned_entry_id) {
137            (Some(checkpoint), Some(pruned_entry_id)) => Some(Self {
138                entry_id: checkpoint.entry_id.max(pruned_entry_id),
139                metadata_entry_id: if is_metric_engine {
140                    Some(
141                        checkpoint
142                            .metadata_entry_id
143                            .unwrap_or_default()
144                            .max(pruned_entry_id),
145                    )
146                } else {
147                    checkpoint.metadata_entry_id
148                },
149            }),
150            (None, Some(pruned_entry_id)) => Some(Self {
151                entry_id: pruned_entry_id,
152                metadata_entry_id: is_metric_engine.then_some(pruned_entry_id),
153            }),
154            (checkpoint, None) => checkpoint,
155        }
156    }
157}
158
159impl TopicRegionValue {
160    /// Creates a new [`TopicRegionValue`] with the given checkpoint.
161    pub fn new(checkpoint: Option<ReplayCheckpoint>) -> Self {
162        Self { checkpoint }
163    }
164
165    /// Returns the minimum entry id of the region.
166    ///
167    /// If the metadata entry id is not set, it returns the entry id.
168    pub fn min_entry_id(&self) -> Option<u64> {
169        match self.checkpoint {
170            Some(ReplayCheckpoint {
171                entry_id,
172                metadata_entry_id,
173            }) => match metadata_entry_id {
174                Some(metadata_entry_id) => Some(entry_id.min(metadata_entry_id)),
175                None => Some(entry_id),
176            },
177            None => None,
178        }
179    }
180}
181
182fn topic_region_decoder(value: &KeyValue) -> Result<(TopicRegionKey<'_>, TopicRegionValue)> {
183    let key = TopicRegionKey::from_bytes(&value.key)?;
184    let value = if value.value.is_empty() {
185        TopicRegionValue::default()
186    } else {
187        TopicRegionValue::try_from_raw_value(&value.value)?
188    };
189    Ok((key, value))
190}
191
192/// Manages map of topics and regions in kvbackend.
193pub struct TopicRegionManager {
194    kv_backend: KvBackendRef,
195}
196
197impl TopicRegionManager {
198    pub fn new(kv_backend: KvBackendRef) -> Self {
199        Self { kv_backend }
200    }
201
202    pub async fn put(&self, key: TopicRegionKey<'_>) -> Result<()> {
203        let put_req = PutRequest {
204            key: key.to_bytes(),
205            value: vec![],
206            prev_kv: false,
207        };
208        self.kv_backend.put(put_req).await?;
209        Ok(())
210    }
211
212    pub async fn batch_get(
213        &self,
214        keys: Vec<TopicRegionKey<'_>>,
215    ) -> Result<HashMap<RegionId, TopicRegionValue>> {
216        let raw_keys = keys.iter().map(|key| key.to_bytes()).collect::<Vec<_>>();
217        let req = BatchGetRequest { keys: raw_keys };
218        let resp = self.kv_backend.batch_get(req).await?;
219
220        let v = resp
221            .kvs
222            .into_iter()
223            .map(|kv| topic_region_decoder(&kv).map(|(key, value)| (key.region_id, value)))
224            .collect::<Result<HashMap<_, _>>>()?;
225
226        Ok(v)
227    }
228
229    pub async fn get(&self, key: TopicRegionKey<'_>) -> Result<Option<TopicRegionValue>> {
230        let key_bytes = key.to_bytes();
231        let resp = self.kv_backend.get(&key_bytes).await?;
232        let value = resp
233            .map(|kv| topic_region_decoder(&kv).map(|(_, value)| value))
234            .transpose()?;
235
236        Ok(value)
237    }
238
239    pub async fn batch_put(
240        &self,
241        keys: &[(TopicRegionKey<'_>, Option<TopicRegionValue>)],
242    ) -> Result<()> {
243        let req = BatchPutRequest {
244            kvs: keys
245                .iter()
246                .map(|(key, value)| {
247                    let value = value
248                        .map(|v| v.try_as_raw_value())
249                        .transpose()?
250                        .unwrap_or_default();
251
252                    Ok(KeyValue {
253                        key: key.to_bytes(),
254                        value,
255                    })
256                })
257                .collect::<Result<Vec<_>>>()?,
258            prev_kv: false,
259        };
260        self.kv_backend.batch_put(req).await?;
261        Ok(())
262    }
263
264    /// Build a create topic region mapping transaction. It only executes while the primary keys comparing successes.
265    pub fn build_create_txn(
266        &self,
267        table_id: TableId,
268        region_wal_options: &RegionWalOptions,
269    ) -> Result<Txn> {
270        let topic_region_mapping = self.get_topic_region_mapping(table_id, region_wal_options);
271        let topic_region_keys = topic_region_mapping
272            .iter()
273            .map(|(region_id, topic)| TopicRegionKey::new(*region_id, topic))
274            .collect::<Vec<_>>();
275        let operations = topic_region_keys
276            .into_iter()
277            .map(|key| TxnOp::Put(key.to_bytes(), vec![]))
278            .collect::<Vec<_>>();
279        Ok(Txn::new().and_then(operations))
280    }
281
282    /// Build a update topic region mapping transaction.
283    pub fn build_update_txn(
284        &self,
285        table_id: TableId,
286        old_region_wal_options: &RegionWalOptions,
287        new_region_wal_options: &RegionWalOptions,
288    ) -> Result<Txn> {
289        let old_mapping = self.get_topic_region_mapping(table_id, old_region_wal_options);
290        let new_mapping = self.get_topic_region_mapping(table_id, new_region_wal_options);
291
292        // Convert to HashMap for easier lookup: RegionId -> Topic
293        let old_map: HashMap<RegionId, &str> = old_mapping.into_iter().collect();
294        let new_map: HashMap<RegionId, &str> = new_mapping.into_iter().collect();
295        let mut ops = Vec::new();
296
297        // Check for deletes (in old but not in new, or topic changed)
298        for (region_id, old_topic) in &old_map {
299            match new_map.get(region_id) {
300                Some(new_topic) if *new_topic == *old_topic => {
301                    // Same topic, do nothing (preserve checkpoint)
302                }
303                _ => {
304                    // Removed or topic changed -> Delete old
305                    let key = TopicRegionKey::new(*region_id, old_topic);
306                    ops.push(TxnOp::Delete(key.to_bytes()));
307                }
308            }
309        }
310
311        // Check for adds (in new but not in old, or topic changed)
312        for (region_id, new_topic) in &new_map {
313            match old_map.get(region_id) {
314                Some(old_topic) if *old_topic == *new_topic => {
315                    // Same topic, already handled (do nothing)
316                }
317                _ => {
318                    // New or topic changed -> Put new
319                    let key = TopicRegionKey::new(*region_id, new_topic);
320                    // Initialize with empty value (default TopicRegionValue)
321                    ops.push(TxnOp::Put(key.to_bytes(), vec![]));
322                }
323            }
324        }
325
326        Ok(Txn::new().and_then(ops))
327    }
328
329    /// Returns the map of [`RegionId`] to their corresponding topic [`TopicRegionValue`].
330    pub async fn regions(&self, topic: &str) -> Result<HashMap<RegionId, TopicRegionValue>> {
331        let prefix = TopicRegionKey::range_topic_key(topic);
332        let req = RangeRequest::new().with_prefix(prefix.as_bytes());
333        let resp = self.kv_backend.range(req).await?;
334        let region_ids = resp
335            .kvs
336            .iter()
337            .map(topic_region_decoder)
338            .collect::<Result<Vec<_>>>()?;
339        Ok(region_ids
340            .into_iter()
341            .map(|(key, value)| (key.region_id, value))
342            .collect())
343    }
344
345    pub async fn delete(&self, key: TopicRegionKey<'_>) -> Result<()> {
346        let raw_key = key.to_bytes();
347        self.kv_backend.delete(&raw_key, false).await?;
348        Ok(())
349    }
350
351    pub async fn batch_delete(&self, keys: Vec<TopicRegionKey<'_>>) -> Result<()> {
352        let raw_keys = keys.iter().map(|key| key.to_bytes()).collect::<Vec<_>>();
353        let req = BatchDeleteRequest {
354            keys: raw_keys,
355            prev_kv: false,
356        };
357        self.kv_backend.batch_delete(req).await?;
358        Ok(())
359    }
360
361    /// Retrieves a mapping of [`RegionId`]s to their corresponding topics name
362    /// based on the provided table ID and WAL options.
363    ///
364    /// # Returns
365    /// A vector of tuples, where each tuple contains a [`RegionId`] and its corresponding topic name.
366    pub fn get_topic_region_mapping<'a>(
367        &self,
368        table_id: TableId,
369        region_wal_options: &'a RegionWalOptions,
370    ) -> Vec<(RegionId, &'a str)> {
371        region_wal_options
372            .keys()
373            .filter_map(
374                |region_number| match region_wal_options.get(region_number) {
375                    Some(WalOptions::Kafka(kafka)) => {
376                        let region_id = RegionId::new(table_id, *region_number);
377                        Some((region_id, kafka.topic.as_str()))
378                    }
379                    Some(WalOptions::RaftEngine) => None,
380                    Some(WalOptions::Noop) => None,
381                    Some(WalOptions::ObjectStore(_)) => None,
382                    None => None,
383                },
384            )
385            .collect::<Vec<_>>()
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use std::sync::Arc;
392
393    use common_wal::options::KafkaWalOptions;
394
395    use super::*;
396    use crate::kv_backend::memory::MemoryKvBackend;
397
398    #[test]
399    fn test_merge_checkpoint_with_topic_pruned_entry_id_missing_pruned() {
400        let checkpoint = Some(ReplayCheckpoint::new(10, None));
401
402        assert_eq!(
403            ReplayCheckpoint::merge_with_topic_pruned_entry_id(checkpoint, None, true),
404            checkpoint
405        );
406    }
407
408    #[test]
409    fn test_merge_checkpoint_with_topic_pruned_entry_id_creates_checkpoint() {
410        assert_eq!(
411            ReplayCheckpoint::merge_with_topic_pruned_entry_id(None, Some(10), true),
412            Some(ReplayCheckpoint::new(10, Some(10)))
413        );
414    }
415
416    #[test]
417    fn test_merge_checkpoint_with_topic_pruned_entry_id_updates_both_ids() {
418        let checkpoint = ReplayCheckpoint::new(10, Some(5));
419
420        assert_eq!(
421            ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(20), true),
422            Some(ReplayCheckpoint::new(20, Some(20)))
423        );
424    }
425
426    #[test]
427    fn test_merge_checkpoint_with_topic_pruned_entry_id_preserves_larger_ids() {
428        let checkpoint = ReplayCheckpoint::new(30, Some(40));
429
430        assert_eq!(
431            ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(20), true),
432            Some(checkpoint)
433        );
434    }
435
436    #[test]
437    fn test_merge_checkpoint_with_topic_pruned_entry_id_for_mito() {
438        assert_eq!(
439            ReplayCheckpoint::merge_with_topic_pruned_entry_id(None, Some(10), false),
440            Some(ReplayCheckpoint::new(10, None))
441        );
442
443        let checkpoint = ReplayCheckpoint::new(5, Some(8));
444        assert_eq!(
445            ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(10), false),
446            Some(ReplayCheckpoint::new(10, Some(8)))
447        );
448    }
449
450    #[tokio::test]
451    async fn test_topic_region_manager() {
452        let kv_backend = Arc::new(MemoryKvBackend::default());
453        let manager = TopicRegionManager::new(kv_backend.clone());
454
455        let topics = (0..16).map(|i| format!("topic_{}", i)).collect::<Vec<_>>();
456        let keys = (0..64)
457            .map(|i| {
458                (
459                    TopicRegionKey::new(RegionId::from_u64(i), &topics[(i % 16) as usize]),
460                    None,
461                )
462            })
463            .collect::<Vec<_>>();
464
465        manager.batch_put(&keys).await.unwrap();
466        let mut key_values = manager
467            .regions(&topics[0])
468            .await
469            .unwrap()
470            .into_keys()
471            .collect::<Vec<_>>();
472        let expected = keys
473            .iter()
474            .filter_map(|(key, _)| {
475                if key.topic == topics[0] {
476                    Some(key.region_id)
477                } else {
478                    None
479                }
480            })
481            .collect::<Vec<_>>();
482        key_values.sort_by_key(|id| id.as_u64());
483        assert_eq!(key_values, expected);
484
485        let key = TopicRegionKey::new(RegionId::from_u64(0), "topic_0");
486        manager.delete(key.clone()).await.unwrap();
487        let mut key_values = manager
488            .regions(&topics[0])
489            .await
490            .unwrap()
491            .into_keys()
492            .collect::<Vec<_>>();
493        let expected = keys
494            .iter()
495            .filter_map(|(key, _)| {
496                if key.topic == topics[0] && key.region_id != RegionId::from_u64(0) {
497                    Some(key.region_id)
498                } else {
499                    None
500                }
501            })
502            .collect::<Vec<_>>();
503        key_values.sort_by_key(|id| id.as_u64());
504        assert_eq!(key_values, expected);
505    }
506
507    #[test]
508    fn test_topic_region_map() {
509        let kv_backend = Arc::new(MemoryKvBackend::default());
510        let manager = TopicRegionManager::new(kv_backend.clone());
511
512        let table_id = 1;
513        let region_wal_options = (0..64)
514            .map(|i| {
515                let region_number = i;
516                let wal_options = if i % 2 == 0 {
517                    WalOptions::Kafka(KafkaWalOptions::new(format!("topic_{}", i)))
518                } else {
519                    WalOptions::RaftEngine
520                };
521                (region_number, wal_options)
522            })
523            .collect::<HashMap<_, _>>();
524
525        let mut topic_region_mapping =
526            manager.get_topic_region_mapping(table_id, &region_wal_options);
527        let mut expected = (0..64)
528            .filter_map(|i| {
529                if i % 2 == 0 {
530                    Some((RegionId::new(table_id, i), format!("topic_{}", i)))
531                } else {
532                    None
533                }
534            })
535            .collect::<Vec<_>>();
536        topic_region_mapping.sort_by_key(|(region_id, _)| region_id.as_u64());
537        let topic_region_map = topic_region_mapping
538            .iter()
539            .map(|(region_id, topic)| (*region_id, topic.to_string()))
540            .collect::<Vec<_>>();
541        expected.sort_by_key(|(region_id, _)| region_id.as_u64());
542        assert_eq!(topic_region_map, expected);
543    }
544
545    #[test]
546    fn test_topic_region_key_is_match() {
547        let key = "__topic_region/6f153a64-7fac-4cf6-8b0b-a7967dd73879_2/4410931412992";
548        let topic_region_key = TopicRegionKey::try_from(key).unwrap();
549        assert_eq!(
550            topic_region_key.topic,
551            "6f153a64-7fac-4cf6-8b0b-a7967dd73879_2"
552        );
553        assert_eq!(
554            topic_region_key.region_id,
555            RegionId::from_u64(4410931412992)
556        );
557    }
558
559    #[test]
560    fn test_build_create_txn() {
561        let kv_backend = Arc::new(MemoryKvBackend::default());
562        let manager = TopicRegionManager::new(kv_backend.clone());
563        let table_id = 1;
564        let region_wal_options = vec![
565            (
566                0,
567                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
568            ),
569            (
570                1,
571                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())),
572            ),
573            (2, WalOptions::RaftEngine), // Should be ignored
574        ]
575        .into_iter()
576        .collect::<HashMap<_, _>>();
577
578        let txn = manager
579            .build_create_txn(table_id, &region_wal_options)
580            .unwrap();
581
582        // Verify the transaction contains correct operations
583        // Should create mappings for region 0 and 1, but not region 2 (RaftEngine)
584        let ops = txn.req().success.clone();
585        assert_eq!(ops.len(), 2);
586
587        let keys: Vec<_> = ops
588            .iter()
589            .filter_map(|op| {
590                if let TxnOp::Put(key, _) = op {
591                    TopicRegionKey::from_bytes(key).ok()
592                } else {
593                    None
594                }
595            })
596            .collect();
597
598        assert_eq!(keys.len(), 2);
599        let region_ids: Vec<_> = keys.iter().map(|k| k.region_id).collect();
600        assert!(region_ids.contains(&RegionId::new(table_id, 0)));
601        assert!(region_ids.contains(&RegionId::new(table_id, 1)));
602        assert!(!region_ids.contains(&RegionId::new(table_id, 2)));
603
604        // Verify topics are correct
605        for key in keys {
606            match key.region_id.region_number() {
607                0 => assert_eq!(key.topic, "topic_0"),
608                1 => assert_eq!(key.topic, "topic_1"),
609                _ => panic!("Unexpected region number"),
610            }
611        }
612    }
613
614    #[test]
615    fn test_build_update_txn_add_new_region() {
616        let kv_backend = Arc::new(MemoryKvBackend::default());
617        let manager = TopicRegionManager::new(kv_backend.clone());
618        let table_id = 1;
619        let old_region_wal_options = vec![(
620            0,
621            WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
622        )]
623        .into_iter()
624        .collect::<HashMap<_, _>>();
625        let new_region_wal_options = vec![
626            (
627                0,
628                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
629            ),
630            (
631                1,
632                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())),
633            ),
634        ]
635        .into_iter()
636        .collect::<HashMap<_, _>>();
637        let txn = manager
638            .build_update_txn(table_id, &old_region_wal_options, &new_region_wal_options)
639            .unwrap();
640        let ops = txn.req().success.clone();
641        // Should only have Put for new region 1 (region 0 unchanged)
642        assert_eq!(ops.len(), 1);
643        if let TxnOp::Put(key, _) = &ops[0] {
644            let topic_key = TopicRegionKey::from_bytes(key).unwrap();
645            assert_eq!(topic_key.region_id, RegionId::new(table_id, 1));
646            assert_eq!(topic_key.topic, "topic_1");
647        } else {
648            panic!("Expected Put operation");
649        }
650    }
651
652    #[test]
653    fn test_build_update_txn_remove_region() {
654        let kv_backend = Arc::new(MemoryKvBackend::default());
655        let manager = TopicRegionManager::new(kv_backend.clone());
656        let table_id = 1;
657        let old_region_wal_options = vec![
658            (
659                0,
660                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
661            ),
662            (
663                1,
664                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())),
665            ),
666        ]
667        .into_iter()
668        .collect::<HashMap<_, _>>();
669        let new_region_wal_options = vec![(
670            0,
671            WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
672        )]
673        .into_iter()
674        .collect::<HashMap<_, _>>();
675        let txn = manager
676            .build_update_txn(table_id, &old_region_wal_options, &new_region_wal_options)
677            .unwrap();
678        let ops = txn.req().success.clone();
679        // Should only have Delete for removed region 1 (region 0 unchanged)
680        assert_eq!(ops.len(), 1);
681        match &ops[0] {
682            TxnOp::Delete(key) => {
683                let topic_key = TopicRegionKey::from_bytes(key).unwrap();
684                assert_eq!(topic_key.region_id, RegionId::new(table_id, 1));
685                assert_eq!(topic_key.topic, "topic_1");
686            }
687            TxnOp::Put(_, _) | TxnOp::Get(_) => {
688                panic!("Expected Delete operation");
689            }
690        }
691    }
692
693    #[test]
694    fn test_build_update_txn_change_topic() {
695        let kv_backend = Arc::new(MemoryKvBackend::default());
696        let manager = TopicRegionManager::new(kv_backend.clone());
697        let table_id = 1;
698        let old_region_wal_options = vec![(
699            0,
700            WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
701        )]
702        .into_iter()
703        .collect::<HashMap<_, _>>();
704        let new_region_wal_options = vec![(
705            0,
706            WalOptions::Kafka(KafkaWalOptions::new("topic_0_new".to_string())),
707        )]
708        .into_iter()
709        .collect::<HashMap<_, _>>();
710        let txn = manager
711            .build_update_txn(table_id, &old_region_wal_options, &new_region_wal_options)
712            .unwrap();
713        let ops = txn.req().success.clone();
714        // Should have Delete for old topic and Put for new topic
715        assert_eq!(ops.len(), 2);
716
717        let mut delete_found = false;
718        let mut put_found = false;
719        for op in ops {
720            match op {
721                TxnOp::Delete(key) => {
722                    let topic_key = TopicRegionKey::from_bytes(&key).unwrap();
723                    assert_eq!(topic_key.region_id, RegionId::new(table_id, 0));
724                    assert_eq!(topic_key.topic, "topic_0");
725                    delete_found = true;
726                }
727                TxnOp::Put(key, _) => {
728                    let topic_key = TopicRegionKey::from_bytes(&key).unwrap();
729                    assert_eq!(topic_key.region_id, RegionId::new(table_id, 0));
730                    assert_eq!(topic_key.topic, "topic_0_new");
731                    put_found = true;
732                }
733                TxnOp::Get(_) => {
734                    // Get operations shouldn't appear in this context
735                    panic!("Unexpected Get operation in update transaction");
736                }
737            }
738        }
739        assert!(delete_found, "Expected Delete operation for old topic");
740        assert!(put_found, "Expected Put operation for new topic");
741    }
742
743    #[test]
744    fn test_build_update_txn_no_change() {
745        let kv_backend = Arc::new(MemoryKvBackend::default());
746        let manager = TopicRegionManager::new(kv_backend.clone());
747        let table_id = 1;
748        let region_wal_options = vec![
749            (
750                0,
751                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
752            ),
753            (
754                1,
755                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())),
756            ),
757        ]
758        .into_iter()
759        .collect::<HashMap<_, _>>();
760        let txn = manager
761            .build_update_txn(table_id, &region_wal_options, &region_wal_options)
762            .unwrap();
763        // Should have no operations when nothing changes (preserves checkpoint)
764        let ops = txn.req().success.clone();
765        assert_eq!(ops.len(), 0);
766    }
767
768    #[test]
769    fn test_build_update_txn_mixed_scenarios() {
770        let kv_backend = Arc::new(MemoryKvBackend::default());
771        let manager = TopicRegionManager::new(kv_backend.clone());
772        let table_id = 1;
773        let old_region_wal_options = vec![
774            (
775                0,
776                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
777            ),
778            (
779                1,
780                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())),
781            ),
782            (
783                2,
784                WalOptions::Kafka(KafkaWalOptions::new("topic_2".to_string())),
785            ),
786        ]
787        .into_iter()
788        .collect::<HashMap<_, _>>();
789        let new_region_wal_options = vec![
790            (
791                0,
792                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())), // Unchanged
793            ),
794            (
795                1,
796                WalOptions::Kafka(KafkaWalOptions::new("topic_1_new".to_string())), // Topic changed
797            ),
798            // Region 2 removed
799            (
800                3,
801                WalOptions::Kafka(KafkaWalOptions::new("topic_3".to_string())), // New region
802            ),
803        ]
804        .into_iter()
805        .collect::<HashMap<_, _>>();
806        let txn = manager
807            .build_update_txn(table_id, &old_region_wal_options, &new_region_wal_options)
808            .unwrap();
809
810        let ops = txn.req().success.clone();
811        // Should have:
812        // - Delete for region 2 (removed)
813        // - Delete for region 1 old topic (topic changed)
814        // - Put for region 1 new topic (topic changed)
815        // - Put for region 3 (new)
816        // Region 0 unchanged, so no operation
817        assert_eq!(ops.len(), 4);
818
819        let mut delete_ops = 0;
820        let mut put_ops = 0;
821        let mut delete_region_2 = false;
822        let mut delete_region_1_old = false;
823        let mut put_region_1_new = false;
824        let mut put_region_3 = false;
825
826        for op in ops {
827            match op {
828                TxnOp::Delete(key) => {
829                    delete_ops += 1;
830                    let topic_key = TopicRegionKey::from_bytes(&key).unwrap();
831                    match topic_key.region_id.region_number() {
832                        1 => {
833                            assert_eq!(topic_key.topic, "topic_1");
834                            delete_region_1_old = true;
835                        }
836                        2 => {
837                            assert_eq!(topic_key.topic, "topic_2");
838                            delete_region_2 = true;
839                        }
840                        _ => panic!("Unexpected delete operation for region"),
841                    }
842                }
843                TxnOp::Put(key, _) => {
844                    put_ops += 1;
845                    let topic_key: TopicRegionKey<'_> = TopicRegionKey::from_bytes(&key).unwrap();
846                    match topic_key.region_id.region_number() {
847                        1 => {
848                            assert_eq!(topic_key.topic, "topic_1_new");
849                            put_region_1_new = true;
850                        }
851                        3 => {
852                            assert_eq!(topic_key.topic, "topic_3");
853                            put_region_3 = true;
854                        }
855                        _ => panic!("Unexpected put operation for region"),
856                    }
857                }
858                TxnOp::Get(_) => {
859                    panic!("Unexpected Get operation in update transaction");
860                }
861            }
862        }
863
864        assert_eq!(delete_ops, 2);
865        assert_eq!(put_ops, 2);
866        assert!(delete_region_2, "Expected delete for removed region 2");
867        assert!(
868            delete_region_1_old,
869            "Expected delete for region 1 old topic"
870        );
871        assert!(put_region_1_new, "Expected put for region 1 new topic");
872        assert!(put_region_3, "Expected put for new region 3");
873    }
874
875    #[test]
876    fn test_build_update_txn_with_raft_engine() {
877        let kv_backend = Arc::new(MemoryKvBackend::default());
878        let manager = TopicRegionManager::new(kv_backend.clone());
879        let table_id = 1;
880        let old_region_wal_options = vec![
881            (
882                0,
883                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
884            ),
885            (1, WalOptions::RaftEngine), // Should be ignored
886        ]
887        .into_iter()
888        .collect::<HashMap<_, _>>();
889        let new_region_wal_options = vec![
890            (
891                0,
892                WalOptions::Kafka(KafkaWalOptions::new("topic_0".to_string())),
893            ),
894            (
895                1,
896                WalOptions::Kafka(KafkaWalOptions::new("topic_1".to_string())), // Changed from RaftEngine to Kafka
897            ),
898        ]
899        .into_iter()
900        .collect::<HashMap<_, _>>();
901        let txn = manager
902            .build_update_txn(table_id, &old_region_wal_options, &new_region_wal_options)
903            .unwrap();
904        let ops = txn.req().success.clone();
905        // Should only have Put for region 1 (new Kafka topic)
906        // Region 0 unchanged, so no operation
907        // Region 1 was RaftEngine before (not tracked), so only Put needed
908        assert_eq!(ops.len(), 1);
909        match &ops[0] {
910            TxnOp::Put(key, _) => {
911                let topic_key = TopicRegionKey::from_bytes(key).unwrap();
912                assert_eq!(topic_key.region_id, RegionId::new(table_id, 1));
913                assert_eq!(topic_key.topic, "topic_1");
914            }
915            TxnOp::Delete(_) | TxnOp::Get(_) => {
916                panic!("Expected Put operation for new Kafka region");
917            }
918        }
919    }
920}
common_meta/key/topic_region.rs

common_meta/key/
topic_region.rs
