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::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
49/// Default interval between processing-state polls.
50pub const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
51
52/// Default time after which a still-processing upload fails.
53pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 300_000;
54
55/// Provider options of uploads (`provider_options["google"]`).
56#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
57#[serde(rename_all = "camelCase")]
58pub struct GoogleFilesOptions {
59    /// Display name stored with the file (defaults to `filename`).
60    #[serde(default)]
61    pub display_name: Option<String>,
62    /// Poll interval in milliseconds while the file is `PROCESSING`.
63    #[serde(default)]
64    pub poll_interval_ms: Option<u64>,
65    /// Poll timeout in milliseconds.
66    #[serde(default)]
67    pub poll_timeout_ms: Option<u64>,
68}
69
70/// A file resource.
71#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
72#[serde(rename_all = "camelCase")]
73pub struct GoogleFile {
74    /// Resource name (`files/abc-123`).
75    pub name: String,
76    /// Display name.
77    #[serde(default)]
78    pub display_name: Option<String>,
79    /// Media type.
80    #[serde(default)]
81    pub mime_type: Option<String>,
82    /// Size in bytes.
83    #[serde(default, deserialize_with = "deserialize_count")]
84    pub size_bytes: Option<u64>,
85    /// Creation time (RFC 3339).
86    #[serde(default)]
87    pub create_time: Option<String>,
88    /// Update time (RFC 3339).
89    #[serde(default)]
90    pub update_time: Option<String>,
91    /// Expiration time (RFC 3339).
92    #[serde(default)]
93    pub expiration_time: Option<String>,
94    /// SHA-256 hash (base64).
95    #[serde(default)]
96    pub sha256_hash: Option<String>,
97    /// URI usable as `fileData.fileUri`.
98    #[serde(default)]
99    pub uri: Option<String>,
100    /// `PROCESSING`, `ACTIVE` or `FAILED`.
101    #[serde(default)]
102    pub state: Option<String>,
103}
104
105/// Upload responses wrap the file in `{file}`; `files.get` returns it bare.
106#[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/// A byte upload through the resumable protocol.
122#[derive(Debug, Clone)]
123pub struct UploadRequest {
124    /// Content.
125    pub data: Bytes,
126    /// Media type sent as `X-Goog-Upload-Header-Content-Type`.
127    pub media_type: String,
128    /// Display name.
129    pub display_name: Option<String>,
130    /// Poll interval while `PROCESSING`.
131    pub poll_interval: Duration,
132    /// Poll timeout.
133    pub poll_timeout: Duration,
134    /// Additional request headers.
135    pub headers: Headers,
136    /// Cancellation token.
137    pub cancellation: CancellationToken,
138}
139
140impl UploadRequest {
141    /// Creates a request with the default poll settings.
142    #[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
179/// Normalizes a provider reference (a `files/...` name, a bare id or the
180/// full `uri`) to the resource name.
181fn 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/// Files service backed by the Files API.
192#[derive(Debug, Clone)]
193pub struct GoogleFiles {
194    config: SharedConfig,
195    provider: ProviderId,
196}
197
198impl GoogleFiles {
199    /// Creates the service.
200    #[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    /// Maps a file resource to the specification result.
217    #[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    /// Fetches the resource `name` (`files/...`).
249    ///
250    /// # Errors
251    ///
252    /// Returns the API error of the request.
253    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    /// Uploads bytes through the resumable protocol and waits until the file
275    /// leaves the `PROCESSING` state.
276    ///
277    /// # Errors
278    ///
279    /// Returns [`ProviderError::InvalidResponseData`] when the upload session
280    /// returns no upload URL, [`ProviderError::ApiCall`] when processing
281    /// fails or times out, and [`ProviderError::Cancelled`] when the token
282    /// fires while polling.
283    #[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}