Skip to main content

common_meta/key/
flow.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
15pub mod flow_info;
16pub(crate) mod flow_name;
17pub mod flow_route;
18pub mod flow_state;
19mod flownode_addr_helper;
20pub(crate) mod flownode_flow;
21pub(crate) mod table_flow;
22use std::collections::BTreeMap;
23use std::ops::Deref;
24use std::sync::Arc;
25
26use common_telemetry::info;
27use flow_route::{FlowRouteKey, FlowRouteManager, FlowRouteValue};
28use snafu::{OptionExt, ensure};
29use table_flow::TableFlowValue;
30
31use self::flow_info::{FlowInfoKey, FlowInfoValue};
32use self::flow_name::FlowNameKey;
33use self::flownode_flow::FlownodeFlowKey;
34use self::table_flow::TableFlowKey;
35use crate::ensure_values;
36use crate::error::{self, Result};
37use crate::key::flow::flow_info::FlowInfoManager;
38use crate::key::flow::flow_name::FlowNameManager;
39use crate::key::flow::flow_state::FlowStateManager;
40use crate::key::flow::flownode_flow::FlownodeFlowManager;
41pub use crate::key::flow::table_flow::{TableFlowManager, TableFlowManagerRef};
42use crate::key::txn_helper::TxnOpGetResponseSet;
43use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey};
44use crate::kv_backend::KvBackendRef;
45use crate::kv_backend::txn::Txn;
46use crate::rpc::store::BatchDeleteRequest;
47
48pub const FLOW_KEY_PREFIX: &str = "__flow";
49
50/// Returns a flow metadata key under the `__flow` scope.
51pub fn scoped_flow_key(inner_key: &str) -> String {
52    format!("{FLOW_KEY_PREFIX}/{inner_key}")
53}
54
55/// Returns a flow metadata key prefix under the `__flow` scope.
56pub fn scoped_flow_key_prefix(inner_prefix: &str) -> String {
57    scoped_flow_key(inner_prefix)
58}
59
60/// Returns the full prefix of flow info keys.
61pub fn flow_info_key_prefix() -> String {
62    scoped_flow_key_prefix(flow_info::FLOW_INFO_KEY_PREFIX)
63}
64
65/// Returns the full prefix of flow name keys.
66pub fn flow_name_key_prefix() -> String {
67    scoped_flow_key_prefix(flow_name::FLOW_NAME_KEY_PREFIX)
68}
69
70/// Returns the full prefix of flow route keys.
71pub fn flow_route_key_prefix() -> String {
72    scoped_flow_key_prefix(flow_route::FLOW_ROUTE_KEY_PREFIX)
73}
74
75/// Returns the full prefix of table flow keys.
76pub fn table_flow_key_prefix() -> String {
77    scoped_flow_key_prefix(table_flow::TABLE_FLOW_KEY_PREFIX)
78}
79
80/// Returns the full prefix of flownode flow keys.
81pub fn flownode_flow_key_prefix() -> String {
82    scoped_flow_key_prefix(flownode_flow::FLOWNODE_FLOW_KEY_PREFIX)
83}
84
85/// Returns the full key of flow state.
86pub fn flow_state_full_key() -> String {
87    scoped_flow_key(flow_state::FLOW_STATE_KEY)
88}
89
90/// The key of `__flow/` scope.
91#[derive(Debug, Clone, PartialEq)]
92pub struct FlowScoped<T> {
93    inner: T,
94}
95
96impl<T> Deref for FlowScoped<T> {
97    type Target = T;
98
99    fn deref(&self) -> &Self::Target {
100        &self.inner
101    }
102}
103
104impl<T> FlowScoped<T> {
105    const PREFIX: &'static str = "__flow/";
106
107    /// Returns a new [FlowScoped] key.
108    pub fn new(inner: T) -> FlowScoped<T> {
109        Self { inner }
110    }
111}
112
113impl<'a, T: MetadataKey<'a, T>> MetadataKey<'a, FlowScoped<T>> for FlowScoped<T> {
114    fn to_bytes(&self) -> Vec<u8> {
115        let prefix = FlowScoped::<T>::PREFIX.as_bytes();
116        let inner = self.inner.to_bytes();
117        let mut bytes = Vec::with_capacity(prefix.len() + inner.len());
118        bytes.extend(prefix);
119        bytes.extend(inner);
120        bytes
121    }
122
123    fn from_bytes(bytes: &'a [u8]) -> Result<FlowScoped<T>> {
124        let prefix = FlowScoped::<T>::PREFIX.as_bytes();
125        ensure!(
126            bytes.starts_with(prefix),
127            error::MismatchPrefixSnafu {
128                prefix: String::from_utf8_lossy(prefix),
129                key: String::from_utf8_lossy(bytes),
130            }
131        );
132        let inner = T::from_bytes(&bytes[prefix.len()..])?;
133        Ok(FlowScoped { inner })
134    }
135}
136
137pub type FlowMetadataManagerRef = Arc<FlowMetadataManager>;
138
139/// The manager of metadata, provides ability to:
140/// - Create metadata of the flow.
141/// - Retrieve metadata of the flow.
142/// - Delete metadata of the flow.
143pub struct FlowMetadataManager {
144    flow_info_manager: FlowInfoManager,
145    flow_route_manager: FlowRouteManager,
146    flownode_flow_manager: FlownodeFlowManager,
147    table_flow_manager: TableFlowManager,
148    flow_name_manager: FlowNameManager,
149    /// only metasrv have access to itself's memory backend, so for other case it should be None
150    flow_state_manager: Option<FlowStateManager>,
151    kv_backend: KvBackendRef,
152}
153
154impl FlowMetadataManager {
155    /// Returns a new [`FlowMetadataManager`].
156    pub fn new(kv_backend: KvBackendRef) -> Self {
157        Self {
158            flow_info_manager: FlowInfoManager::new(kv_backend.clone()),
159            flow_route_manager: FlowRouteManager::new(kv_backend.clone()),
160            flow_name_manager: FlowNameManager::new(kv_backend.clone()),
161            flownode_flow_manager: FlownodeFlowManager::new(kv_backend.clone()),
162            table_flow_manager: TableFlowManager::new(kv_backend.clone()),
163            flow_state_manager: None,
164            kv_backend,
165        }
166    }
167
168    /// Returns the [`FlowNameManager`].
169    pub fn flow_name_manager(&self) -> &FlowNameManager {
170        &self.flow_name_manager
171    }
172
173    pub fn flow_state_manager(&self) -> Option<&FlowStateManager> {
174        self.flow_state_manager.as_ref()
175    }
176
177    /// Returns the [`FlowInfoManager`].
178    pub fn flow_info_manager(&self) -> &FlowInfoManager {
179        &self.flow_info_manager
180    }
181
182    /// Returns the [`FlowRouteManager`].
183    pub fn flow_route_manager(&self) -> &FlowRouteManager {
184        &self.flow_route_manager
185    }
186
187    /// Returns the [`FlownodeFlowManager`].
188    pub fn flownode_flow_manager(&self) -> &FlownodeFlowManager {
189        &self.flownode_flow_manager
190    }
191
192    /// Returns the [`TableFlowManager`].
193    pub fn table_flow_manager(&self) -> &TableFlowManager {
194        &self.table_flow_manager
195    }
196
197    /// Returns a best-effort mapping from flow partition to flownode address.
198    pub async fn flownode_addrs(
199        &self,
200        flow_id: FlowId,
201    ) -> Result<BTreeMap<FlowPartitionId, String>> {
202        let routes = self.flow_route_manager.routes(flow_id).await?;
203
204        Ok(routes
205            .into_iter()
206            .filter_map(|(key, route)| {
207                let addr = route.peer.addr;
208                (!addr.is_empty()).then_some((key.partition_id(), addr))
209            })
210            .collect())
211    }
212
213    /// Creates metadata for flow and returns an error if different metadata exists.
214    pub async fn create_flow_metadata(
215        &self,
216        flow_id: FlowId,
217        flow_info: FlowInfoValue,
218        flow_routes: Vec<(FlowPartitionId, FlowRouteValue)>,
219    ) -> Result<()> {
220        let (create_flow_flow_name_txn, on_create_flow_flow_name_failure) = self
221            .flow_name_manager
222            .build_create_txn(&flow_info.catalog_name, &flow_info.flow_name, flow_id)?;
223
224        let (create_flow_txn, on_create_flow_failure) = self
225            .flow_info_manager
226            .build_create_txn(flow_id, &flow_info)?;
227
228        let create_flow_routes_txn = self
229            .flow_route_manager
230            .build_create_txn(flow_id, flow_routes.clone())?;
231
232        let create_flownode_flow_txn = self
233            .flownode_flow_manager
234            .build_create_txn(flow_id, flow_info.flownode_ids().clone());
235
236        let create_table_flow_txn = self.table_flow_manager.build_create_txn(
237            flow_id,
238            flow_routes
239                .into_iter()
240                .map(|(partition_id, route)| (partition_id, TableFlowValue { peer: route.peer }))
241                .collect(),
242            flow_info.source_table_ids(),
243        )?;
244
245        let txn = Txn::merge_all(vec![
246            create_flow_flow_name_txn,
247            create_flow_txn,
248            create_flow_routes_txn,
249            create_flownode_flow_txn,
250            create_table_flow_txn,
251        ]);
252        info!(
253            "Creating flow {}.{}({}), with {} txn operations",
254            flow_info.catalog_name,
255            flow_info.flow_name,
256            flow_id,
257            txn.max_operations()
258        );
259
260        let mut resp = self.kv_backend.txn(txn).await?;
261        if !resp.succeeded {
262            let mut set = TxnOpGetResponseSet::from(&mut resp.responses);
263            let remote_flow_flow_name =
264                on_create_flow_flow_name_failure(&mut set)?.with_context(|| {
265                    error::UnexpectedSnafu {
266                        err_msg: format!(
267                    "Reads the empty flow name in comparing operation of the creating flow, flow_id: {flow_id}"
268                ),
269                    }
270                })?;
271
272            if remote_flow_flow_name.flow_id() != flow_id {
273                info!(
274                    "Trying to create flow {}.{}({}), but flow({}) already exists",
275                    flow_info.catalog_name,
276                    flow_info.flow_name,
277                    flow_id,
278                    remote_flow_flow_name.flow_id()
279                );
280
281                return error::FlowAlreadyExistsSnafu {
282                    flow_name: format!("{}.{}", flow_info.catalog_name, flow_info.flow_name),
283                }
284                .fail();
285            }
286
287            let remote_flow =
288                on_create_flow_failure(&mut set)?.with_context(|| error::UnexpectedSnafu {
289                    err_msg: format!(
290                        "Reads the empty flow in comparing operation of creating flow, flow_id: {flow_id}"
291                    ),
292                })?;
293            let op_name = "creating flow";
294            ensure_values!(*remote_flow, flow_info, op_name);
295        }
296
297        Ok(())
298    }
299
300    /// Update metadata for flow and returns an error if old metadata IS NOT exists.
301    pub async fn update_flow_metadata(
302        &self,
303        flow_id: FlowId,
304        current_flow_info: &DeserializedValueWithBytes<FlowInfoValue>,
305        new_flow_info: &FlowInfoValue,
306        flow_routes: Vec<(FlowPartitionId, FlowRouteValue)>,
307    ) -> Result<()> {
308        let (update_flow_flow_name_txn, on_create_flow_flow_name_failure) =
309            self.flow_name_manager.build_update_txn(
310                &new_flow_info.catalog_name,
311                &new_flow_info.flow_name,
312                flow_id,
313            )?;
314
315        let (update_flow_txn, on_create_flow_failure) =
316            self.flow_info_manager
317                .build_update_txn(flow_id, current_flow_info, new_flow_info)?;
318
319        let update_flow_routes_txn = self.flow_route_manager.build_update_txn(
320            flow_id,
321            current_flow_info,
322            flow_routes.clone(),
323        )?;
324
325        let update_flownode_flow_txn = self.flownode_flow_manager.build_update_txn(
326            flow_id,
327            current_flow_info,
328            new_flow_info.flownode_ids().clone(),
329        );
330
331        let update_table_flow_txn = self.table_flow_manager.build_update_txn(
332            flow_id,
333            current_flow_info,
334            flow_routes
335                .into_iter()
336                .map(|(partition_id, route)| (partition_id, TableFlowValue { peer: route.peer }))
337                .collect(),
338            new_flow_info.source_table_ids(),
339        )?;
340
341        let txn = Txn::merge_all(vec![
342            update_flow_flow_name_txn,
343            update_flow_txn,
344            update_flow_routes_txn,
345            update_flownode_flow_txn,
346            update_table_flow_txn,
347        ]);
348        info!(
349            "Creating flow {}.{}({}), with {} txn operations",
350            new_flow_info.catalog_name,
351            new_flow_info.flow_name,
352            flow_id,
353            txn.max_operations()
354        );
355
356        let mut resp = self.kv_backend.txn(txn).await?;
357        if !resp.succeeded {
358            let mut set = TxnOpGetResponseSet::from(&mut resp.responses);
359            let remote_flow_flow_name =
360                on_create_flow_flow_name_failure(&mut set)?.with_context(|| {
361                    error::UnexpectedSnafu {
362                        err_msg: format!(
363                        "Reads the empty flow name in comparing operation of the updating flow, flow_id: {flow_id}"
364                    ),
365                    }
366                })?;
367
368            if remote_flow_flow_name.flow_id() != flow_id {
369                info!(
370                    "Trying to updating flow {}.{}({}), but flow({}) already exists with a different flow id",
371                    new_flow_info.catalog_name,
372                    new_flow_info.flow_name,
373                    flow_id,
374                    remote_flow_flow_name.flow_id()
375                );
376
377                return error::UnexpectedSnafu {
378                    err_msg: format!(
379                        "Reads different flow id when updating flow({2}.{3}), prev flow id = {0}, updating with flow id = {1}",
380                        remote_flow_flow_name.flow_id(),
381                        flow_id,
382                        new_flow_info.catalog_name,
383                        new_flow_info.flow_name,
384                    ),
385                }.fail();
386            }
387
388            let remote_flow =
389                on_create_flow_failure(&mut set)?.with_context(|| error::UnexpectedSnafu {
390                    err_msg: format!(
391                        "Reads the empty flow in comparing operation of the updating flow, flow_id: {flow_id}"
392                    ),
393                })?;
394            let op_name = "updating flow";
395            ensure_values!(*remote_flow, new_flow_info.clone(), op_name);
396        }
397
398        Ok(())
399    }
400
401    fn flow_metadata_keys(&self, flow_id: FlowId, flow_value: &FlowInfoValue) -> Vec<Vec<u8>> {
402        let source_table_ids = flow_value.source_table_ids();
403        let mut keys =
404            Vec::with_capacity(2 + flow_value.flownode_ids.len() * (source_table_ids.len() + 2));
405        // Builds flow name key
406        let flow_name = FlowNameKey::new(&flow_value.catalog_name, &flow_value.flow_name);
407        keys.push(flow_name.to_bytes());
408
409        // Builds flow value key
410        let flow_info_key = FlowInfoKey::new(flow_id);
411        keys.push(flow_info_key.to_bytes());
412
413        // Builds flownode flow keys & table flow keys
414        flow_value
415            .flownode_ids
416            .iter()
417            .for_each(|(&partition_id, &flownode_id)| {
418                keys.push(FlownodeFlowKey::new(flownode_id, flow_id, partition_id).to_bytes());
419                keys.push(FlowRouteKey::new(flow_id, partition_id).to_bytes());
420                source_table_ids.iter().for_each(|&table_id| {
421                    keys.push(
422                        TableFlowKey::new(table_id, flownode_id, flow_id, partition_id).to_bytes(),
423                    );
424                })
425            });
426        keys
427    }
428
429    /// Deletes metadata for table **permanently**.
430    pub async fn destroy_flow_metadata(
431        &self,
432        flow_id: FlowId,
433        flow_value: &FlowInfoValue,
434    ) -> Result<()> {
435        let keys = self.flow_metadata_keys(flow_id, flow_value);
436        let _ = self
437            .kv_backend
438            .batch_delete(BatchDeleteRequest::new().with_keys(keys))
439            .await?;
440        Ok(())
441    }
442}
443
444impl std::fmt::Debug for FlowMetadataManager {
445    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
446        f.debug_struct("FlowMetadataManager").finish()
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use std::assert_matches;
453    use std::collections::BTreeMap;
454    use std::sync::Arc;
455
456    use futures::TryStreamExt;
457    use table::metadata::TableId;
458    use table::table_name::TableName;
459
460    use super::*;
461    use crate::FlownodeId;
462    use crate::key::flow::flow_info::FlowStatus;
463    use crate::key::flow::table_flow::TableFlowKey;
464    use crate::key::node_address::{NodeAddressKey, NodeAddressValue};
465    use crate::key::{FlowPartitionId, MetadataValue};
466    use crate::kv_backend::KvBackend;
467    use crate::kv_backend::memory::MemoryKvBackend;
468    use crate::peer::Peer;
469    use crate::rpc::store::PutRequest;
470
471    #[derive(Debug)]
472    struct MockKey {
473        inner: Vec<u8>,
474    }
475
476    impl<'a> MetadataKey<'a, MockKey> for MockKey {
477        fn to_bytes(&self) -> Vec<u8> {
478            self.inner.clone()
479        }
480
481        fn from_bytes(bytes: &'a [u8]) -> Result<MockKey> {
482            Ok(MockKey {
483                inner: bytes.to_vec(),
484            })
485        }
486    }
487
488    #[test]
489    fn test_flow_scoped_to_bytes() {
490        let key = FlowScoped::new(MockKey {
491            inner: b"hi".to_vec(),
492        });
493        assert_eq!(b"__flow/hi".to_vec(), key.to_bytes());
494    }
495
496    #[test]
497    fn test_flow_scoped_from_bytes() {
498        let bytes = b"__flow/hi";
499        let key = FlowScoped::<MockKey>::from_bytes(bytes).unwrap();
500        assert_eq!(key.inner.inner, b"hi".to_vec());
501    }
502
503    #[test]
504    fn test_flow_scoped_from_bytes_mismatch() {
505        let bytes = b"__table/hi";
506        let err = FlowScoped::<MockKey>::from_bytes(bytes).unwrap_err();
507        assert_matches!(err, error::Error::MismatchPrefix { .. });
508    }
509
510    fn test_flow_info_value(
511        flow_name: &str,
512        flownode_ids: BTreeMap<FlowPartitionId, FlownodeId>,
513        source_table_ids: Vec<TableId>,
514    ) -> FlowInfoValue {
515        let catalog_name = "greptime";
516        let sink_table_name = TableName {
517            catalog_name: catalog_name.to_string(),
518            schema_name: "my_schema".to_string(),
519            table_name: "sink_table".to_string(),
520        };
521        FlowInfoValue {
522            catalog_name: catalog_name.to_string(),
523            query_context: None,
524            flow_name: flow_name.to_string(),
525            source_table_ids,
526            all_source_table_names: vec![],
527            unresolved_source_table_names: vec![],
528            sink_table_name,
529            flownode_ids,
530            raw_sql: "raw".to_string(),
531            expire_after: Some(300),
532            eval_interval_secs: None,
533            comment: "hi".to_string(),
534            options: Default::default(),
535            status: FlowStatus::Active,
536            created_time: chrono::Utc::now(),
537            updated_time: chrono::Utc::now(),
538        }
539    }
540
541    #[tokio::test]
542    async fn test_create_flow_metadata() {
543        let mem_kv = Arc::new(MemoryKvBackend::default());
544        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
545        let flow_id = 10;
546        let flow_value = test_flow_info_value(
547            "flow",
548            [(0, 1u64), (1, 2u64)].into(),
549            vec![1024, 1025, 1026],
550        );
551        let flow_routes = vec![
552            (
553                1u32,
554                FlowRouteValue {
555                    peer: Peer::empty(1),
556                },
557            ),
558            (
559                2,
560                FlowRouteValue {
561                    peer: Peer::empty(2),
562                },
563            ),
564        ];
565        flow_metadata_manager
566            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
567            .await
568            .unwrap();
569        // Creates again.
570        flow_metadata_manager
571            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
572            .await
573            .unwrap();
574        let got = flow_metadata_manager
575            .flow_info_manager()
576            .get(flow_id)
577            .await
578            .unwrap()
579            .unwrap();
580        let routes = flow_metadata_manager
581            .flow_route_manager()
582            .routes(flow_id)
583            .await
584            .unwrap();
585        assert_eq!(
586            routes,
587            vec![
588                (
589                    FlowRouteKey::new(flow_id, 1),
590                    FlowRouteValue {
591                        peer: Peer::empty(1),
592                    },
593                ),
594                (
595                    FlowRouteKey::new(flow_id, 2),
596                    FlowRouteValue {
597                        peer: Peer::empty(2),
598                    },
599                ),
600            ]
601        );
602        assert_eq!(got, flow_value);
603        let flows = flow_metadata_manager
604            .flownode_flow_manager()
605            .flows(1)
606            .try_collect::<Vec<_>>()
607            .await
608            .unwrap();
609        assert_eq!(flows, vec![(flow_id, 0)]);
610        for table_id in [1024, 1025, 1026] {
611            let nodes = flow_metadata_manager
612                .table_flow_manager()
613                .flows(table_id)
614                .await
615                .unwrap();
616            assert_eq!(
617                nodes,
618                vec![
619                    (
620                        TableFlowKey::new(table_id, 1, flow_id, 1),
621                        TableFlowValue {
622                            peer: Peer::empty(1)
623                        }
624                    ),
625                    (
626                        TableFlowKey::new(table_id, 2, flow_id, 2),
627                        TableFlowValue {
628                            peer: Peer::empty(2)
629                        }
630                    )
631                ]
632            );
633        }
634    }
635
636    #[tokio::test]
637    async fn test_flownode_addrs_remaps_to_latest_address() {
638        let mem_kv = Arc::new(MemoryKvBackend::default());
639        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
640        let flow_id = 10;
641        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
642        let flow_routes = vec![(
643            0u32,
644            FlowRouteValue {
645                peer: Peer::new(1, "old-addr"),
646            },
647        )];
648
649        flow_metadata_manager
650            .create_flow_metadata(flow_id, flow_value, flow_routes)
651            .await
652            .unwrap();
653
654        mem_kv
655            .put(PutRequest {
656                key: NodeAddressKey::with_flownode(1).to_bytes(),
657                value: NodeAddressValue::new(Peer::new(1, "new-addr"))
658                    .try_as_raw_value()
659                    .unwrap(),
660                ..Default::default()
661            })
662            .await
663            .unwrap();
664
665        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
666        assert_eq!(addrs, BTreeMap::from([(0, "new-addr".to_string())]));
667    }
668
669    #[tokio::test]
670    async fn test_flownode_addrs_falls_back_to_route_address() {
671        let mem_kv = Arc::new(MemoryKvBackend::default());
672        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
673        let flow_id = 10;
674        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
675        let flow_routes = vec![(
676            0u32,
677            FlowRouteValue {
678                peer: Peer::new(1, "route-addr"),
679            },
680        )];
681
682        flow_metadata_manager
683            .create_flow_metadata(flow_id, flow_value, flow_routes)
684            .await
685            .unwrap();
686
687        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
688        assert_eq!(addrs, BTreeMap::from([(0, "route-addr".to_string())]));
689    }
690
691    #[tokio::test]
692    async fn test_flownode_addrs_skips_empty_addresses() {
693        let mem_kv = Arc::new(MemoryKvBackend::default());
694        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
695        let flow_id = 10;
696        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
697        let flow_routes = vec![(
698            0u32,
699            FlowRouteValue {
700                peer: Peer::empty(1),
701            },
702        )];
703
704        flow_metadata_manager
705            .create_flow_metadata(flow_id, flow_value, flow_routes)
706            .await
707            .unwrap();
708
709        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
710        assert!(addrs.is_empty());
711    }
712
713    #[tokio::test]
714    async fn test_create_flow_metadata_flow_exists_err() {
715        let mem_kv = Arc::new(MemoryKvBackend::default());
716        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
717        let flow_id = 10;
718        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
719        let flow_routes = vec![
720            (
721                1u32,
722                FlowRouteValue {
723                    peer: Peer::empty(1),
724                },
725            ),
726            (
727                2,
728                FlowRouteValue {
729                    peer: Peer::empty(2),
730                },
731            ),
732        ];
733        flow_metadata_manager
734            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
735            .await
736            .unwrap();
737        // Creates again
738        let err = flow_metadata_manager
739            .create_flow_metadata(flow_id + 1, flow_value, flow_routes.clone())
740            .await
741            .unwrap_err();
742        assert_matches!(err, error::Error::FlowAlreadyExists { .. });
743    }
744
745    #[tokio::test]
746    async fn test_create_flow_metadata_unexpected_err() {
747        let mem_kv = Arc::new(MemoryKvBackend::default());
748        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
749        let flow_id = 10;
750        let catalog_name = "greptime";
751        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
752        let flow_routes = vec![
753            (
754                1u32,
755                FlowRouteValue {
756                    peer: Peer::empty(1),
757                },
758            ),
759            (
760                2,
761                FlowRouteValue {
762                    peer: Peer::empty(2),
763                },
764            ),
765        ];
766        flow_metadata_manager
767            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
768            .await
769            .unwrap();
770        // Creates again.
771        let another_sink_table_name = TableName {
772            catalog_name: catalog_name.to_string(),
773            schema_name: "my_schema".to_string(),
774            table_name: "another_sink_table".to_string(),
775        };
776        let flow_value = FlowInfoValue {
777            catalog_name: "greptime".to_string(),
778            query_context: None,
779            flow_name: "flow".to_string(),
780            source_table_ids: vec![1024, 1025, 1026],
781            all_source_table_names: vec![],
782            unresolved_source_table_names: vec![],
783            sink_table_name: another_sink_table_name,
784            flownode_ids: [(0, 1u64)].into(),
785            raw_sql: "raw".to_string(),
786            expire_after: Some(300),
787            eval_interval_secs: None,
788            comment: "hi".to_string(),
789            options: Default::default(),
790            status: FlowStatus::Active,
791            created_time: chrono::Utc::now(),
792            updated_time: chrono::Utc::now(),
793        };
794        let err = flow_metadata_manager
795            .create_flow_metadata(flow_id, flow_value, flow_routes.clone())
796            .await
797            .unwrap_err();
798        assert!(err.to_string().contains("Reads the different value"));
799    }
800
801    #[tokio::test]
802    async fn test_destroy_flow_metadata() {
803        let mem_kv = Arc::new(MemoryKvBackend::default());
804        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
805        let flow_id = 10;
806        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
807        let flow_routes = vec![(
808            0u32,
809            FlowRouteValue {
810                peer: Peer::empty(1),
811            },
812        )];
813        flow_metadata_manager
814            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
815            .await
816            .unwrap();
817
818        flow_metadata_manager
819            .destroy_flow_metadata(flow_id, &flow_value)
820            .await
821            .unwrap();
822        // Destroys again
823        flow_metadata_manager
824            .destroy_flow_metadata(flow_id, &flow_value)
825            .await
826            .unwrap();
827        // Ensures all keys are deleted
828        assert!(mem_kv.is_empty())
829    }
830
831    #[tokio::test]
832    async fn test_update_flow_metadata() {
833        let mem_kv = Arc::new(MemoryKvBackend::default());
834        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
835        let flow_id = 10;
836        let flow_value = test_flow_info_value(
837            "flow",
838            [(0, 1u64), (1, 2u64)].into(),
839            vec![1024, 1025, 1026],
840        );
841        let flow_routes = vec![
842            (
843                1u32,
844                FlowRouteValue {
845                    peer: Peer::empty(1),
846                },
847            ),
848            (
849                2,
850                FlowRouteValue {
851                    peer: Peer::empty(2),
852                },
853            ),
854        ];
855        flow_metadata_manager
856            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
857            .await
858            .unwrap();
859
860        let new_flow_value = {
861            let mut tmp = flow_value.clone();
862            tmp.raw_sql = "new".to_string();
863            tmp
864        };
865
866        // Update flow instead
867        flow_metadata_manager
868            .update_flow_metadata(
869                flow_id,
870                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
871                &new_flow_value,
872                flow_routes.clone(),
873            )
874            .await
875            .unwrap();
876
877        let got = flow_metadata_manager
878            .flow_info_manager()
879            .get(flow_id)
880            .await
881            .unwrap()
882            .unwrap();
883        let routes = flow_metadata_manager
884            .flow_route_manager()
885            .routes(flow_id)
886            .await
887            .unwrap();
888        assert_eq!(
889            routes,
890            vec![
891                (
892                    FlowRouteKey::new(flow_id, 1),
893                    FlowRouteValue {
894                        peer: Peer::empty(1),
895                    },
896                ),
897                (
898                    FlowRouteKey::new(flow_id, 2),
899                    FlowRouteValue {
900                        peer: Peer::empty(2),
901                    },
902                ),
903            ]
904        );
905        assert_eq!(got, new_flow_value);
906        let flows = flow_metadata_manager
907            .flownode_flow_manager()
908            .flows(1)
909            .try_collect::<Vec<_>>()
910            .await
911            .unwrap();
912        assert_eq!(flows, vec![(flow_id, 0)]);
913        for table_id in [1024, 1025, 1026] {
914            let nodes = flow_metadata_manager
915                .table_flow_manager()
916                .flows(table_id)
917                .await
918                .unwrap();
919            assert_eq!(
920                nodes,
921                vec![
922                    (
923                        TableFlowKey::new(table_id, 1, flow_id, 1),
924                        TableFlowValue {
925                            peer: Peer::empty(1)
926                        }
927                    ),
928                    (
929                        TableFlowKey::new(table_id, 2, flow_id, 2),
930                        TableFlowValue {
931                            peer: Peer::empty(2)
932                        }
933                    )
934                ]
935            );
936        }
937    }
938
939    #[tokio::test]
940    async fn test_update_flow_metadata_diff_flownode() {
941        let mem_kv = Arc::new(MemoryKvBackend::default());
942        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
943        let flow_id = 10;
944        let flow_value = test_flow_info_value(
945            "flow",
946            [(0u32, 1u64), (1u32, 2u64)].into(),
947            vec![1024, 1025, 1026],
948        );
949        let flow_routes = vec![
950            (
951                0u32,
952                FlowRouteValue {
953                    peer: Peer::empty(1),
954                },
955            ),
956            (
957                1,
958                FlowRouteValue {
959                    peer: Peer::empty(2),
960                },
961            ),
962        ];
963        flow_metadata_manager
964            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
965            .await
966            .unwrap();
967
968        let new_flow_value = {
969            let mut tmp = flow_value.clone();
970            tmp.raw_sql = "new".to_string();
971            // move to different flownodes
972            tmp.flownode_ids = [(0, 3u64), (1, 4u64)].into();
973            tmp
974        };
975        let new_flow_routes = vec![
976            (
977                0u32,
978                FlowRouteValue {
979                    peer: Peer::empty(3),
980                },
981            ),
982            (
983                1,
984                FlowRouteValue {
985                    peer: Peer::empty(4),
986                },
987            ),
988        ];
989
990        // Update flow instead
991        flow_metadata_manager
992            .update_flow_metadata(
993                flow_id,
994                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
995                &new_flow_value,
996                new_flow_routes.clone(),
997            )
998            .await
999            .unwrap();
1000
1001        let got = flow_metadata_manager
1002            .flow_info_manager()
1003            .get(flow_id)
1004            .await
1005            .unwrap()
1006            .unwrap();
1007        let routes = flow_metadata_manager
1008            .flow_route_manager()
1009            .routes(flow_id)
1010            .await
1011            .unwrap();
1012        assert_eq!(
1013            routes,
1014            vec![
1015                (
1016                    FlowRouteKey::new(flow_id, 0),
1017                    FlowRouteValue {
1018                        peer: Peer::empty(3),
1019                    },
1020                ),
1021                (
1022                    FlowRouteKey::new(flow_id, 1),
1023                    FlowRouteValue {
1024                        peer: Peer::empty(4),
1025                    },
1026                ),
1027            ]
1028        );
1029        assert_eq!(got, new_flow_value);
1030
1031        let flows = flow_metadata_manager
1032            .flownode_flow_manager()
1033            .flows(1)
1034            .try_collect::<Vec<_>>()
1035            .await
1036            .unwrap();
1037        // should moved to different flownode
1038        assert_eq!(flows, vec![]);
1039
1040        let flows = flow_metadata_manager
1041            .flownode_flow_manager()
1042            .flows(3)
1043            .try_collect::<Vec<_>>()
1044            .await
1045            .unwrap();
1046        assert_eq!(flows, vec![(flow_id, 0)]);
1047
1048        for table_id in [1024, 1025, 1026] {
1049            let nodes = flow_metadata_manager
1050                .table_flow_manager()
1051                .flows(table_id)
1052                .await
1053                .unwrap();
1054            assert_eq!(
1055                nodes,
1056                vec![
1057                    (
1058                        TableFlowKey::new(table_id, 3, flow_id, 0),
1059                        TableFlowValue {
1060                            peer: Peer::empty(3)
1061                        }
1062                    ),
1063                    (
1064                        TableFlowKey::new(table_id, 4, flow_id, 1),
1065                        TableFlowValue {
1066                            peer: Peer::empty(4)
1067                        }
1068                    )
1069                ]
1070            );
1071        }
1072    }
1073
1074    #[tokio::test]
1075    async fn test_update_flow_metadata_flow_replace_diff_id_err() {
1076        let mem_kv = Arc::new(MemoryKvBackend::default());
1077        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
1078        let flow_id = 10;
1079        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
1080        let flow_routes = vec![
1081            (
1082                1u32,
1083                FlowRouteValue {
1084                    peer: Peer::empty(1),
1085                },
1086            ),
1087            (
1088                2,
1089                FlowRouteValue {
1090                    peer: Peer::empty(2),
1091                },
1092            ),
1093        ];
1094        flow_metadata_manager
1095            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
1096            .await
1097            .unwrap();
1098        // update again with same flow id
1099        flow_metadata_manager
1100            .update_flow_metadata(
1101                flow_id,
1102                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1103                &flow_value,
1104                flow_routes.clone(),
1105            )
1106            .await
1107            .unwrap();
1108        // update again with wrong flow id, expected error
1109        let err = flow_metadata_manager
1110            .update_flow_metadata(
1111                flow_id + 1,
1112                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1113                &flow_value,
1114                flow_routes,
1115            )
1116            .await
1117            .unwrap_err();
1118        assert_matches!(err, error::Error::Unexpected { .. });
1119        assert!(
1120            err.to_string()
1121                .contains("Reads different flow id when updating flow")
1122        );
1123    }
1124
1125    #[tokio::test]
1126    async fn test_update_flow_metadata_unexpected_err_prev_value_diff() {
1127        let mem_kv = Arc::new(MemoryKvBackend::default());
1128        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
1129        let flow_id = 10;
1130        let catalog_name = "greptime";
1131        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
1132        let flow_routes = vec![
1133            (
1134                1u32,
1135                FlowRouteValue {
1136                    peer: Peer::empty(1),
1137                },
1138            ),
1139            (
1140                2,
1141                FlowRouteValue {
1142                    peer: Peer::empty(2),
1143                },
1144            ),
1145        ];
1146        flow_metadata_manager
1147            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
1148            .await
1149            .unwrap();
1150        // Creates again.
1151        let another_sink_table_name = TableName {
1152            catalog_name: catalog_name.to_string(),
1153            schema_name: "my_schema".to_string(),
1154            table_name: "another_sink_table".to_string(),
1155        };
1156        let flow_value = FlowInfoValue {
1157            catalog_name: "greptime".to_string(),
1158            query_context: None,
1159            flow_name: "flow".to_string(),
1160            source_table_ids: vec![1024, 1025, 1026],
1161            all_source_table_names: vec![],
1162            unresolved_source_table_names: vec![],
1163            sink_table_name: another_sink_table_name,
1164            flownode_ids: [(0, 1u64)].into(),
1165            raw_sql: "raw".to_string(),
1166            expire_after: Some(300),
1167            eval_interval_secs: None,
1168            comment: "hi".to_string(),
1169            options: Default::default(),
1170            status: FlowStatus::Active,
1171            created_time: chrono::Utc::now(),
1172            updated_time: chrono::Utc::now(),
1173        };
1174        let err = flow_metadata_manager
1175            .update_flow_metadata(
1176                flow_id,
1177                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1178                &flow_value,
1179                flow_routes.clone(),
1180            )
1181            .await
1182            .unwrap_err();
1183        assert!(
1184            err.to_string().contains("Reads the different value"),
1185            "error: {:?}",
1186            err
1187        );
1188    }
1189}