micromegas_telemetry/
blob_storage.rs

1use anyhow::Result;
2use futures::StreamExt;
3use futures::stream;
4use object_store::prefix::PrefixStore;
5use object_store::{ObjectStore, ObjectStoreExt, path::Path};
6use std::sync::Arc;
7
8/// Parses an object-store URI into the raw store + root prefix, feeding
9/// `object_store` the process env vars lowercased (its expected option keys).
10/// The single home for the `parse_url_opts(url, env-vars-lowercased)` idiom.
11pub fn parse_object_store_url(uri: &str) -> Result<(Arc<dyn ObjectStore>, Path)> {
12    parse_object_store_url_parsed(&url::Url::parse(uri)?)
13}
14
15/// Like `parse_object_store_url` but takes an already-parsed URL, for callers
16/// that also need the `url::Url` themselves and would otherwise parse it twice.
17pub fn parse_object_store_url_parsed(url: &url::Url) -> Result<(Arc<dyn ObjectStore>, Path)> {
18    let (store, prefix) =
19        object_store::parse_url_opts(url, std::env::vars().map(|(k, v)| (k.to_lowercase(), v)))?;
20    Ok((Arc::new(store), prefix))
21}
22
23/// A client for interacting with blob storage.
24///
25/// This struct wraps an `ObjectStore` and prefixes all paths with a root path,
26/// providing a convenient way to interact with a specific "folder" within the blob storage.
27#[derive(Debug)]
28pub struct BlobStorage {
29    blob_store: Arc<dyn ObjectStore>,
30}
31
32impl BlobStorage {
33    /// Creates a new `BlobStorage` instance.
34    pub fn new(blob_store: Arc<dyn ObjectStore>, blob_store_root: Path) -> Self {
35        Self {
36            blob_store: Arc::new(PrefixStore::new(blob_store, blob_store_root)),
37        }
38    }
39
40    /// Connects to a blob storage service using the provided URL.
41    pub fn connect(object_store_url: &str) -> Result<Self> {
42        Self::connect_with_layer(object_store_url, |s| s)
43    }
44
45    /// Parses an object store URL into the raw (unwrapped, un-prefixed) store and
46    /// the lake root path, using the same env-var-derived options as `connect` /
47    /// `connect_with_layer`. Exposed so callers that need to apply their own layer
48    /// (e.g. the object cache client) and also need the root prefix (e.g. to key a
49    /// write-time cache warm the same way `PrefixStore` keys a demand read) don't
50    /// have to re-parse the URL or duplicate the env-var lowercasing.
51    pub fn parse_url_opts(object_store_url: &str) -> Result<(Arc<dyn ObjectStore>, Path)> {
52        parse_object_store_url(object_store_url)
53    }
54
55    /// Connects to a blob storage service and applies a layer to the raw store before
56    /// wrapping it in `PrefixStore`. The layer receives the full-bucket store so its
57    /// keys are bucket-relative (including the lake root prefix).
58    pub fn connect_with_layer(
59        object_store_url: &str,
60        layer: impl FnOnce(Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore>,
61    ) -> Result<Self> {
62        let (blob_store, blob_store_root) = Self::parse_url_opts(object_store_url)?;
63        let layered = layer(blob_store);
64        Ok(Self {
65            blob_store: Arc::new(PrefixStore::new(layered, blob_store_root)),
66        })
67    }
68
69    /// Returns a shared reference to the inner `ObjectStore`.
70    pub fn inner(&self) -> Arc<dyn ObjectStore> {
71        self.blob_store.clone()
72    }
73
74    /// Puts a blob into storage at the specified path.
75    pub async fn put(&self, obj_path: &str, buffer: bytes::Bytes) -> Result<()> {
76        self.blob_store
77            .put(&Path::from(obj_path), buffer.into())
78            .await?;
79        Ok(())
80    }
81
82    /// Reads a blob from storage at the specified path.
83    pub async fn read_blob(&self, obj_path: &str) -> Result<bytes::Bytes> {
84        let get_result = self.blob_store.get(&Path::from(obj_path)).await?;
85        Ok(get_result.bytes().await?)
86    }
87
88    /// Deletes a blob from storage at the specified path.
89    pub async fn delete(&self, obj_path: &str) -> Result<()> {
90        self.blob_store.delete(&Path::from(obj_path)).await?;
91        Ok(())
92    }
93
94    /// Probes blob storage reachability and credentials by fetching the first
95    /// page of a bucket listing (a single request; does not enumerate the bucket).
96    pub async fn probe(&self) -> anyhow::Result<()> {
97        match self.blob_store.list(None).next().await {
98            Some(Ok(_)) | None => Ok(()),
99            Some(Err(e)) => Err(e.into()),
100        }
101    }
102
103    /// Deletes a batch of blobs from storage.
104    pub async fn delete_batch(&self, objects: &[String]) -> Result<()> {
105        let paths: Vec<_> = objects
106            .iter()
107            .map(|obj_path| Ok(Path::from(obj_path.as_str())))
108            .collect();
109        let path_stream = stream::iter(paths);
110        let mut stream = self.blob_store.delete_stream(Box::pin(path_stream));
111        while let Some(res) = stream.next().await {
112            if let Err(e) = res {
113                match e {
114                    object_store::Error::NotFound { path: _, source: _ } => Ok(()),
115                    ref _other_error => Err(e),
116                }?
117            }
118        }
119        Ok(())
120    }
121}
micromegas_telemetry/blob_storage.rs

micromegas_telemetry/
blob_storage.rs
