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
36#[cfg(feature = "opendal-s3")]
37use crate::s3::CustomAwsCredentialLoader;
38use crate::{OpenDalClientConfig, OpenDalStorage};
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    let client_config = OpenDalClientConfig::from_properties(props)?;
90    match scheme {
91        #[cfg(feature = "opendal-s3")]
92        "s3" => {
93            let config = crate::s3::s3_config_parse(props.clone())?;
94            Ok(OpenDalStorage::S3 {
95                config: Arc::new(config),
96                customized_credential_load: customized_credential_load.clone(),
97                client_config,
98            })
99        }
100        #[cfg(feature = "opendal-gcs")]
101        "gcs" => {
102            let config = crate::gcs::gcs_config_parse(props.clone())?;
103            Ok(OpenDalStorage::Gcs {
104                config: Arc::new(config),
105                client_config,
106            })
107        }
108        #[cfg(feature = "opendal-oss")]
109        "oss" => {
110            let config = crate::oss::oss_config_parse(props.clone())?;
111            Ok(OpenDalStorage::Oss {
112                config: Arc::new(config),
113                client_config,
114            })
115        }
116        #[cfg(feature = "opendal-azdls")]
117        "azdls" => {
118            let config = crate::azdls::azdls_config_parse(props.clone())?;
119            Ok(OpenDalStorage::Azdls {
120                config: Arc::new(config),
121                client_config,
122            })
123        }
124        #[cfg(feature = "opendal-fs")]
125        "file" => Ok(OpenDalStorage::LocalFs { client_config }),
126        #[cfg(feature = "opendal-memory")]
127        "memory" => Ok(OpenDalStorage::Memory {
128            operator: crate::memory::memory_config_build()?,
129            client_config,
130        }),
131        #[cfg(feature = "opendal-hf")]
132        "hf" => {
133            let config = crate::hf::hf_config_parse(props.clone())?;
134            Ok(OpenDalStorage::Hf {
135                config: Arc::new(config),
136                client_config,
137            })
138        }
139        unsupported => Err(Error::new(
140            ErrorKind::FeatureUnsupported,
141            format!("Unsupported storage scheme: {unsupported}"),
142        )),
143    }
144}
145
146/// A resolving storage factory that creates [`OpenDalResolvingStorage`] instances.
147///
148/// This factory accepts paths from any supported storage system and dynamically
149/// delegates operations to the appropriate [`OpenDalStorage`] variant based on
150/// the path scheme.
151///
152/// # Serialization
153///
154/// Serialization fails when a custom S3 credential loader is configured because the loader holds
155/// process-local state that cannot be reconstructed in another process. Construct the factory
156/// without a custom loader before serializing it.
157///
158/// # Example
159///
160/// ```rust,ignore
161/// use std::sync::Arc;
162/// use iceberg::io::FileIOBuilder;
163/// use iceberg_storage_opendal::OpenDalResolvingStorageFactory;
164///
165/// let factory = OpenDalResolvingStorageFactory::new();
166/// let file_io = FileIOBuilder::new(Arc::new(factory))
167///     .with_prop("s3.region", "us-east-1")
168///     .build();
169/// ```
170#[derive(Clone, Debug, Serialize, Deserialize)]
171pub struct OpenDalResolvingStorageFactory {
172    /// Custom AWS credential loader for S3 storage.
173    #[cfg(feature = "opendal-s3")]
174    #[serde(
175        skip_deserializing,
176        skip_serializing_if = "Option::is_none",
177        serialize_with = "crate::serialize_custom_credential_loader"
178    )]
179    customized_credential_load: Option<CustomAwsCredentialLoader>,
180}
181
182impl Default for OpenDalResolvingStorageFactory {
183    fn default() -> Self {
184        Self::new()
185    }
186}
187
188impl OpenDalResolvingStorageFactory {
189    /// Create a new resolving storage factory.
190    pub fn new() -> Self {
191        Self {
192            #[cfg(feature = "opendal-s3")]
193            customized_credential_load: None,
194        }
195    }
196
197    /// Set a custom AWS credential loader for S3 storage.
198    #[cfg(feature = "opendal-s3")]
199    pub fn with_s3_credential_loader(mut self, loader: CustomAwsCredentialLoader) -> Self {
200        self.customized_credential_load = Some(loader);
201        self
202    }
203}
204
205#[typetag::serde]
206impl StorageFactory for OpenDalResolvingStorageFactory {
207    fn build(&self, config: &StorageConfig) -> Result<Arc<dyn Storage>> {
208        Ok(Arc::new(OpenDalResolvingStorage {
209            props: config.props().clone(),
210            storages: RwLock::new(HashMap::new()),
211            #[cfg(feature = "opendal-s3")]
212            customized_credential_load: self.customized_credential_load.clone(),
213        }))
214    }
215}
216
217/// A resolving storage that auto-detects the scheme from a path and delegates
218/// to the appropriate [`OpenDalStorage`] variant.
219///
220/// Sub-storages are lazily created on first use for each scheme and cached
221/// for subsequent operations. Scheme aliases like `s3`/`s3a`/`s3n` map to
222/// the same canonical scheme, so they share a storage instance.
223#[derive(Debug, Serialize, Deserialize)]
224pub struct OpenDalResolvingStorage {
225    /// Configuration properties shared across all backends.
226    props: HashMap<String, String>,
227    /// Cache of canonical scheme to storage mappings.
228    #[serde(skip, default)]
229    storages: RwLock<HashMap<&'static str, Arc<OpenDalStorage>>>,
230    /// Custom AWS credential loader for S3 storage.
231    #[cfg(feature = "opendal-s3")]
232    #[serde(skip)]
233    customized_credential_load: Option<CustomAwsCredentialLoader>,
234}
235
236impl OpenDalResolvingStorage {
237    /// Resolve the storage for the given path by extracting the canonical scheme and
238    /// returning the cached or newly-created [`OpenDalStorage`].
239    fn resolve(&self, path: &str) -> Result<Arc<OpenDalStorage>> {
240        let scheme = extract_scheme(path)?;
241
242        // Fast path: check read lock first.
243        {
244            let cache = self
245                .storages
246                .read()
247                .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
248            if let Some(storage) = cache.get(&scheme) {
249                return Ok(storage.clone());
250            }
251        }
252
253        // Slow path: build and insert under write lock.
254        let mut cache = self
255            .storages
256            .write()
257            .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
258
259        // Double-check after acquiring write lock.
260        if let Some(storage) = cache.get(&scheme) {
261            return Ok(storage.clone());
262        }
263
264        let storage = build_storage_for_scheme(
265            scheme,
266            &self.props,
267            #[cfg(feature = "opendal-s3")]
268            &self.customized_credential_load,
269        )?;
270        let storage = Arc::new(storage);
271        cache.insert(scheme, storage.clone());
272        Ok(storage)
273    }
274}
275
276#[async_trait]
277#[typetag::serde]
278impl Storage for OpenDalResolvingStorage {
279    async fn exists(&self, path: &str) -> Result<bool> {
280        self.resolve(path)?.exists(path).await
281    }
282
283    async fn metadata(&self, path: &str) -> Result<FileMetadata> {
284        self.resolve(path)?.metadata(path).await
285    }
286
287    async fn read(&self, path: &str) -> Result<Bytes> {
288        self.resolve(path)?.read(path).await
289    }
290
291    async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
292        self.resolve(path)?.reader(path).await
293    }
294
295    async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
296        self.resolve(path)?.write(path, bs).await
297    }
298
299    async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
300        self.resolve(path)?.writer(path).await
301    }
302
303    async fn delete(&self, path: &str) -> Result<()> {
304        self.resolve(path)?.delete(path).await
305    }
306
307    async fn delete_prefix(&self, path: &str) -> Result<()> {
308        self.resolve(path)?.delete_prefix(path).await
309    }
310
311    async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> {
312        // Group paths by canonical scheme so each resolved storage receives a batch,
313        // avoiding repeated operator creation per path.
314        let mut grouped: HashMap<&'static str, Vec<String>> = HashMap::new();
315        while let Some(path) = paths.next().await {
316            let scheme = extract_scheme(&path)?;
317            grouped.entry(scheme).or_default().push(path);
318        }
319
320        for (_, paths) in grouped {
321            let storage = self.resolve(&paths[0])?;
322            storage
323                .delete_stream(futures::stream::iter(paths).boxed())
324                .await?;
325        }
326        Ok(())
327    }
328
329    fn new_input(&self, path: &str) -> Result<InputFile> {
330        Ok(InputFile::new(
331            Arc::new(self.resolve(path)?.as_ref().clone()),
332            path.to_string(),
333        ))
334    }
335
336    fn new_output(&self, path: &str) -> Result<OutputFile> {
337        Ok(OutputFile::new(
338            Arc::new(self.resolve(path)?.as_ref().clone()),
339            path.to_string(),
340        ))
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::time::Duration;
347
348    use super::*;
349    use crate::OPENDAL_IO_TIMEOUT_MS;
350
351    #[cfg(feature = "opendal-s3")]
352    #[derive(Debug)]
353    struct EmptyCredentialLoader;
354
355    #[cfg(feature = "opendal-s3")]
356    impl crate::s3::ProvideCredential for EmptyCredentialLoader {
357        type Credential = crate::s3::AwsCredential;
358
359        async fn provide_credential(
360            &self,
361            _ctx: &reqsign_core::Context,
362        ) -> reqsign_core::Result<Option<Self::Credential>> {
363            Ok(None)
364        }
365    }
366
367    #[cfg(feature = "opendal-s3")]
368    #[test]
369    fn test_custom_credential_loader_serialization_fails() {
370        let factory = OpenDalResolvingStorageFactory::new()
371            .with_s3_credential_loader(CustomAwsCredentialLoader::new(EmptyCredentialLoader));
372        let file_io = iceberg::io::FileIOBuilder::new(Arc::new(factory)).build();
373
374        let err = file_io.serialize_all().unwrap_err();
375        assert!(
376            err.to_string()
377                .contains("custom AWS credential loaders cannot be serialized")
378        );
379    }
380
381    /// Builds a resolving storage with empty props, suitable for `resolve()`
382    /// calls that don't actually hit any backend.
383    fn empty_resolving_storage() -> OpenDalResolvingStorage {
384        OpenDalResolvingStorage {
385            props: HashMap::new(),
386            storages: RwLock::new(HashMap::new()),
387            #[cfg(feature = "opendal-s3")]
388            customized_credential_load: None,
389        }
390    }
391
392    #[test]
393    fn test_resolve_propagates_io_timeout() {
394        let mut storage = empty_resolving_storage();
395        storage
396            .props
397            .insert(OPENDAL_IO_TIMEOUT_MS.to_string(), "45000".to_string());
398
399        let paths: &[&str] = &[
400            #[cfg(feature = "opendal-memory")]
401            "memory:/key",
402            #[cfg(feature = "opendal-fs")]
403            "file:/key",
404            #[cfg(feature = "opendal-s3")]
405            "s3://bucket/key",
406            #[cfg(feature = "opendal-gcs")]
407            "gs://bucket/key",
408            #[cfg(feature = "opendal-oss")]
409            "oss://bucket/key",
410            #[cfg(feature = "opendal-azdls")]
411            "abfss://myfs@myaccount.dfs.core.windows.net/key",
412            #[cfg(feature = "opendal-hf")]
413            "hf://datasets/user/repo/key",
414        ];
415        for path in paths {
416            let resolved = storage.resolve(path).unwrap();
417            assert_eq!(
418                resolved.client_config().io_timeout(),
419                Duration::from_secs(45),
420                "{path}"
421            );
422        }
423    }
424
425    #[cfg(feature = "opendal-s3")]
426    #[test]
427    fn test_resolve_s3_aliases_share_instance() {
428        let storage = empty_resolving_storage();
429
430        // All three S3-family schemes must collapse to a single cached
431        // `Arc<OpenDalStorage>` so that catalogs handing the resolver a mix
432        // of `s3://`, `s3a://`, `s3n://` paths don't rebuild operators.
433        let a = storage.resolve("s3://bucket/key").unwrap();
434        let b = storage.resolve("s3a://bucket/key").unwrap();
435        let c = storage.resolve("s3n://bucket/key").unwrap();
436
437        assert!(Arc::ptr_eq(&a, &b), "s3 and s3a should share one instance");
438        assert!(Arc::ptr_eq(&a, &c), "s3 and s3n should share one instance");
439    }
440
441    #[cfg(feature = "opendal-azdls")]
442    #[test]
443    fn test_resolve_azdls_aliases_share_instance() {
444        let storage = empty_resolving_storage();
445
446        let path_for = |scheme: &str| {
447            format!("{scheme}://myfs@myaccount.dfs.core.windows.net/path/to/file.parquet")
448        };
449
450        // All Azure schemes collapse onto one cached instance.
451        let abfss = storage.resolve(&path_for("abfss")).unwrap();
452        let abfs = storage.resolve(&path_for("abfs")).unwrap();
453
454        assert!(
455            Arc::ptr_eq(&abfss, &abfs),
456            "abfss and abfs should share one instance"
457        );
458    }
459}
iceberg_storage_opendal/resolving.rs

iceberg_storage_opendal/
resolving.rs
