Skip to main content

ferrin_google/
files.rs

1//! Files API: resumable upload with processing poll, metadata and delete.
2
3use std::time::Duration;
4use std::time::Instant;
5
6use bytes::Bytes;
7use bytes::BytesMut;
8use ferrin_provider_util::http::HttpRequest;
9use ferrin_provider_util::http::RequestBody;
10use ferrin_provider_util::http::ResponseHandlers;
11use ferrin_provider_util::http::delete;
12use ferrin_provider_util::http::get;
13use ferrin_provider_util::http::json_response_handler;
14use ferrin_provider_util::http::post_json;
15use ferrin_provider_util::http::send;
16use ferrin_provider_util::http::text_response_handler;
17use ferrin_provider_util::secure_url::validate_url;
18use ferrin_spec::Headers;
19use ferrin_spec::JsonObject;
20use ferrin_spec::JsonValue;
21use ferrin_spec::MediaType;
22use ferrin_spec::ProviderId;
23use ferrin_spec::ProviderReference;
24use ferrin_spec::error::ApiCallError;
25use ferrin_spec::error::InvalidArgumentError;
26use ferrin_spec::error::InvalidResponseDataError;
27use ferrin_spec::error::ProviderError;
28use ferrin_spec::files::DeleteFileResult;
29use ferrin_spec::files::FileMetadataResult;
30use ferrin_spec::files::FileReferenceOptions;
31use ferrin_spec::files::Files;
32use ferrin_spec::files::UploadData;
33use ferrin_spec::files::UploadFileOptions;
34use ferrin_spec::files::UploadFileResult;
35use futures_util::StreamExt;
36use futures_util::future::Either;
37use futures_util::future::select;
38use serde::Deserialize;
39use serde_json::json;
40use tokio_util::sync::CancellationToken;
41use url::Url;
42
43use crate::api_types::deserialize_count;
44use crate::config::CANONICAL_OPTIONS_KEY;
45use crate::config::SharedConfig;
46use crate::config::UPLOAD_PATH;
47use crate::convert_prompt::resolve_reference;
48use crate::error::failed_response_handler;
49use crate::options::parse_merged;
50use crate::output::OutputMapper;
51
52/// Default interval between processing-state polls.
53pub const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
54
55/// Default time after which a still-processing upload fails.
56pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 300_000;
57
58/// Provider options of uploads (`provider_options["google"]`).
59#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
60#[serde(rename_all = "camelCase")]
61pub struct GoogleFilesOptions {
62    /// Display name stored with the file (defaults to `filename`).
63    #[serde(default)]
64    pub display_name: Option<String>,
65    /// Poll interval in milliseconds while the file is `PROCESSING`.
66    #[serde(default)]
67    pub poll_interval_ms: Option<u64>,
68    /// Poll timeout in milliseconds.
69    #[serde(default)]
70    pub poll_timeout_ms: Option<u64>,
71}
72
73/// A file resource.
74#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub struct GoogleFile {
77    /// Resource name (`files/abc-123`).
78    pub name: String,
79    /// Display name.
80    #[serde(default)]
81    pub display_name: Option<String>,
82    /// Media type.
83    #[serde(default)]
84    pub mime_type: Option<String>,
85    /// Size in bytes.
86    #[serde(default, deserialize_with = "deserialize_count")]
87    pub size_bytes: Option<u64>,
88    /// Creation time (RFC 3339).
89    #[serde(default)]
90    pub create_time: Option<String>,
91    /// Update time (RFC 3339).
92    #[serde(default)]
93    pub update_time: Option<String>,
94    /// Expiration time (RFC 3339).
95    #[serde(default)]
96    pub expiration_time: Option<String>,
97    /// SHA-256 hash (base64).
98    #[serde(default)]
99    pub sha256_hash: Option<String>,
100    /// URI usable as `fileData.fileUri`.
101    #[serde(default)]
102    pub uri: Option<String>,
103    /// `PROCESSING`, `ACTIVE` or `FAILED`.
104    #[serde(default)]
105    pub state: Option<String>,
106}
107
108/// Upload responses wrap the file in `{file}`; `files.get` returns it bare.
109#[derive(Debug, Deserialize)]
110#[serde(untagged)]
111enum FileResponse {
112    Envelope { file: GoogleFile },
113    Bare(GoogleFile),
114}
115
116impl FileResponse {
117    fn into_file(self) -> GoogleFile {
118        match self {
119            Self::Envelope { file } | Self::Bare(file) => file,
120        }
121    }
122}
123
124/// A byte upload through the resumable protocol.
125#[derive(Debug, Clone)]
126pub struct UploadRequest {
127    /// Content.
128    pub data: Bytes,
129    /// Media type sent as `X-Goog-Upload-Header-Content-Type`.
130    pub media_type: String,
131    /// Display name.
132    pub display_name: Option<String>,
133    /// Poll interval while `PROCESSING`.
134    pub poll_interval: Duration,
135    /// Poll timeout.
136    pub poll_timeout: Duration,
137    /// Additional request headers.
138    pub headers: Headers,
139    /// Cancellation token.
140    pub cancellation: CancellationToken,
141}
142
143impl UploadRequest {
144    /// Creates a request with the default poll settings.
145    #[must_use]
146    pub fn new(data: Bytes, media_type: impl Into<String>) -> Self {
147        Self {
148            data,
149            media_type: media_type.into(),
150            display_name: None,
151            poll_interval: Duration::from_millis(DEFAULT_POLL_INTERVAL_MS),
152            poll_timeout: Duration::from_millis(DEFAULT_POLL_TIMEOUT_MS),
153            headers: Headers::new(),
154            cancellation: CancellationToken::new(),
155        }
156    }
157}
158
159async fn collect(data: UploadData) -> Result<Bytes, ProviderError> {
160    match data {
161        UploadData::Bytes(bytes) => Ok(bytes),
162        UploadData::Text(text) => Ok(Bytes::from(text)),
163        UploadData::Stream(stream) => {
164            let mut buffer = BytesMut::new();
165            let mut stream = stream;
166            while let Some(chunk) = stream.next().await {
167                buffer.extend_from_slice(&chunk?);
168            }
169            Ok(buffer.freeze())
170        }
171        #[allow(unreachable_patterns, reason = "UploadData is non-exhaustive")]
172        _ => Err(ProviderError::unsupported("upload data type")),
173    }
174}
175
176fn parse_time(value: Option<&str>) -> Option<chrono::DateTime<chrono::Utc>> {
177    value
178        .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
179        .map(|time| time.with_timezone(&chrono::Utc))
180}
181
182/// Normalizes a provider reference (a `files/...` name, a bare id or the
183/// full `uri`) to the resource name.
184fn file_name(reference: &str) -> String {
185    if let Some((_, rest)) = reference.rsplit_once("/files/") {
186        return format!("files/{rest}");
187    }
188    if reference.starts_with("files/") {
189        return reference.to_owned();
190    }
191    format!("files/{reference}")
192}
193
194/// Files service backed by the Files API.
195#[derive(Debug, Clone)]
196pub struct GoogleFiles {
197    config: SharedConfig,
198    provider: ProviderId,
199}
200
201impl GoogleFiles {
202    /// Creates the service.
203    #[must_use]
204    pub fn new(config: SharedConfig) -> Self {
205        Self {
206            provider: ProviderId::new(config.name.clone()),
207            config,
208        }
209    }
210
211    fn reference(&self, file: &GoogleFile) -> ProviderReference {
212        let value = file.uri.clone().unwrap_or_else(|| file.name.clone());
213        let mut reference = ProviderReference::new();
214        reference.insert(CANONICAL_OPTIONS_KEY.to_owned(), value.clone());
215        reference.insert(self.config.name.clone(), value);
216        reference
217    }
218
219    /// Maps a file resource to the specification result.
220    #[must_use]
221    pub fn to_result(&self, file: &GoogleFile) -> UploadFileResult {
222        let mut meta = JsonObject::new();
223        let string =
224            |value: &Option<String>| value.clone().map_or(JsonValue::Null, JsonValue::from);
225        meta.insert("name".to_owned(), JsonValue::from(file.name.clone()));
226        meta.insert("displayName".to_owned(), string(&file.display_name));
227        meta.insert("mimeType".to_owned(), string(&file.mime_type));
228        meta.insert(
229            "sizeBytes".to_owned(),
230            file.size_bytes
231                .map_or(JsonValue::Null, |size| JsonValue::from(size.to_string())),
232        );
233        meta.insert("state".to_owned(), string(&file.state));
234        meta.insert("uri".to_owned(), string(&file.uri));
235        for (key, value) in [
236            ("createTime", &file.create_time),
237            ("updateTime", &file.update_time),
238            ("expirationTime", &file.expiration_time),
239            ("sha256Hash", &file.sha256_hash),
240        ] {
241            if let Some(value) = value {
242                meta.insert(key.to_owned(), json!(value));
243            }
244        }
245        let mapper = OutputMapper::new(self.config.clone(), Default::default());
246        UploadFileResult {
247            provider_reference: self.reference(file),
248            media_type: file.mime_type.as_deref().map(MediaType::new),
249            filename: file.display_name.clone(),
250            byte_size: file.size_bytes,
251            created_at: parse_time(file.create_time.as_deref()),
252            expires_at: parse_time(file.expiration_time.as_deref()),
253            provider_metadata: Some(mapper.metadata(meta)),
254            warnings: Vec::new(),
255        }
256    }
257
258    fn resource_url(&self, name: &str) -> Result<Url, ProviderError> {
259        let mut url = self.config.base_url.clone();
260        let mut segments = url.path_segments_mut().map_err(|()| {
261            InvalidArgumentError::new("base_url", "file base URL cannot contain path segments")
262        })?;
263        segments.pop_if_empty();
264        let id = if let Some(id) = name.strip_prefix("files/").filter(|id| !id.contains('/')) {
265            segments.push("files");
266            id
267        } else {
268            name
269        };
270        // URL parsers normalize dot segments even when encoded once.
271        segments.push(match id {
272            "." => "%2E",
273            ".." => "%2E%2E",
274            id => id,
275        });
276        drop(segments);
277        Ok(url)
278    }
279
280    /// Fetches the resource `name` (`files/...`).
281    ///
282    /// # Errors
283    ///
284    /// Returns the API error of the request.
285    pub async fn fetch_file(
286        &self,
287        name: &str,
288        headers: &Headers,
289        cancellation: CancellationToken,
290    ) -> Result<GoogleFile, ProviderError> {
291        let handlers = ResponseHandlers::new(
292            json_response_handler::<FileResponse>(),
293            failed_response_handler(),
294        );
295        let response = get(
296            self.config.transport.as_ref(),
297            self.resource_url(name)?,
298            self.config.headers(headers)?,
299            &handlers,
300            cancellation,
301        )
302        .await?;
303        Ok(response.value.into_file())
304    }
305
306    /// Uploads bytes through the resumable protocol and waits until the file
307    /// leaves the `PROCESSING` state.
308    ///
309    /// # Errors
310    ///
311    /// Returns [`ProviderError::InvalidResponseData`] when the upload session
312    /// returns no upload URL, [`ProviderError::ApiCall`] when processing
313    /// fails or times out, and [`ProviderError::Cancelled`] when the token
314    /// fires while polling.
315    #[tracing::instrument(skip_all, fields(media_type = %request.media_type, bytes = request.data.len()))]
316    pub async fn upload_bytes(&self, request: UploadRequest) -> Result<GoogleFile, ProviderError> {
317        let start_url = self.config.origin_url(UPLOAD_PATH);
318        let start_headers = self
319            .config
320            .headers(&request.headers)?
321            .with("x-goog-upload-protocol", "resumable")
322            .with("x-goog-upload-command", "start")
323            .with(
324                "x-goog-upload-header-content-length",
325                &request.data.len().to_string(),
326            )
327            .with("x-goog-upload-header-content-type", &request.media_type);
328        let mut file = JsonObject::new();
329        if let Some(name) = &request.display_name {
330            file.insert("display_name".to_owned(), JsonValue::from(name.as_str()));
331        }
332        let start_handlers =
333            ResponseHandlers::new(text_response_handler(), failed_response_handler());
334        let started = post_json(
335            self.config.transport.as_ref(),
336            start_url,
337            start_headers,
338            &json!({"file": file}),
339            &start_handlers,
340            request.cancellation.clone(),
341        )
342        .await?;
343        let upload_url = started
344            .response_headers
345            .get_str("x-goog-upload-url")
346            .and_then(|value| Url::parse(value).ok())
347            .ok_or_else(|| {
348                ProviderError::InvalidResponseData(Box::new(InvalidResponseDataError::new(
349                    "google did not return a resumable upload URL",
350                    JsonValue::Null,
351                )))
352            })?;
353        let validated = validate_url(&upload_url, &self.config.url_policy)
354            .await
355            .map_err(|error| {
356                InvalidResponseDataError::new(
357                    format!("google returned an unsafe upload URL: {error}"),
358                    JsonValue::Null,
359                )
360            })?;
361        let mut finalize_headers = if upload_url.origin() == self.config.base_url.origin()
362            || self.config.url_policy.is_credentialed(&upload_url)
363        {
364            self.config.unauthenticated_headers(&request.headers)
365        } else {
366            Headers::new().with_user_agent_suffix([crate::config::USER_AGENT])
367        };
368        finalize_headers.remove(crate::config::API_KEY_HEADER);
369        let finalize_headers = finalize_headers
370            .with("x-goog-upload-offset", "0")
371            .with("x-goog-upload-command", "upload, finalize");
372        let handlers = ResponseHandlers::new(
373            json_response_handler::<FileResponse>()
374                .with_max_bytes(self.config.url_policy.max_body_bytes),
375            failed_response_handler().with_max_bytes(self.config.url_policy.max_body_bytes),
376        );
377        let finalize = HttpRequest::post(upload_url.clone())
378            .with_headers(finalize_headers)
379            .with_body(RequestBody::Bytes {
380                content_type: request.media_type.clone(),
381                data: request.data,
382            })
383            .with_cancellation(request.cancellation.clone())
384            .with_pinned_addresses(validated.addresses);
385        let uploaded = send(self.config.transport.as_ref(), finalize, None, &handlers).await?;
386        let mut file = uploaded.value.into_file();
387        let started_at = Instant::now();
388        while file.state.as_deref() == Some("PROCESSING") {
389            if started_at.elapsed() > request.poll_timeout {
390                return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
391                    format!(
392                        "file processing timed out after {}ms",
393                        request.poll_timeout.as_millis()
394                    ),
395                    upload_url,
396                ))));
397            }
398            let sleep = Box::pin(tokio::time::sleep(request.poll_interval));
399            let cancelled = Box::pin(request.cancellation.cancelled());
400            if let Either::Right(_) = select(sleep, cancelled).await {
401                return Err(ProviderError::Cancelled);
402            }
403            file = self
404                .fetch_file(&file.name, &request.headers, request.cancellation.clone())
405                .await?;
406        }
407        if file.state.as_deref() == Some("FAILED") {
408            return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
409                format!("file processing failed for {}", file.name),
410                upload_url,
411            ))));
412        }
413        Ok(file)
414    }
415}
416
417impl Files for GoogleFiles {
418    fn provider(&self) -> &ProviderId {
419        &self.provider
420    }
421
422    async fn upload_file(
423        &self,
424        options: UploadFileOptions,
425    ) -> Result<UploadFileResult, ProviderError> {
426        let google = parse_merged::<GoogleFilesOptions>(
427            &self.config,
428            &options.provider_options,
429            |mut canonical, custom| {
430                if custom.display_name.is_some() {
431                    canonical.display_name = custom.display_name;
432                }
433                if custom.poll_interval_ms.is_some() {
434                    canonical.poll_interval_ms = custom.poll_interval_ms;
435                }
436                if custom.poll_timeout_ms.is_some() {
437                    canonical.poll_timeout_ms = custom.poll_timeout_ms;
438                }
439                canonical
440            },
441        )?;
442        if google.poll_interval_ms == Some(0) || google.poll_timeout_ms == Some(0) {
443            return Err(InvalidArgumentError::new(
444                "provider_options",
445                "file polling intervals and timeouts must be positive",
446            )
447            .into());
448        }
449        let ignored_filename = options.filename.is_some();
450        let data = match select(
451            Box::pin(collect(options.data)),
452            Box::pin(options.cancellation.cancelled()),
453        )
454        .await
455        {
456            Either::Left((data, _)) => data?,
457            Either::Right(_) => return Err(ProviderError::Cancelled),
458        };
459        let mut request = UploadRequest::new(data, options.media_type.as_str());
460        request.display_name = google.display_name;
461        if let Some(interval) = google.poll_interval_ms {
462            request.poll_interval = Duration::from_millis(interval);
463        }
464        if let Some(timeout) = google.poll_timeout_ms {
465            request.poll_timeout = Duration::from_millis(timeout);
466        }
467        request.headers = options.headers;
468        request.cancellation = options.cancellation;
469        let file = self.upload_bytes(request).await?;
470        let mut result = self.to_result(&file);
471        if ignored_filename {
472            result
473                .warnings
474                .push(ferrin_spec::Warning::unsupported("filename"));
475        }
476        if result.media_type.is_none() {
477            result.media_type = Some(options.media_type);
478        }
479        Ok(result)
480    }
481
482    fn supports_get_file_metadata(&self) -> bool {
483        true
484    }
485
486    async fn get_file_metadata(
487        &self,
488        options: FileReferenceOptions,
489    ) -> Result<FileMetadataResult, ProviderError> {
490        let name = file_name(resolve_reference(&self.config, &options.file)?);
491        let file = self
492            .fetch_file(&name, &options.headers, options.cancellation)
493            .await?;
494        Ok(self.to_result(&file))
495    }
496
497    fn supports_delete_file(&self) -> bool {
498        true
499    }
500
501    async fn delete_file(
502        &self,
503        options: FileReferenceOptions,
504    ) -> Result<DeleteFileResult, ProviderError> {
505        let name = file_name(resolve_reference(&self.config, &options.file)?);
506        let handlers = ResponseHandlers::new(text_response_handler(), failed_response_handler());
507        delete(
508            self.config.transport.as_ref(),
509            self.resource_url(&name)?,
510            self.config.headers(&options.headers)?,
511            &handlers,
512            options.cancellation,
513        )
514        .await?;
515        Ok(DeleteFileResult {
516            provider_reference: options.file,
517            deleted: true,
518            provider_metadata: None,
519            warnings: Vec::new(),
520        })
521    }
522}