1use 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
52pub const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
54
55pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 300_000;
57
58#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
60#[serde(rename_all = "camelCase")]
61pub struct GoogleFilesOptions {
62 #[serde(default)]
64 pub display_name: Option<String>,
65 #[serde(default)]
67 pub poll_interval_ms: Option<u64>,
68 #[serde(default)]
70 pub poll_timeout_ms: Option<u64>,
71}
72
73#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub struct GoogleFile {
77 pub name: String,
79 #[serde(default)]
81 pub display_name: Option<String>,
82 #[serde(default)]
84 pub mime_type: Option<String>,
85 #[serde(default, deserialize_with = "deserialize_count")]
87 pub size_bytes: Option<u64>,
88 #[serde(default)]
90 pub create_time: Option<String>,
91 #[serde(default)]
93 pub update_time: Option<String>,
94 #[serde(default)]
96 pub expiration_time: Option<String>,
97 #[serde(default)]
99 pub sha256_hash: Option<String>,
100 #[serde(default)]
102 pub uri: Option<String>,
103 #[serde(default)]
105 pub state: Option<String>,
106}
107
108#[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#[derive(Debug, Clone)]
126pub struct UploadRequest {
127 pub data: Bytes,
129 pub media_type: String,
131 pub display_name: Option<String>,
133 pub poll_interval: Duration,
135 pub poll_timeout: Duration,
137 pub headers: Headers,
139 pub cancellation: CancellationToken,
141}
142
143impl UploadRequest {
144 #[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
182fn 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#[derive(Debug, Clone)]
196pub struct GoogleFiles {
197 config: SharedConfig,
198 provider: ProviderId,
199}
200
201impl GoogleFiles {
202 #[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 #[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 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 #[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}