Skip to main content

headless_lms_utils/file_store/
google_cloud_file_store.rs

1use std::path::{Path, PathBuf};
2
3use crate::prelude::*;
4use async_trait::async_trait;
5use bytes::Bytes;
6use futures::StreamExt;
7use google_cloud_auth::credentials::{
8    Builder as CredentialsBuilder, Credentials,
9    service_account::Builder as ServiceAccountCredentialsBuilder,
10};
11use google_cloud_auth::signer::Signer;
12use google_cloud_storage::builder::storage::SignedUrlBuilder;
13use google_cloud_storage::client::{Storage, StorageControl};
14use serde_json::Value;
15use std::{env, io};
16
17use super::{
18    FileStore, GenericPayload, attachment_content_disposition, generate_cache_folder_dir,
19    path_to_str,
20};
21
22pub struct GoogleCloudFileStore {
23    bucket_name: String,
24    storage_client: Storage,
25    control_client: StorageControl,
26    signer: Signer,
27    pub cache_files_path: PathBuf,
28}
29
30impl GoogleCloudFileStore {
31    /// Creates a Google Cloud file store using ADC-backed SDK clients.
32    #[instrument]
33    pub async fn new(bucket_name: String) -> UtilResult<Self> {
34        let (credentials, signer) = Self::build_credentials_and_signer()?;
35        let storage_client = match credentials.clone() {
36            Some(credentials) => Storage::builder()
37                .with_credentials(credentials)
38                .build()
39                .await
40                .map_err(Self::map_init_error)?,
41            None => Storage::builder()
42                .build()
43                .await
44                .map_err(Self::map_init_error)?,
45        };
46        let control_client = match credentials {
47            Some(credentials) => StorageControl::builder()
48                .with_credentials(credentials)
49                .build()
50                .await
51                .map_err(Self::map_init_error)?,
52            None => StorageControl::builder()
53                .build()
54                .await
55                .map_err(Self::map_init_error)?,
56        };
57        let cache_files_path = generate_cache_folder_dir()?;
58
59        Ok(Self {
60            bucket_name,
61            storage_client,
62            control_client,
63            signer,
64            cache_files_path,
65        })
66    }
67
68    /// Converts a raw bucket id into the GCS API resource name format.
69    fn bucket_resource_name(&self) -> String {
70        format!("projects/_/buckets/{}", self.bucket_name)
71    }
72
73    /// Builds Google Cloud credentials while preserving the legacy `cloud-storage`
74    /// crate's support for JSON credentials supplied directly in an environment variable.
75    fn build_credentials_and_signer() -> UtilResult<(Option<Credentials>, Signer)> {
76        if let Some(credentials_json) = Self::credentials_json_from_env()? {
77            let credentials = ServiceAccountCredentialsBuilder::new(credentials_json.clone())
78                .build()
79                .map_err(Self::map_init_error)?;
80            let signer = ServiceAccountCredentialsBuilder::new(credentials_json)
81                .build_signer()
82                .map_err(Self::map_init_error)?;
83            return Ok((Some(credentials), signer));
84        }
85
86        let signer = CredentialsBuilder::default()
87            .build_signer()
88            .map_err(Self::map_init_error)?;
89        Ok((None, signer))
90    }
91
92    fn credentials_json_from_env() -> UtilResult<Option<Value>> {
93        let raw_json = env::var("GOOGLE_APPLICATION_CREDENTIALS_JSON")
94            .or_else(|_| env::var("SERVICE_ACCOUNT_JSON"))
95            .ok();
96        raw_json
97            .map(|value| serde_json::from_str(&value).map_err(Self::map_init_error))
98            .transpose()
99    }
100
101    /// Maps SDK initialization errors into the cloud storage util error.
102    fn map_init_error<E: std::error::Error + Send + Sync + 'static>(error: E) -> UtilError {
103        UtilError::new(
104            UtilErrorType::CloudStorage,
105            error.to_string(),
106            Some(error.into()),
107        )
108    }
109}
110
111impl GoogleCloudFileStore {
112    /// `content_disposition` overrides the header the signed url's response carries.
113    async fn signed_get_url(
114        &self,
115        path: &Path,
116        content_disposition: Option<String>,
117    ) -> UtilResult<String> {
118        let object_name = path_to_str(path)?;
119        // Keep old behavior: verify object exists before returning a URL.
120        self.control_client
121            .get_object()
122            .set_bucket(self.bucket_resource_name())
123            .set_object(object_name)
124            .send()
125            .await?;
126        let mut builder = SignedUrlBuilder::for_object(self.bucket_resource_name(), object_name)
127            .with_method(google_cloud_storage::http::Method::GET)
128            .with_expiration(std::time::Duration::from_secs(300));
129        if let Some(content_disposition) = content_disposition {
130            builder = builder.with_query_param("response-content-disposition", content_disposition);
131        }
132        let url = builder
133            .sign_with(&self.signer)
134            .await
135            .map_err(|e| util_err!(CloudStorage, e.to_string(), e))?;
136        Ok(url)
137    }
138}
139
140#[async_trait(?Send)]
141impl FileStore for GoogleCloudFileStore {
142    async fn upload(&self, path: &Path, file: Vec<u8>, mime_type: &str) -> UtilResult<()> {
143        self.storage_client
144            .write_object(
145                self.bucket_resource_name(),
146                path_to_str(path)?,
147                Bytes::from(file),
148            )
149            .set_content_type(mime_type)
150            .send_buffered()
151            .await?;
152        Ok(())
153    }
154
155    async fn download(&self, path: &Path) -> UtilResult<Vec<u8>> {
156        let mut reader = self
157            .storage_client
158            .read_object(self.bucket_resource_name(), path_to_str(path)?)
159            .send()
160            .await?;
161        let mut out = Vec::new();
162        while let Some(chunk) = reader.next().await.transpose()? {
163            out.extend_from_slice(&chunk);
164        }
165        Ok(out)
166    }
167
168    async fn get_direct_download_url(&self, path: &Path) -> UtilResult<String> {
169        self.signed_get_url(path, None).await
170    }
171
172    async fn get_direct_attachment_download_url(
173        &self,
174        path: &Path,
175        file_name: &str,
176    ) -> UtilResult<String> {
177        self.signed_get_url(path, Some(attachment_content_disposition(file_name)))
178            .await
179    }
180
181    async fn delete(&self, path: &Path) -> UtilResult<()> {
182        self.control_client
183            .delete_object()
184            .set_bucket(self.bucket_resource_name())
185            .set_object(path_to_str(path)?)
186            .send()
187            .await?;
188        Ok(())
189    }
190
191    async fn upload_stream(
192        &self,
193        path: &Path,
194        mut contents: GenericPayload,
195        mime_type: &str,
196    ) -> UtilResult<()> {
197        let mut collected = Vec::new();
198        while let Some(chunk) = contents.next().await {
199            let chunk = chunk
200                .map_err(|e| UtilError::new(UtilErrorType::CloudStorage, e.to_string(), None))?;
201            collected.extend_from_slice(&chunk);
202        }
203        self.storage_client
204            .write_object(
205                self.bucket_resource_name(),
206                path_to_str(path)?,
207                Bytes::from(collected),
208            )
209            .set_content_type(mime_type)
210            .send_buffered()
211            .await?;
212        Ok(())
213    }
214
215    async fn download_stream(
216        &self,
217        path: &Path,
218    ) -> UtilResult<Box<dyn futures::Stream<Item = std::io::Result<bytes::Bytes>>>> {
219        let mut reader = self
220            .storage_client
221            .read_object(self.bucket_resource_name(), path_to_str(path)?)
222            .send()
223            .await?;
224        let stream = async_stream::stream! {
225            while let Some(item) = reader.next().await {
226                yield item.map_err(io::Error::other);
227            }
228        };
229        Ok(Box::new(stream))
230    }
231
232    fn get_cache_files_folder_path(&self) -> UtilResult<&Path> {
233        Ok(&self.cache_files_path)
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use std::sync::Mutex;
241
242    static ENV_MUTEX: Mutex<()> = Mutex::new(());
243
244    struct EnvRestore {
245        google_application_credentials_json: Option<String>,
246        service_account_json: Option<String>,
247    }
248
249    impl EnvRestore {
250        fn capture() -> Self {
251            Self {
252                google_application_credentials_json: env::var(
253                    "GOOGLE_APPLICATION_CREDENTIALS_JSON",
254                )
255                .ok(),
256                service_account_json: env::var("SERVICE_ACCOUNT_JSON").ok(),
257            }
258        }
259    }
260
261    impl Drop for EnvRestore {
262        fn drop(&mut self) {
263            unsafe {
264                match &self.google_application_credentials_json {
265                    Some(value) => env::set_var("GOOGLE_APPLICATION_CREDENTIALS_JSON", value),
266                    None => env::remove_var("GOOGLE_APPLICATION_CREDENTIALS_JSON"),
267                }
268                match &self.service_account_json {
269                    Some(value) => env::set_var("SERVICE_ACCOUNT_JSON", value),
270                    None => env::remove_var("SERVICE_ACCOUNT_JSON"),
271                }
272            }
273        }
274    }
275
276    #[tokio::test]
277    async fn uses_google_application_credentials_json_contents() {
278        let signer = {
279            let _guard = ENV_MUTEX.lock().expect("environment mutex was poisoned");
280            let _restore = EnvRestore::capture();
281            unsafe {
282                env::set_var(
283                    "GOOGLE_APPLICATION_CREDENTIALS_JSON",
284                    r#"{
285                        "type": "service_account",
286                        "client_email": "json-creds@test-project.iam.example.com",
287                        "private_key_id": "test-private-key-id",
288                        "private_key": "not-a-real-private-key",
289                        "project_id": "test-project-id"
290                    }"#,
291                );
292                env::remove_var("SERVICE_ACCOUNT_JSON");
293            }
294
295            GoogleCloudFileStore::build_credentials_and_signer()
296                .expect("raw JSON credentials should build a signer")
297                .1
298        };
299
300        let client_email = signer
301            .client_email()
302            .await
303            .expect("service account signer should expose client email without IAM API");
304        assert_eq!(client_email, "json-creds@test-project.iam.example.com");
305    }
306}