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