cli/metadata/
snapshot.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::path::{Component, Path, PathBuf};
16
17use async_trait::async_trait;
18use clap::{Parser, Subcommand};
19use common_error::ext::BoxedError;
20use common_meta::snapshot::MetadataSnapshotManager;
21use object_store::ObjectStore;
22
23use crate::Tool;
24use crate::common::{ObjectStoreConfig, StoreConfig, new_fs_object_store};
25use crate::utils::resolve_relative_path_with_current_dir;
26
27/// Subcommand for metadata snapshot operations, including saving snapshots, restoring from snapshots, and viewing snapshot information.
28#[derive(Subcommand)]
29pub enum SnapshotCommand {
30    /// Save a snapshot of the current metadata state to a specified location.
31    Save(SaveCommand),
32    /// Restore metadata from a snapshot.
33    Restore(RestoreCommand),
34    /// Explore metadata from a snapshot.
35    Info(InfoCommand),
36}
37
38impl SnapshotCommand {
39    pub async fn build(&self) -> Result<Box<dyn Tool>, BoxedError> {
40        match self {
41            SnapshotCommand::Save(cmd) => Ok(Box::new(cmd.build().await?)),
42            SnapshotCommand::Restore(cmd) => Ok(Box::new(cmd.build().await?)),
43            SnapshotCommand::Info(cmd) => Ok(Box::new(cmd.build().await?)),
44        }
45    }
46}
47
48/// Export metadata snapshot tool.
49/// This tool is used to export metadata snapshot from etcd, pg or mysql.
50/// It will dump the metadata snapshot to local file or s3 bucket.
51/// The snapshot file will be in binary format.
52#[derive(Debug, Default, Parser)]
53pub struct SaveCommand {
54    /// The store configuration.
55    #[clap(flatten)]
56    store: StoreConfig,
57    /// The object store configuration.
58    #[clap(flatten)]
59    object_store: ObjectStoreConfig,
60    /// The path of the target snapshot file.
61    #[clap(
62        long,
63        default_value = "metadata_snapshot.metadata.fb",
64        alias = "file_name"
65    )]
66    file_path: String,
67    /// Specifies the root directory used for I/O operations.
68    #[clap(long, default_value = "/", alias = "output_dir")]
69    dir: String,
70}
71
72impl SaveCommand {
73    async fn build(&self) -> Result<MetaSnapshotTool, BoxedError> {
74        let kvbackend = self.store.build().await?;
75        let (object_store, file_path) = build_object_store_and_resolve_file_path(
76            self.object_store.clone(),
77            &self.dir,
78            &self.file_path,
79        )?;
80        let tool = MetaSnapshotTool {
81            inner: MetadataSnapshotManager::new(kvbackend, object_store),
82            file_path,
83        };
84        Ok(tool)
85    }
86}
87
88struct MetaSnapshotTool {
89    inner: MetadataSnapshotManager,
90    file_path: String,
91}
92
93#[async_trait]
94impl Tool for MetaSnapshotTool {
95    async fn do_work(&self) -> std::result::Result<(), BoxedError> {
96        self.inner
97            .dump(&self.file_path)
98            .await
99            .map_err(BoxedError::new)?;
100        Ok(())
101    }
102}
103
104/// Restore metadata from a snapshot file.
105///
106/// This command restores the metadata state from a previously saved snapshot.
107/// The snapshot can be loaded from either a local file system or an S3 bucket,
108/// depending on the provided configuration.
109#[derive(Debug, Default, Parser)]
110pub struct RestoreCommand {
111    /// The store configuration.
112    #[clap(flatten)]
113    store: StoreConfig,
114    /// The object store config.
115    #[clap(flatten)]
116    object_store: ObjectStoreConfig,
117    /// The path of the target snapshot file.
118    #[clap(
119        long,
120        default_value = "metadata_snapshot.metadata.fb",
121        alias = "file_name"
122    )]
123    file_path: String,
124    /// Specifies the root directory used for I/O operations.
125    #[clap(long, default_value = "/", alias = "input_dir")]
126    dir: String,
127    #[clap(long, default_value = "false")]
128    force: bool,
129}
130
131impl RestoreCommand {
132    async fn build(&self) -> Result<MetaRestoreTool, BoxedError> {
133        let kvbackend = self.store.build().await?;
134        let (object_store, file_path) = build_object_store_and_resolve_file_path(
135            self.object_store.clone(),
136            &self.dir,
137            &self.file_path,
138        )
139        .map_err(BoxedError::new)?;
140        let tool = MetaRestoreTool::new(
141            MetadataSnapshotManager::new(kvbackend, object_store),
142            file_path,
143            self.force,
144        );
145        Ok(tool)
146    }
147}
148
149struct MetaRestoreTool {
150    inner: MetadataSnapshotManager,
151    file_path: String,
152    force: bool,
153}
154
155impl MetaRestoreTool {
156    pub fn new(inner: MetadataSnapshotManager, file_path: String, force: bool) -> Self {
157        Self {
158            inner,
159            file_path,
160            force,
161        }
162    }
163}
164
165#[async_trait]
166impl Tool for MetaRestoreTool {
167    async fn do_work(&self) -> std::result::Result<(), BoxedError> {
168        let clean = self
169            .inner
170            .check_target_source_clean()
171            .await
172            .map_err(BoxedError::new)?;
173        if clean {
174            common_telemetry::info!(
175                "The target source is clean, we will restore the metadata snapshot."
176            );
177            self.inner
178                .restore(&self.file_path)
179                .await
180                .map_err(BoxedError::new)?;
181            Ok(())
182        } else if !self.force {
183            common_telemetry::warn!(
184                "The target source is not clean, if you want to restore the metadata snapshot forcefully, please use --force option."
185            );
186            Ok(())
187        } else {
188            common_telemetry::info!(
189                "The target source is not clean, We will restore the metadata snapshot with --force."
190            );
191            self.inner
192                .restore(&self.file_path)
193                .await
194                .map_err(BoxedError::new)?;
195            Ok(())
196        }
197    }
198}
199
200/// Explore metadata from a snapshot file.
201///
202/// This command allows filtering the metadata by a specific key and limiting the number of results.
203/// It prints the filtered metadata to the console.
204#[derive(Debug, Default, Parser)]
205pub struct InfoCommand {
206    /// The object store config.
207    #[clap(flatten)]
208    object_store: ObjectStoreConfig,
209    /// The path of the target snapshot file.
210    #[clap(
211        long,
212        default_value = "metadata_snapshot.metadata.fb",
213        alias = "file_name"
214    )]
215    file_path: String,
216    /// Specifies the root directory used for I/O operations.
217    #[clap(long, default_value = "/", alias = "input_dir")]
218    dir: String,
219    /// The query string to filter the metadata.
220    #[clap(long, default_value = "*")]
221    inspect_key: String,
222    /// The limit of the metadata to query.
223    #[clap(long)]
224    limit: Option<usize>,
225}
226
227struct MetaInfoTool {
228    inner: ObjectStore,
229    file_path: String,
230    inspect_key: String,
231    limit: Option<usize>,
232}
233
234#[async_trait]
235impl Tool for MetaInfoTool {
236    async fn do_work(&self) -> std::result::Result<(), BoxedError> {
237        let result = MetadataSnapshotManager::info(
238            &self.inner,
239            &self.file_path,
240            &self.inspect_key,
241            self.limit,
242        )
243        .await
244        .map_err(BoxedError::new)?;
245        for item in result {
246            println!("{}", item);
247        }
248        Ok(())
249    }
250}
251
252impl InfoCommand {
253    async fn build(&self) -> Result<MetaInfoTool, BoxedError> {
254        let (object_store, file_path) = build_object_store_and_resolve_file_path(
255            self.object_store.clone(),
256            &self.dir,
257            &self.file_path,
258        )?;
259        let tool = MetaInfoTool {
260            inner: object_store,
261            file_path,
262            inspect_key: self.inspect_key.clone(),
263            limit: self.limit,
264        };
265        Ok(tool)
266    }
267}
268
269/// Builds the object store and resolves the file path.
270fn build_object_store_and_resolve_file_path(
271    object_store: ObjectStoreConfig,
272    fs_root: &str,
273    file_path: &str,
274) -> Result<(ObjectStore, String), BoxedError> {
275    if let Some(object_store) = object_store.build().map_err(BoxedError::new)? {
276        return Ok((object_store, file_path.to_string()));
277    }
278
279    let absolute_path =
280        resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)?;
281    let (root, key) = split_fs_snapshot_path(fs_root, &absolute_path);
282    Ok((new_fs_object_store(&root)?, key))
283}
284
285/// Separates a local absolute path into an operator root and a relative key.
286fn split_fs_snapshot_path(fs_root: &str, absolute_path: &str) -> (String, String) {
287    let path = Path::new(absolute_path);
288    // OpenDAL keys cannot contain a Windows drive/UNC prefix or a root directory.
289    // Keep the filesystem anchor in the operator root, not in the object key.
290    let anchor: PathBuf = path
291        .components()
292        .take_while(|c| matches!(c, Component::Prefix(_) | Component::RootDir))
293        .collect();
294    let key = path
295        .components()
296        .filter(|c| !matches!(c, Component::Prefix(_) | Component::RootDir))
297        .map(|c| c.as_os_str().to_string_lossy())
298        .collect::<Vec<_>>()
299        .join("/");
300    // Preserve the existing custom-root semantics: cwd-resolved keys are rooted
301    // under --dir. The default root follows the input path's drive on Windows.
302    let root = if fs_root == "/" {
303        anchor.to_string_lossy()
304    } else {
305        fs_root.into()
306    };
307    (root.into_owned(), key)
308}
309
310#[cfg(test)]
311mod tests {
312    use std::env;
313    use std::sync::Arc;
314    use std::time::Duration;
315
316    use clap::Parser;
317    use common_meta::kv_backend::KvBackend;
318    use common_meta::kv_backend::memory::MemoryKvBackend;
319    use common_meta::rpc::store::PutRequest;
320    use object_store::ObjectStore;
321
322    use super::*;
323    use crate::metadata::snapshot::RestoreCommand;
324
325    fn create_raftengine_url(path: &std::path::Path) -> String {
326        let mut path = path.to_string_lossy().replace('\\', "/");
327        if !path.starts_with('/') {
328            path = format!("/{}", path);
329        }
330        format!("raftengine://{}", path)
331    }
332
333    #[tokio::test]
334    async fn test_cmd_resolve_file_path() {
335        let current_dir = env::current_dir().unwrap();
336        for file_path in [
337            "metadata_snapshot.metadata.fb".to_string(),
338            current_dir
339                .join("metadata_snapshot.metadata.fb")
340                .display()
341                .to_string(),
342        ] {
343            let cmd = RestoreCommand::parse_from([
344                "",
345                "--file_name",
346                &file_path,
347                "--backend",
348                "memory-store",
349                "--store-addrs",
350                "memory://",
351            ]);
352            let tool = cmd.build().await.unwrap();
353            assert!(!Path::new(&tool.file_path).is_absolute());
354            assert!(
355                Path::new(&tool.file_path)
356                    .components()
357                    .all(|c| matches!(c, Component::Normal(_)))
358            );
359            let expected = current_dir
360                .components()
361                .filter_map(|c| match c {
362                    Component::Normal(part) => Some(part.to_str().unwrap()),
363                    _ => None,
364                })
365                .chain(["metadata_snapshot.metadata.fb"])
366                .collect::<Vec<_>>()
367                .join("/");
368            assert_eq!(tool.file_path, expected);
369        }
370    }
371
372    #[cfg(windows)]
373    #[test]
374    fn test_split_windows_snapshot_path() {
375        for (path, root, key) in [
376            (
377                r"C:\backup\snapshot.metadata.fb",
378                r"C:\",
379                "backup/snapshot.metadata.fb",
380            ),
381            (
382                r"D:\备份 dir\snapshot.metadata.fb",
383                r"D:\",
384                "备份 dir/snapshot.metadata.fb",
385            ),
386            (
387                r"\\server\share\backup\snapshot.metadata.fb",
388                r"\\server\share\",
389                "backup/snapshot.metadata.fb",
390            ),
391            (
392                r"\\?\C:\backup\snapshot.metadata.fb",
393                r"\\?\C:\",
394                "backup/snapshot.metadata.fb",
395            ),
396        ] {
397            assert_eq!(
398                split_fs_snapshot_path("/", path),
399                (root.to_string(), key.to_string())
400            );
401            assert_eq!(
402                split_fs_snapshot_path(r"D:\root", path),
403                (r"D:\root".to_string(), key.to_string())
404            );
405        }
406    }
407
408    #[tokio::test]
409    async fn test_fs_snapshot_paths() {
410        let temp_dir = tempfile::tempdir().unwrap();
411        for name in ["snapshot.metadata.fb", "nested dir/快照.metadata.fb"] {
412            let path = temp_dir.path().join(name);
413            let (store, key) = build_object_store_and_resolve_file_path(
414                ObjectStoreConfig::default(),
415                "/",
416                path.to_str().unwrap(),
417            )
418            .unwrap();
419            assert!(!path.exists());
420            store.write(&key, "snapshot").await.unwrap();
421            assert_eq!(std::fs::read(&path).unwrap(), b"snapshot");
422            assert_eq!(store.read(&key).await.unwrap().to_bytes(), &b"snapshot"[..]);
423        }
424    }
425
426    #[tokio::test]
427    async fn test_fs_snapshot_custom_root() {
428        let temp_dir = tempfile::tempdir().unwrap();
429        let (store, key) = build_object_store_and_resolve_file_path(
430            ObjectStoreConfig::default(),
431            temp_dir.path().to_str().unwrap(),
432            "snapshot.metadata.fb",
433        )
434        .unwrap();
435        let current_dir = env::current_dir().unwrap();
436        let expected = current_dir
437            .components()
438            .filter_map(|c| match c {
439                Component::Normal(part) => Some(part.to_str().unwrap()),
440                _ => None,
441            })
442            .chain(["snapshot.metadata.fb"])
443            .collect::<Vec<_>>()
444            .join("/");
445        assert_eq!(key, expected);
446        store.write(&key, "snapshot").await.unwrap();
447        assert_eq!(
448            std::fs::read(temp_dir.path().join(expected)).unwrap(),
449            b"snapshot"
450        );
451    }
452
453    #[tokio::test]
454    async fn test_fs_snapshot_rejects_parent_traversal() {
455        let temp_dir = tempfile::tempdir().unwrap();
456        let path = temp_dir.path().join("../snapshot.metadata.fb");
457        let (store, key) = build_object_store_and_resolve_file_path(
458            ObjectStoreConfig::default(),
459            "/",
460            path.to_str().unwrap(),
461        )
462        .unwrap();
463        assert!(key.split('/').any(|part| part == ".."));
464        let err = store.write(&key, "snapshot").await.unwrap_err();
465        assert_eq!(err.kind(), object_store::ErrorKind::NotFound);
466        assert!(err.to_string().contains("path escapes the configured root"));
467    }
468
469    #[test]
470    fn test_remote_snapshot_key_unchanged() {
471        let cmd = RestoreCommand::parse_from([
472            "",
473            "--backend",
474            "memory-store",
475            "--store-addrs",
476            "memory://",
477            "--s3",
478            "--s3-bucket",
479            "snapshots",
480            "--s3-region",
481            "us-east-1",
482        ]);
483        let key = "nested/snapshot.metadata.fb";
484        let (store, resolved) =
485            build_object_store_and_resolve_file_path(cmd.object_store, "/ignored", key).unwrap();
486        assert_eq!(store.info().scheme(), object_store::services::S3_SCHEME);
487        assert_eq!(resolved, key);
488    }
489
490    async fn setup_backup_file(object_store: ObjectStore, file_path: &str) {
491        let kv_backend = Arc::new(MemoryKvBackend::default());
492        let manager = MetadataSnapshotManager::new(kv_backend.clone(), object_store);
493        // Put some data into the kv backend
494        kv_backend
495            .put(
496                PutRequest::new()
497                    .with_key(b"test".to_vec())
498                    .with_value(b"test".to_vec()),
499            )
500            .await
501            .unwrap();
502        manager.dump(file_path).await.unwrap();
503    }
504
505    #[tokio::test]
506    async fn test_restore_raft_engine_store() {
507        common_telemetry::init_default_ut_logging();
508        let temp_dir = tempfile::tempdir().unwrap();
509        let root = temp_dir.path().display().to_string();
510        let object_store = new_fs_object_store(&root).unwrap();
511        setup_backup_file(object_store, "/backup/metadata_snapshot.metadata.fb").await;
512        let metadata_path = temp_dir.path().join("metadata");
513        {
514            let cmd = RestoreCommand::parse_from([
515                "",
516                "--file_name",
517                temp_dir
518                    .path()
519                    .join("backup")
520                    .join("metadata_snapshot.metadata.fb")
521                    .to_str()
522                    .unwrap(),
523                "--backend",
524                "raft-engine-store",
525                "--store-addrs",
526                &create_raftengine_url(&metadata_path),
527            ]);
528            let tool = cmd.build().await.unwrap();
529            tool.do_work().await.unwrap();
530        }
531        // Waits for the raft engine release the file lock.
532        tokio::time::sleep(Duration::from_secs(1)).await;
533        let kv = standalone::build_metadata_kvbackend(
534            metadata_path.display().to_string(),
535            Default::default(),
536        )
537        .unwrap();
538
539        let value = kv.get(b"test").await.unwrap().unwrap().value;
540        assert_eq!(value, b"test");
541    }
542
543    #[tokio::test]
544    async fn test_save_raft_engine_store() {
545        common_telemetry::init_default_ut_logging();
546        let temp_dir = tempfile::tempdir().unwrap();
547        let root = temp_dir.path().display().to_string();
548        let metadata_path = temp_dir.path().join("metadata");
549        {
550            let kv = standalone::build_metadata_kvbackend(
551                metadata_path.to_string_lossy().to_string(),
552                Default::default(),
553            )
554            .unwrap();
555            kv.put(
556                PutRequest::new()
557                    .with_key(b"test".to_vec())
558                    .with_value(b"test".to_vec()),
559            )
560            .await
561            .unwrap();
562        }
563        // Waits for the raft engine release the file lock.
564        tokio::time::sleep(Duration::from_secs(1)).await;
565        {
566            let cmd = SaveCommand::parse_from([
567                "",
568                "--file_name",
569                temp_dir
570                    .path()
571                    .join("backup")
572                    .join("metadata_snapshot.metadata.fb")
573                    .to_str()
574                    .unwrap(),
575                "--backend",
576                "raft-engine-store",
577                "--store-addrs",
578                &create_raftengine_url(&metadata_path),
579            ]);
580            let tool = cmd.build().await.unwrap();
581            tool.do_work().await.unwrap();
582        }
583
584        // Reads the snapshot file from the object store.
585        let object_store = new_fs_object_store(&root).unwrap();
586        let kv_backend = Arc::new(MemoryKvBackend::default());
587        let manager = MetadataSnapshotManager::new(kv_backend.clone(), object_store);
588        manager
589            .restore("/backup/metadata_snapshot.metadata.fb")
590            .await
591            .unwrap();
592        let value = kv_backend.get(b"test").await.unwrap().unwrap().value;
593        assert_eq!(value, b"test");
594    }
595
596    #[test]
597    fn test_path() {
598        let path = "C:\\Users\\user\\AppData\\Local\\Temp\\.tmpuPiVuB\\metadata";
599        let path = Path::new(path);
600        let url = create_raftengine_url(path);
601        assert_eq!(
602            url,
603            "raftengine:///C:/Users/user/AppData/Local/Temp/.tmpuPiVuB/metadata"
604        );
605        let url = url::Url::parse(&url).unwrap();
606        assert_eq!(
607            url.path(),
608            "/C:/Users/user/AppData/Local/Temp/.tmpuPiVuB/metadata"
609        );
610    }
611}
cli/metadata/snapshot.rs

cli/metadata/
snapshot.rs
