micromegas_telemetry/

blob_storage.rs

use anyhow::Result;
use futures::StreamExt;
use futures::stream;
use object_store::prefix::PrefixStore;
use object_store::{ObjectStore, ObjectStoreExt, PutMode, PutOptions, path::Path};
use std::sync::Arc;

/// Outcome of a create-only write.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PutIfAbsent {
    /// The object did not exist and now holds the supplied bytes.
    Created,
    /// The object already existed; it was left untouched and the bytes were discarded.
    AlreadyExists,
}

/// Parses an object-store URI into the raw store + root prefix, feeding
/// `object_store` the process env vars as options (it normalizes the keys).
/// The single home for the `parse_url_opts(url, env vars)` idiom.
pub fn parse_object_store_url(uri: &str) -> Result<(Arc<dyn ObjectStore>, Path)> {
    parse_object_store_url_parsed(&url::Url::parse(uri)?)
}

/// Like `parse_object_store_url` but takes an already-parsed URL, for callers
/// that also need the `url::Url` themselves and would otherwise parse it twice.
pub fn parse_object_store_url_parsed(url: &url::Url) -> Result<(Arc<dyn ObjectStore>, Path)> {
    let (store, prefix) = object_store::parse_url_opts(url, std::env::vars())?;
    Ok((Arc::new(store), prefix))
}

/// A client for interacting with blob storage.
///
/// This struct wraps an `ObjectStore` and prefixes all paths with a root path,
/// providing a convenient way to interact with a specific "folder" within the blob storage.
#[derive(Debug)]
pub struct BlobStorage {
    blob_store: Arc<dyn ObjectStore>,
}

impl BlobStorage {
    /// Creates a new `BlobStorage` instance.
    pub fn new(blob_store: Arc<dyn ObjectStore>, blob_store_root: Path) -> Self {
        Self {
            blob_store: Arc::new(PrefixStore::new(blob_store, blob_store_root)),
        }
    }

    /// Connects to a blob storage service using the provided URL.
    pub fn connect(object_store_url: &str) -> Result<Self> {
        Self::connect_with_layer(object_store_url, |s| s)
    }

    /// Parses an object store URL into the raw (unwrapped, un-prefixed) store and
    /// the lake root path, using the same env-var-derived options as `connect` /
    /// `connect_with_layer`. Exposed so callers that need to apply their own layer
    /// (e.g. the object cache client) and also need the root prefix (e.g. to key a
    /// write-time cache warm the same way `PrefixStore` keys a demand read) don't
    /// have to re-parse the URL or duplicate the env-var option plumbing.
    pub fn parse_url_opts(object_store_url: &str) -> Result<(Arc<dyn ObjectStore>, Path)> {
        parse_object_store_url(object_store_url)
    }

    /// Connects to a blob storage service and applies a layer to the raw store before
    /// wrapping it in `PrefixStore`. The layer receives the full-bucket store so its
    /// keys are bucket-relative (including the lake root prefix).
    pub fn connect_with_layer(
        object_store_url: &str,
        layer: impl FnOnce(Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore>,
    ) -> Result<Self> {
        let (blob_store, blob_store_root) = Self::parse_url_opts(object_store_url)?;
        let layered = layer(blob_store);
        Ok(Self {
            blob_store: Arc::new(PrefixStore::new(layered, blob_store_root)),
        })
    }

    /// Returns a shared reference to the inner `ObjectStore`.
    pub fn inner(&self) -> Arc<dyn ObjectStore> {
        self.blob_store.clone()
    }

    /// Puts a blob into storage at the specified path.
    pub async fn put(&self, obj_path: &str, buffer: bytes::Bytes) -> Result<()> {
        self.blob_store
            .put(&Path::from(obj_path), buffer.into())
            .await?;
        Ok(())
    }

    /// Writes a blob only if the key is absent, leaving an existing object untouched.
    ///
    /// The lake's object keys are write-once and content-addressed — the range cache keys
    /// carry no etag or generation and are never invalidated (see `range_cache`'s module
    /// docs) — and this is what enforces that: a colliding write is reported, never applied.
    pub async fn put_if_absent(&self, obj_path: &str, buffer: bytes::Bytes) -> Result<PutIfAbsent> {
        match self
            .blob_store
            .put_opts(
                &Path::from(obj_path),
                buffer.into(),
                PutOptions::from(PutMode::Create),
            )
            .await
        {
            Ok(_) => Ok(PutIfAbsent::Created),
            Err(object_store::Error::AlreadyExists { .. }) => Ok(PutIfAbsent::AlreadyExists),
            Err(object_store::Error::NotImplemented { .. }) => Err(anyhow::anyhow!(
                "object store does not support conditional put (PutMode::Create), which \
                 create-only block writes require; if this is an S3-compatible store, check \
                 that it is not configured with aws_conditional_put=disabled"
            )),
            Err(e) => Err(e.into()),
        }
    }

    /// Reads a blob from storage at the specified path.
    pub async fn read_blob(&self, obj_path: &str) -> Result<bytes::Bytes> {
        let get_result = self.blob_store.get(&Path::from(obj_path)).await?;
        Ok(get_result.bytes().await?)
    }

    /// Deletes a blob from storage at the specified path.
    pub async fn delete(&self, obj_path: &str) -> Result<()> {
        self.blob_store.delete(&Path::from(obj_path)).await?;
        Ok(())
    }

    /// Probes blob storage reachability and credentials by fetching the first
    /// page of a bucket listing (a single request; does not enumerate the bucket).
    pub async fn probe(&self) -> anyhow::Result<()> {
        match self.blob_store.list(None).next().await {
            Some(Ok(_)) | None => Ok(()),
            Some(Err(e)) => Err(e.into()),
        }
    }

    /// Deletes a batch of blobs from storage.
    pub async fn delete_batch(&self, objects: &[String]) -> Result<()> {
        let paths: Vec<_> = objects
            .iter()
            .map(|obj_path| Ok(Path::from(obj_path.as_str())))
            .collect();
        let path_stream = stream::iter(paths);
        let mut stream = self.blob_store.delete_stream(Box::pin(path_stream));
        while let Some(res) = stream.next().await {
            if let Err(e) = res {
                match e {
                    object_store::Error::NotFound { path: _, source: _ } => Ok(()),
                    ref _other_error => Err(e),
                }?
            }
        }
        Ok(())
    }
}