headless_lms_utils/file_store/
google_cloud_file_store.rs1use 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 #[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 fn bucket_resource_name(&self) -> String {
70 format!("projects/_/buckets/{}", self.bucket_name)
71 }
72
73 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 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 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 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}