athena_rs 2.12.1

Database gateway API
Documentation
use actix_web::HttpResponse;
use aws_config::BehaviorVersion;
use aws_sdk_s3::Client as S3Client;
use aws_sdk_s3::config::{Credentials, Region};
use moka::future::Cache;
use once_cell::sync::Lazy;
use std::path::PathBuf;
use std::time::Duration;

use crate::api::response::bad_request;

pub const PRESIGN_URL_EXPIRY_SECS: u64 = 3600;

static S3_CLIENT_CACHE: Lazy<Cache<String, S3Client>> = Lazy::new(|| {
    Cache::builder()
        .max_capacity(256)
        .time_to_live(Duration::from_secs(900))
        .build()
});

pub struct TempPathCleanup {
    path: PathBuf,
}

impl TempPathCleanup {
    pub fn new(path: PathBuf) -> Self {
        Self { path }
    }
}

impl Drop for TempPathCleanup {
    fn drop(&mut self) {
        let _ = std::fs::remove_file(&self.path);
    }
}

pub fn ensure_bucket_present(bucket: &str) -> Result<(), HttpResponse> {
    if bucket.trim().is_empty() {
        return Err(bad_request("Bucket is required", "bucket field is empty"));
    }
    Ok(())
}

pub async fn build_s3_client(
    endpoint: &str,
    region: &str,
    bucket: &str,
    access_key_id: &str,
    secret_key: &str,
) -> S3Client {
    let cache_key: String = format!(
        "{}|{}|{}|{}",
        endpoint.trim(),
        region.trim(),
        bucket.trim(),
        access_key_id.trim()
    );

    if let Some(client) = S3_CLIENT_CACHE.get(&cache_key).await {
        return client;
    }

    let aws_region: Region = Region::new(region.to_string());
    let credentials: Credentials =
        Credentials::new(access_key_id, secret_key, None, None, "athena-storage");

    let mut builder: aws_config::ConfigLoader = aws_config::defaults(BehaviorVersion::latest())
        .region(aws_region)
        .credentials_provider(credentials);

    let normalized_endpoint: String = endpoint.trim().to_string();
    if !normalized_endpoint.is_empty() && !normalized_endpoint.contains("amazonaws.com") {
        builder = builder.endpoint_url(&normalized_endpoint);
    }

    let aws_config: aws_config::SdkConfig = builder.load().await;
    let mut s3_builder: aws_sdk_s3::config::Builder =
        aws_sdk_s3::config::Builder::from(&aws_config);

    if !normalized_endpoint.is_empty() && !normalized_endpoint.contains("amazonaws.com") {
        s3_builder = s3_builder.force_path_style(true);
    }

    let client: S3Client = S3Client::from_conf(s3_builder.build());
    S3_CLIENT_CACHE.insert(cache_key, client.clone()).await;
    client
}