common_meta/key/
node_address.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::fmt::Display;
16
17use api::v1::meta::Role;
18use serde::{Deserialize, Serialize};
19use snafu::OptionExt;
20
21use crate::error::{InvalidMetadataSnafu, Result};
22use crate::key::{MetadataKey, NODE_ADDRESS_PATTERN, NODE_ADDRESS_PREFIX};
23use crate::peer::Peer;
24
25/// The key stores node address.
26///
27/// The layout: `__node_address/{role}/{node_id}`
28#[derive(Debug, PartialEq)]
29pub struct NodeAddressKey {
30    pub role: Role,
31    pub node_id: u64,
32}
33
34impl NodeAddressKey {
35    pub fn new(role: Role, node_id: u64) -> Self {
36        Self { role, node_id }
37    }
38
39    pub fn with_datanode(node_id: u64) -> Self {
40        Self::new(Role::Datanode, node_id)
41    }
42
43    pub fn with_flownode(node_id: u64) -> Self {
44        Self::new(Role::Flownode, node_id)
45    }
46}
47
48#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
49pub struct NodeAddressValue {
50    pub peer: Peer,
51}
52
53impl NodeAddressValue {
54    pub fn new(peer: Peer) -> Self {
55        Self { peer }
56    }
57}
58
59impl MetadataKey<'_, NodeAddressKey> for NodeAddressKey {
60    fn to_bytes(&self) -> Vec<u8> {
61        self.to_string().into_bytes()
62    }
63
64    fn from_bytes(bytes: &[u8]) -> Result<NodeAddressKey> {
65        let key = std::str::from_utf8(bytes).map_err(|e| {
66            InvalidMetadataSnafu {
67                err_msg: format!(
68                    "NodeAddressKey '{}' is not a valid UTF8 string: {e}",
69                    String::from_utf8_lossy(bytes)
70                ),
71            }
72            .build()
73        })?;
74        let captures = NODE_ADDRESS_PATTERN
75            .captures(key)
76            .context(InvalidMetadataSnafu {
77                err_msg: format!("Invalid NodeAddressKey '{key}'"),
78            })?;
79        // Safety: pass the regex check above
80        let role = captures[1].parse::<i32>().unwrap();
81        let role = Role::try_from(role).map_err(|_| {
82            InvalidMetadataSnafu {
83                err_msg: format!("Invalid Role value: {role}"),
84            }
85            .build()
86        })?;
87        let node_id = captures[2].parse::<u64>().unwrap();
88        Ok(NodeAddressKey::new(role, node_id))
89    }
90}
91
92impl Display for NodeAddressKey {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        write!(
95            f,
96            "{}/{}/{}",
97            NODE_ADDRESS_PREFIX, self.role as i32, self.node_id
98        )
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn test_node_address_key() {
108        let key = NodeAddressKey::new(Role::Datanode, 1);
109        let bytes = key.to_bytes();
110        let key2 = NodeAddressKey::from_bytes(&bytes).unwrap();
111        assert_eq!(key, key2);
112
113        let key = NodeAddressKey::new(Role::Flownode, 3);
114        let bytes = key.to_bytes();
115        let key2 = NodeAddressKey::from_bytes(&bytes).unwrap();
116        assert_eq!(key, key2);
117    }
118}
common_meta/key/node_address.rs

common_meta/key/
node_address.rs
