Skip to main content

iota_storage/object_store/http/
gcs.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{fmt, sync::Arc};
6
7use anyhow::Result;
8use async_trait::async_trait;
9use bytes::Bytes;
10use iota_config::object_storage_config::{CONNECT_TIMEOUT, TRANSFER_STALL_TIMEOUT};
11use object_store::{GetResult, path::Path};
12use percent_encoding::{NON_ALPHANUMERIC, percent_encode, utf8_percent_encode};
13use reqwest::{Client, ClientBuilder};
14
15use crate::object_store::{
16    ObjectStoreGetExt, collect_get_result_with_progress,
17    http::{DEFAULT_USER_AGENT, exists, get, size},
18};
19
20#[derive(Debug)]
21struct GoogleCloudStorageClient {
22    client: Client,
23    bucket_name_encoded: String,
24}
25
26impl GoogleCloudStorageClient {
27    pub fn new(bucket: &str) -> Result<Self> {
28        let _ = rustls::crypto::ring::default_provider().install_default();
29        let mut builder = ClientBuilder::new().pool_idle_timeout(None);
30        builder = builder
31            .user_agent(DEFAULT_USER_AGENT)
32            .connect_timeout(CONNECT_TIMEOUT)
33            .read_timeout(TRANSFER_STALL_TIMEOUT);
34        let client = builder.https_only(false).build()?;
35        let bucket_name_encoded = percent_encode(bucket.as_bytes(), NON_ALPHANUMERIC).to_string();
36
37        Ok(Self {
38            client,
39            bucket_name_encoded,
40        })
41    }
42
43    async fn get(&self, path: &Path) -> Result<GetResult> {
44        let url = self.object_url(path);
45        get(&url, "gcs", path, &self.client).await
46    }
47
48    async fn exists(&self, path: &Path) -> Result<bool> {
49        let url = self.object_url(path);
50        exists(&url, "gcs", path, &self.client).await
51    }
52
53    async fn size(&self, path: &Path) -> Result<u64> {
54        let url = self.object_url(path);
55        size(&url, "gcs", path, &self.client).await
56    }
57
58    fn object_url(&self, path: &Path) -> String {
59        let encoded = utf8_percent_encode(path.as_ref(), NON_ALPHANUMERIC);
60        format!(
61            "https://storage.googleapis.com/{}/{}",
62            self.bucket_name_encoded, encoded
63        )
64    }
65}
66
67/// Interface for [Google Cloud Storage](https://cloud.google.com/storage/).
68#[derive(Debug)]
69pub struct GoogleCloudStorage {
70    client: Arc<GoogleCloudStorageClient>,
71}
72
73impl GoogleCloudStorage {
74    pub fn new(bucket: &str) -> Result<Self> {
75        let gcs_client = GoogleCloudStorageClient::new(bucket)?;
76        Ok(GoogleCloudStorage {
77            client: Arc::new(gcs_client),
78        })
79    }
80}
81
82impl fmt::Display for GoogleCloudStorage {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        write!(f, "gcs:{}", self.client.bucket_name_encoded)
85    }
86}
87
88#[async_trait]
89impl ObjectStoreGetExt for GoogleCloudStorage {
90    async fn get_bytes(&self, location: &Path) -> Result<Bytes> {
91        let result = self.client.get(location).await?;
92        let bytes = result.bytes().await?;
93        Ok(bytes)
94    }
95
96    async fn get_bytes_with_progress(
97        &self,
98        location: &Path,
99        on_bytes: &(dyn Fn(u64) + Send + Sync),
100    ) -> Result<Bytes> {
101        let result = self.client.get(location).await?;
102        collect_get_result_with_progress(result, location, on_bytes).await
103    }
104
105    async fn exists(&self, location: &Path) -> Result<bool> {
106        self.client.exists(location).await
107    }
108
109    async fn object_size(&self, location: &Path) -> Result<u64> {
110        self.client.size(location).await
111    }
112}