lumen-server 0.1.0

use std::{sync::Arc, time::Duration};

use super::{ArtifactRef, ArtifactStore, ArtifactWrite, BoxFuture, ServiceError, ServiceResult};

pub trait PresignedUrlResolver: Send + Sync {
    fn resolve<'a>(&'a self, artifact: &'a ArtifactWrite) -> BoxFuture<'a, ServiceResult<String>>;
}

#[derive(Debug, Clone)]
pub struct StaticPresignedUrlResolver {
    upload_url: String,
}

impl StaticPresignedUrlResolver {
    pub fn new(upload_url: impl Into<String>) -> Self {
        Self {
            upload_url: upload_url.into(),
        }
    }
}

impl PresignedUrlResolver for StaticPresignedUrlResolver {
    fn resolve<'a>(&'a self, _artifact: &'a ArtifactWrite) -> BoxFuture<'a, ServiceResult<String>> {
        Box::pin(async move { Ok(self.upload_url.clone()) })
    }
}

#[derive(Clone)]
pub struct PresignedUrlArtifactStore<R> {
    resolver: Arc<R>,
    artifact_uri_prefix: String,
    bearer_token: Option<String>,
    timeout: Duration,
}

impl<R> PresignedUrlArtifactStore<R>
where
    R: PresignedUrlResolver,
{
    pub fn new(resolver: R, artifact_uri_prefix: impl Into<String>) -> Self {
        Self {
            resolver: Arc::new(resolver),
            artifact_uri_prefix: artifact_uri_prefix.into(),
            bearer_token: None,
            timeout: Duration::from_secs(60),
        }
    }

    pub fn with_bearer_token(mut self, bearer_token: impl Into<String>) -> Self {
        self.bearer_token = Some(bearer_token.into());
        self
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }
}

impl<R> ArtifactStore for PresignedUrlArtifactStore<R>
where
    R: PresignedUrlResolver + 'static,
{
    fn put<'a>(&'a self, artifact: ArtifactWrite) -> BoxFuture<'a, ServiceResult<ArtifactRef>> {
        Box::pin(async move {
            let upload_url = self.resolver.resolve(&artifact).await?;
            let mut request = reqwest::Client::new()
                .put(&upload_url)
                .header(reqwest::header::CONTENT_TYPE, artifact.content_type)
                .body(artifact.bytes.clone())
                .timeout(self.timeout);
            if let Some(token) = self.bearer_token.as_deref() {
                request = request.bearer_auth(token);
            }
            let response = request.send().await.map_err(|err| ServiceError {
                code: "artifact_upload_failed",
                message: err.to_string(),
                retryable: true,
            })?;
            if !response.status().is_success() {
                return Err(ServiceError {
                    code: "artifact_upload_failed",
                    message: format!("artifact upload returned status {}", response.status()),
                    retryable: response.status().is_server_error()
                        || response.status().as_u16() == 429,
                });
            }

            let id = artifact.job_id.0;
            Ok(ArtifactRef {
                id: id.clone(),
                content_type: artifact.content_type,
                bytes: artifact.bytes.len(),
                uri: format!("{}/{}", self.artifact_uri_prefix.trim_end_matches('/'), id),
            })
        })
    }
}