Skip to main content

floe_core/io/storage/
target.rs

1use crate::{config, io, ConfigError, FloeResult};
2
3#[derive(Debug, Clone)]
4pub enum Target {
5    Local {
6        storage: String,
7        uri: String,
8        base_path: String,
9    },
10    S3 {
11        storage: String,
12        uri: String,
13        bucket: String,
14        base_key: String,
15    },
16    Adls {
17        storage: String,
18        uri: String,
19        account: String,
20        container: String,
21        base_path: String,
22    },
23    Gcs {
24        storage: String,
25        uri: String,
26        bucket: String,
27        base_key: String,
28    },
29}
30
31impl Target {
32    pub fn from_resolved(resolved: &config::ResolvedPath) -> FloeResult<Self> {
33        if let Some(path) = &resolved.local_path {
34            return Ok(Target::Local {
35                storage: resolved.storage.clone(),
36                uri: resolved.uri.clone(),
37                base_path: path.display().to_string(),
38            });
39        }
40        if resolved.uri.starts_with("s3://") {
41            let location = io::storage::s3::parse_s3_uri(&resolved.uri)?;
42            return Ok(Target::S3 {
43                storage: resolved.storage.clone(),
44                uri: resolved.uri.clone(),
45                bucket: location.bucket,
46                base_key: location.key,
47            });
48        }
49        if resolved.uri.starts_with("abfs://") {
50            let location = io::storage::adls::parse_adls_uri(&resolved.uri)?;
51            return Ok(Target::Adls {
52                storage: resolved.storage.clone(),
53                uri: resolved.uri.clone(),
54                account: location.account,
55                container: location.container,
56                base_path: location.path,
57            });
58        }
59        if resolved.uri.starts_with("gs://") {
60            let location = io::storage::gcs::parse_gcs_uri(&resolved.uri)?;
61            return Ok(Target::Gcs {
62                storage: resolved.storage.clone(),
63                uri: resolved.uri.clone(),
64                bucket: location.bucket,
65                base_key: location.key,
66            });
67        }
68        Err(Box::new(ConfigError(format!(
69            "unsupported storage uri: {}",
70            resolved.uri
71        ))))
72    }
73
74    pub fn storage(&self) -> &str {
75        match self {
76            Target::Local { storage, .. }
77            | Target::S3 { storage, .. }
78            | Target::Adls { storage, .. }
79            | Target::Gcs { storage, .. } => storage.as_str(),
80        }
81    }
82
83    pub fn target_uri(&self) -> &str {
84        match self {
85            Target::Local { uri, .. }
86            | Target::S3 { uri, .. }
87            | Target::Adls { uri, .. }
88            | Target::Gcs { uri, .. } => uri.as_str(),
89        }
90    }
91
92    pub fn s3_parts(&self) -> Option<(&str, &str)> {
93        match self {
94            Target::S3 {
95                bucket, base_key, ..
96            } => Some((bucket.as_str(), base_key.as_str())),
97            Target::Local { .. } => None,
98            Target::Adls { .. } | Target::Gcs { .. } => None,
99        }
100    }
101
102    pub fn gcs_parts(&self) -> Option<(&str, &str)> {
103        match self {
104            Target::Gcs {
105                bucket, base_key, ..
106            } => Some((bucket.as_str(), base_key.as_str())),
107            Target::Local { .. } | Target::S3 { .. } | Target::Adls { .. } => None,
108        }
109    }
110
111    pub fn adls_parts(&self) -> Option<(&str, &str, &str)> {
112        match self {
113            Target::Adls {
114                container,
115                account,
116                base_path,
117                ..
118            } => Some((container.as_str(), account.as_str(), base_path.as_str())),
119            Target::Local { .. } | Target::S3 { .. } | Target::Gcs { .. } => None,
120        }
121    }
122}