1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! Functions to export data from a store

use std::path::PathBuf;

use anyhow::Context;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use tracing::trace;

use crate::{
    format::collection::Collection,
    store::{BaoBlobSize, ExportFormat, ExportMode, MapEntry, Store as BaoStore},
    util::progress::{IdGenerator, ProgressSender},
    Hash,
};

/// Export a hash to the local file system.
///
/// This exports a single hash, or a collection `recursive` is true, from the `db` store to the
/// local filesystem. Depending on `mode` the data is either copied or reflinked (if possible).
///
/// Progress is reported as [`ExportProgress`] through a [`ProgressSender`]. Note that the
/// [`ExportProgress::AllDone`] event is not emitted from here, but left to an upper layer to send,
/// if desired.
pub async fn export<D: BaoStore>(
    db: &D,
    hash: Hash,
    outpath: PathBuf,
    format: ExportFormat,
    mode: ExportMode,
    progress: impl ProgressSender<Msg = ExportProgress> + IdGenerator,
) -> anyhow::Result<()> {
    match format {
        ExportFormat::Blob => export_blob(db, hash, outpath, mode, progress).await,
        ExportFormat::Collection => export_collection(db, hash, outpath, mode, progress).await,
    }
}

/// Export all entries of a collection, recursively, to files on the local filesystem.
pub async fn export_collection<D: BaoStore>(
    db: &D,
    hash: Hash,
    outpath: PathBuf,
    mode: ExportMode,
    progress: impl ProgressSender<Msg = ExportProgress> + IdGenerator,
) -> anyhow::Result<()> {
    tokio::fs::create_dir_all(&outpath).await?;
    let collection = Collection::load_db(db, &hash).await?;
    for (name, hash) in collection.into_iter() {
        #[allow(clippy::needless_borrow)]
        let path = outpath.join(pathbuf_from_name(&name));
        export_blob(db, hash, path, mode, progress.clone()).await?;
    }
    Ok(())
}

/// Export a single blob to a file on the local filesystem.
pub async fn export_blob<D: BaoStore>(
    db: &D,
    hash: Hash,
    outpath: PathBuf,
    mode: ExportMode,
    progress: impl ProgressSender<Msg = ExportProgress> + IdGenerator,
) -> anyhow::Result<()> {
    if let Some(parent) = outpath.parent() {
        tokio::fs::create_dir_all(parent).await?;
    }
    trace!("exporting blob {} to {}", hash, outpath.display());
    let id = progress.new_id();
    let entry = db.get(&hash).await?.context("entry not there")?;
    progress
        .send(ExportProgress::Found {
            id,
            hash,
            outpath: outpath.clone(),
            size: entry.size(),
            meta: None,
        })
        .await?;
    let progress1 = progress.clone();
    db.export(
        hash,
        outpath,
        mode,
        Box::new(move |offset| Ok(progress1.try_send(ExportProgress::Progress { id, offset })?)),
    )
    .await?;
    progress.send(ExportProgress::Done { id }).await?;
    Ok(())
}

/// Progress events for an export operation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ExportProgress {
    /// The download part is done for this id, we are now exporting the data
    /// to the specified out path.
    Found {
        /// Unique id of the entry.
        id: u64,
        /// The hash of the entry.
        hash: Hash,
        /// The size of the entry in bytes.
        size: BaoBlobSize,
        /// The path to the file where the data is exported.
        outpath: PathBuf,
        /// Operation-specific metadata.
        meta: Option<Bytes>,
    },
    /// We have made progress exporting the data.
    ///
    /// This is only sent for large blobs.
    Progress {
        /// Unique id of the entry that is being exported.
        id: u64,
        /// The offset of the progress, in bytes.
        offset: u64,
    },
    /// We finished exporting a blob
    Done {
        /// Unique id of the entry that is being exported.
        id: u64,
    },
    /// We are done with the whole operation.
    AllDone,
    /// We got an error and need to abort.
    Abort(serde_error::Error),
}

fn pathbuf_from_name(name: &str) -> PathBuf {
    let mut path = PathBuf::new();
    for part in name.split('/') {
        path.push(part);
    }
    path
}