iceberg_storage_opendal/
resolving.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Resolving storage that auto-detects the scheme from a path and delegates
19//! to the appropriate [`OpenDalStorage`] variant.
20
21use std::collections::HashMap;
22use std::sync::{Arc, RwLock};
23
24use async_trait::async_trait;
25use bytes::Bytes;
26use futures::StreamExt;
27use futures::stream::BoxStream;
28use iceberg::io::{
29    FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageConfig,
30    StorageFactory,
31};
32use iceberg::{Error, ErrorKind, Result};
33use serde::{Deserialize, Serialize};
34use url::Url;
35
36use crate::OpenDalStorage;
37#[cfg(feature = "opendal-s3")]
38use crate::s3::CustomAwsCredentialLoader;
39
40/// Schemes supported by OpenDalResolvingStorage
41pub const SCHEME_MEMORY: &str = "memory";
42pub const SCHEME_FILE: &str = "file";
43pub const SCHEME_S3: &str = "s3";
44pub const SCHEME_S3A: &str = "s3a";
45pub const SCHEME_S3N: &str = "s3n";
46pub const SCHEME_GS: &str = "gs";
47pub const SCHEME_GCS: &str = "gcs";
48pub const SCHEME_OSS: &str = "oss";
49pub const SCHEME_ABFSS: &str = "abfss";
50pub const SCHEME_ABFS: &str = "abfs";
51pub const SCHEME_WASBS: &str = "wasbs";
52pub const SCHEME_WASB: &str = "wasb";
53pub const SCHEME_HF: &str = "hf";
54
55/// Parse a URL scheme string.
56fn parse_scheme(scheme: &str) -> Result<&'static str> {
57    match scheme {
58        SCHEME_MEMORY => Ok("memory"),
59        SCHEME_FILE | "" => Ok("file"),
60        SCHEME_S3 | SCHEME_S3A | SCHEME_S3N => Ok("s3"),
61        SCHEME_GS | SCHEME_GCS => Ok("gcs"),
62        SCHEME_OSS => Ok("oss"),
63        SCHEME_ABFSS | SCHEME_ABFS | SCHEME_WASBS | SCHEME_WASB => Ok("azdls"),
64        SCHEME_HF => Ok("hf"),
65        s => Err(Error::new(
66            ErrorKind::FeatureUnsupported,
67            format!("Unsupported storage scheme: {s}"),
68        )),
69    }
70}
71
72/// Extract the scheme from a path URL.
73fn extract_scheme(path: &str) -> Result<&'static str> {
74    let url = Url::parse(path).map_err(|e| {
75        Error::new(
76            ErrorKind::DataInvalid,
77            format!("Invalid path: {path}, failed to parse URL: {e}"),
78        )
79    })?;
80    parse_scheme(url.scheme())
81}
82
83/// Build an [`OpenDalStorage`] variant for the given scheme and config properties.
84fn build_storage_for_scheme(
85    scheme: &'static str,
86    props: &HashMap<String, String>,
87    #[cfg(feature = "opendal-s3")] customized_credential_load: &Option<CustomAwsCredentialLoader>,
88) -> Result<OpenDalStorage> {
89    match scheme {
90        #[cfg(feature = "opendal-s3")]
91        "s3" => {
92            let config = crate::s3::s3_config_parse(props.clone())?;
93            Ok(OpenDalStorage::S3 {
94                config: Arc::new(config),
95                customized_credential_load: customized_credential_load.clone(),
96            })
97        }
98        #[cfg(feature = "opendal-gcs")]
99        "gcs" => {
100            let config = crate::gcs::gcs_config_parse(props.clone())?;
101            Ok(OpenDalStorage::Gcs {
102                config: Arc::new(config),
103            })
104        }
105        #[cfg(feature = "opendal-oss")]
106        "oss" => {
107            let config = crate::oss::oss_config_parse(props.clone())?;
108            Ok(OpenDalStorage::Oss {
109                config: Arc::new(config),
110            })
111        }
112        #[cfg(feature = "opendal-azdls")]
113        "azdls" => {
114            let config = crate::azdls::azdls_config_parse(props.clone())?;
115            Ok(OpenDalStorage::Azdls {
116                config: Arc::new(config),
117            })
118        }
119        #[cfg(feature = "opendal-fs")]
120        "file" => Ok(OpenDalStorage::LocalFs),
121        #[cfg(feature = "opendal-memory")]
122        "memory" => Ok(OpenDalStorage::Memory(crate::memory::memory_config_build()?)),
123        #[cfg(feature = "opendal-hf")]
124        "hf" => {
125            let config = crate::hf::hf_config_parse(props.clone())?;
126            Ok(OpenDalStorage::Hf {
127                config: Arc::new(config),
128            })
129        }
130        unsupported => Err(Error::new(
131            ErrorKind::FeatureUnsupported,
132            format!("Unsupported storage scheme: {unsupported}"),
133        )),
134    }
135}
136
137/// A resolving storage factory that creates [`OpenDalResolvingStorage`] instances.
138///
139/// This factory accepts paths from any supported storage system and dynamically
140/// delegates operations to the appropriate [`OpenDalStorage`] variant based on
141/// the path scheme.
142///
143/// # Serialization
144///
145/// Serialization fails when a custom S3 credential loader is configured because the loader holds
146/// process-local state that cannot be reconstructed in another process. Construct the factory
147/// without a custom loader before serializing it.
148///
149/// # Example
150///
151/// ```rust,ignore
152/// use std::sync::Arc;
153/// use iceberg::io::FileIOBuilder;
154/// use iceberg_storage_opendal::OpenDalResolvingStorageFactory;
155///
156/// let factory = OpenDalResolvingStorageFactory::new();
157/// let file_io = FileIOBuilder::new(Arc::new(factory))
158///     .with_prop("s3.region", "us-east-1")
159///     .build();
160/// ```
161#[derive(Clone, Debug, Serialize, Deserialize)]
162pub struct OpenDalResolvingStorageFactory {
163    /// Custom AWS credential loader for S3 storage.
164    #[cfg(feature = "opendal-s3")]
165    #[serde(
166        skip_deserializing,
167        skip_serializing_if = "Option::is_none",
168        serialize_with = "crate::serialize_custom_credential_loader"
169    )]
170    customized_credential_load: Option<CustomAwsCredentialLoader>,
171}
172
173impl Default for OpenDalResolvingStorageFactory {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179impl OpenDalResolvingStorageFactory {
180    /// Create a new resolving storage factory.
181    pub fn new() -> Self {
182        Self {
183            #[cfg(feature = "opendal-s3")]
184            customized_credential_load: None,
185        }
186    }
187
188    /// Set a custom AWS credential loader for S3 storage.
189    #[cfg(feature = "opendal-s3")]
190    pub fn with_s3_credential_loader(mut self, loader: CustomAwsCredentialLoader) -> Self {
191        self.customized_credential_load = Some(loader);
192        self
193    }
194}
195
196#[typetag::serde]
197impl StorageFactory for OpenDalResolvingStorageFactory {
198    fn build(&self, config: &StorageConfig) -> Result<Arc<dyn Storage>> {
199        Ok(Arc::new(OpenDalResolvingStorage {
200            props: config.props().clone(),
201            storages: RwLock::new(HashMap::new()),
202            #[cfg(feature = "opendal-s3")]
203            customized_credential_load: self.customized_credential_load.clone(),
204        }))
205    }
206}
207
208/// A resolving storage that auto-detects the scheme from a path and delegates
209/// to the appropriate [`OpenDalStorage`] variant.
210///
211/// Sub-storages are lazily created on first use for each scheme and cached
212/// for subsequent operations. Scheme aliases like `s3`/`s3a`/`s3n` map to
213/// the same canonical scheme, so they share a storage instance.
214#[derive(Debug, Serialize, Deserialize)]
215pub struct OpenDalResolvingStorage {
216    /// Configuration properties shared across all backends.
217    props: HashMap<String, String>,
218    /// Cache of canonical scheme to storage mappings.
219    #[serde(skip, default)]
220    storages: RwLock<HashMap<&'static str, Arc<OpenDalStorage>>>,
221    /// Custom AWS credential loader for S3 storage.
222    #[cfg(feature = "opendal-s3")]
223    #[serde(skip)]
224    customized_credential_load: Option<CustomAwsCredentialLoader>,
225}
226
227impl OpenDalResolvingStorage {
228    /// Resolve the storage for the given path by extracting the canonical scheme and
229    /// returning the cached or newly-created [`OpenDalStorage`].
230    fn resolve(&self, path: &str) -> Result<Arc<OpenDalStorage>> {
231        let scheme = extract_scheme(path)?;
232
233        // Fast path: check read lock first.
234        {
235            let cache = self
236                .storages
237                .read()
238                .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
239            if let Some(storage) = cache.get(&scheme) {
240                return Ok(storage.clone());
241            }
242        }
243
244        // Slow path: build and insert under write lock.
245        let mut cache = self
246            .storages
247            .write()
248            .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
249
250        // Double-check after acquiring write lock.
251        if let Some(storage) = cache.get(&scheme) {
252            return Ok(storage.clone());
253        }
254
255        let storage = build_storage_for_scheme(
256            scheme,
257            &self.props,
258            #[cfg(feature = "opendal-s3")]
259            &self.customized_credential_load,
260        )?;
261        let storage = Arc::new(storage);
262        cache.insert(scheme, storage.clone());
263        Ok(storage)
264    }
265}
266
267#[async_trait]
268#[typetag::serde]
269impl Storage for OpenDalResolvingStorage {
270    async fn exists(&self, path: &str) -> Result<bool> {
271        self.resolve(path)?.exists(path).await
272    }
273
274    async fn metadata(&self, path: &str) -> Result<FileMetadata> {
275        self.resolve(path)?.metadata(path).await
276    }
277
278    async fn read(&self, path: &str) -> Result<Bytes> {
279        self.resolve(path)?.read(path).await
280    }
281
282    async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
283        self.resolve(path)?.reader(path).await
284    }
285
286    async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
287        self.resolve(path)?.write(path, bs).await
288    }
289
290    async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
291        self.resolve(path)?.writer(path).await
292    }
293
294    async fn delete(&self, path: &str) -> Result<()> {
295        self.resolve(path)?.delete(path).await
296    }
297
298    async fn delete_prefix(&self, path: &str) -> Result<()> {
299        self.resolve(path)?.delete_prefix(path).await
300    }
301
302    async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> {
303        // Group paths by canonical scheme so each resolved storage receives a batch,
304        // avoiding repeated operator creation per path.
305        let mut grouped: HashMap<&'static str, Vec<String>> = HashMap::new();
306        while let Some(path) = paths.next().await {
307            let scheme = extract_scheme(&path)?;
308            grouped.entry(scheme).or_default().push(path);
309        }
310
311        for (_, paths) in grouped {
312            let storage = self.resolve(&paths[0])?;
313            storage
314                .delete_stream(futures::stream::iter(paths).boxed())
315                .await?;
316        }
317        Ok(())
318    }
319
320    fn new_input(&self, path: &str) -> Result<InputFile> {
321        Ok(InputFile::new(
322            Arc::new(self.resolve(path)?.as_ref().clone()),
323            path.to_string(),
324        ))
325    }
326
327    fn new_output(&self, path: &str) -> Result<OutputFile> {
328        Ok(OutputFile::new(
329            Arc::new(self.resolve(path)?.as_ref().clone()),
330            path.to_string(),
331        ))
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    #[cfg(feature = "opendal-s3")]
340    #[derive(Debug)]
341    struct EmptyCredentialLoader;
342
343    #[cfg(feature = "opendal-s3")]
344    impl crate::s3::ProvideCredential for EmptyCredentialLoader {
345        type Credential = crate::s3::AwsCredential;
346
347        async fn provide_credential(
348            &self,
349            _ctx: &reqsign_core::Context,
350        ) -> reqsign_core::Result<Option<Self::Credential>> {
351            Ok(None)
352        }
353    }
354
355    #[cfg(feature = "opendal-s3")]
356    #[test]
357    fn test_custom_credential_loader_serialization_fails() {
358        let factory = OpenDalResolvingStorageFactory::new()
359            .with_s3_credential_loader(CustomAwsCredentialLoader::new(EmptyCredentialLoader));
360        let file_io = iceberg::io::FileIOBuilder::new(Arc::new(factory)).build();
361
362        let err = file_io.serialize_all().unwrap_err();
363        assert!(
364            err.to_string()
365                .contains("custom AWS credential loaders cannot be serialized")
366        );
367    }
368
369    /// Builds a resolving storage with empty props, suitable for `resolve()`
370    /// calls that don't actually hit any backend.
371    fn empty_resolving_storage() -> OpenDalResolvingStorage {
372        OpenDalResolvingStorage {
373            props: HashMap::new(),
374            storages: RwLock::new(HashMap::new()),
375            #[cfg(feature = "opendal-s3")]
376            customized_credential_load: None,
377        }
378    }
379
380    #[cfg(feature = "opendal-s3")]
381    #[test]
382    fn test_resolve_s3_aliases_share_instance() {
383        let storage = empty_resolving_storage();
384
385        // All three S3-family schemes must collapse to a single cached
386        // `Arc<OpenDalStorage>` so that catalogs handing the resolver a mix
387        // of `s3://`, `s3a://`, `s3n://` paths don't rebuild operators.
388        let a = storage.resolve("s3://bucket/key").unwrap();
389        let b = storage.resolve("s3a://bucket/key").unwrap();
390        let c = storage.resolve("s3n://bucket/key").unwrap();
391
392        assert!(Arc::ptr_eq(&a, &b), "s3 and s3a should share one instance");
393        assert!(Arc::ptr_eq(&a, &c), "s3 and s3n should share one instance");
394    }
395
396    #[cfg(feature = "opendal-azdls")]
397    #[test]
398    fn test_resolve_azdls_aliases_share_instance() {
399        let storage = empty_resolving_storage();
400
401        let path_for = |scheme: &str| {
402            format!("{scheme}://myfs@myaccount.dfs.core.windows.net/path/to/file.parquet")
403        };
404
405        // All Azure schemes collapse onto one cached instance.
406        let abfss = storage.resolve(&path_for("abfss")).unwrap();
407        let abfs = storage.resolve(&path_for("abfs")).unwrap();
408
409        assert!(
410            Arc::ptr_eq(&abfss, &abfs),
411            "abfss and abfs should share one instance"
412        );
413    }
414}
iceberg_storage_opendal/resolving.rs

iceberg_storage_opendal/
resolving.rs
