arraylake/

lib.rs

use std::{
    collections::{HashMap, HashSet},
    num::{NonZeroU16, NonZeroUsize},
    sync::Arc,
};

use anyhow::{Context, bail};
use async_trait::async_trait;
use chrono::{DateTime, TimeDelta, Utc};
use icechunk::{
    Storage,
    config::{
        GcsBearerCredential, GcsCredentials, GcsCredentialsFetcher, S3Credentials,
        S3CredentialsFetcher, S3Options, S3StaticCredentials,
    },
    format::SnapshotId,
    new_s3_storage,
    ops::gc::{Action, ExpireResult, GCConfig, GCSummary},
    refs::Ref,
    storage::{new_gcs_storage, new_r2_storage},
};
use serde::{Deserialize, Serialize};
use tracing::{debug, info, instrument};
use url::Url;

#[derive(Debug, Deserialize, PartialEq, Eq)]
#[serde(untagged)]
pub enum ALBucketCredentials {
    S3 {
        aws_access_key_id: String,
        aws_secret_access_key: String,
        aws_session_token: Option<String>,
        expiration: Option<DateTime<Utc>>,
    },
    GCS {
        access_token: String,
        principal: String,
        expiration: Option<DateTime<Utc>>,
    },
}

#[derive(Debug, Deserialize, PartialEq, Eq)]
#[serde(untagged)]
pub enum ALBucketExtraConfigValue {
    B(bool),
    S(String),
}

#[derive(Debug, Deserialize, PartialEq, Eq)]
pub enum ALBucketPlatform {
    #[serde(rename = "s3")]
    S3,
    #[serde(rename = "s3-compatible")]
    S3Compatible,
    #[serde(rename = "minio")]
    Minio,
    #[serde(rename = "gs")]
    GS,
}

#[derive(Debug, Deserialize, PartialEq, Eq)]
#[serde(tag = "method")]
pub enum ALBucketAuthConfig {
    #[serde(rename = "aws_customer_managed_role")]
    #[serde(alias = "customer_managed_role")]
    AWSCustomerManagedRole {
        external_customer_id: String,
        external_role_name: String,
        shared_secret: String, // This is a secret string, not used by the client
    },

    #[serde(rename = "gcp_customer_managed_role")]
    GCPCustomerManagedRole { target_service_account: String },

    #[serde(rename = "r2_customer_managed_role")]
    R2CustomerManagedRole {
        external_accound_id: String,
        account_api_token: String, // This is a secret string, not used by the client
        parent_access_key_id: String, // This is a secret string, not used by the client
        duration_seconds: i32,     // 3600 seconds by default
    },

    #[serde(rename = "hmac")]
    Hmac { access_key_id: String, secret_access_key: String },

    #[serde(rename = "anonymous")]
    Anonymous,
}

#[derive(Debug, Deserialize, PartialEq, Eq)]
pub struct ALBucketInfo {
    pub nickname: String,
    pub platform: ALBucketPlatform,
    pub name: String,
    pub prefix: String,
    pub is_default: bool,
    pub extra_config: HashMap<String, ALBucketExtraConfigValue>,
    pub auth_config: ALBucketAuthConfig,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct OptimizationConfig {
    pub gc_config: Option<ALGCConfig>,
    pub expiration_config: Option<ALExpirationConfig>,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct ALRepoInfo {
    pub name: String,
    pub org: String,
    pub bucket: Option<ALBucketInfo>,
    pub prefix: String,
    pub kind: String,
    pub optimization_config: Option<OptimizationConfig>,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct ALGCConfig {
    pub extra_gc_roots: HashSet<String>,
    pub dangling_chunks: Option<f64>,
    pub dangling_manifests: Option<f64>,
    pub dangling_attributes: Option<f64>,
    pub dangling_transaction_logs: Option<f64>,
    pub dangling_snapshots: Option<f64>,
    pub enabled: bool,
}

impl TryFrom<ALGCConfig> for GCConfig {
    type Error = anyhow::Error;

    fn try_from(value: ALGCConfig) -> Result<Self, Self::Error> {
        fn to_action(now: DateTime<Utc>, seconds: Option<f64>) -> Action {
            seconds
                .map(|sec| {
                    Action::DeleteIfCreatedBefore(
                        now - TimeDelta::seconds(sec.round() as i64),
                    )
                })
                .unwrap_or(Action::Keep)
        }

        let now = Utc::now();
        let config = GCConfig::new(
            value
                .extra_gc_roots
                .into_iter()
                .map(|s| {
                    SnapshotId::try_from(s.as_str()).map_err(|error| {
                        anyhow::anyhow!(
                            "Cannot parse a SnapshotId from <{}>: {}",
                            s,
                            error
                        )
                    })
                })
                .collect::<Result<_, _>>()?,
            to_action(now, value.dangling_chunks),
            to_action(now, value.dangling_manifests),
            to_action(now, value.dangling_attributes),
            to_action(now, value.dangling_transaction_logs),
            to_action(now, value.dangling_snapshots),
            NonZeroU16::new(1_000).unwrap(),
            NonZeroUsize::new(5 * 1_024 * 1_024 * 1_024).unwrap(), // 5GB
            NonZeroU16::new(500).unwrap(),
            false,
        );
        Ok(config)
    }
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct ALExpirationConfig {
    pub expire_versions_older_than: f64,
    pub expire_every: Option<f64>,
    pub enabled: bool,
}

#[derive(Debug)]
pub struct ALClient {
    http: reqwest::Client,
    uri: Url,
    token: String,
}

// We implement default to satisfy the serialization requirements or the credential refreshers
impl Default for ALClient {
    fn default() -> Self {
        Self {
            http: Default::default(),
            uri: Url::parse("https://api.earthmover.io").unwrap(),
            token: Default::default(),
        }
    }
}

impl ALClient {
    pub fn new(service_uri: Option<String>, token: String) -> anyhow::Result<Self> {
        let uri = service_uri.unwrap_or("https://api.earthmover.io".to_string());
        let uri = Url::parse(uri.as_str())?;
        Ok(ALClient { token, uri, http: reqwest::Client::new() })
    }

    #[instrument(skip(self))]
    pub async fn get_credentials(
        &self,
        org: &str,
        repo: &str,
    ) -> anyhow::Result<ALBucketCredentials> {
        info!(repo_name = repo, org_name = org, "Getting credentials");
        let path = format!("repos/{org}/{repo}/bucket-credentials");
        let url = self.make_url(path.as_str())?;
        let res = self
            .http
            .get(url)
            .bearer_auth(self.token.clone())
            .send()
            .await
            .context("Fetching credentials")?
            .json::<ALBucketCredentials>()
            .await
            .context("Parsing credentials")?;
        Ok(res)
    }

    #[instrument(skip(self))]
    pub async fn get_repo_info(
        &self,
        org: &str,
        repo: &str,
    ) -> anyhow::Result<ALRepoInfo> {
        info!(repo_name = repo, org_name = org, "Getting repository info");
        let path = format!("repos/{org}/{repo}");
        let url = self.make_url(path.as_str())?;
        let res = self
            .http
            .get(url)
            .bearer_auth(self.token.clone())
            .send()
            .await
            .context("Fetching repo info")?
            .json::<ALRepoInfo>()
            .await
            .context("Parsing repo info")?;
        Ok(res)
    }

    #[instrument(skip(self))]
    pub async fn get_storage_for_repo(
        self: &Arc<Self>,
        repo: &ALRepoInfo,
    ) -> anyhow::Result<Arc<dyn Storage>> {
        info!(?repo, "Getting storage for repo");
        let Some(bucket) = &repo.bucket else { bail!("Bucket info is missing") };

        if !can_we_get_bucket_credentials(bucket) {
            bail!("We don't have credentials to the bucket")
        }
        let credential_refresher = ALCredentialsFetcher {
            al: Arc::clone(self),
            repo: repo.name.clone(),
            org: repo.org.clone(),
        };
        match bucket.platform {
            ALBucketPlatform::S3 => get_s3_storage(repo, credential_refresher),
            ALBucketPlatform::S3Compatible => {
                get_s3_compatible_storage(repo, credential_refresher)
            }
            ALBucketPlatform::Minio => {
                get_s3_compatible_storage(repo, credential_refresher)
            }
            ALBucketPlatform::GS => get_gcs_storage(repo, credential_refresher).await,
        }
    }

    #[instrument(skip(self))]
    pub async fn store_gc_results(
        &self,
        org: &str,
        repo: &str,
        job_run_id: Option<&str>,
        result: GCSummary,
    ) -> anyhow::Result<()> {
        info!(?repo, "Storing gc result");
        let path = format!("repos/icechunk/{org}/{repo}/gc_results");
        let url = self.make_url(path.as_str())?;
        let mut body: HashMap<String, serde_json::Value> = [
            ("bytes_deleted".to_string(), serde_json::json!(result.bytes_deleted)),
            ("chunks_deleted".to_string(), serde_json::json!(result.chunks_deleted)),
            (
                "manifests_deleted".to_string(),
                serde_json::json!(result.manifests_deleted),
            ),
            (
                "snapshots_deleted".to_string(),
                serde_json::json!(result.snapshots_deleted),
            ),
            (
                "attributes_deleted".to_string(),
                serde_json::json!(result.attributes_deleted),
            ),
            (
                "transaction_logs_deleted".to_string(),
                serde_json::json!(result.transaction_logs_deleted),
            ),
        ]
        .into_iter()
        .collect();

        if let Some(id) = job_run_id {
            body.insert("job_run_id".to_string(), serde_json::json!(id));
        }

        self.http
            .post(url)
            .bearer_auth(self.token.clone())
            .json(&body)
            .send()
            .await
            .context("Uploading gc results")?
            .error_for_status()?;

        Ok(())
    }

    #[instrument(skip(self))]
    pub async fn store_expiration_results(
        &self,
        org: &str,
        repo: &str,
        job_run_id: Option<&str>,
        result: ExpireResult,
    ) -> anyhow::Result<()> {
        info!(?repo, "Storing expiration result");
        let path = format!("repos/icechunk/{org}/{repo}/expiration_results");
        let url = self.make_url(path.as_str())?;
        let mut body: HashMap<String, serde_json::Value> = [
            (
                "released_snapshots".to_string(),
                serde_json::json!(
                    result
                        .released_snapshots
                        .into_iter()
                        .map(|r| r.to_string())
                        .collect::<HashSet<_>>()
                ),
            ),
            (
                "edited_snapshots".to_string(),
                serde_json::json!(
                    result
                        .edited_snapshots
                        .into_iter()
                        .map(|r| r.to_string())
                        .collect::<HashSet<_>>()
                ),
            ),
            (
                "deleted_tags".to_string(),
                serde_json::json!(
                    result
                        .deleted_refs
                        .iter()
                        .filter_map(|r| match r {
                            Ref::Tag(name) => Some(name.clone()),
                            Ref::Branch(_) => None,
                        })
                        .collect::<HashSet<_>>()
                ),
            ),
            (
                "deleted_branches".to_string(),
                serde_json::json!(
                    result
                        .deleted_refs
                        .iter()
                        .filter_map(|r| match r {
                            Ref::Tag(_) => None,
                            Ref::Branch(name) => Some(name.clone()),
                        })
                        .collect::<HashSet<_>>()
                ),
            ),
        ]
        .into_iter()
        .collect();

        if let Some(id) = job_run_id {
            body.insert("job_run_id".to_string(), serde_json::json!(id));
        }

        self.http
            .post(url)
            .bearer_auth(self.token.clone())
            .json(&body)
            .send()
            .await
            .context("Uploading expiration results")?
            .error_for_status()?;

        Ok(())
    }

    fn make_url(&self, path: &str) -> anyhow::Result<String> {
        Ok(self.uri.join(path)?.as_str().to_string())
    }
}

fn use_delegated_credentials(bucket: &ALBucketInfo) -> bool {
    matches!(
        bucket.auth_config,
        ALBucketAuthConfig::AWSCustomerManagedRole { .. }
            | ALBucketAuthConfig::GCPCustomerManagedRole { .. }
            | ALBucketAuthConfig::R2CustomerManagedRole { .. }
    ) && matches!(
        bucket.platform,
        ALBucketPlatform::GS | ALBucketPlatform::S3 | ALBucketPlatform::S3Compatible
    )
}

fn use_hmac_credentials(bucket: &ALBucketInfo) -> bool {
    matches!(bucket.auth_config, ALBucketAuthConfig::Hmac { .. })
}

fn can_we_get_bucket_credentials(bucket: &ALBucketInfo) -> bool {
    use_delegated_credentials(bucket) || use_hmac_credentials(bucket)
}

#[derive(Debug, Serialize, Deserialize)]
struct ALCredentialsFetcher {
    #[serde(skip)]
    al: Arc<ALClient>,
    repo: String,
    org: String,
}

#[async_trait]
#[typetag::serde]
impl GcsCredentialsFetcher for ALCredentialsFetcher {
    async fn get(&self) -> Result<GcsBearerCredential, String> {
        debug!("Refreshing GCS credentials");
        match self.al.get_credentials(self.org.as_str(), self.repo.as_str()).await {
            Ok(ALBucketCredentials::GCS { access_token, expiration, .. }) => {
                Ok(GcsBearerCredential {
                    bearer: access_token,
                    expires_after: expiration,
                })
            }
            Ok(_) => Err("Invalid credential type for GCS bucket".to_string()),
            Err(_) => Err("Cannot get credentials for GCS bucket".to_string()),
        }
    }
}

#[async_trait]
#[typetag::serde]
impl S3CredentialsFetcher for ALCredentialsFetcher {
    async fn get(&self) -> Result<S3StaticCredentials, String> {
        debug!("Refreshing S3 credentials");
        match self.al.get_credentials(self.org.as_str(), self.repo.as_str()).await {
            Ok(ALBucketCredentials::S3 {
                aws_access_key_id,
                aws_secret_access_key,
                aws_session_token,
                expiration,
            }) => Ok(S3StaticCredentials {
                access_key_id: aws_access_key_id,
                secret_access_key: aws_secret_access_key,
                session_token: aws_session_token,
                expires_after: expiration,
            }),
            Ok(_) => Err("Invalid credential type for S3 bucket".to_string()),
            Err(_) => Err("Cannot get credentials for S3 bucket".to_string()),
        }
    }
}

#[instrument(skip(fetcher))]
async fn get_gcs_storage(
    repo: &ALRepoInfo,
    fetcher: ALCredentialsFetcher,
) -> anyhow::Result<Arc<dyn Storage>> {
    use ALBucketAuthConfig::*;

    let Some(bucket) = &repo.bucket else { bail!("Bucket info is missing") };

    let credentials = match &bucket.auth_config {
        GCPCustomerManagedRole { .. } => GcsCredentials::Refreshable(Arc::new(fetcher)),
        Anonymous => GcsCredentials::Anonymous,
        _ => bail!("Cannot create credentials for the bucket"),
    };
    new_gcs_storage(
        bucket.name.clone(),
        Some(repo.prefix.clone()), // already includes the bucket prefix
        Some(credentials),
        None, // FIXME: should we pass something here?
    )
    .await
    .context("Creating GCS storage")
}

#[instrument(skip(fetcher))]
fn get_s3_storage(
    repo: &ALRepoInfo,
    fetcher: ALCredentialsFetcher,
) -> anyhow::Result<Arc<dyn Storage>> {
    use ALBucketAuthConfig::*;

    let Some(bucket) = &repo.bucket else { bail!("Bucket info is missing") };

    let credentials = match &bucket.auth_config {
        AWSCustomerManagedRole { .. } => S3Credentials::Refreshable(Arc::new(fetcher)),
        Hmac { access_key_id, secret_access_key } => {
            S3Credentials::Static(S3StaticCredentials {
                access_key_id: access_key_id.clone(),
                secret_access_key: secret_access_key.clone(),
                session_token: None,
                expires_after: None,
            })
        }
        Anonymous => S3Credentials::Anonymous,
        _ => bail!("Cannot create credentials for the bucket"),
    };

    let region = match bucket.extra_config.get("region_name") {
        Some(ALBucketExtraConfigValue::S(region)) => Some(region.clone()),
        _ => None,
    };
    let options = S3Options {
        region,
        endpoint_url: None,
        anonymous: false,
        allow_http: false,
        force_path_style: false,
        network_stream_timeout_seconds: None,
        requester_pays: false,
    };

    new_s3_storage(
        options,
        bucket.name.clone(),
        Some(repo.prefix.clone()), // already includes the bucket prefix
        Some(credentials),
    )
    .context("Creating S3 storage")
}

#[instrument(skip(fetcher))]
fn get_s3_compatible_storage(
    repo: &ALRepoInfo,
    fetcher: ALCredentialsFetcher,
) -> anyhow::Result<Arc<dyn Storage>> {
    use ALBucketAuthConfig::*;
    let Some(bucket) = repo.bucket.as_ref() else {
        bail!("Bucket not found");
    };
    let credentials = match &bucket.auth_config {
        AWSCustomerManagedRole { .. } => S3Credentials::Refreshable(Arc::new(fetcher)),
        R2CustomerManagedRole { .. } => S3Credentials::Refreshable(Arc::new(fetcher)),
        Hmac { access_key_id, secret_access_key } => {
            S3Credentials::Static(S3StaticCredentials {
                access_key_id: access_key_id.clone(),
                secret_access_key: secret_access_key.clone(),
                session_token: None,
                expires_after: None,
            })
        }
        Anonymous => S3Credentials::Anonymous,
        _ => bail!("Cannot create credentials for the bucket"),
    };
    // FIXME: add tigris, minio and R2
    let region = match bucket.extra_config.get("region_name") {
        Some(ALBucketExtraConfigValue::S(region)) => Some(region.clone()),
        _ => None,
    };
    let endpoint_url = match bucket.extra_config.get("endpoint_url") {
        Some(ALBucketExtraConfigValue::S(ep)) => Some(ep.clone()),
        _ => None,
    };
    let allow_http = match bucket.extra_config.get("use_ssl") {
        Some(ALBucketExtraConfigValue::B(b)) => *b,
        _ => false,
    };
    let force_path_style = allow_http; // this is what the Python AL client does

    // If its an r2 bucket, use r2 storage
    let options = S3Options {
        region,
        endpoint_url: endpoint_url.clone(),
        anonymous: false,
        allow_http,
        force_path_style,
        network_stream_timeout_seconds: None,
        requester_pays: false,
    };

    if let Some(endpoint_url) = endpoint_url
        && endpoint_url.contains("r2.cloudflarestorage.com")
    {
        new_r2_storage(
            options,
            Some(bucket.name.clone()),
            Some(repo.prefix.clone()),
            None,
            Some(credentials),
        )
        .context("Creating R2 storage")
    } else {
        new_s3_storage(
            options,
            bucket.name.clone(),
            Some(repo.prefix.clone()), // already includes the bucket prefix
            Some(credentials),
        )
        .context("Creating S3 storage")
    }
}

#[cfg(test)]
#[allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use chrono::{TimeZone, Utc};

    use super::{
        ALBucketAuthConfig, ALBucketCredentials, ALBucketExtraConfigValue, ALBucketInfo,
        ALBucketPlatform, ALRepoInfo,
    };

    #[test]
    fn test_parsing_repo_info() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"
            {
            "_id": "id",
            "org": "earthmover-demos",
            "name": "repo",
            "created": "2025-04-23T16:31:30+00:00",
            "updated": "2025-04-23T16:31:30+00:00",
            "description": null,
            "metadata": null,
            "created_by": "abcdef",
            "visibility": "PRIVATE",
            "bucket": {
                "id": "some-bucket-id",
                "nickname": "nickname",
                "platform": "s3",
                "name": "name",
                "created": "2024-04-17T23:00:43+00:00",
                "updated": "2025-04-08T13:28:28+00:00",
                "created_by": "somebody@example.com",
                "prefix": "bucket_prefix",
                "extra_config": {
                    "region_name": "us-east-1",
                    "anonymous": false
                },
                "is_default": true,
                "auth_config": {
                    "method": "aws_customer_managed_role",
                    "external_customer_id": "1234",
                    "external_role_name": "role",
                    "shared_secret": "secret"
                }
            },
            "status": {
                "mode": "online",
                "message": "new repo creation",
                "initiated_by": {
                    "principal_id": "principal",
                    "system_id": null
                },
                "estimated_end_time": null
            },
            "kind": "icechunk",
            "prefix": "repo_prefix",
            "gc_config": null,
            "expiration_config": null
            }
        "#;

        let info: ALRepoInfo = serde_json::from_str(json)?;
        let extra_config = [
            (
                "region_name".to_string(),
                ALBucketExtraConfigValue::S("us-east-1".to_string()),
            ),
            ("anonymous".to_string(), ALBucketExtraConfigValue::B(false)),
        ]
        .into_iter()
        .collect();
        let auth_config = ALBucketAuthConfig::AWSCustomerManagedRole {
            external_customer_id: "1234".to_string(),
            external_role_name: "role".to_string(),
            shared_secret: "secret".to_string(),
        };
        assert_eq!(
            info,
            ALRepoInfo {
                org: "earthmover-demos".to_string(),
                name: "repo".to_string(),
                bucket: Some(ALBucketInfo {
                    nickname: "nickname".to_string(),
                    platform: ALBucketPlatform::S3,
                    name: "name".to_string(),
                    prefix: "bucket_prefix".to_string(),
                    is_default: true,
                    extra_config,
                    auth_config,
                }),
                prefix: "repo_prefix".to_string(),
                optimization_config: None,
                kind: "icechunk".to_string(),
            }
        );

        Ok(())
    }

    #[test]
    fn test_parsing_credentials() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"
            {
                "aws_access_key_id": "access",
                "aws_secret_access_key": "secret",
                "aws_session_token": "token",
                "expiration": "2025-05-13T19:11:50Z"
            }
        "#;

        let expiration = Utc.with_ymd_and_hms(2025, 5, 13, 19, 11, 50).single().unwrap();
        let creds: ALBucketCredentials = serde_json::from_str(json)?;
        assert_eq!(
            creds,
            ALBucketCredentials::S3 {
                aws_access_key_id: "access".to_string(),
                aws_secret_access_key: "secret".to_string(),
                aws_session_token: Some("token".to_string()),
                expiration: Some(expiration),
            }
        );
        Ok(())
    }
}