1use std::time::Duration;
4use std::time::Instant;
5
6use bytes::Bytes;
7use bytes::BytesMut;
8use ferrin_provider_util::http::ResponseHandlers;
9use ferrin_provider_util::http::delete;
10use ferrin_provider_util::http::get;
11use ferrin_provider_util::http::json_response_handler;
12use ferrin_provider_util::http::post_bytes;
13use ferrin_provider_util::http::post_json;
14use ferrin_provider_util::http::text_response_handler;
15use ferrin_spec::Headers;
16use ferrin_spec::JsonObject;
17use ferrin_spec::JsonValue;
18use ferrin_spec::MediaType;
19use ferrin_spec::ProviderId;
20use ferrin_spec::ProviderReference;
21use ferrin_spec::error::ApiCallError;
22use ferrin_spec::error::InvalidResponseDataError;
23use ferrin_spec::error::ProviderError;
24use ferrin_spec::files::DeleteFileResult;
25use ferrin_spec::files::FileMetadataResult;
26use ferrin_spec::files::FileReferenceOptions;
27use ferrin_spec::files::Files;
28use ferrin_spec::files::UploadData;
29use ferrin_spec::files::UploadFileOptions;
30use ferrin_spec::files::UploadFileResult;
31use futures_util::StreamExt;
32use futures_util::future::Either;
33use futures_util::future::select;
34use serde::Deserialize;
35use serde_json::json;
36use tokio_util::sync::CancellationToken;
37use url::Url;
38
39use crate::api_types::deserialize_count;
40use crate::config::CANONICAL_OPTIONS_KEY;
41use crate::config::SharedConfig;
42use crate::config::UPLOAD_PATH;
43use crate::convert_prompt::resolve_reference;
44use crate::error::failed_response_handler;
45use crate::options::part_options;
46use crate::options::read_options;
47use crate::output::OutputMapper;
48
49pub const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
51
52pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 300_000;
54
55#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
57#[serde(rename_all = "camelCase")]
58pub struct GoogleFilesOptions {
59 #[serde(default)]
61 pub display_name: Option<String>,
62 #[serde(default)]
64 pub poll_interval_ms: Option<u64>,
65 #[serde(default)]
67 pub poll_timeout_ms: Option<u64>,
68}
69
70#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
72#[serde(rename_all = "camelCase")]
73pub struct GoogleFile {
74 pub name: String,
76 #[serde(default)]
78 pub display_name: Option<String>,
79 #[serde(default)]
81 pub mime_type: Option<String>,
82 #[serde(default, deserialize_with = "deserialize_count")]
84 pub size_bytes: Option<u64>,
85 #[serde(default)]
87 pub create_time: Option<String>,
88 #[serde(default)]
90 pub update_time: Option<String>,
91 #[serde(default)]
93 pub expiration_time: Option<String>,
94 #[serde(default)]
96 pub sha256_hash: Option<String>,
97 #[serde(default)]
99 pub uri: Option<String>,
100 #[serde(default)]
102 pub state: Option<String>,
103}
104
105#[derive(Debug, Deserialize)]
107#[serde(untagged)]
108enum FileResponse {
109 Envelope { file: GoogleFile },
110 Bare(GoogleFile),
111}
112
113impl FileResponse {
114 fn into_file(self) -> GoogleFile {
115 match self {
116 Self::Envelope { file } | Self::Bare(file) => file,
117 }
118 }
119}
120
121#[derive(Debug, Clone)]
123pub struct UploadRequest {
124 pub data: Bytes,
126 pub media_type: String,
128 pub display_name: Option<String>,
130 pub poll_interval: Duration,
132 pub poll_timeout: Duration,
134 pub headers: Headers,
136 pub cancellation: CancellationToken,
138}
139
140impl UploadRequest {
141 #[must_use]
143 pub fn new(data: Bytes, media_type: impl Into<String>) -> Self {
144 Self {
145 data,
146 media_type: media_type.into(),
147 display_name: None,
148 poll_interval: Duration::from_millis(DEFAULT_POLL_INTERVAL_MS),
149 poll_timeout: Duration::from_millis(DEFAULT_POLL_TIMEOUT_MS),
150 headers: Headers::new(),
151 cancellation: CancellationToken::new(),
152 }
153 }
154}
155
156async fn collect(data: UploadData) -> Result<Bytes, ProviderError> {
157 match data {
158 UploadData::Bytes(bytes) => Ok(bytes),
159 UploadData::Text(text) => Ok(Bytes::from(text)),
160 UploadData::Stream(stream) => {
161 let mut buffer = BytesMut::new();
162 let mut stream = stream;
163 while let Some(chunk) = stream.next().await {
164 buffer.extend_from_slice(&chunk?);
165 }
166 Ok(buffer.freeze())
167 }
168 #[allow(unreachable_patterns, reason = "UploadData is non-exhaustive")]
169 _ => Err(ProviderError::unsupported("upload data type")),
170 }
171}
172
173fn parse_time(value: Option<&str>) -> Option<chrono::DateTime<chrono::Utc>> {
174 value
175 .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
176 .map(|time| time.with_timezone(&chrono::Utc))
177}
178
179fn file_name(reference: &str) -> String {
182 if let Some((_, rest)) = reference.rsplit_once("/files/") {
183 return format!("files/{rest}");
184 }
185 if reference.starts_with("files/") {
186 return reference.to_owned();
187 }
188 format!("files/{reference}")
189}
190
191#[derive(Debug, Clone)]
193pub struct GoogleFiles {
194 config: SharedConfig,
195 provider: ProviderId,
196}
197
198impl GoogleFiles {
199 #[must_use]
201 pub fn new(config: SharedConfig) -> Self {
202 Self {
203 provider: ProviderId::new(config.name.clone()),
204 config,
205 }
206 }
207
208 fn reference(&self, file: &GoogleFile) -> ProviderReference {
209 let value = file.uri.clone().unwrap_or_else(|| file.name.clone());
210 let mut reference = ProviderReference::new();
211 reference.insert(CANONICAL_OPTIONS_KEY.to_owned(), value.clone());
212 reference.insert(self.config.name.clone(), value);
213 reference
214 }
215
216 #[must_use]
218 pub fn to_result(&self, file: &GoogleFile) -> UploadFileResult {
219 let mut meta = JsonObject::new();
220 let string =
221 |value: &Option<String>| value.clone().map_or(JsonValue::Null, JsonValue::from);
222 meta.insert("name".to_owned(), JsonValue::from(file.name.clone()));
223 meta.insert("displayName".to_owned(), string(&file.display_name));
224 meta.insert("mimeType".to_owned(), string(&file.mime_type));
225 meta.insert(
226 "sizeBytes".to_owned(),
227 file.size_bytes.map_or(JsonValue::Null, JsonValue::from),
228 );
229 meta.insert("state".to_owned(), string(&file.state));
230 meta.insert("uri".to_owned(), string(&file.uri));
231 meta.insert("createTime".to_owned(), string(&file.create_time));
232 meta.insert("updateTime".to_owned(), string(&file.update_time));
233 meta.insert("expirationTime".to_owned(), string(&file.expiration_time));
234 meta.insert("sha256Hash".to_owned(), string(&file.sha256_hash));
235 let mapper = OutputMapper::new(self.config.clone(), Default::default());
236 UploadFileResult {
237 provider_reference: self.reference(file),
238 media_type: file.mime_type.as_deref().map(MediaType::new),
239 filename: file.display_name.clone(),
240 byte_size: file.size_bytes,
241 created_at: parse_time(file.create_time.as_deref()),
242 expires_at: parse_time(file.expiration_time.as_deref()),
243 provider_metadata: Some(mapper.metadata(meta)),
244 warnings: Vec::new(),
245 }
246 }
247
248 pub async fn fetch_file(
254 &self,
255 name: &str,
256 headers: &Headers,
257 cancellation: CancellationToken,
258 ) -> Result<GoogleFile, ProviderError> {
259 let handlers = ResponseHandlers::new(
260 json_response_handler::<FileResponse>(),
261 failed_response_handler(),
262 );
263 let response = get(
264 self.config.transport.as_ref(),
265 self.config.url(name),
266 self.config.headers(headers)?,
267 &handlers,
268 cancellation,
269 )
270 .await?;
271 Ok(response.value.into_file())
272 }
273
274 #[tracing::instrument(skip_all, fields(media_type = %request.media_type, bytes = request.data.len()))]
284 pub async fn upload_bytes(&self, request: UploadRequest) -> Result<GoogleFile, ProviderError> {
285 let start_url = self.config.origin_url(UPLOAD_PATH);
286 let start_headers = self
287 .config
288 .headers(&request.headers)?
289 .with("x-goog-upload-protocol", "resumable")
290 .with("x-goog-upload-command", "start")
291 .with(
292 "x-goog-upload-header-content-length",
293 &request.data.len().to_string(),
294 )
295 .with("x-goog-upload-header-content-type", &request.media_type);
296 let mut file = JsonObject::new();
297 if let Some(name) = &request.display_name {
298 file.insert("display_name".to_owned(), JsonValue::from(name.as_str()));
299 }
300 let start_handlers =
301 ResponseHandlers::new(text_response_handler(), failed_response_handler());
302 let started = post_json(
303 self.config.transport.as_ref(),
304 start_url,
305 start_headers,
306 &json!({"file": file}),
307 &start_handlers,
308 request.cancellation.clone(),
309 )
310 .await?;
311 let upload_url = started
312 .response_headers
313 .get_str("x-goog-upload-url")
314 .and_then(|value| Url::parse(value).ok())
315 .ok_or_else(|| {
316 ProviderError::InvalidResponseData(Box::new(InvalidResponseDataError::new(
317 "google did not return a resumable upload URL",
318 JsonValue::Null,
319 )))
320 })?;
321 let finalize_headers = self
322 .config
323 .unauthenticated_headers(&request.headers)
324 .with("x-goog-upload-offset", "0")
325 .with("x-goog-upload-command", "upload, finalize");
326 let handlers = ResponseHandlers::new(
327 json_response_handler::<FileResponse>(),
328 failed_response_handler(),
329 );
330 let uploaded = post_bytes(
331 self.config.transport.as_ref(),
332 upload_url.clone(),
333 finalize_headers,
334 &request.media_type,
335 request.data,
336 &handlers,
337 request.cancellation.clone(),
338 )
339 .await?;
340 let mut file = uploaded.value.into_file();
341 let started_at = Instant::now();
342 while file.state.as_deref() == Some("PROCESSING") {
343 if started_at.elapsed() > request.poll_timeout {
344 return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
345 format!(
346 "file processing timed out after {}ms",
347 request.poll_timeout.as_millis()
348 ),
349 upload_url,
350 ))));
351 }
352 let sleep = Box::pin(tokio::time::sleep(request.poll_interval));
353 let cancelled = Box::pin(request.cancellation.cancelled());
354 if let Either::Right(_) = select(sleep, cancelled).await {
355 return Err(ProviderError::Cancelled);
356 }
357 file = self
358 .fetch_file(&file.name, &request.headers, request.cancellation.clone())
359 .await?;
360 }
361 if file.state.as_deref() == Some("FAILED") {
362 return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
363 format!("file processing failed for {}", file.name),
364 upload_url,
365 ))));
366 }
367 Ok(file)
368 }
369}
370
371impl Files for GoogleFiles {
372 fn provider(&self) -> &ProviderId {
373 &self.provider
374 }
375
376 async fn upload_file(
377 &self,
378 options: UploadFileOptions,
379 ) -> Result<UploadFileResult, ProviderError> {
380 let google: GoogleFilesOptions =
381 read_options(part_options(&self.config, Some(&options.provider_options)))
382 .unwrap_or_default();
383 let data = collect(options.data).await?;
384 let mut request = UploadRequest::new(data, options.media_type.as_str());
385 request.display_name = google.display_name.or(options.filename);
386 if let Some(interval) = google.poll_interval_ms {
387 request.poll_interval = Duration::from_millis(interval);
388 }
389 if let Some(timeout) = google.poll_timeout_ms {
390 request.poll_timeout = Duration::from_millis(timeout);
391 }
392 request.headers = options.headers;
393 request.cancellation = options.cancellation;
394 let file = self.upload_bytes(request).await?;
395 Ok(self.to_result(&file))
396 }
397
398 fn supports_get_file_metadata(&self) -> bool {
399 true
400 }
401
402 async fn get_file_metadata(
403 &self,
404 options: FileReferenceOptions,
405 ) -> Result<FileMetadataResult, ProviderError> {
406 let name = file_name(resolve_reference(&self.config, &options.file)?);
407 let file = self
408 .fetch_file(&name, &options.headers, options.cancellation)
409 .await?;
410 Ok(self.to_result(&file))
411 }
412
413 fn supports_delete_file(&self) -> bool {
414 true
415 }
416
417 async fn delete_file(
418 &self,
419 options: FileReferenceOptions,
420 ) -> Result<DeleteFileResult, ProviderError> {
421 let name = file_name(resolve_reference(&self.config, &options.file)?);
422 let handlers = ResponseHandlers::new(text_response_handler(), failed_response_handler());
423 delete(
424 self.config.transport.as_ref(),
425 self.config.url(&name),
426 self.config.headers(&options.headers)?,
427 &handlers,
428 options.cancellation,
429 )
430 .await?;
431 Ok(DeleteFileResult {
432 provider_reference: options.file,
433 deleted: true,
434 provider_metadata: None,
435 warnings: Vec::new(),
436 })
437 }
438}