sk_core/
external_storage.rs

1/// We use the [object_store](https://docs.rs/object_store/latest/object_store/index.html) crate to
2/// enable reading/writing from the three major cloud providers (AWS, Azure, GCP), as well as
3/// to/from a local filesystem or an in-memory store.  Supposedly HTTP with WebDAV is supported as
4/// well but that is completely untested.
5///
6/// The reader will load credentials from the environment to communicate with the cloud provider,
7/// as follows (other auth mechanisms _may_ work as well but are currently untested):
8///
9/// ### AWS
10///
11/// Set the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables, and pass in a
12/// URL like `s3://bucket/path/to/resource`.
13///
14/// ### Azure
15///
16/// Set the `AZURE_STORAGE_ACCOUNT_NAME` and `AZURE_STORAGE_ACCOUNT_KEY` environment variables, and
17/// pass in a URL like `azure://container/path/to/resources` (do not include the storage acocunt
18/// name in the URL).
19///
20/// ### GCP
21///
22/// Set the `GOOGLE_SERVICE_ACCOUNT` environment variable to the path for your service account JSON
23/// file (if you're running inside a container, you'll need that file injected as well).  Pass in a
24/// URL like `gs://bucket/path/to/resource`.
25use std::path::{
26    absolute,
27    PathBuf,
28};
29
30use anyhow::anyhow;
31use async_trait::async_trait;
32use bytes::Bytes;
33#[cfg(feature = "mock")]
34use mockall::automock;
35use object_store::path::Path;
36use object_store::{
37    DynObjectStore,
38    ObjectStoreScheme,
39    PutPayload,
40};
41use reqwest::Url;
42
43use crate::errors::*;
44
45#[cfg_attr(feature = "mock", automock)]
46#[async_trait]
47pub trait ObjectStoreWrapper {
48    fn scheme(&self) -> ObjectStoreScheme;
49    async fn put(&self, data: Bytes) -> EmptyResult;
50    async fn get(&self) -> anyhow::Result<Bytes>;
51}
52
53#[derive(Debug)]
54pub struct SkObjectStore {
55    scheme: ObjectStoreScheme,
56    store: Box<DynObjectStore>,
57    path: Path,
58}
59
60impl SkObjectStore {
61    pub fn new(path_str: &str) -> anyhow::Result<SkObjectStore> {
62        let (scheme, path) = parse_path(path_str)?;
63        let store: Box<DynObjectStore> = match scheme {
64            ObjectStoreScheme::Local => Box::new(object_store::local::LocalFileSystem::new()),
65            ObjectStoreScheme::Memory => Box::new(object_store::memory::InMemory::new()),
66            ObjectStoreScheme::AmazonS3 => {
67                Box::new(object_store::aws::AmazonS3Builder::from_env().with_url(path_str).build()?)
68            },
69            ObjectStoreScheme::MicrosoftAzure => Box::new(
70                object_store::azure::MicrosoftAzureBuilder::from_env()
71                    .with_url(path_str)
72                    .build()?,
73            ),
74            ObjectStoreScheme::GoogleCloudStorage => Box::new(
75                object_store::gcp::GoogleCloudStorageBuilder::from_env()
76                    .with_url(path_str)
77                    .build()?,
78            ),
79            ObjectStoreScheme::Http => Box::new(object_store::http::HttpBuilder::new().with_url(path_str).build()?),
80            _ => unimplemented!(),
81        };
82
83        Ok(SkObjectStore { scheme, store, path })
84    }
85}
86
87#[async_trait]
88impl ObjectStoreWrapper for SkObjectStore {
89    fn scheme(&self) -> ObjectStoreScheme {
90        self.scheme.clone()
91    }
92
93    async fn put(&self, data: Bytes) -> EmptyResult {
94        let payload = PutPayload::from_bytes(data);
95        self.store.put(&self.path, payload).await?;
96        Ok(())
97    }
98
99    async fn get(&self) -> anyhow::Result<Bytes> {
100        Ok(self.store.get(&self.path).await?.bytes().await?)
101    }
102}
103
104fn parse_path(path_str: &str) -> anyhow::Result<(ObjectStoreScheme, Path)> {
105    let url = match Url::parse(path_str) {
106        Err(url::ParseError::RelativeUrlWithoutBase) => {
107            let path = absolute(PathBuf::from(path_str))?;
108            Url::from_file_path(path).map_err(|e| anyhow!("could not create URL from file path: {e:?}"))?
109        },
110        res => res?,
111    };
112
113    Ok(ObjectStoreScheme::parse(&url)?)
114}
115
116#[cfg(test)]
117#[cfg_attr(coverage, coverage(off))]
118mod test {
119    use sk_testutils::*;
120
121    use super::*;
122
123    #[rstest]
124    fn test_new_sk_object_store_invalid() {
125        let _ = SkObjectStore::new("oracle3://foo/bar").unwrap_err();
126    }
127
128    #[rstest]
129    fn test_new_sk_object_store() {
130        let store = SkObjectStore::new("s3://foo/bar").unwrap();
131        assert_eq!(store.scheme(), ObjectStoreScheme::AmazonS3);
132    }
133
134    #[rstest]
135    #[case::with_base("file:///tmp/foo")]
136    #[case::without_base("/tmp/foo")]
137    #[case::relative("foo")]
138    fn test_new_sk_object_store_local_path(#[case] path: &str) {
139        let store = SkObjectStore::new(path).unwrap();
140        assert_eq!(store.scheme(), ObjectStoreScheme::Local);
141    }
142}