pile/crates/pile-dataset/src/serve/extract.rs

use axum::{
	Json,
	extract::{Query, RawQuery, State},
	http::{StatusCode, header},
	response::{IntoResponse, Response},
};
use pile_config::{Label, objectpath::ObjectPath};
use pile_value::{extract::traits::ExtractState, value::PileValue};
use serde::Deserialize;
use std::{sync::Arc, time::Instant};
use tracing::debug;
use utoipa::ToSchema;

use crate::Datasets;

#[derive(Deserialize, ToSchema)]
pub struct ExtractQuery {
	source: String,
	key: String,

	#[serde(default)]
	download: bool,
	name: Option<String>,
}

/// Extract a specific field from an item's metadata.
/// Multiple `path` parameters may be provided; the first non-null result is returned.
#[utoipa::path(
	get,
	path = "/extract",
	params(
		("source" = String, Query, description = "Source label"),
		("key" = String, Query, description = "Item key"),
		("path" = String, Query, description = "Object path (e.g. $.flac.title); repeat for fallbacks"),
		("name" = Option<String>, Query, description = "Downloaded filename; defaults to the last segment of the key"),
	),
	responses(
		(status = 200, description = "Field value as JSON"),
		(status = 400, description = "Invalid source label or path"),
		(status = 404, description = "Item or field not found"),
		(status = 500, description = "Internal server error"),
	)
)]
pub async fn get_extract(
	State(state): State<Arc<Datasets>>,
	Query(params): Query<ExtractQuery>,
	RawQuery(raw_query): RawQuery,
) -> Response {
	let start = Instant::now();

	let label = match Label::try_from(params.source.clone()) {
		Ok(l) => l,
		Err(e) => return (StatusCode::BAD_REQUEST, format!("{e:?}")).into_response(),
	};

	// Collect all `path` query params in order (supports repeated ?path=...&path=...)
	let raw = raw_query.as_deref().unwrap_or("");
	let paths: Vec<ObjectPath> = {
		let mut result = Vec::new();
		for part in raw.split('&') {
			if let Some((k, v)) = part.split_once('=')
				&& k == "path"
			{
				match v.parse::<ObjectPath>() {
					Ok(p) => result.push(p),
					Err(e) => {
						return (StatusCode::BAD_REQUEST, format!("{e:?}")).into_response();
					}
				}
			}
		}
		result
	};

	if paths.is_empty() {
		return (StatusCode::BAD_REQUEST, "Missing `path` query parameter").into_response();
	}

	debug!(
		message = "Serving /extract",
		source = params.source,
		key = params.key,
		paths = paths.len(),
	);

	let Some(item) = state.get(&label, &params.key).await else {
		return StatusCode::NOT_FOUND.into_response();
	};

	let extract_state = ExtractState { ignore_mime: false };
	let item = PileValue::Item(item);

	// Try each path in order, returning the first non-null result
	let mut value = None;
	for path in &paths {
		match item.query(&extract_state, path).await {
			Ok(Some(PileValue::Null)) | Ok(None) => continue,
			Ok(Some(v)) => {
				value = Some(v);
				break;
			}
			Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, format!("{e:?}")).into_response(),
		}
	}

	let Some(value) = value else {
		return StatusCode::NOT_FOUND.into_response();
	};

	debug!(
		message = "Served /extract",
		source = params.source,
		key = params.key,
		time_ms = start.elapsed().as_millis()
	);

	let disposition_type = if params.download {
		"attachment"
	} else {
		"inline"
	};
	let file_name = params.name.unwrap_or_else(|| {
		params
			.key
			.rsplit('/')
			.next()
			.unwrap_or(&params.key)
			.to_owned()
	});
	let disposition = format!("{disposition_type}; filename=\"{file_name}\"");

	match value {
		PileValue::String(s) => (
			StatusCode::OK,
			[
				(header::CONTENT_TYPE, "text/plain".to_owned()),
				(header::CONTENT_DISPOSITION, disposition),
			],
			s.to_string(),
		)
			.into_response(),
		PileValue::Blob { mime, bytes } => (
			StatusCode::OK,
			[
				(header::CONTENT_TYPE, mime.to_string()),
				(header::CONTENT_DISPOSITION, disposition),
			],
			bytes.as_ref().clone(),
		)
			.into_response(),
		_ => match value.to_json(&extract_state).await {
			Ok(json) => (
				StatusCode::OK,
				[(header::CONTENT_DISPOSITION, disposition)],
				Json(json),
			)
				.into_response(),
			Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, format!("{e:?}")).into_response(),
		},
	}
}