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::InvalidResponseDataError;
26use ferrin_spec::error::ProviderError;
27use ferrin_spec::files::DeleteFileResult;
28use ferrin_spec::files::FileMetadataResult;
29use ferrin_spec::files::FileReferenceOptions;
30use ferrin_spec::files::Files;
31use ferrin_spec::files::UploadData;
32use ferrin_spec::files::UploadFileOptions;
33use ferrin_spec::files::UploadFileResult;
34use futures_util::StreamExt;
35use futures_util::future::Either;
36use futures_util::future::select;
37use serde::Deserialize;
38use serde_json::json;
39use tokio_util::sync::CancellationToken;
40use url::Url;
41
42use crate::api_types::deserialize_count;
43use crate::config::CANONICAL_OPTIONS_KEY;
44use crate::config::SharedConfig;
45use crate::config::UPLOAD_PATH;
46use crate::convert_prompt::resolve_reference;
47use crate::error::failed_response_handler;
48use crate::options::part_options;
49use crate::options::read_options;
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.map_or(JsonValue::Null, JsonValue::from),
231        );
232        meta.insert("state".to_owned(), string(&file.state));
233        meta.insert("uri".to_owned(), string(&file.uri));
234        meta.insert("createTime".to_owned(), string(&file.create_time));
235        meta.insert("updateTime".to_owned(), string(&file.update_time));
236        meta.insert("expirationTime".to_owned(), string(&file.expiration_time));
237        meta.insert("sha256Hash".to_owned(), string(&file.sha256_hash));
238        let mapper = OutputMapper::new(self.config.clone(), Default::default());
239        UploadFileResult {
240            provider_reference: self.reference(file),
241            media_type: file.mime_type.as_deref().map(MediaType::new),
242            filename: file.display_name.clone(),
243            byte_size: file.size_bytes,
244            created_at: parse_time(file.create_time.as_deref()),
245            expires_at: parse_time(file.expiration_time.as_deref()),
246            provider_metadata: Some(mapper.metadata(meta)),
247            warnings: Vec::new(),
248        }
249    }
250
251    /// Fetches the resource `name` (`files/...`).
252    ///
253    /// # Errors
254    ///
255    /// Returns the API error of the request.
256    pub async fn fetch_file(
257        &self,
258        name: &str,
259        headers: &Headers,
260        cancellation: CancellationToken,
261    ) -> Result<GoogleFile, ProviderError> {
262        let handlers = ResponseHandlers::new(
263            json_response_handler::<FileResponse>(),
264            failed_response_handler(),
265        );
266        let response = get(
267            self.config.transport.as_ref(),
268            self.config.url(name),
269            self.config.headers(headers)?,
270            &handlers,
271            cancellation,
272        )
273        .await?;
274        Ok(response.value.into_file())
275    }
276
277    /// Uploads bytes through the resumable protocol and waits until the file
278    /// leaves the `PROCESSING` state.
279    ///
280    /// # Errors
281    ///
282    /// Returns [`ProviderError::InvalidResponseData`] when the upload session
283    /// returns no upload URL, [`ProviderError::ApiCall`] when processing
284    /// fails or times out, and [`ProviderError::Cancelled`] when the token
285    /// fires while polling.
286    #[tracing::instrument(skip_all, fields(media_type = %request.media_type, bytes = request.data.len()))]
287    pub async fn upload_bytes(&self, request: UploadRequest) -> Result<GoogleFile, ProviderError> {
288        let start_url = self.config.origin_url(UPLOAD_PATH);
289        let start_headers = self
290            .config
291            .headers(&request.headers)?
292            .with("x-goog-upload-protocol", "resumable")
293            .with("x-goog-upload-command", "start")
294            .with(
295                "x-goog-upload-header-content-length",
296                &request.data.len().to_string(),
297            )
298            .with("x-goog-upload-header-content-type", &request.media_type);
299        let mut file = JsonObject::new();
300        if let Some(name) = &request.display_name {
301            file.insert("display_name".to_owned(), JsonValue::from(name.as_str()));
302        }
303        let start_handlers =
304            ResponseHandlers::new(text_response_handler(), failed_response_handler());
305        let started = post_json(
306            self.config.transport.as_ref(),
307            start_url,
308            start_headers,
309            &json!({"file": file}),
310            &start_handlers,
311            request.cancellation.clone(),
312        )
313        .await?;
314        let upload_url = started
315            .response_headers
316            .get_str("x-goog-upload-url")
317            .and_then(|value| Url::parse(value).ok())
318            .ok_or_else(|| {
319                ProviderError::InvalidResponseData(Box::new(InvalidResponseDataError::new(
320                    "google did not return a resumable upload URL",
321                    JsonValue::Null,
322                )))
323            })?;
324        let validated = validate_url(&upload_url, &self.config.url_policy)
325            .await
326            .map_err(|error| {
327                InvalidResponseDataError::new(
328                    format!("google returned an unsafe upload URL: {error}"),
329                    JsonValue::Null,
330                )
331            })?;
332        let mut finalize_headers = if upload_url.origin() == self.config.base_url.origin()
333            || self.config.url_policy.is_credentialed(&upload_url)
334        {
335            self.config.unauthenticated_headers(&request.headers)
336        } else {
337            Headers::new().with_user_agent_suffix([crate::config::USER_AGENT])
338        };
339        finalize_headers.remove(crate::config::API_KEY_HEADER);
340        let finalize_headers = finalize_headers
341            .with("x-goog-upload-offset", "0")
342            .with("x-goog-upload-command", "upload, finalize");
343        let handlers = ResponseHandlers::new(
344            json_response_handler::<FileResponse>()
345                .with_max_bytes(self.config.url_policy.max_body_bytes),
346            failed_response_handler().with_max_bytes(self.config.url_policy.max_body_bytes),
347        );
348        let finalize = HttpRequest::post(upload_url.clone())
349            .with_headers(finalize_headers)
350            .with_body(RequestBody::Bytes {
351                content_type: request.media_type.clone(),
352                data: request.data,
353            })
354            .with_cancellation(request.cancellation.clone())
355            .with_pinned_addresses(validated.addresses);
356        let uploaded = send(self.config.transport.as_ref(), finalize, None, &handlers).await?;
357        let mut file = uploaded.value.into_file();
358        let started_at = Instant::now();
359        while file.state.as_deref() == Some("PROCESSING") {
360            if started_at.elapsed() > request.poll_timeout {
361                return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
362                    format!(
363                        "file processing timed out after {}ms",
364                        request.poll_timeout.as_millis()
365                    ),
366                    upload_url,
367                ))));
368            }
369            let sleep = Box::pin(tokio::time::sleep(request.poll_interval));
370            let cancelled = Box::pin(request.cancellation.cancelled());
371            if let Either::Right(_) = select(sleep, cancelled).await {
372                return Err(ProviderError::Cancelled);
373            }
374            file = self
375                .fetch_file(&file.name, &request.headers, request.cancellation.clone())
376                .await?;
377        }
378        if file.state.as_deref() == Some("FAILED") {
379            return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
380                format!("file processing failed for {}", file.name),
381                upload_url,
382            ))));
383        }
384        Ok(file)
385    }
386}
387
388impl Files for GoogleFiles {
389    fn provider(&self) -> &ProviderId {
390        &self.provider
391    }
392
393    async fn upload_file(
394        &self,
395        options: UploadFileOptions,
396    ) -> Result<UploadFileResult, ProviderError> {
397        let google: GoogleFilesOptions =
398            read_options(part_options(&self.config, Some(&options.provider_options)))
399                .unwrap_or_default();
400        let data = collect(options.data).await?;
401        let mut request = UploadRequest::new(data, options.media_type.as_str());
402        request.display_name = google.display_name.or(options.filename);
403        if let Some(interval) = google.poll_interval_ms {
404            request.poll_interval = Duration::from_millis(interval);
405        }
406        if let Some(timeout) = google.poll_timeout_ms {
407            request.poll_timeout = Duration::from_millis(timeout);
408        }
409        request.headers = options.headers;
410        request.cancellation = options.cancellation;
411        let file = self.upload_bytes(request).await?;
412        Ok(self.to_result(&file))
413    }
414
415    fn supports_get_file_metadata(&self) -> bool {
416        true
417    }
418
419    async fn get_file_metadata(
420        &self,
421        options: FileReferenceOptions,
422    ) -> Result<FileMetadataResult, ProviderError> {
423        let name = file_name(resolve_reference(&self.config, &options.file)?);
424        let file = self
425            .fetch_file(&name, &options.headers, options.cancellation)
426            .await?;
427        Ok(self.to_result(&file))
428    }
429
430    fn supports_delete_file(&self) -> bool {
431        true
432    }
433
434    async fn delete_file(
435        &self,
436        options: FileReferenceOptions,
437    ) -> Result<DeleteFileResult, ProviderError> {
438        let name = file_name(resolve_reference(&self.config, &options.file)?);
439        let handlers = ResponseHandlers::new(text_response_handler(), failed_response_handler());
440        delete(
441            self.config.transport.as_ref(),
442            self.config.url(&name),
443            self.config.headers(&options.headers)?,
444            &handlers,
445            options.cancellation,
446        )
447        .await?;
448        Ok(DeleteFileResult {
449            provider_reference: options.file,
450            deleted: true,
451            provider_metadata: None,
452            warnings: Vec::new(),
453        })
454    }
455}