Skip to main content

common_datasource/object_store/
oss.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::collections::HashMap;
16
17use object_store::ObjectStore;
18use object_store::services::Oss;
19use object_store::util::{with_instrument_layers, with_retry_layers};
20use snafu::ResultExt;
21
22use crate::error::{self, Result};
23
24const BUCKET: &str = "bucket";
25const ENDPOINT: &str = "endpoint";
26const ACCESS_KEY_ID: &str = "access_key_id";
27const ACCESS_KEY_SECRET: &str = "access_key_secret";
28const ROOT: &str = "root";
29const ALLOW_ANONYMOUS: &str = "allow_anonymous";
30const SKIP_SIGNATURE: &str = "skip_signature";
31
32/// Check if the key is supported in OSS configuration.
33pub fn is_supported_in_oss(key: &str) -> bool {
34    [
35        ROOT,
36        ALLOW_ANONYMOUS,
37        SKIP_SIGNATURE,
38        BUCKET,
39        ENDPOINT,
40        ACCESS_KEY_ID,
41        ACCESS_KEY_SECRET,
42    ]
43    .contains(&key)
44}
45
46/// Build an OSS backend using the provided bucket, root, and connection parameters.
47pub fn build_oss_backend(
48    bucket: &str,
49    root: &str,
50    connection: &HashMap<String, String>,
51) -> Result<ObjectStore> {
52    let mut builder = Oss::default().bucket(bucket).root(root);
53
54    if let Some(endpoint) = connection.get(ENDPOINT) {
55        builder = builder.endpoint(endpoint);
56    }
57
58    if let Some(access_key_id) = connection.get(ACCESS_KEY_ID) {
59        builder = builder.access_key_id(access_key_id);
60    }
61
62    if let Some(access_key_secret) = connection.get(ACCESS_KEY_SECRET) {
63        builder = builder.access_key_secret(access_key_secret);
64    }
65
66    if let Some((key, value)) = connection
67        .get(SKIP_SIGNATURE)
68        .map(|value| (SKIP_SIGNATURE, value))
69        .or_else(|| {
70            connection
71                .get(ALLOW_ANONYMOUS)
72                .map(|value| (ALLOW_ANONYMOUS, value))
73        })
74    {
75        let skip_signature = value.as_str().parse::<bool>().map_err(|e| {
76            error::InvalidConnectionSnafu {
77                msg: format!("failed to parse the option {}={}, {}", key, value, e),
78            }
79            .build()
80        })?;
81        if skip_signature {
82            builder = builder.skip_signature();
83        }
84    }
85
86    let object_store = ObjectStore::new(builder)
87        .context(error::BuildBackendSnafu)?
88        .finish();
89    Ok(with_instrument_layers(
90        with_retry_layers(object_store),
91        true,
92    ))
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_is_supported_in_oss() {
101        assert!(is_supported_in_oss(ROOT));
102        assert!(is_supported_in_oss(ALLOW_ANONYMOUS));
103        assert!(is_supported_in_oss(SKIP_SIGNATURE));
104        assert!(is_supported_in_oss(BUCKET));
105        assert!(is_supported_in_oss(ENDPOINT));
106        assert!(is_supported_in_oss(ACCESS_KEY_ID));
107        assert!(is_supported_in_oss(ACCESS_KEY_SECRET));
108        assert!(!is_supported_in_oss("foo"));
109        assert!(!is_supported_in_oss("BAR"));
110    }
111
112    #[test]
113    fn test_build_oss_backend_all_fields_valid() {
114        let mut connection = HashMap::new();
115        connection.insert(
116            ENDPOINT.to_string(),
117            "http://oss-ap-southeast-1.aliyuncs.com".to_string(),
118        );
119        connection.insert(ACCESS_KEY_ID.to_string(), "key_id".to_string());
120        connection.insert(ACCESS_KEY_SECRET.to_string(), "key_secret".to_string());
121        connection.insert(ALLOW_ANONYMOUS.to_string(), "true".to_string());
122
123        let result = build_oss_backend("my-bucket", "my-root", &connection);
124        assert!(result.is_ok());
125    }
126}