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            eval_schedule: None,
539        }
540    }
541
542    #[tokio::test]
543    async fn test_create_flow_metadata() {
544        let mem_kv = Arc::new(MemoryKvBackend::default());
545        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
546        let flow_id = 10;
547        let flow_value = test_flow_info_value(
548            "flow",
549            [(0, 1u64), (1, 2u64)].into(),
550            vec![1024, 1025, 1026],
551        );
552        let flow_routes = vec![
553            (
554                1u32,
555                FlowRouteValue {
556                    peer: Peer::empty(1),
557                },
558            ),
559            (
560                2,
561                FlowRouteValue {
562                    peer: Peer::empty(2),
563                },
564            ),
565        ];
566        flow_metadata_manager
567            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
568            .await
569            .unwrap();
570        // Creates again.
571        flow_metadata_manager
572            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
573            .await
574            .unwrap();
575        let got = flow_metadata_manager
576            .flow_info_manager()
577            .get(flow_id)
578            .await
579            .unwrap()
580            .unwrap();
581        let routes = flow_metadata_manager
582            .flow_route_manager()
583            .routes(flow_id)
584            .await
585            .unwrap();
586        assert_eq!(
587            routes,
588            vec![
589                (
590                    FlowRouteKey::new(flow_id, 1),
591                    FlowRouteValue {
592                        peer: Peer::empty(1),
593                    },
594                ),
595                (
596                    FlowRouteKey::new(flow_id, 2),
597                    FlowRouteValue {
598                        peer: Peer::empty(2),
599                    },
600                ),
601            ]
602        );
603        assert_eq!(got, flow_value);
604        let flows = flow_metadata_manager
605            .flownode_flow_manager()
606            .flows(1)
607            .try_collect::<Vec<_>>()
608            .await
609            .unwrap();
610        assert_eq!(flows, vec![(flow_id, 0)]);
611        for table_id in [1024, 1025, 1026] {
612            let nodes = flow_metadata_manager
613                .table_flow_manager()
614                .flows(table_id)
615                .await
616                .unwrap();
617            assert_eq!(
618                nodes,
619                vec![
620                    (
621                        TableFlowKey::new(table_id, 1, flow_id, 1),
622                        TableFlowValue {
623                            peer: Peer::empty(1)
624                        }
625                    ),
626                    (
627                        TableFlowKey::new(table_id, 2, flow_id, 2),
628                        TableFlowValue {
629                            peer: Peer::empty(2)
630                        }
631                    )
632                ]
633            );
634        }
635    }
636
637    #[tokio::test]
638    async fn test_flownode_addrs_remaps_to_latest_address() {
639        let mem_kv = Arc::new(MemoryKvBackend::default());
640        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
641        let flow_id = 10;
642        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
643        let flow_routes = vec![(
644            0u32,
645            FlowRouteValue {
646                peer: Peer::new(1, "old-addr"),
647            },
648        )];
649
650        flow_metadata_manager
651            .create_flow_metadata(flow_id, flow_value, flow_routes)
652            .await
653            .unwrap();
654
655        mem_kv
656            .put(PutRequest {
657                key: NodeAddressKey::with_flownode(1).to_bytes(),
658                value: NodeAddressValue::new(Peer::new(1, "new-addr"))
659                    .try_as_raw_value()
660                    .unwrap(),
661                ..Default::default()
662            })
663            .await
664            .unwrap();
665
666        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
667        assert_eq!(addrs, BTreeMap::from([(0, "new-addr".to_string())]));
668    }
669
670    #[tokio::test]
671    async fn test_flownode_addrs_falls_back_to_route_address() {
672        let mem_kv = Arc::new(MemoryKvBackend::default());
673        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
674        let flow_id = 10;
675        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
676        let flow_routes = vec![(
677            0u32,
678            FlowRouteValue {
679                peer: Peer::new(1, "route-addr"),
680            },
681        )];
682
683        flow_metadata_manager
684            .create_flow_metadata(flow_id, flow_value, flow_routes)
685            .await
686            .unwrap();
687
688        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
689        assert_eq!(addrs, BTreeMap::from([(0, "route-addr".to_string())]));
690    }
691
692    #[tokio::test]
693    async fn test_flownode_addrs_skips_empty_addresses() {
694        let mem_kv = Arc::new(MemoryKvBackend::default());
695        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
696        let flow_id = 10;
697        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024]);
698        let flow_routes = vec![(
699            0u32,
700            FlowRouteValue {
701                peer: Peer::empty(1),
702            },
703        )];
704
705        flow_metadata_manager
706            .create_flow_metadata(flow_id, flow_value, flow_routes)
707            .await
708            .unwrap();
709
710        let addrs = flow_metadata_manager.flownode_addrs(flow_id).await.unwrap();
711        assert!(addrs.is_empty());
712    }
713
714    #[tokio::test]
715    async fn test_create_flow_metadata_flow_exists_err() {
716        let mem_kv = Arc::new(MemoryKvBackend::default());
717        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
718        let flow_id = 10;
719        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
720        let flow_routes = vec![
721            (
722                1u32,
723                FlowRouteValue {
724                    peer: Peer::empty(1),
725                },
726            ),
727            (
728                2,
729                FlowRouteValue {
730                    peer: Peer::empty(2),
731                },
732            ),
733        ];
734        flow_metadata_manager
735            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
736            .await
737            .unwrap();
738        // Creates again
739        let err = flow_metadata_manager
740            .create_flow_metadata(flow_id + 1, flow_value, flow_routes.clone())
741            .await
742            .unwrap_err();
743        assert_matches!(err, error::Error::FlowAlreadyExists { .. });
744    }
745
746    #[tokio::test]
747    async fn test_create_flow_metadata_unexpected_err() {
748        let mem_kv = Arc::new(MemoryKvBackend::default());
749        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
750        let flow_id = 10;
751        let catalog_name = "greptime";
752        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
753        let flow_routes = vec![
754            (
755                1u32,
756                FlowRouteValue {
757                    peer: Peer::empty(1),
758                },
759            ),
760            (
761                2,
762                FlowRouteValue {
763                    peer: Peer::empty(2),
764                },
765            ),
766        ];
767        flow_metadata_manager
768            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
769            .await
770            .unwrap();
771        // Creates again.
772        let another_sink_table_name = TableName {
773            catalog_name: catalog_name.to_string(),
774            schema_name: "my_schema".to_string(),
775            table_name: "another_sink_table".to_string(),
776        };
777        let flow_value = FlowInfoValue {
778            catalog_name: "greptime".to_string(),
779            query_context: None,
780            flow_name: "flow".to_string(),
781            source_table_ids: vec![1024, 1025, 1026],
782            all_source_table_names: vec![],
783            unresolved_source_table_names: vec![],
784            sink_table_name: another_sink_table_name,
785            flownode_ids: [(0, 1u64)].into(),
786            raw_sql: "raw".to_string(),
787            expire_after: Some(300),
788            eval_interval_secs: None,
789            comment: "hi".to_string(),
790            options: Default::default(),
791            status: FlowStatus::Active,
792            created_time: chrono::Utc::now(),
793            updated_time: chrono::Utc::now(),
794            eval_schedule: None,
795        };
796        let err = flow_metadata_manager
797            .create_flow_metadata(flow_id, flow_value, flow_routes.clone())
798            .await
799            .unwrap_err();
800        assert!(err.to_string().contains("Reads the different value"));
801    }
802
803    #[tokio::test]
804    async fn test_destroy_flow_metadata() {
805        let mem_kv = Arc::new(MemoryKvBackend::default());
806        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
807        let flow_id = 10;
808        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
809        let flow_routes = vec![(
810            0u32,
811            FlowRouteValue {
812                peer: Peer::empty(1),
813            },
814        )];
815        flow_metadata_manager
816            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
817            .await
818            .unwrap();
819
820        flow_metadata_manager
821            .destroy_flow_metadata(flow_id, &flow_value)
822            .await
823            .unwrap();
824        // Destroys again
825        flow_metadata_manager
826            .destroy_flow_metadata(flow_id, &flow_value)
827            .await
828            .unwrap();
829        // Ensures all keys are deleted
830        assert!(mem_kv.is_empty())
831    }
832
833    #[tokio::test]
834    async fn test_update_flow_metadata() {
835        let mem_kv = Arc::new(MemoryKvBackend::default());
836        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
837        let flow_id = 10;
838        let flow_value = test_flow_info_value(
839            "flow",
840            [(0, 1u64), (1, 2u64)].into(),
841            vec![1024, 1025, 1026],
842        );
843        let flow_routes = vec![
844            (
845                1u32,
846                FlowRouteValue {
847                    peer: Peer::empty(1),
848                },
849            ),
850            (
851                2,
852                FlowRouteValue {
853                    peer: Peer::empty(2),
854                },
855            ),
856        ];
857        flow_metadata_manager
858            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
859            .await
860            .unwrap();
861
862        let new_flow_value = {
863            let mut tmp = flow_value.clone();
864            tmp.raw_sql = "new".to_string();
865            tmp
866        };
867
868        // Update flow instead
869        flow_metadata_manager
870            .update_flow_metadata(
871                flow_id,
872                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
873                &new_flow_value,
874                flow_routes.clone(),
875            )
876            .await
877            .unwrap();
878
879        let got = flow_metadata_manager
880            .flow_info_manager()
881            .get(flow_id)
882            .await
883            .unwrap()
884            .unwrap();
885        let routes = flow_metadata_manager
886            .flow_route_manager()
887            .routes(flow_id)
888            .await
889            .unwrap();
890        assert_eq!(
891            routes,
892            vec![
893                (
894                    FlowRouteKey::new(flow_id, 1),
895                    FlowRouteValue {
896                        peer: Peer::empty(1),
897                    },
898                ),
899                (
900                    FlowRouteKey::new(flow_id, 2),
901                    FlowRouteValue {
902                        peer: Peer::empty(2),
903                    },
904                ),
905            ]
906        );
907        assert_eq!(got, new_flow_value);
908        let flows = flow_metadata_manager
909            .flownode_flow_manager()
910            .flows(1)
911            .try_collect::<Vec<_>>()
912            .await
913            .unwrap();
914        assert_eq!(flows, vec![(flow_id, 0)]);
915        for table_id in [1024, 1025, 1026] {
916            let nodes = flow_metadata_manager
917                .table_flow_manager()
918                .flows(table_id)
919                .await
920                .unwrap();
921            assert_eq!(
922                nodes,
923                vec![
924                    (
925                        TableFlowKey::new(table_id, 1, flow_id, 1),
926                        TableFlowValue {
927                            peer: Peer::empty(1)
928                        }
929                    ),
930                    (
931                        TableFlowKey::new(table_id, 2, flow_id, 2),
932                        TableFlowValue {
933                            peer: Peer::empty(2)
934                        }
935                    )
936                ]
937            );
938        }
939    }
940
941    #[tokio::test]
942    async fn test_update_flow_metadata_diff_flownode() {
943        let mem_kv = Arc::new(MemoryKvBackend::default());
944        let flow_metadata_manager = FlowMetadataManager::new(mem_kv.clone());
945        let flow_id = 10;
946        let flow_value = test_flow_info_value(
947            "flow",
948            [(0u32, 1u64), (1u32, 2u64)].into(),
949            vec![1024, 1025, 1026],
950        );
951        let flow_routes = vec![
952            (
953                0u32,
954                FlowRouteValue {
955                    peer: Peer::empty(1),
956                },
957            ),
958            (
959                1,
960                FlowRouteValue {
961                    peer: Peer::empty(2),
962                },
963            ),
964        ];
965        flow_metadata_manager
966            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
967            .await
968            .unwrap();
969
970        let new_flow_value = {
971            let mut tmp = flow_value.clone();
972            tmp.raw_sql = "new".to_string();
973            // move to different flownodes
974            tmp.flownode_ids = [(0, 3u64), (1, 4u64)].into();
975            tmp
976        };
977        let new_flow_routes = vec![
978            (
979                0u32,
980                FlowRouteValue {
981                    peer: Peer::empty(3),
982                },
983            ),
984            (
985                1,
986                FlowRouteValue {
987                    peer: Peer::empty(4),
988                },
989            ),
990        ];
991
992        // Update flow instead
993        flow_metadata_manager
994            .update_flow_metadata(
995                flow_id,
996                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
997                &new_flow_value,
998                new_flow_routes.clone(),
999            )
1000            .await
1001            .unwrap();
1002
1003        let got = flow_metadata_manager
1004            .flow_info_manager()
1005            .get(flow_id)
1006            .await
1007            .unwrap()
1008            .unwrap();
1009        let routes = flow_metadata_manager
1010            .flow_route_manager()
1011            .routes(flow_id)
1012            .await
1013            .unwrap();
1014        assert_eq!(
1015            routes,
1016            vec![
1017                (
1018                    FlowRouteKey::new(flow_id, 0),
1019                    FlowRouteValue {
1020                        peer: Peer::empty(3),
1021                    },
1022                ),
1023                (
1024                    FlowRouteKey::new(flow_id, 1),
1025                    FlowRouteValue {
1026                        peer: Peer::empty(4),
1027                    },
1028                ),
1029            ]
1030        );
1031        assert_eq!(got, new_flow_value);
1032
1033        let flows = flow_metadata_manager
1034            .flownode_flow_manager()
1035            .flows(1)
1036            .try_collect::<Vec<_>>()
1037            .await
1038            .unwrap();
1039        // should moved to different flownode
1040        assert_eq!(flows, vec![]);
1041
1042        let flows = flow_metadata_manager
1043            .flownode_flow_manager()
1044            .flows(3)
1045            .try_collect::<Vec<_>>()
1046            .await
1047            .unwrap();
1048        assert_eq!(flows, vec![(flow_id, 0)]);
1049
1050        for table_id in [1024, 1025, 1026] {
1051            let nodes = flow_metadata_manager
1052                .table_flow_manager()
1053                .flows(table_id)
1054                .await
1055                .unwrap();
1056            assert_eq!(
1057                nodes,
1058                vec![
1059                    (
1060                        TableFlowKey::new(table_id, 3, flow_id, 0),
1061                        TableFlowValue {
1062                            peer: Peer::empty(3)
1063                        }
1064                    ),
1065                    (
1066                        TableFlowKey::new(table_id, 4, flow_id, 1),
1067                        TableFlowValue {
1068                            peer: Peer::empty(4)
1069                        }
1070                    )
1071                ]
1072            );
1073        }
1074    }
1075
1076    #[tokio::test]
1077    async fn test_update_flow_metadata_flow_replace_diff_id_err() {
1078        let mem_kv = Arc::new(MemoryKvBackend::default());
1079        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
1080        let flow_id = 10;
1081        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
1082        let flow_routes = vec![
1083            (
1084                1u32,
1085                FlowRouteValue {
1086                    peer: Peer::empty(1),
1087                },
1088            ),
1089            (
1090                2,
1091                FlowRouteValue {
1092                    peer: Peer::empty(2),
1093                },
1094            ),
1095        ];
1096        flow_metadata_manager
1097            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
1098            .await
1099            .unwrap();
1100        // update again with same flow id
1101        flow_metadata_manager
1102            .update_flow_metadata(
1103                flow_id,
1104                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1105                &flow_value,
1106                flow_routes.clone(),
1107            )
1108            .await
1109            .unwrap();
1110        // update again with wrong flow id, expected error
1111        let err = flow_metadata_manager
1112            .update_flow_metadata(
1113                flow_id + 1,
1114                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1115                &flow_value,
1116                flow_routes,
1117            )
1118            .await
1119            .unwrap_err();
1120        assert_matches!(err, error::Error::Unexpected { .. });
1121        assert!(
1122            err.to_string()
1123                .contains("Reads different flow id when updating flow")
1124        );
1125    }
1126
1127    #[tokio::test]
1128    async fn test_update_flow_metadata_unexpected_err_prev_value_diff() {
1129        let mem_kv = Arc::new(MemoryKvBackend::default());
1130        let flow_metadata_manager = FlowMetadataManager::new(mem_kv);
1131        let flow_id = 10;
1132        let catalog_name = "greptime";
1133        let flow_value = test_flow_info_value("flow", [(0, 1u64)].into(), vec![1024, 1025, 1026]);
1134        let flow_routes = vec![
1135            (
1136                1u32,
1137                FlowRouteValue {
1138                    peer: Peer::empty(1),
1139                },
1140            ),
1141            (
1142                2,
1143                FlowRouteValue {
1144                    peer: Peer::empty(2),
1145                },
1146            ),
1147        ];
1148        flow_metadata_manager
1149            .create_flow_metadata(flow_id, flow_value.clone(), flow_routes.clone())
1150            .await
1151            .unwrap();
1152        // Creates again.
1153        let another_sink_table_name = TableName {
1154            catalog_name: catalog_name.to_string(),
1155            schema_name: "my_schema".to_string(),
1156            table_name: "another_sink_table".to_string(),
1157        };
1158        let flow_value = FlowInfoValue {
1159            catalog_name: "greptime".to_string(),
1160            query_context: None,
1161            flow_name: "flow".to_string(),
1162            source_table_ids: vec![1024, 1025, 1026],
1163            all_source_table_names: vec![],
1164            unresolved_source_table_names: vec![],
1165            sink_table_name: another_sink_table_name,
1166            flownode_ids: [(0, 1u64)].into(),
1167            raw_sql: "raw".to_string(),
1168            expire_after: Some(300),
1169            eval_interval_secs: None,
1170            comment: "hi".to_string(),
1171            options: Default::default(),
1172            status: FlowStatus::Active,
1173            created_time: chrono::Utc::now(),
1174            updated_time: chrono::Utc::now(),
1175            eval_schedule: None,
1176        };
1177        let err = flow_metadata_manager
1178            .update_flow_metadata(
1179                flow_id,
1180                &DeserializedValueWithBytes::from_inner(flow_value.clone()),
1181                &flow_value,
1182                flow_routes.clone(),
1183            )
1184            .await
1185            .unwrap_err();
1186        assert!(
1187            err.to_string().contains("Reads the different value"),
1188            "error: {:?}",
1189            err
1190        );
1191    }
1192}