cmd/datanode/
tool_util.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
15#[cfg(feature = "dev-tools")]
16use std::fs::File;
17use std::path::Path;
18use std::sync::Arc;
19
20use common_wal::config::DatanodeWalConfig;
21use datanode::config::RegionEngineConfig;
22use datanode::store;
23use mito2::config::MitoConfig;
24use object_store::ObjectStore;
25#[cfg(feature = "dev-tools")]
26use parquet::basic::Compression;
27use parquet::file::metadata::{KeyValue, ParquetMetaData};
28#[cfg(feature = "dev-tools")]
29use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
30use snafu::OptionExt;
31#[cfg(feature = "dev-tools")]
32use snafu::ResultExt;
33use store_api::metadata::{RegionMetadata, RegionMetadataRef};
34use store_api::region_request::PathType;
35use store_api::storage::{FileId, RegionId};
36
37use crate::datanode::{StorageConfig, StorageConfigWrapper};
38use crate::error;
39
40pub(crate) fn parse_config(
41    config_path: &Path,
42) -> error::Result<(StorageConfig, MitoConfig, DatanodeWalConfig)> {
43    let cfg_str = std::fs::read_to_string(config_path).map_err(|e| {
44        error::IllegalConfigSnafu {
45            msg: format!("failed to read config {}: {e}", config_path.display()),
46        }
47        .build()
48    })?;
49
50    let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| {
51        error::IllegalConfigSnafu {
52            msg: format!("failed to parse config {}: {e}", config_path.display()),
53        }
54        .build()
55    })?;
56
57    let wal_config = store_cfg.wal;
58    let storage_config = store_cfg.storage;
59    let mito_engine_config = store_cfg
60        .region_engine
61        .into_iter()
62        .find_map(|config| match config {
63            RegionEngineConfig::Mito(mito) => Some(mito),
64            _ => None,
65        })
66        .with_context(|| error::IllegalConfigSnafu {
67            msg: format!("Engine config not found in {:?}", config_path),
68        })?;
69
70    Ok((storage_config, mito_engine_config, wal_config))
71}
72
73pub(crate) async fn build_object_store(config: &StorageConfig) -> error::Result<ObjectStore> {
74    store::new_object_store(config.store.clone(), &config.data_home)
75        .await
76        .map_err(|e| {
77            error::IllegalConfigSnafu {
78                msg: format!("Failed to build object store: {e:?}"),
79            }
80            .build()
81        })
82}
83
84pub(crate) fn extract_region_metadata(
85    file_path: &str,
86    metadata: &ParquetMetaData,
87) -> error::Result<RegionMetadataRef> {
88    let key_values: Option<&Vec<KeyValue>> = metadata.file_metadata().key_value_metadata();
89    let Some(key_values) = key_values else {
90        return Err(error::IllegalConfigSnafu {
91            msg: format!("{file_path}: missing parquet key_value metadata"),
92        }
93        .build());
94    };
95    let json = key_values
96        .iter()
97        .find(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY)
98        .and_then(|key_value| key_value.value.as_ref())
99        .ok_or_else(|| {
100            error::IllegalConfigSnafu {
101                msg: format!(
102                    "{file_path}: key {} not found or empty",
103                    mito2::sst::parquet::PARQUET_METADATA_KEY
104                ),
105            }
106            .build()
107        })?;
108    let region = RegionMetadata::from_json(json).map_err(|e| {
109        error::IllegalConfigSnafu {
110            msg: format!("invalid region metadata json: {e}"),
111        }
112        .build()
113    })?;
114    Ok(Arc::new(region))
115}
116
117pub(crate) fn parse_region_id(value: &str) -> error::Result<RegionId> {
118    if let Some((table_id, region_number)) = value.split_once(':') {
119        let table_id = table_id.parse().map_err(|e| {
120            error::IllegalConfigSnafu {
121                msg: format!("invalid table_id in region_id '{value}': {e}"),
122            }
123            .build()
124        })?;
125        let region_number = region_number.parse().map_err(|e| {
126            error::IllegalConfigSnafu {
127                msg: format!("invalid region_num in region_id '{value}': {e}"),
128            }
129            .build()
130        })?;
131        Ok(RegionId::new(table_id, region_number))
132    } else {
133        value.parse().map(RegionId::from_u64).map_err(|e| {
134            error::IllegalConfigSnafu {
135                msg: format!("invalid region_id '{value}': {e}"),
136            }
137            .build()
138        })
139    }
140}
141
142pub(crate) fn parse_file_id(value: &str) -> error::Result<FileId> {
143    FileId::parse_str(value).map_err(|e| {
144        error::IllegalConfigSnafu {
145            msg: format!("invalid file_id '{value}': {e}"),
146        }
147        .build()
148    })
149}
150
151pub(crate) fn parse_path_type(value: &str) -> error::Result<PathType> {
152    match value.to_lowercase().as_str() {
153        "bare" => Ok(PathType::Bare),
154        "data" => Ok(PathType::Data),
155        "metadata" => Ok(PathType::Metadata),
156        _ => Err(error::IllegalConfigSnafu {
157            msg: format!("invalid path_type '{value}', expected: bare, data, metadata"),
158        }
159        .build()),
160    }
161}
162
163pub(crate) fn format_bytes(bytes: u64) -> String {
164    const KIB: u64 = 1024;
165    const MIB: u64 = 1024 * KIB;
166    const GIB: u64 = 1024 * MIB;
167    if bytes >= GIB {
168        format!("{:.2} GiB", bytes as f64 / GIB as f64)
169    } else if bytes >= MIB {
170        format!("{:.2} MiB", bytes as f64 / MIB as f64)
171    } else if bytes >= KIB {
172        format!("{:.2} KiB", bytes as f64 / KIB as f64)
173    } else {
174        format!("{bytes} B")
175    }
176}
177
178pub(crate) fn max_row_group_uncompressed_size(metadata: &ParquetMetaData) -> u64 {
179    metadata
180        .row_groups()
181        .iter()
182        .map(|row_group| {
183            row_group
184                .columns()
185                .iter()
186                .map(|column| column.uncompressed_size() as u64)
187                .sum::<u64>()
188        })
189        .max()
190        .unwrap_or(0)
191}
192
193#[cfg(feature = "dev-tools")]
194pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result<ParquetMetaData> {
195    let file = File::open(path).context(error::FileIoSnafu)?;
196    ParquetMetaDataReader::new()
197        .with_page_index_policy(PageIndexPolicy::Optional)
198        .parse_and_finish(&file)
199        .map_err(|e| {
200            error::IllegalConfigSnafu {
201                msg: format!("read parquet metadata failed for {}: {e}", path.display()),
202            }
203            .build()
204        })
205}
206
207#[cfg(feature = "dev-tools")]
208pub(crate) fn compression_name(compression: Compression) -> &'static str {
209    match compression {
210        Compression::UNCOMPRESSED => "uncompressed",
211        Compression::SNAPPY => "snappy",
212        Compression::GZIP(_) => "gzip",
213        Compression::LZO => "lzo",
214        Compression::BROTLI(_) => "brotli",
215        Compression::LZ4 => "lz4",
216        Compression::ZSTD(_) => "zstd",
217        Compression::LZ4_RAW => "lz4-raw",
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    #[test]
226    fn test_parse_region_and_path_type() {
227        assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7));
228        assert_eq!(
229            parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(),
230            RegionId::new(1, 2)
231        );
232        assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare);
233        assert_eq!(parse_path_type("data").unwrap(), PathType::Data);
234        assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata);
235    }
236}
cmd/datanode/tool_util.rs

cmd/datanode/
tool_util.rs
